feat(tenancy): harden shared cloud runtime boundaries

This commit is contained in:
Junyan Qin
2026-07-20 01:47:42 +08:00
parent 3f01ffe63b
commit 627eb6b8ef
121 changed files with 18142 additions and 5615 deletions
+10 -2
View File
@@ -44,7 +44,9 @@ jobs:
runs-on: ubuntu-latest
services:
postgres:
image: postgres:16
# Release migration 0013 installs the pgvector extension in the shared
# business database; CI must exercise the same extension availability.
image: pgvector/pgvector:pg16
env:
POSTGRES_USER: langbot
POSTGRES_PASSWORD: langbot
@@ -75,4 +77,10 @@ jobs:
- name: Run PostgreSQL migration tests
env:
TEST_POSTGRES_URL: postgresql+asyncpg://langbot:langbot@localhost:5432/langbot_test
run: uv run pytest tests/integration/persistence/test_migrations_postgres.py -q --tb=short
run: >-
uv run pytest
tests/integration/persistence/test_migrations_postgres.py
tests/integration/persistence/test_pgvector_postgres.py
tests/integration/persistence/test_release_migration_postgres.py
tests/integration/persistence/test_plugin_identity_migration.py
-q --tb=short
+33 -5
View File
@@ -12,9 +12,13 @@ verification gates. Exact commands and observed results are recorded in the
- [x] Unrelated untracked files in either repository remain untouched.
- [x] Open-source startup cannot enable SaaS multi-workspace through edition flags or unsigned configuration.
## SaaS-only release gates
## SaaS activation gates
These items intentionally remain incomplete. The feature branch delivers the Core isolation kernel, not the closed SaaS product or a production Cloud v2 deployment.
These items intentionally remain incomplete. Some require additional Core
transaction/cutover primitives and others require the closed Control Plane or
deployment. The feature branch delivers the Core isolation kernel, not the
closed SaaS product or a production Cloud v2 deployment. Checked implementation
items later in this document do not supersede these gates.
- [ ] The closed Control Plane owns the global Account, Workspace, Membership, and Invitation directory.
- [ ] The closed placement service issues monotonic generations and leases for projected Workspaces.
@@ -22,10 +26,17 @@ These items intentionally remain incomplete. The feature branch delivers the Cor
- [ ] Tenant database writes hold a generation-aware shared transaction fence through commit, while placement cutovers take the exclusive fence.
- [ ] Business writes and non-transactional side effects use a generation-stamped outbox or equivalent publish fence.
- [ ] Durable object references survive a placement-generation change through stable published keys or an explicitly atomic key/reference migration.
- [ ] SaaS cells enforce tenant-safe egress and SSRF controls for Webhooks, providers, MCP servers, and every tenant-configurable outbound URL.
- [ ] The SaaS runtime pools enforce tenant-safe egress and SSRF controls for Webhooks, providers, MCP servers, and every tenant-configurable outbound URL.
- [ ] Entitlement checks, usage aggregation, subscription lifecycle, and billing are implemented in the closed Control Plane.
- [ ] OAuth state and directory projection use an atomic shared store suitable for horizontally scaled SaaS services.
- [ ] A greenfield Cloud v2 deployment is designed and validated independently of the legacy Space deployment scheme.
- [ ] The Plugin Runtime deployment provides delegated cgroup v2 and tenant-safe egress; the shared profile refuses to run without hard CPU, memory, and PID limits.
- [ ] The Box deployment provides an operator-owned quota provider that proves hard byte and inode limits for Workspace, Skill, root, tmp, and home storage.
- [ ] Core and Box Runtime mount the same durable volume and pass the authenticated marker challenge during startup and reconnect.
- [ ] Production provisions distinct migrator/runtime credentials and runs the implemented same-host/port/database release command as a one-shot Job, with tested orchestration retry, backup, and rollback procedures.
- [ ] Production PostgreSQL uses a dedicated cluster/endpoint, or a tested HBA/proxy policy proves the cluster-wide runtime credential can connect only to the target business database.
- [ ] Any future direct-migrator/pooler-runtime endpoint split is admitted only by a migrator-owned, runtime-read-only database cluster identity that the runtime role cannot spoof.
- [ ] Legacy pgvector migration failure and retry integration paths prove exact source-table RLS/FORCE restoration; the non-superuser, non-`BYPASSRLS` success path is already covered below.
- [ ] Multi-workspace is enabled in SaaS only after all closed Control Plane, deployment, and security gates pass.
## 1. Persistence foundation
@@ -50,6 +61,15 @@ These items intentionally remain incomplete. The feature branch delivers the Cor
- [x] SQLite destructive boundaries create verified, revision-aware backups and atomically restore after failure.
- [x] Migration can resume safely after interruption.
- [x] New installs and upgraded installs produce the same tenancy-kernel schema.
- [x] The first Cloud release pins migrator and runtime sessions to `public` with `current_schemas(false)` containing only that business schema; runtime-role/database `search_path` overrides are rejected.
- [x] Cloud sessions require `session_replication_role=origin`, `row_security=on`, and `lo_compat_privileges=off`; every persistent `pg_db_role_setting` applicable to the runtime role or current business database is rejected.
- [x] The release migrator grants the runtime role exact business-table DML, `alembic_version` read-only access, and business-sequence `USAGE/SELECT`, with no `WITH GRANT OPTION` or non-business object grants.
- [x] Every Cloud runtime startup revalidates the login role, current user, schema, effective/direct ACLs, ownership, memberships in all directions, column ACLs, routines, extensions, foreign objects, parameter ACLs, and other-schema access before serving traffic.
- [x] The business database requires `vector`, permits only `plpgsql`/`vector` extensions, forbids runtime extension ownership, and contains no foreign data wrapper, foreign server, or user mapping.
- [x] The runtime role and `PUBLIC` have no explicit routine or parameter ACL; the runtime owns no routine and cannot effectively execute any `SECURITY DEFINER` routine, including extension-owned routines.
- [x] PostgreSQL's default `PUBLIC TEMP` is documented and tested as a dedicated-business-database v1 compatibility exception; the migrator never grants `TEMP` directly to the runtime role.
- [x] Legacy pgvector migration succeeds as a non-superuser, non-`BYPASSRLS` source-table owner and restores mixed source-table RLS/FORCE states exactly.
- [ ] Legacy pgvector migration still needs explicit failure-and-retry integration coverage before SaaS activation.
## 2. Authentication and authorization
@@ -138,7 +158,10 @@ Each row type must have a non-null Workspace UUID, scoped indexes, scoped unique
- [x] Plugin installation and configuration are Workspace scoped.
- [x] Runtime control actions carry trusted Workspace binding and placement generation.
- [x] A plugin process or supervisor never serves multiple Workspaces in SaaS mode.
- [x] The Plugin Runtime supervisor is instance-scoped and intentionally serves multiple Workspaces.
- [x] Every plugin process is bound to exactly one Workspace, installation, generation, revision, and verified artifact digest.
- [x] Same-digest plugin code may be cached once, while worker processes and writable data remain isolated.
- [x] Same-digest plugin dependencies are prepared once in a Runtime-owned immutable environment and mounted read-only into each isolated worker; dependency failure is surfaced before launch and recorded per installation without blocking other desired-state recovery.
- [x] Host API derives Workspace from the connection, installation, and trusted action context, not plugin input.
- [x] Plugin get_bots, models, tools, vector, RAG, configuration, and messaging calls are scoped.
- [x] Plugin Workspace storage no longer uses owner default.
@@ -159,6 +182,11 @@ Each row type must have a non-null Workspace UUID, scoped indexes, scoped unique
- [x] Same-named Box sessions and processes cannot collide across Workspaces or placement generations.
- [x] Box relay and process I/O reject or retire stale generations.
- [x] External paths and privileged mounts cannot be supplied by an untrusted plugin.
- [x] Cloud attachment host I/O uses query UUIDs and link-free dirfd operations with bounded inode traversal.
- [x] Cloud Skill package paths are Runtime-owned, Workspace-scoped, read-only mounts; Python env/cache stays tenant-writable.
- [x] Skill ZIP preview/install rejects path escape, links, non-regular files, duplicate entries, excessive compression ratio, entry count, per-file size, and total size.
- [x] Cloud Box startup proves Core and Runtime see the same durable volume.
- [x] Cloud Box readiness fails until hard Workspace, Skill, ephemeral-storage, and inode quota capabilities are available.
## 6. SDK and protocol
@@ -223,7 +251,7 @@ Each row type must have a non-null Workspace UUID, scoped indexes, scoped unique
- [x] LangBot integration tests pass.
- [x] Frontend lint completes without errors and the production build passes.
- [x] SDK focused and full relevant tests pass.
- [x] Local SDK is installed into LangBot from the exact pushed SDK commit and cross-repo tests pass with no sync.
- [x] LangBot is pinned to the exact pushed SDK commit and cross-repo tests pass against that revision.
## 9. Real browser E2E
+37 -5
View File
@@ -35,12 +35,14 @@ This log records implementation choices made while delivering the Workspace arch
- Decision: Resource lookups always include Workspace UUID. A guessed UUID belonging to another Workspace returns 404; a visible resource with insufficient same-Workspace permission returns 403.
- Reason: This avoids leaking resource existence across tenants while preserving actionable same-tenant errors.
### Plugin Runtime binds to exactly one Workspace
### Plugin Runtime is shared; every plugin process is single-Workspace
- Decision: A trusted LangBot control connection binds a Plugin Runtime to one `instance_uuid`, `workspace_uuid`, `placement_generation`, and optional installation. Repeating the same binding is idempotent; rebinding is rejected. Plugin-supplied context fields are stripped.
- Reason: One untrusted plugin process must never become a cross-Workspace router.
- Compatibility: Older SDK payloads without context still deserialize, but Workspace storage and tenant Host APIs fail explicitly until a trusted binding is established.
- Startup fence: The Runtime does not launch or register plugins until `SET_RUNTIME_CONFIG` establishes the trusted Workspace binding. A cloud or multi-Workspace connector must be constructed for one explicit projected Workspace and generation; it never falls back to a migration-created local compatibility Workspace.
- Decision: One instance-scoped Plugin Runtime control plane serves all Workspaces in the logical LangBot instance. Each plugin installation still launches as its own nsjail worker with an immutable binding containing `instance_uuid`, `workspace_uuid`, `placement_generation`, `installation_uuid`, `runtime_revision`, and verified artifact digest. A worker never routes actions for another Workspace or installation, and plugin-supplied scope fields are stripped.
- Isolation: Plugin code is mounted read-only. Home, tmp, and data paths are installation-scoped; process, file-descriptor, file-size, CPU, memory, and PID limits come only from `data/config.yaml` (including native environment overrides), never from a plugin manifest. Cloud requires nsjail and delegated cgroup v2 hard limits or fails closed.
- Cost boundary: Identical verified package bytes share one digest-addressed code cache. A dependency environment is keyed by the artifact and requirements digests, Python ABI, Runtime version, and installer schema, then atomically published read-only for reuse. Installations and processes are not merged, even for the same plugin and version. Registration creates database desired state only; a worker is launched only for an enabled installation.
- Recovery: PostgreSQL installation desired state and durable binary storage are authoritative. Runtime reconnect performs an instance-wide full reconciliation, removes stale workers, and can replay a verified package after Runtime-local cache loss. Dependency preparation failure is recorded per installation with `dependency_prepare_failed`; it prevents that worker launch without blocking recovery of other desired installations, and the same revision can be retried.
- Compatibility: Older SDK payloads and legacy `data/plugins` remain an OSS-only bridge. Shared mode requires complete bindings and rejects incomplete context.
- Reason: Sharing the supervisor and immutable code cache removes per-Workspace service cost without turning an untrusted plugin process into a cross-tenant router.
### Invitation delivery does not require SMTP
@@ -176,3 +178,33 @@ This log records implementation choices made while delivering the Workspace arch
- Decision: Unhandled HTTP and webhook failures return a stable `internal_error` response and request ID, and expose that ID in `X-Request-Id`; the detailed exception is retained only in server logs correlated by the same ID. Explicit domain and validation errors keep their documented status and code.
- Secret boundary: Shared sanitization removes URL user information and masks sensitive query parameters before provider or MCP configuration is serialized, logged, or shown to a reader. Masked placeholders can be round-tripped by an authorized manager without replacing the stored secret.
- Reason: Tenant isolation is incomplete if framework exceptions, connection URLs, or configuration reads can export credentials across otherwise authorized interfaces.
### Box is one shared control plane with one admitted sandbox per Workspace
- Decision: The logical instance has one shared Box Runtime. A closed entitlement adapter projects generic `managed_sandbox` capability and `managed_sandbox_sessions` limit; Core and Runtime never branch on a plan name. An eligible Workspace receives at most one persistent logical `global` session, while each ordinary command remains a one-shot nsjail process. Managed processes and network are disabled in the first Cloud release.
- Storage boundary: Core and Runtime prove they see the same durable volume with an authenticated random-marker challenge. Attachments use opaque query UUID directories and link-free dirfd operations. Skill packages remain in the Runtime-owned Workspace store and enter a sandbox only as a read-only logical-name mount; Python environments and caches live in the tenant's writable Workspace.
- Resource boundary: Cloud readiness requires cgroup v2 plus hard byte and inode limits for Workspace files, Skill storage, root, tmp, and home. The existing directory scan is only a compatibility soft check. Plain nsjail reports these storage capabilities as unavailable, so Cloud Box intentionally cannot start until the greenfield deployment supplies and verifies a real quota provider.
- Archive boundary: Skill ZIP processing is bounded by compressed input, entry count, per-entry size, total uncompressed size, and compression ratio, and rejects links, non-regular entries, duplicates, and path escape before streaming extraction.
- Reason: A shared supervisor removes per-Workspace services and idle control-plane cost, but storage and process admission must still fail closed at the untrusted execution boundary.
### Cloud business data and vectors share one PostgreSQL schema
- Decision: SaaS uses one PostgreSQL business database and shared schema. Every tenant row has an explicit Workspace key; application scope is the first boundary and precise `ENABLE` plus `FORCE ROW LEVEL SECURITY` policies are the second. The Cloud runtime role must be non-owner and have neither superuser nor `BYPASSRLS`.
- Vector boundary: pgvector is the Cloud default in the same business database. Vectors use `(workspace_uuid, knowledge_base_uuid, vector_id)` identity, an untyped vector column with explicit checked dimension, and release-created partial expression indexes for the enabled dimensions. Cloud never falls back to Chroma or performs vector DDL at runtime.
- Transaction boundary: A tenant UoW binds `SET LOCAL` and SQL to one transaction. Long-running pipeline and streaming MCP execution carry a trusted transaction-free tenant scope; each database helper opens a short scoped transaction, avoiding a held pool connection during LLM or network waits. Detached tasks start only after commit and create their own short UoW; rollback cancels them.
- Schema boundary: The first release has exactly one business schema, `public`. Both migrator and runtime sessions must report `current_schema() = 'public'` and `current_schemas(false) = ARRAY['public']`; the runtime role and business database must not carry a `search_path` override. Runtime startup validates this before using the prepared schema and reruns the complete catalog and privilege validation on every process start; it never runs DDL.
- Session boundary: Both Cloud modes require `session_replication_role = 'origin'`, `row_security = 'on'`, and `lo_compat_privileges = 'off'`. Every persistent setting applicable to the runtime role or current business database in `pg_db_role_setting` is rejected, even if its present value appears safe; tenant context remains transaction-local application state rather than a persistent role/database override.
- Grant boundary: The migrator grants the runtime role direct `CONNECT` on the dedicated business database and `USAGE` on `public`; exact `SELECT, INSERT, UPDATE, DELETE` on every allowlisted business table; `SELECT` only on `alembic_version`; and exact `USAGE, SELECT` on business-owned sequences. It grants neither `CREATE`, `TRUNCATE`, `REFERENCES`, `TRIGGER`, sequence `UPDATE`, nor any privilege with `WITH GRANT OPTION`, and grants nothing on other relations or schemas.
- Role boundary: The runtime identity is a `LOGIN` role with no superuser, `BYPASSRLS`, `CREATEDB`, `CREATEROLE`, or replication attribute; no role membership in any direction, including acting as grantor; no ownership of the business database, `public` schema, relations, sequences, routines, or extensions; no column ACLs; and no use, create, or ownership in another non-system schema. Neither the runtime role nor `PUBLIC` may have an explicit routine or parameter ACL, and the runtime role may not effectively execute any `SECURITY DEFINER` routine, including an extension-owned one. PostgreSQL's default `TEMP` privilege inherited from `PUBLIC` is an explicit first-release compatibility decision for this dedicated business database, not a direct runtime-role grant.
- Catalog boundary: The business database must contain `vector` and may contain no extension other than `plpgsql` and `vector`; the runtime role owns neither. It contains no foreign data wrapper, foreign server, or user mapping. These checks remove catalog-level escape paths without forbidding the ordinary implicit execution of non-`SECURITY DEFINER` built-in routines.
- Migration boundary: In the first release, the migrator and runtime URLs must name the same normalized PostgreSQL host, port, and database while using different roles. The migrator owns the application schema, establishes the exact allowlist above, and validates both required access and every prohibited escalation path before releasing the advisory lock. An exact Alembic head, RLS checks, and pgvector table/index/constraint validation remain mandatory; concurrent Jobs fail explicitly and are retried by orchestration.
- Deployment boundary: PostgreSQL roles are cluster-wide, while the in-database audit proves only the target business database contract. SaaS production must therefore use a dedicated PostgreSQL cluster or endpoint that exposes only this business database to the runtime credential, or enforce and test an HBA/proxy policy proving that the credential cannot connect to any other database. This external connectivity proof is still an incomplete SaaS activation gate.
- Endpoint evolution: A future deployment may use a direct endpoint for migrations and a pooler endpoint for runtime traffic. That topology may relax literal host/port equality only after both endpoints are proven to reach the same database through a database-internal, migrator-owned cluster identity that the runtime role can read but cannot create, alter, or spoof.
- Legacy pgvector boundary: Revision 0013 records the exact `ENABLE` and `FORCE ROW LEVEL SECURITY` state of each RLS-protected source table, temporarily suspends those source policies as their table owner inside the migration transaction, and restores every table to its recorded state in `finally`. The migrator does not require superuser or `BYPASSRLS` for this data move.
- Activation gate: The shared schema, pgvector adapter, and database-local runtime audit are implemented, but the external cluster/endpoint or HBA/proxy connectivity proof remains deployment work. Ordinary business writes also do not yet hold a generation-aware fence through commit; a generation-stamped outbox (or equivalent atomic publish fence) and stable durable-object references across generation cutover remain required before SaaS activation.
- Reason: Sharing one database and pool keeps marginal Workspace cost low, while transaction-local context and RLS prevent that shared storage from becoming shared authority.
### stdio MCP has an independent deployment gate
- Decision: `mcp.stdio.enabled` is independent of Box availability and entitlement. OSS defaults it on for compatibility; Cloud requires it off at bootstrap and enforces the same gate on create, update, test, startup loading, and final runtime execution.
- Reason: Treating Box availability as stdio permission would silently create another persistent `mcp-shared` sandbox for each Workspace and bypass the one-sandbox subscription and cost boundary.
@@ -1,12 +1,13 @@
# Cloud v2 多租户架构决策与待决策项
状态:`PARTIALLY_DECIDED — MVP target confirmed, implementation pending`
状态:`DECIDED — Core isolation kernel implemented; SaaS activation gates remain`
创建日期:2026-07-19
最近更新:2026-07-19
最近更新:2026-07-20
本文记录 Cloud v2 多租户架构中已经确认的首期决策、明确淘汰的方案和仍需在后续阶段决定的扩展项。
“已决定”表示实现目标已经确定,不表示代码已经完成。正式实现时还需要把结论同步到
[implementation-decisions.md](./implementation-decisions.md)、主架构、实施清单、配置模板和协议设计。
本文同时记录实现状态。“实现完成”仅指开源 Core/SDK 的隔离内核和 fail-closed 门禁,
不表示闭源 Control Plane、计费或 Cloud v2 部署已经可上线。最终实现选择同步记录在
[implementation-decisions.md](./implementation-decisions.md),剩余发布门禁记录在实施清单和验证报告。
## 0. 已确认的 SaaS 拓扑前提
@@ -25,12 +26,12 @@
### 0.1 本轮确认的首期决策
| 编号 | 结论 | 首期状态 |
| ----- | -------------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------- |
| D-001 | 一个共享 Plugin Runtime 控制面;每个运行中的 plugin installation 独占一个 nsjail 子进程;只有 digest 相同且已验证的代码/依赖 artifact 可以只读共享 | `DECIDED — implementation pending` |
| D-002 | 一个共享 Box RuntimeCloud 固定使用 nsjail;符合套餐的 Workspace 最多一个持久 `global` 逻辑 sandbox,普通执行按需启动 nsjail 进程 | `DECIDED — implementation pending` |
| D-003 | SaaS 业务数据使用 PostgreSQL shared schema、应用层作用域和 RLS 双重隔离;pgvector 使用同一 PostgreSQL,作为 SaaS 默认向量后端 | `DECIDED — implementation pending` |
| D-004 | stdio MCP 与 Box availability 解耦;Cloud v2 首期强制关闭 stdio MCP,避免为每个 Workspace 创建额外的 `mcp-shared` persistent sandbox | `DECIDED — implementation pending` |
| 编号 | 结论 | 首期状态 |
| ----- | --------------------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------- |
| D-001 | 一个共享 Plugin Runtime 控制面;每个运行中的 plugin installation 独占一个 nsjail 子进程;只有 digest 相同且已验证的代码 artifact 可以只读共享 | `IMPLEMENTED — real Linux/egress deployment gate pending` |
| D-002 | 一个共享 Box RuntimeCloud 固定使用 nsjail;符合套餐的 Workspace 最多一个持久 `global` 逻辑 sandbox,普通执行按需启动 nsjail 进程 | `IMPLEMENTED FAIL-CLOSED — hard filesystem quota provider pending` |
| D-003 | SaaS 业务数据使用 PostgreSQL shared schema、应用层作用域和 RLS 双重隔离;pgvector 使用同一 PostgreSQL,作为 SaaS 默认向量后端 | `PARTIALLY IMPLEMENTED — transaction/outbox/deployment gates remain` |
| D-004 | stdio MCP 与 Box availability 解耦;Cloud v2 首期强制关闭 stdio MCP,避免为每个 Workspace 创建额外的 `mcp-shared` persistent sandbox | `IMPLEMENTED` |
Workspace 的具体创建、释放、数据导出和单 Workspace 恢复机制不在本轮决定;本文只保证这些后续能力不会改变稳定的
`workspace_uuid`,也不会要求重建租户专属部署。
@@ -81,21 +82,18 @@ Core 不能继承 Plugin Runtime 所需的 nsjail/cgroup 权限。Box Runtime
## 3. D-001Plugin Runtime 多租户控制面
状态:`DECIDED — MVP target confirmed, implementation pending`
状态:`IMPLEMENTED — Cloud deployment verification pending`
### 3.1 当前实现事实
### 3.1 已实现的基础
- SDK `RuntimeContext` 和 Core `PluginRuntimeConnector` 当前都绑定一个 WorkspaceRuntime 只有一个全局 PluginManager
- 插件已经是直接子进程,但当前由 `asyncio.create_subprocess_exec` 启动,没有 nsjail 或同等级内核隔离
- 插件包目录当前按 author/name 组织,依赖安装会修改 Runtime 的全局 Python 环境,不能直接作为共享多租户写入边界
- 生产 `run-plugin` 路径仍会读取插件目录中的 `.env`;公开 SaaS 不能让 artifact 自行注入 Runtime secret
- 当前控制协议已有 `SET_RUNTIME_CONFIG`,可以由 Core 把验证后的实例级 worker policy 下发给 Runtime
但该控制 handler 当前仍要求单 Workspace context;目标协议必须先解除控制连接的 Workspace 绑定。
- LangBot 的配置加载器原生支持把 `A__B__C` 环境变量映射到 `data/config.yaml` 的嵌套字段;
数值和布尔字段必须先出现在配置模板中,才能稳定保留类型
- 现有镜像和 Box backend 已包含 nsjail、mount namespace、private `/proc`、rlimit 和 cgroup v2 相关实现,可复用隔离参数生成逻辑。
- 当前 installation identity 及持久化模型还没有完整的随机 `installation_uuid``artifact_digest``runtime_revision`
和 desired-state 字段;它们是目标模型,不是现有能力。
- Plugin Runtime 控制连接只绑定稳定实例身份;一个 Supervisor 通过完整 installation binding 管理多个 Workspace
- 每个 enabled installation 使用独立 nsjail worker;代码只读,home/tmp/data 私有,shared profile 不读取 artifact `.env`
- 实例级 `PluginWorkerPolicy` 由 Core 的 `data/config.yaml` 下发,支持原生环境变量覆写;manifest 不能覆盖
- `installation_uuid``artifact_digest``runtime_revision` 已持久化并进入 desired-state、注册、Host API 和 generation/revision fence
- 已验证 `.lbpkg` 先进入 Workspace-scoped durable binary storageRuntime 本地缓存丢失后可由 Core replay。
- 相同 digest 的代码和 Runtime 准备的只读依赖环境可以共享,但 worker、运行时写入、配置和数据不合并;
dependency preparation 在启动 worker 前完成,失败会进入明确的 installation failed 状态。
- Cloud shared profile 强制 Linux nsjail`plugin.worker.require_hard_limits=true` 时 cgroup v2 delegation 不可用会启动失败
### 3.2 已确认的不变量
@@ -132,7 +130,7 @@ Core 不能继承 Plugin Runtime 所需的 nsjail/cgroup 权限。Box Runtime
```text
data/plugin-runtime/
├── artifacts/sha256/<artifact_digest>/code/ # digest 校验后只读共享
├── envs/<environment_digest>/ # 可选,只读共享依赖环境
├── environments/sha256/<environment_digest>/ # 原子发布、只读共享依赖环境
└── installations/<installation_uuid>/
├── home/ # 私有可写
├── tmp/ # 私有可写、可清理
@@ -140,8 +138,10 @@ data/plugin-runtime/
```
- artifact 只有在内容摘要和完整性校验一致时才允许共享,不能只凭 author/name/version 复用目录;
cache 接受哪些签名/来源以及如何撤销和 GC,仍是实现前需要补充的供应链规则
- artifact 与共享依赖环境以只读 mount 进入 nsjailinstallation 的 home/tmp/data 使用独立可写 mount。
cache 接受签名/来源撤销和 GC 规范属于后续发布规则,不改变本轮基于已验证 digest 的只读共享边界
- artifact 与按环境摘要构建的共享依赖环境以只读 mount 进入 nsjailinstallation 的 home/tmp/data 使用独立可写 mount。
环境摘要包含 artifact、requirements、Python ABI、Runtime 版本和 installer schema。依赖只能从已验证 artifact 的 PEP 508 声明构建,
index/trusted-host 只由实例配置控制;构建在独立 nsjail 的临时路径中完成并在成功后原子发布,失败或并发安装不能留下可见半成品。
- 插件 cwd 可以是其私有 mount namespace 内的只读 `/plugin`,不要求为每个 installation 复制代码;
必须私有的是 home/tmp/data 等所有可写路径。
- nsjail 必须启用 mount、PID、IPC、UTS 和 private `/proc` 等必要 namespace,插件不能枚举或 signal 其他插件及 Runtime 进程,
@@ -164,6 +164,7 @@ plugin:
max_pids: 128
max_open_files: 256
max_file_size_mb: 512
require_hard_limits: true # Cloud; OSS defaults false
```
配置文件路径为 `data/config.yaml`,沿用现有原生环境变量覆写:
@@ -173,6 +174,7 @@ plugin:
- `PLUGIN__WORKER__MAX_PIDS`
- `PLUGIN__WORKER__MAX_OPEN_FILES`
- `PLUGIN__WORKER__MAX_FILE_SIZE_MB`
- `PLUGIN__WORKER__REQUIRE_HARD_LIMITS`
Core 启动时校验配置并通过现有 `SET_RUNTIME_CONFIG` 下发不可变 `PluginWorkerPolicy`
Runtime 不读取另一份环境变量配置,避免两个配置源不一致。CPU、内存和 PID 使用 cgroup 硬限制,
@@ -203,22 +205,23 @@ installation data 的总空间硬配额需要 filesystem project quota 或独立
- 修改 manifest 不能改变任何资源上限。
- 旧 generation/revision 的回调、消息、副作用和存储访问全部失败关闭。
- Runtime 重启能从业务 desired state 恢复,不依赖本地进程表作为权威真相。
- requirements 中存在 Runtime 基础镜像未预装的包时,Supervisor 仍能先完成共享依赖环境准备再启动 worker;
安装失败不会留下持续重启的半启动进程,也不会影响同 digest 已就绪环境的其他 installation。
## 4. D-002Box 多租户控制面和套餐边界
状态:`DECIDED — MVP target confirmed, implementation pending`
状态:`IMPLEMENTED FAIL-CLOSED — production quota provider pending`
### 4.1 当前实现事实
### 4.1 已实现的基础
- Box 已经按 `instance_uuid + workspace_uuid` 派生持久 namespace,并按 generation 派生运行 namespace
单个 server/control connection 已有服务多个 Workspace 的数据结构基础
- 当前 nsjail backend 的 session 是“逻辑 session + 共享 host workspace 目录”。每次普通 `exec` 都启动一个 one-shot nsjail 进程,
命令结束后进程退出;managed process 才会保持对应 nsjail 进程运行
- `persistent=True` 会让 session 跳过 300 秒 TTL reaper 和普通 shutdown 清理,但 Runtime 重启不会恢复内存中的 session 或运行进程
- 现有 `{global}` 强制 session ID 只覆盖部分 Agent 执行入口,不能单独保证所有入口最多一个 session,也不会自动设置 `persistent=True`
- 当前没有 `max_sessions_per_workspace`,调用其他入口仍可能使用不同逻辑 session ID,必须在 Box admission 层补最终门禁
- nsjail 文件机制不是对象存储同步算法,而是把 Box Runtime 容器中的 Workspace host path bind mount 到 sandbox 的 `/workspace`
- 当前 nsjail 在不可写/不可用的 cgroup v2 环境会降级并告警,无法保证共享 SaaS 的硬内存隔离。
- 共享 Box 控制连接可服务多个 Workspace;所有操作绑定 instance、Workspace generationRuntime namespace 由可信 context 派生。
- 短期 `SandboxAdmissionGrant`、revision tombstone 和原子 session admission 强制每个合资格 Workspace 最多一个 `global` persistent sessionmanaged process 固定为零
- Core 与 Runtime 使用认证 host-control challenge 校验同一个 durable volume,而不是比较路径字符串;不一致时启动和重连失败。
- Cloud skill 只传逻辑名称;Runtime 从 Workspace-scoped store 解析只读包路径,Python env/cache 写入租户自己的 `/workspace/.skill-envs`
- ZIP 安装限制压缩输入、条目、单项、总解压量和压缩比,采用流式解压并拒绝 link、非普通文件、重复项和路径逃逸
- 附件 host path 使用 query UUID 和 dirfd/openat/O_NOFOLLOWCloud replica 启动不再全局清理其他请求目录,遍历和删除有 inode 预算
- grant-enforced readiness 强制 cgroup、namespace、mount、共享卷、Workspace hard quota、Skill hard quota、ephemeral storage 和 inode quota 全部被证明
普通 nsjail backend 对尚未实现的硬磁盘能力明确返回 false,因此当前 Cloud Box 会按设计拒绝启动,直到新部署提供真实 quota provider
### 4.2 首期套餐与 entitlement 模型
@@ -307,23 +310,50 @@ installation data 的总空间硬配额需要 filesystem project quota 或独立
- 非 Pro、entitlement 缺失/过期及伪造 plan 的 API 直调全部失败关闭。
- TTL 不回收 persistent sessionRuntime 重启后进程和临时目录失效,但 `/workspace` 保留并能在下一次使用时安全重建。
- 两个 Workspace 的文件、进程、session、attach token 和 generation 完全隔离;network/managed-process 请求在首期失败关闭。
- Core 与 Box Runtime 使用同一路径的共享持久卷;filesystem quota 在 Box Runtime 写入点真实生效,现有目录扫描不被当作硬配额
- cgroup hard limit 不可用时 Cloud Box Runtime readiness 失败
- Core 与 Box Runtime 通过随机 marker challenge 证明同一共享持久卷;只配置相同路径字符串不算通过
- Workspace、Skill store、ephemeral root/tmp/home 的 byte quota 与 inode quota 在写入点真实生效;现有目录扫描不被当作硬配额
- cgroup 或任一硬存储能力不可用时 Cloud Box Runtime readiness 失败;普通 nsjail 因此不会被误当成 production-ready provider。
## 5. D-003SaaS PostgreSQL 与 pgvector
状态:`DECIDED — MVP target confirmed, implementation pending`
状态:`PARTIALLY IMPLEMENTED — shared schema/pgvector complete; SaaS transaction and deployment gates remain`
当前实现仍有明确差距:Cloud 模板默认 Chromapgvector adapter 使用独立 engine、全局 `id` 和固定 `vector(1536)`
没有持久化 `workspace_uuid/knowledge_base_uuid`,并会在应用启动时执行 `CREATE EXTENSION/create_all`
当前 persistence helper 也不能保证 `SET LOCAL` 与业务查询处于同一 transaction。以下均是目标约束,不是现状描述。
当前分支已实现 PostgreSQL shared schema、transaction-local scope、FORCE RLS、Cloud runtime 非 DDL 模式、
同业务数据库 pgvector、显式向量维度和 tenant-scoped vector 主键。一次性 migrator 使用独立凭据、advisory lock
负责建立并校验 runtime role 的最小权限,并完成全量 schema 验证;
普通业务写入贯穿 commit 的 generation-aware fence、与外部副作用同事务的 outbox,以及 generation cutover 后稳定的 durable object 引用尚未实现。
这些 Core 事务原语与生产 Job、凭据发放、备份和回滚流程,以及 runtime credential 的跨 database 连接隔离证明一起,
都是 Cloud v2 的 SaaS activation gate。
### 5.1 已确认的数据库边界
- PostgreSQL 是 SaaS 的业务数据库,不把它扩展成通用 Runtime coordinator、Box session directory 或新控制面数据库。
- M0 使用一个 PostgreSQL business database/shared schema 承载全部 Workspace。创建 Workspace 不创建 database、schema、role 或专属连接池。
- 首版 migrator URL 和 runtime URL 必须归一化到同一 host、port 和 database,但必须使用不同 role。
未来如果 migrator 使用 direct endpoint、runtime 使用 pooler endpoint,只能在两个端点都验证同一个数据库内部 cluster identity 后放宽 host/port 相等。
该 identity 由 migrator 所有并固定,runtime role 只能读取,不能创建、修改或伪造。
- 首版唯一业务 schema 固定为 `public`。migrator 和 runtime 连接都必须满足
`current_schema() = 'public'``current_schemas(false) = ARRAY['public']`;禁止 runtime role 级和 business database 级 `search_path` 覆写。
- migrator 和 runtime session 的安全值固定为 `session_replication_role=origin``row_security=on`
`lo_compat_privileges=off`。runtime role 或当前 business database 作用域内只要存在任意 `pg_db_role_setting` 持久化设置就失败关闭,
即使该设置当前看似等于安全值也不接受;tenant context 只能通过事务内 `SET LOCAL` 建立。
- 业务行显式携带 `workspace_uuid`Repository/Service 的应用层 scope 是第一道边界,PostgreSQL RLS 是第二道边界。
- 应用角色不得拥有 `BYPASSRLS`,关键租户表使用 `FORCE ROW LEVEL SECURITY`migration/repair/audit 使用独立受控角色。
- runtime role 的直接 ACL 固定为:business database 的 `CONNECT``public``USAGE`、全部 allowlisted business table 的
`SELECT/INSERT/UPDATE/DELETE``alembic_version` 的只读 `SELECT`,以及业务表自有 sequence 的 `USAGE/SELECT`
不授予 database/schema `CREATE`、table `TRUNCATE/REFERENCES/TRIGGER`、sequence `UPDATE`、其他对象权限或任何 `WITH GRANT OPTION`
- runtime role 必须是 `LOGIN`,但不得具有 superuser、`BYPASSRLS``CREATEDB``CREATEROLE` 或 replication 属性;
不得在 role membership 中以 granted role、member 或 grantor 任一方向出现;不得拥有 database、schema、table、view、sequence、routine 或 extension
不得持有 column ACL,也不得使用、创建或拥有其他非系统 schema。
- business database 必须安装 `vector`,且 extension catalog 只允许 `plpgsql``vector`;不得存在 FDW、foreign server 或 user mapping。
runtime role 和 `PUBLIC` 都不得有显式 routine ACL 或 parameter `SET/ALTER SYSTEM` ACLruntime role 不得有效执行任何
`SECURITY DEFINER` routine,包括被 allowlisted extension 收编的 routine。普通非 `SECURITY DEFINER` 内建函数的隐式执行权限不在此禁令内。
- PostgreSQL 新 database 默认向 `PUBLIC` 提供的 `TEMP` 是首版在专用业务 database 上明确接受的兼容性决定,
不是 migrator 对 runtime role 的直接 grant;首版不得据此把业务 database 与不受信任工作负载混用。
migrator 在释放 advisory lock 前建立并校验上述精确 allowlistCloud runtime 每次启动都必须重新完成 schema、身份、有效权限和 catalog 负向校验,发现 drift 立即失败关闭。
- PostgreSQL role 是 cluster-wide identity,当前 database 内的 catalog audit 不能证明同一 credential 无法连接 cluster 中的其他 database。
SaaS 生产环境必须使用仅向该 credential 暴露目标 business database 的专用 PostgreSQL cluster/endpoint
或通过已验证的 HBA/proxy policy 证明该 credential 只能连接目标 business database;这项部署隔离仍是未完成的 activation gate。
- 关键租户表使用 `FORCE ROW LEVEL SECURITY`migration/repair/audit 使用独立受控 migrator role。
- 每个租户事务通过 `SET LOCAL` 设置 tenant context,并由统一 `TenantUnitOfWork` 保证设置 context 和业务查询使用同一事务/连接。
禁止使用连接级 session variable 或 `search_path`,避免连接池、PgBouncer、异常回滚和后台任务串租户。
- 一个 `TenantUnitOfWork` 只能访问一个 Workspace。业务写入与对应 business outbox 在同一事务中提交;
@@ -344,6 +374,9 @@ installation data 的总空间硬配额需要 filesystem project quota 或独立
`CHECK (vector_dims(embedding) = embedding_dimension)` 校验;release migration 为允许的维度创建带 dimension predicate 的 expression/partial ANN index。
知识库/model 元数据必须选择已启用维度,写入和查询 mismatch 或未启用维度时失败关闭,不能截断、补齐、退化为无界扫描或换后端。
- `vector` extension、表、索引和 RLS 由 release migration 创建。应用进程不在启动时执行 DDL。
- 0013 如需搬迁 legacy pgvector 数据,migrator 作为源表 owner 先记录每个受保护源表的 `ENABLE/FORCE RLS` 状态,
仅在同一 migration transaction 内临时暂停 RLS,并在 `finally` 中精确恢复各表原状态。
该流程不依赖 superuser 或 `BYPASSRLS`,也不允许在迁移事务外留下已禁用的 RLS。
### 5.3 候选拓扑与未来演进
@@ -357,6 +390,7 @@ installation data 的总空间硬配额需要 filesystem project quota 或独立
M0 不提前增加始终返回 `primary` 的 resolver、shard router 或 shard binding。
P1 的 resolver、映射、在线迁移、连接池预算、shard-affine replica 和 dedicated shard 细节等到出现容量、地域或合规需求时再设计。
在此之前,direct endpoint 与 pooler endpoint 分离只能通过数据库内部、runtime 不可伪造的 cluster identity 开启,不使用 DNS 名、数据库名或配置声明代替。
### 5.4 备份与生命周期边界
@@ -373,6 +407,15 @@ P1 的 resolver、映射、在线迁移、连接池预算、shard-affine replica
- 故意遗漏应用层 Workspace filter 时,RLS 仍阻止跨租户读写。
- 连接池/事务池复用、异常回滚、并发请求和后台任务不会残留 tenant context;如部署 PgBouncer,也必须覆盖 transaction pooling。
- migration 对 shared schema 只执行一次,不产生 Workspace 级 schema drift;应用启动角色不能执行 DDL。
- 首版拒绝 host、port 或 database 不同的 migrator/runtime URL,并拒绝相同 rolemigrator/runtime 都只解析到 `public`
migrator 在迁移后完成精确 table/sequence/`alembic_version` ACL grant 和正反向 role 校验,runtime 每次启动重新校验。
- runtime role 没有任一方向的 role membership、`WITH GRANT OPTION``search_path` 覆写、对象所有权、其他 schema 访问或非业务对象权限;
专用业务 database 上可继承 PostgreSQL 默认 `PUBLIC TEMP`,但 runtime role 没有直接 `TEMP` ACL。
- migrator/runtime session GUC 保持 `session_replication_role=origin``row_security=on``lo_compat_privileges=off`
runtime role/当前 database 没有任何 `pg_db_role_setting`extension 仅为 `plpgsql/vector` 且 runtime 不拥有 extension
database 中没有 FDW/server/user mapping、runtime 或 `PUBLIC` 显式 routine/parameter ACL、runtime-owned routine 或 runtime 可执行的 `SECURITY DEFINER` routine。
- 生产 deployment 证明 cluster-wide runtime credential 只能连接目标 business database;专用 cluster/endpoint 或 HBA/proxy 隔离未经验证前不得启用 SaaS。
- legacy pgvector 搬迁在非 superuser、非 `BYPASSRLS` 的 table-owner migrator 下可成功,成功、异常和重试路径都精确恢复所有源表的 RLS/FORCE 状态。
- 业务写入和对应 business outbox 在同一事务内具备可证明的提交顺序;外部 generation/fencing token 校验失败时不产生写入。
- pgvector 使用真实 PostgreSQL 集成测试覆盖:两个 Workspace 使用相同 `vector_id`、猜测其他 Workspace ID、
故意遗漏 scope、连接复用、CRUD 和后台任务,全部不能越权。
@@ -381,13 +424,13 @@ P1 的 resolver、映射、在线迁移、连接池预算、shard-affine replica
## 6. D-004stdio MCP 独立开关
状态:`DECIDED — MVP target confirmed, implementation pending`
状态:`IMPLEMENTED`
### 6.1 当前问题
### 6.1 已修复的原问题
- 当前 stdio MCP 的启用条件只检查 transport 为 `stdio` 且 Box available,没有独立 feature gate。
- 所有 stdio MCP 当前使用固定的 `mcp-shared` 逻辑 session,并强制 `persistent=True`
- 如果多租户 Cloud 通过 `box.enabled` 开放能力,每个配置 stdio MCP 的 Workspace 都会额外保留一个 persistent sandbox
- 修复前,stdio MCP 的启用条件只检查 transport 为 `stdio` 且 Box available,没有独立 feature gate。
- 修复前,所有 stdio MCP 使用固定的 `mcp-shared` 逻辑 session,并强制 `persistent=True`
- 在该旧逻辑下,如果多租户 Cloud 通过 `box.enabled` 开放能力,每个配置 stdio MCP 的 Workspace 都会额外保留一个 persistent sandbox
绕过“每 Workspace 最多一个 managed `global` sandbox”的成本和套餐边界。
### 6.2 首期决定
+1 -1
View File
@@ -70,7 +70,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@a1544b6b38a37ba72e3284f2836618144f0742c1",
"langbot-plugin @ git+https://github.com/langbot-app/langbot-plugin-sdk.git@044536f3720a2a8424d92f78934b9623da9f7f1d",
"asyncpg>=0.30.0",
"line-bot-sdk>=3.19.0",
"matrix-nio>=0.25.2",
+39 -17
View File
@@ -20,8 +20,7 @@ asciiart = r"""
"""
async def main_entry(loop: asyncio.AbstractEventLoop):
"""Main entry point for LangBot"""
def _build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description='LangBot')
parser.add_argument(
'--standalone-runtime',
@@ -36,7 +35,20 @@ async def main_entry(loop: asyncio.AbstractEventLoop):
default=False,
)
parser.add_argument('--debug', action='store_true', help='Debug mode / 调试模式', default=False)
args = parser.parse_args()
subparsers = parser.add_subparsers(dest='command')
migrate_parser = subparsers.add_parser('migrate', help='Run an operator-only database migration')
migrate_parser.add_argument(
'--cloud',
action='store_true',
required=True,
help='Migrate and validate the Cloud PostgreSQL business database',
)
return parser
async def main_entry(loop: asyncio.AbstractEventLoop):
"""Main entry point for LangBot"""
args = _build_parser().parse_args()
if args.standalone_runtime:
from langbot.pkg.utils import platform
@@ -55,22 +67,26 @@ async def main_entry(loop: asyncio.AbstractEventLoop):
print(asciiart)
# Check dependencies
from langbot.pkg.core.bootutils import deps
# A release migration is a deterministic one-shot deployment Job. It must
# fail with the current image when a dependency is absent, never mutate its
# environment and ask an orchestrator to restart it.
if args.command != 'migrate':
from langbot.pkg.core.bootutils import deps
missing_deps = await deps.check_deps()
missing_deps = await deps.check_deps()
if missing_deps:
print('以下依赖包未安装,将自动安装,请完成后重启程序:')
print(
'These dependencies are missing, they will be installed automatically, please restart the program after completion:'
)
for dep in missing_deps:
print('-', dep)
await deps.install_deps(missing_deps)
print('已自动安装缺失的依赖包,请重启程序。')
print('The missing dependencies have been installed automatically, please restart the program.')
sys.exit(0)
if missing_deps:
print('以下依赖包未安装,将自动安装,请完成后重启程序:')
print(
'These dependencies are missing, they will be installed automatically, '
'please restart the program after completion:'
)
for dep in missing_deps:
print('-', dep)
await deps.install_deps(missing_deps)
print('已自动安装缺失的依赖包,请重启程序。')
print('The missing dependencies have been installed automatically, please restart the program.')
sys.exit(0)
# Check configuration files
from langbot.pkg.core.bootutils import files
@@ -83,6 +99,12 @@ async def main_entry(loop: asyncio.AbstractEventLoop):
for file in generated_files:
print('-', file)
if args.command == 'migrate':
from langbot.pkg.persistence.release_migration import run_cloud_release_migration_from_config
await run_cloud_release_migration_from_config(loop)
return
from langbot.pkg.core import boot
await boot.main(loop)
@@ -83,6 +83,7 @@ class RouterGroup(abc.ABC):
rule = self.path + rule
async def handler_error(*args, **kwargs):
request_context: RequestContext | None = None
if auth_type == AuthType.ACCOUNT_TOKEN:
authorization = quart.request.headers.get('Authorization', '')
if not authorization.startswith('Bearer '):
@@ -183,6 +184,17 @@ class RouterGroup(abc.ABC):
return self._auth_error_response(e)
try:
if request_context is not None:
persistence_mgr = getattr(self.ap, 'persistence_mgr', None)
tenant_scope_descriptor = getattr(type(persistence_mgr), 'tenant_scope', None)
if callable(tenant_scope_descriptor):
# Authorization discovery is complete. Carry the
# trusted Workspace identity across the handler, but
# do not reserve a database connection while it waits
# on providers, runtimes, uploads, or streamed clients.
# Services that need atomic writes open a short UoW.
async with persistence_mgr.tenant_scope(request_context.workspace_uuid):
return await f(*args, **kwargs)
return await f(*args, **kwargs)
except Exception as e: # 自动 500
@@ -1,7 +1,9 @@
from __future__ import annotations
from langbot.pkg.utils import constants
from langbot_plugin.box.errors import BoxAdmissionError
from langbot.pkg.cloud.entitlements import EntitlementUnavailableError
from ...authz import Permission
from ...context import RequestContext
from .. import group
@@ -18,7 +20,10 @@ class BoxRouterGroup(group.RouterGroup):
permission=Permission.RESOURCE_VIEW,
)
async def _(request_context: RequestContext) -> str:
status = await self.ap.box_service.get_status(request_context)
try:
status = await self.ap.box_service.get_status(request_context)
except (BoxAdmissionError, EntitlementUnavailableError) as exc:
return self.http_status(403, 'managed_sandbox_unavailable', str(exc))
status['hidden'] = should_hide_box_runtime_status(constants.edition, status.get('enabled'))
return self.success(data=status)
@@ -29,7 +34,10 @@ class BoxRouterGroup(group.RouterGroup):
permission=Permission.AUDIT_VIEW,
)
async def _(request_context: RequestContext) -> str:
sessions = await self.ap.box_service.get_sessions(request_context)
try:
sessions = await self.ap.box_service.get_sessions(request_context)
except (BoxAdmissionError, EntitlementUnavailableError) as exc:
return self.http_status(403, 'managed_sandbox_unavailable', str(exc))
return self.success(data=sessions)
@self.route(
@@ -39,5 +47,10 @@ class BoxRouterGroup(group.RouterGroup):
permission=Permission.AUDIT_VIEW,
)
async def _(request_context: RequestContext) -> str:
try:
if getattr(self.ap.box_service, 'managed_admission_required', False):
await self.ap.box_service.require_workspace_sandbox(request_context)
except (BoxAdmissionError, EntitlementUnavailableError) as exc:
return self.http_status(403, 'managed_sandbox_unavailable', str(exc))
errors = self.ap.box_service.get_recent_errors(request_context)
return self.success(data=errors)
@@ -13,6 +13,7 @@ import quart
from ....authz import Permission, permissions_for_role, require_permission
from ....context import PrincipalContext, PrincipalType, RequestContext, WorkspaceContext
from ... import group
from ......core.task_boundary import run_in_workspace_uow
from ......platform.sources.websocket_manager import WebSocketScope, ws_connection_manager
logger = logging.getLogger(__name__)
@@ -113,7 +114,11 @@ class WebSocketChatRouterGroup(group.RouterGroup):
require_permission(current_context, Permission.RUNTIME_OPERATE)
async def _get_scoped_adapter(self, request_context: RequestContext, pipeline_uuid: str):
pipeline = await self.ap.pipeline_service.get_pipeline(request_context, pipeline_uuid)
pipeline = await run_in_workspace_uow(
self.ap,
request_context.workspace_uuid,
lambda: self.ap.pipeline_service.get_pipeline(request_context, pipeline_uuid),
)
if pipeline is None:
return None
proxy_bot = await self.ap.platform_mgr.get_websocket_proxy_bot(request_context)
@@ -16,11 +16,13 @@ import posixpath
import sqlalchemy
from .....core import taskmgr
from .....core.task_boundary import run_in_workspace_uow
from .....entity.persistence import plugin as persistence_plugin
from ...authz import Permission
from ...context import ExecutionContext, RequestContext
from .. import group
from .....workspace.errors import WorkspaceNotFoundError
from .....plugin.github import validate_github_plugin_install_info
from langbot_plugin.runtime.plugin.mgr import PluginInstallSource
@@ -308,7 +310,11 @@ class PluginsRouterGroup(group.RouterGroup):
):
"""Revalidate a captured task context immediately before Runtime I/O."""
await self.ap.plugin_connector.require_workspace_context(execution_context)
await run_in_workspace_uow(
self.ap,
execution_context.workspace_uuid,
lambda: self.ap.plugin_connector.require_workspace_context(execution_context),
)
return await operation()
async def _require_public_plugin_runtime_context(self) -> ExecutionContext:
@@ -757,27 +763,31 @@ class PluginsRouterGroup(group.RouterGroup):
if limit_error is not None:
return limit_error
data = await quart.request.json
asset_url = data.get('asset_url', '')
owner = data.get('owner', '')
repo = data.get('repo', '')
release_tag = data.get('release_tag', '')
data = await quart.request.json or {}
try:
install_info = validate_github_plugin_install_info(
{
'asset_url': data.get('asset_url'),
'asset_id': data.get('asset_id'),
'release_id': data.get('release_id'),
'owner': data.get('owner'),
'repo': data.get('repo'),
'release_tag': data.get('release_tag'),
'github_url': f'https://github.com/{data.get("owner", "")}/{data.get("repo", "")}',
}
)
except ValueError as exc:
return self.http_status(400, -1, str(exc))
if not asset_url:
return self.http_status(400, -1, 'Missing asset_url parameter')
owner = install_info['owner']
repo = install_info['repo']
release_tag = install_info['release_tag']
execution_context = await self.ap.plugin_connector.require_workspace_context(request_context)
ctx = taskmgr.TaskContext.new()
ctx.metadata['plugin_name'] = f'{owner}/{repo}'
ctx.metadata['install_source'] = 'github'
install_info = {
'asset_url': asset_url,
'owner': owner,
'repo': repo,
'release_tag': release_tag,
'github_url': f'https://github.com/{owner}/{repo}',
}
wrapper = self.ap.task_mgr.create_user_task(
self._run_fenced_plugin_operation(
@@ -25,12 +25,33 @@ class SystemRouterGroup(group.RouterGroup):
account, _ = await self._authenticate_account(authorization.removeprefix('Bearer '))
request_context = await self._resolve_account_context(account, group.AuthType.USER_TOKEN)
if request_context is not None:
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(WorkspaceMetadata).where(
WorkspaceMetadata.workspace_uuid == request_context.workspace_uuid,
WorkspaceMetadata.key.in_(['wizard_status', 'wizard_progress']),
tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
async def load_workspace_metadata():
return await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(
WorkspaceMetadata.key,
WorkspaceMetadata.value,
).where(
WorkspaceMetadata.workspace_uuid == request_context.workspace_uuid,
WorkspaceMetadata.key.in_(['wizard_status', 'wizard_progress']),
)
)
cloud_runtime = (
getattr(getattr(self.ap.persistence_mgr, 'mode', None), 'value', None) == 'cloud_runtime'
)
if cloud_runtime:
if not callable(tenant_uow):
raise RuntimeError('Cloud system metadata requires an explicit tenant UoW')
async with tenant_uow(request_context.workspace_uuid):
result = await load_workspace_metadata()
else:
result = await load_workspace_metadata()
# ``execute_async`` deliberately preserves its historical
# AsyncConnection result shape. Selecting the two fields
# explicitly keeps this reader independent of ORM Session
# scalar semantics inside a tenant UoW.
for row in result:
if row.key == 'wizard_status':
wizard_status = row.value
@@ -44,11 +44,27 @@ class WebhookRouterGroup(group.RouterGroup):
if not hasattr(runtime_bot.adapter, 'handle_unified_webhook'):
return quart.jsonify({'error': 'Adapter does not support unified webhook'}), 501
response = await runtime_bot.adapter.handle_unified_webhook(
bot_uuid=bot_uuid,
path=path,
request=quart.request,
)
async def dispatch():
await self.ap.workspace_service.get_execution_binding(
runtime_bot.workspace_uuid,
expected_generation=runtime_bot.placement_generation,
)
return await runtime_bot.adapter.handle_unified_webhook(
bot_uuid=bot_uuid,
path=path,
request=quart.request,
)
persistence_mgr = self.ap.persistence_mgr
cloud_runtime = getattr(getattr(persistence_mgr, 'mode', None), 'value', None) == 'cloud_runtime'
if cloud_runtime:
tenant_scope = getattr(persistence_mgr, 'tenant_scope', None)
if not callable(tenant_scope):
raise RuntimeError('Cloud webhook dispatch requires an explicit tenant scope')
async with tenant_scope(runtime_bot.workspace_uuid):
response = await dispatch()
else:
response = await dispatch()
return response
+71 -19
View File
@@ -170,37 +170,89 @@ class ApiKeyService:
if not secret.startswith('lbk_'):
return None
secret_hash = self._hash_secret(secret)
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(apikey.ApiKey).where(apikey.ApiKey.key_hash == secret_hash)
)
key = result.first()
if key is None or key.status != apikey.ApiKeyStatus.ACTIVE.value:
current_session = getattr(self.ap.persistence_mgr, 'current_session', lambda: None)
discovery_uow = getattr(self.ap.persistence_mgr, 'api_key_discovery_uow', None)
if current_session() is None and callable(discovery_uow):
async with discovery_uow(secret_hash) as discovery:
key = await discovery.session.scalar(
sqlalchemy.select(apikey.ApiKey).where(apikey.ApiKey.key_hash == secret_hash)
)
else:
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(apikey.ApiKey).where(apikey.ApiKey.key_hash == secret_hash)
)
key = result.first()
if key is None:
return None
discovered_workspace_uuid = key.workspace_uuid
discovered_key_id = key.id
now = self._utcnow()
if key.expires_at is not None and key.expires_at <= now:
return None
raw_scopes = list(key.scopes or [])
async def bind_and_record_use() -> tuple[typing.Any, typing.Any] | None:
# Re-read inside the tenant transaction. A revoke/expiry racing
# discovery must not result in an authenticated identity.
active_session = current_session()
if active_session is not None:
scoped_key = await active_session.scalar(
sqlalchemy.select(apikey.ApiKey).where(
apikey.ApiKey.id == discovered_key_id,
apikey.ApiKey.workspace_uuid == discovered_workspace_uuid,
apikey.ApiKey.key_hash == secret_hash,
)
)
else: # compatibility for isolated service tests
scoped_result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(apikey.ApiKey).where(
apikey.ApiKey.id == discovered_key_id,
apikey.ApiKey.workspace_uuid == discovered_workspace_uuid,
apikey.ApiKey.key_hash == secret_hash,
)
)
scoped_key = scoped_result.first()
if scoped_key is None or scoped_key.status != apikey.ApiKeyStatus.ACTIVE.value:
return None
if scoped_key.expires_at is not None and scoped_key.expires_at <= now:
return None
binding = await self.ap.workspace_service.get_execution_binding(discovered_workspace_uuid)
updated = await self.ap.persistence_mgr.execute_async(
sqlalchemy.update(apikey.ApiKey)
.where(
apikey.ApiKey.id == scoped_key.id,
apikey.ApiKey.workspace_uuid == discovered_workspace_uuid,
apikey.ApiKey.key_hash == secret_hash,
apikey.ApiKey.status == apikey.ApiKeyStatus.ACTIVE.value,
)
.values(last_used_at=now)
.returning(apikey.ApiKey.id)
)
# Authentication and revocation race on this atomic predicate. If
# revoke won, no active row is returned and the stale object read
# above must never become an authenticated identity.
if updated.scalar_one_or_none() is None:
return None
return binding, scoped_key
tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
if current_session() is None and callable(tenant_uow):
async with tenant_uow(discovered_workspace_uuid):
bound = await bind_and_record_use()
else:
bound = await bind_and_record_use()
if bound is None:
return None
binding, scoped_key = bound
raw_scopes = list(scoped_key.scopes or [])
permissions = (
frozenset(permission.value for permission in Permission)
if '*' in raw_scopes
else frozenset(self._normalize_scopes(raw_scopes))
)
binding = await self.ap.workspace_service.get_execution_binding(key.workspace_uuid)
await self.ap.persistence_mgr.execute_async(
sqlalchemy.update(apikey.ApiKey)
.where(
apikey.ApiKey.id == key.id,
apikey.ApiKey.workspace_uuid == key.workspace_uuid,
apikey.ApiKey.key_hash == secret_hash,
)
.values(last_used_at=now)
)
return ApiKeyIdentity(
instance_uuid=binding.instance_uuid,
workspace_uuid=binding.workspace_uuid,
placement_generation=binding.placement_generation,
api_key_uuid=key.uuid,
api_key_uuid=scoped_key.uuid,
permissions=permissions,
)
@@ -1,6 +1,7 @@
from __future__ import annotations
import datetime
import functools
import os
import re
from pathlib import Path
@@ -22,6 +23,25 @@ DEFAULT_LOG_RETENTION_DAYS = 3
UPLOAD_OWNER_TYPES = ('upload_image', 'upload_document', 'upload')
def _workspace_scope(method):
"""Bind maintenance work to a Workspace without spanning external I/O."""
@functools.wraps(method)
async def wrapped(self, context, *args, **kwargs):
workspace_uuid = require_workspace_uuid(context)
persistence_mgr = getattr(self.ap, 'persistence_mgr', None)
tenant_scope = getattr(persistence_mgr, 'tenant_scope', None)
cloud_runtime = getattr(getattr(persistence_mgr, 'mode', None), 'value', None) == 'cloud_runtime'
if cloud_runtime:
if not callable(tenant_scope):
raise RuntimeError('Cloud maintenance requires an explicit tenant scope')
async with tenant_scope(workspace_uuid):
return await method(self, context, *args, **kwargs)
return await method(self, context, *args, **kwargs)
return wrapped
class MaintenanceService:
"""Storage maintenance and diagnostics."""
@@ -30,6 +50,7 @@ class MaintenanceService:
def __init__(self, ap: app.Application) -> None:
self.ap = ap
@_workspace_scope
async def cleanup_expired_files(self, context: ExecutionContext) -> dict[str, int]:
if not isinstance(context, ExecutionContext):
raise WorkspaceRequiredError('Storage cleanup requires an ExecutionContext')
+9 -3
View File
@@ -1,6 +1,5 @@
from __future__ import annotations
import asyncio
import copy
import re
import uuid
@@ -8,6 +7,7 @@ import uuid
import sqlalchemy
from ....core import app, taskmgr
from ....core.task_boundary import create_detached_task
from ....entity.persistence import mcp as persistence_mcp
from ....entity.persistence import plugin as persistence_plugin
from ....provider.tools.loaders.mcp import MCPSessionStatus, RuntimeMCPSession
@@ -249,7 +249,10 @@ class MCPService:
await self.ap.persistence_mgr.execute_async(sqlalchemy.insert(persistence_mcp.MCPServer).values(payload))
created = await self._get_mcp_server_by_uuid_raw(execution_context, payload['uuid'])
if created and self.ap.tool_mgr.mcp_tool_loader:
task = asyncio.create_task(self.ap.tool_mgr.mcp_tool_loader.host_mcp_server(execution_context, created))
task = create_detached_task(
self.ap.tool_mgr.mcp_tool_loader.host_mcp_server(execution_context, created),
after_commit_manager=self.ap.persistence_mgr,
)
self.ap.tool_mgr.mcp_tool_loader._hosted_mcp_tasks.append(task)
return payload['uuid']
@@ -351,7 +354,10 @@ class MCPService:
if old_enable and loader.has_session(execution_context, old_name):
await loader.remove_mcp_server(execution_context, old_name)
if new_enable:
task = asyncio.create_task(loader.host_mcp_server(execution_context, updated))
task = create_detached_task(
loader.host_mcp_server(execution_context, updated),
after_commit_manager=self.ap.persistence_mgr,
)
loader._hosted_mcp_tasks.append(task)
async def delete_mcp_server(self, context: TenantContext, server_uuid: str) -> None:
+138 -107
View File
@@ -2,8 +2,11 @@ from __future__ import annotations
import uuid
import datetime
import functools
import json
import sqlalchemy
from sqlalchemy.dialects import postgresql as postgresql_dialect
from sqlalchemy.dialects import sqlite as sqlite_dialect
from ....core import app
from ....entity.persistence import monitoring as persistence_monitoring
@@ -12,6 +15,21 @@ from ..context import ExecutionContext
from .tenant import TenantContext, require_workspace_uuid
def _workspace_transaction(method):
"""Run an explicit service entrypoint in one Workspace transaction."""
@functools.wraps(method)
async def wrapped(self, context, *args, **kwargs):
workspace_uuid = require_workspace_uuid(context)
tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
if callable(tenant_uow):
async with tenant_uow(workspace_uuid):
return await method(self, context, *args, **kwargs)
return await method(self, context, *args, **kwargs)
return wrapped
class MonitoringService:
"""Monitoring service"""
@@ -49,7 +67,7 @@ class MonitoringService:
Returns:
A dict mapping table name to the number of deleted rows.
"""
self._require_write_context(context)
workspace_uuid = self._require_write_context(context)
if retention_days < 1:
raise ValueError('retention_days must be >= 1')
if batch_size < 1:
@@ -104,17 +122,27 @@ class MonitoringService:
),
]
deleted_counts: dict[str, int] = {}
async def delete_records() -> dict[str, int]:
deleted_counts: dict[str, int] = {}
for table_name, model_cls, ts_column, pk_column in tables_and_columns:
deleted_counts[table_name] = await self._delete_expired_in_batches(
context=context,
model_cls=model_cls,
ts_column=ts_column,
pk_column=pk_column,
cutoff=cutoff,
batch_size=batch_size,
)
return deleted_counts
for table_name, model_cls, ts_column, pk_column in tables_and_columns:
deleted_counts[table_name] = await self._delete_expired_in_batches(
context=context,
model_cls=model_cls,
ts_column=ts_column,
pk_column=pk_column,
cutoff=cutoff,
batch_size=batch_size,
)
tenant_scope = getattr(self.ap.persistence_mgr, 'tenant_scope', None)
if callable(tenant_scope):
# Carry the Workspace across the complete cleanup without holding a
# connection. Each select+delete batch opens and commits its own UoW.
async with tenant_scope(workspace_uuid):
deleted_counts = await delete_records()
else:
deleted_counts = await delete_records()
if sum(deleted_counts.values()) > 0:
await self._release_sqlite_space()
@@ -134,25 +162,36 @@ class MonitoringService:
deleted_total = 0
while True:
select_result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(pk_column)
.where(model_cls.workspace_uuid == workspace_uuid, ts_column < cutoff)
.limit(batch_size)
)
pk_values = list(select_result.scalars().all())
if not pk_values:
break
delete_result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.delete(model_cls).where(
model_cls.workspace_uuid == workspace_uuid,
pk_column.in_(pk_values),
async def delete_batch() -> tuple[int, int]:
select_result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(pk_column)
.where(model_cls.workspace_uuid == workspace_uuid, ts_column < cutoff)
.limit(batch_size)
)
)
deleted = delete_result.rowcount or 0
deleted_total += deleted
pk_values = list(select_result.scalars().all())
if not pk_values:
return 0, 0
if len(pk_values) < batch_size:
delete_result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.delete(model_cls).where(
model_cls.workspace_uuid == workspace_uuid,
pk_column.in_(pk_values),
)
)
return len(pk_values), int(delete_result.rowcount or 0)
tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
if callable(tenant_uow):
async with tenant_uow(workspace_uuid):
selected, deleted = await delete_batch()
else:
selected, deleted = await delete_batch()
deleted_total += deleted
if selected == 0:
break
if selected < batch_size:
break
return deleted_total
@@ -192,22 +231,30 @@ class MonitoringService:
session_id: str | None = None,
):
workspace_uuid = self._require_write_context(context)
context_columns = (
persistence_monitoring.MonitoringMessage.id,
persistence_monitoring.MonitoringMessage.bot_id,
persistence_monitoring.MonitoringMessage.bot_name,
persistence_monitoring.MonitoringMessage.pipeline_id,
persistence_monitoring.MonitoringMessage.pipeline_name,
persistence_monitoring.MonitoringMessage.session_id,
)
if message_id:
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(persistence_monitoring.MonitoringMessage).where(
sqlalchemy.select(*context_columns).where(
persistence_monitoring.MonitoringMessage.workspace_uuid == workspace_uuid,
persistence_monitoring.MonitoringMessage.id == message_id,
)
)
row = result.first()
if row:
return row[0]
return row
if not session_id:
return None
user_query = (
sqlalchemy.select(persistence_monitoring.MonitoringMessage)
sqlalchemy.select(*context_columns)
.where(
sqlalchemy.and_(
persistence_monitoring.MonitoringMessage.session_id == session_id,
@@ -221,10 +268,10 @@ class MonitoringService:
result = await self.ap.persistence_mgr.execute_async(user_query)
row = result.first()
if row:
return row[0]
return row
any_query = (
sqlalchemy.select(persistence_monitoring.MonitoringMessage)
sqlalchemy.select(*context_columns)
.where(
persistence_monitoring.MonitoringMessage.workspace_uuid == workspace_uuid,
persistence_monitoring.MonitoringMessage.session_id == session_id,
@@ -234,10 +281,11 @@ class MonitoringService:
)
result = await self.ap.persistence_mgr.execute_async(any_query)
row = result.first()
return row[0] if row else None
return row
# ========== Recording Methods ==========
@_workspace_transaction
async def record_message(
self,
context: ExecutionContext,
@@ -285,6 +333,7 @@ class MonitoringService:
return message_id
@_workspace_transaction
async def record_llm_call(
self,
context: ExecutionContext,
@@ -331,6 +380,7 @@ class MonitoringService:
return call_id
@_workspace_transaction
async def record_tool_call(
self,
context: ExecutionContext,
@@ -389,6 +439,7 @@ class MonitoringService:
return call_id
@_workspace_transaction
async def record_embedding_call(
self,
context: ExecutionContext,
@@ -432,6 +483,7 @@ class MonitoringService:
return call_id
@_workspace_transaction
async def record_session_start(
self,
context: ExecutionContext,
@@ -466,6 +518,7 @@ class MonitoringService:
sqlalchemy.insert(persistence_monitoring.MonitoringSession).values(session_data)
)
@_workspace_transaction
async def update_session_activity(
self,
context: ExecutionContext,
@@ -503,6 +556,7 @@ class MonitoringService:
# Check if any rows were updated
return result.rowcount > 0
@_workspace_transaction
async def record_error(
self,
context: ExecutionContext,
@@ -540,6 +594,7 @@ class MonitoringService:
return error_id
@_workspace_transaction
async def update_message_status(
self,
context: ExecutionContext,
@@ -1802,81 +1857,57 @@ class MonitoringService:
)
return None
# Check if record with this feedback_id already exists
existing_result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(MonitoringFeedback).where(
MonitoringFeedback.workspace_uuid == workspace_uuid,
MonitoringFeedback.feedback_id == feedback_id,
)
)
existing_row = existing_result.first()
if existing_row:
# UPDATE existing record
existing = existing_row[0] if isinstance(existing_row, tuple) else existing_row
await self.ap.persistence_mgr.execute_async(
sqlalchemy.update(MonitoringFeedback)
.where(
MonitoringFeedback.workspace_uuid == workspace_uuid,
MonitoringFeedback.feedback_id == feedback_id,
)
.values(
timestamp=now,
feedback_type=feedback_type,
feedback_content=feedback_content,
inaccurate_reasons=reasons_json,
bot_id=bot_id or existing.bot_id,
bot_name=bot_name or existing.bot_name,
pipeline_id=pipeline_id or existing.pipeline_id,
pipeline_name=pipeline_name or existing.pipeline_name,
session_id=session_id or existing.session_id,
message_id=message_id or existing.message_id,
stream_id=stream_id or existing.stream_id,
user_id=user_id or existing.user_id,
platform=platform or existing.platform,
)
)
return existing.id
record_data = {
'id': str(uuid.uuid4()),
'workspace_uuid': workspace_uuid,
'timestamp': now,
'feedback_id': feedback_id,
'feedback_type': feedback_type,
'feedback_content': feedback_content,
'inaccurate_reasons': reasons_json,
'bot_id': bot_id,
'bot_name': bot_name,
'pipeline_id': pipeline_id,
'pipeline_name': pipeline_name,
'session_id': session_id,
'message_id': message_id,
'stream_id': stream_id,
'user_id': user_id,
'platform': platform,
}
dialect_name = self.ap.persistence_mgr.get_db_engine().dialect.name
if dialect_name == 'postgresql':
statement = postgresql_dialect.insert(MonitoringFeedback).values(record_data)
elif dialect_name == 'sqlite':
statement = sqlite_dialect.insert(MonitoringFeedback).values(record_data)
else:
# INSERT new record with IntegrityError defense
record_id = str(uuid.uuid4())
record_data = {
'id': record_id,
'workspace_uuid': workspace_uuid,
'timestamp': now,
'feedback_id': feedback_id,
'feedback_type': feedback_type,
'feedback_content': feedback_content,
'inaccurate_reasons': reasons_json,
'bot_id': bot_id,
'bot_name': bot_name,
'pipeline_id': pipeline_id,
'pipeline_name': pipeline_name,
'session_id': session_id,
'message_id': message_id,
'stream_id': stream_id,
'user_id': user_id,
'platform': platform,
}
try:
await self.ap.persistence_mgr.execute_async(sqlalchemy.insert(MonitoringFeedback).values(record_data))
return record_id
except Exception:
# UNIQUE constraint conflict (concurrent feedback for same feedback_id)
await self.ap.persistence_mgr.execute_async(
sqlalchemy.update(MonitoringFeedback)
.where(
MonitoringFeedback.workspace_uuid == workspace_uuid,
MonitoringFeedback.feedback_id == feedback_id,
)
.values(
timestamp=now,
feedback_type=feedback_type,
feedback_content=feedback_content,
inaccurate_reasons=reasons_json,
)
)
return feedback_id
raise RuntimeError(f'Monitoring feedback upsert does not support {dialect_name!r}')
excluded = statement.excluded
def preserve_existing(column):
return sqlalchemy.func.coalesce(sqlalchemy.func.nullif(getattr(excluded, column.key), ''), column)
statement = statement.on_conflict_do_update(
index_elements=[MonitoringFeedback.workspace_uuid, MonitoringFeedback.feedback_id],
set_={
'timestamp': excluded.timestamp,
'feedback_type': excluded.feedback_type,
'feedback_content': excluded.feedback_content,
'inaccurate_reasons': excluded.inaccurate_reasons,
'bot_id': preserve_existing(MonitoringFeedback.bot_id),
'bot_name': preserve_existing(MonitoringFeedback.bot_name),
'pipeline_id': preserve_existing(MonitoringFeedback.pipeline_id),
'pipeline_name': preserve_existing(MonitoringFeedback.pipeline_name),
'session_id': preserve_existing(MonitoringFeedback.session_id),
'message_id': preserve_existing(MonitoringFeedback.message_id),
'stream_id': preserve_existing(MonitoringFeedback.stream_id),
'user_id': preserve_existing(MonitoringFeedback.user_id),
'platform': preserve_existing(MonitoringFeedback.platform),
},
).returning(MonitoringFeedback.id)
result = await self.ap.persistence_mgr.execute_async(statement)
return str(result.scalar_one())
async def get_feedback_stats(
self,
+53 -5
View File
@@ -4,6 +4,7 @@ import io
import inspect
import os
import posixpath
import stat
import zipfile
from typing import Optional
from urllib.parse import quote, unquote, urlparse
@@ -34,6 +35,12 @@ _GITHUB_ASSET_HOSTS = {
'raw.githubusercontent.com',
'codeload.github.com',
}
_MAX_GITHUB_ARCHIVE_BYTES = 10 * 1024 * 1024
_MAX_GITHUB_ARCHIVE_ENTRIES = 4096
_MAX_SKILL_ARCHIVE_FILES = 1024
_MAX_SKILL_FILE_BYTES = 10 * 1024 * 1024
_MAX_SKILL_UNCOMPRESSED_BYTES = 50 * 1024 * 1024
_MAX_SKILL_COMPRESSION_RATIO = 200
class SkillService:
@@ -322,9 +329,22 @@ class SkillService:
async def _download_github_asset(self, asset_url: str) -> bytes:
async with httpx.AsyncClient(follow_redirects=True, timeout=120) as client:
resp = await client.get(asset_url)
resp.raise_for_status()
return resp.content
async with client.stream('GET', asset_url) as resp:
resp.raise_for_status()
content_length = resp.headers.get('content-length')
if content_length is not None:
try:
if int(content_length) > _MAX_GITHUB_ARCHIVE_BYTES:
raise ValueError('GitHub skill archive exceeds the compressed size limit')
except ValueError as exc:
if 'exceeds' in str(exc):
raise
content = bytearray()
async for chunk in resp.aiter_bytes():
content.extend(chunk)
if len(content) > _MAX_GITHUB_ARCHIVE_BYTES:
raise ValueError('GitHub skill archive exceeds the compressed size limit')
return bytes(content)
async def _download_github_skill_directory_as_zip(
self, asset_url: str, *, owner: str, repo: str
@@ -339,7 +359,11 @@ class SkillService:
raise ValueError('GitHub repository archive must be a valid .zip archive') from exc
with source_archive as source_zip:
if len(source_zip.infolist()) > _MAX_GITHUB_ARCHIVE_ENTRIES:
raise ValueError('GitHub repository archive contains too many entries')
skill_entry = self._find_github_skill_archive_entry(source_zip, info['file_path'])
if skill_entry.file_size > _MAX_SKILL_FILE_BYTES:
raise ValueError('GitHub SKILL.md exceeds the file size limit')
try:
skill_md_content = source_zip.read(skill_entry).decode('utf-8')
except UnicodeDecodeError as exc:
@@ -377,6 +401,7 @@ class SkillService:
normalized_source_dir = posixpath.normpath(source_skill_dir)
source_prefix = f'{normalized_source_dir}/'
copied_files = 0
copied_bytes = 0
for member in source_zip.infolist():
normalized_member = posixpath.normpath(member.filename)
@@ -399,10 +424,33 @@ class SkillService:
if member.is_dir():
target_zip.writestr(target_info, b'')
continue
target_zip.writestr(target_info, source_zip.read(member))
if member.flag_bits & 0x1:
raise ValueError('Encrypted GitHub skill archive entries are not supported')
unix_mode = member.external_attr >> 16
if stat.S_IFMT(unix_mode) == stat.S_IFLNK:
raise ValueError(f'GitHub archive contains a symbolic link: {member.filename}')
if member.file_size > _MAX_SKILL_FILE_BYTES:
raise ValueError(f'GitHub skill file exceeds the size limit: {member.filename}')
if member.file_size and member.file_size > max(member.compress_size, 1) * _MAX_SKILL_COMPRESSION_RATIO:
raise ValueError(f'GitHub skill file exceeds the compression-ratio limit: {member.filename}')
copied_files += 1
copied_bytes += member.file_size
if copied_files > _MAX_SKILL_ARCHIVE_FILES:
raise ValueError('GitHub skill directory contains too many files')
if copied_bytes > _MAX_SKILL_UNCOMPRESSED_BYTES:
raise ValueError('GitHub skill directory exceeds the uncompressed size limit')
# Copy in bounded chunks instead of materialising a potentially
# large member in Core memory. The Box Runtime independently
# revalidates the resulting archive before installation.
with source_zip.open(member, 'r') as source_file, target_zip.open(target_info, 'w') as target_file:
remaining = member.file_size
while remaining:
chunk = source_file.read(min(64 * 1024, remaining))
if not chunk:
raise ValueError(f'GitHub skill file is truncated: {member.filename}')
target_file.write(chunk)
remaining -= len(chunk)
if copied_files == 0:
raise ValueError('GitHub skill directory is empty')
+69 -34
View File
@@ -151,10 +151,11 @@ class UserService:
)
async def is_initialized(self) -> bool:
result = await self.ap.persistence_mgr.execute_async(sqlalchemy.select(user.User).limit(1))
result_list = result.all()
return result_list is not None and len(result_list) > 0
account = await self._identity_scalar(
sqlalchemy.select(user.User).limit(1),
f'instance:{self._jwt_identity()[1]}',
)
return account is not None
def _session_factory(self) -> async_sessionmaker[AsyncSession]:
return async_sessionmaker(self.ap.persistence_mgr.get_db_engine(), expire_on_commit=False)
@@ -247,28 +248,24 @@ class UserService:
async def get_user_by_email(self, user_email: str) -> user.User | None:
normalized_email = user_email.strip().casefold()
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(user.User).where(user.User.normalized_email == normalized_email)
return await self._identity_scalar(
sqlalchemy.select(user.User).where(user.User.normalized_email == normalized_email),
f'email:{normalized_email}',
)
result_list = result.all()
return result_list[0] if result_list is not None and len(result_list) > 0 else None
async def get_user_by_uuid(self, account_uuid: str) -> user.User | None:
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(user.User).where(user.User.uuid == account_uuid)
return await self._identity_scalar(
sqlalchemy.select(user.User).where(user.User.uuid == account_uuid),
f'uuid:{account_uuid}',
)
return result.first()
async def get_user_by_space_account_uuid(self, space_account_uuid: str) -> user.User | None:
"""Get user by Space account UUID"""
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(user.User).where(user.User.space_account_uuid == space_account_uuid)
return await self._identity_scalar(
sqlalchemy.select(user.User).where(user.User.space_account_uuid == space_account_uuid),
f'space:{space_account_uuid}',
)
result_list = result.all()
return result_list[0] if result_list is not None and len(result_list) > 0 else None
async def authenticate(self, user_email: str, password: str) -> str | None:
user_obj = await self.get_user_by_email(user_email)
if user_obj is None:
@@ -389,11 +386,13 @@ class UserService:
async def reset_password(self, user_email: str, new_password: str) -> None:
hashed_password = await self._hash_password(new_password)
normalized_email = normalize_email(user_email)
await self.ap.persistence_mgr.execute_async(
await self._identity_execute(
sqlalchemy.update(user.User)
.where(user.User.normalized_email == normalize_email(user_email))
.values(password=hashed_password)
.where(user.User.normalized_email == normalized_email)
.values(password=hashed_password),
f'email:{normalized_email}',
)
async def change_password(self, user_email: str, current_password: str, new_password: str) -> None:
@@ -407,11 +406,13 @@ class UserService:
await self._verify_password(user_obj.password, current_password)
hashed_password = await self._hash_password(new_password)
normalized_email = normalize_email(user_email)
await self.ap.persistence_mgr.execute_async(
await self._identity_execute(
sqlalchemy.update(user.User)
.where(user.User.normalized_email == normalize_email(user_email))
.values(password=hashed_password)
.where(user.User.normalized_email == normalized_email)
.values(password=hashed_password),
f'email:{normalized_email}',
)
# Space user management
@@ -435,7 +436,7 @@ class UserService:
if existing_user:
# Update existing user's tokens
await self.ap.persistence_mgr.execute_async(
await self._identity_execute(
sqlalchemy.update(user.User)
.where(user.User.space_account_uuid == space_account_uuid)
.values(
@@ -443,7 +444,8 @@ class UserService:
space_refresh_token=refresh_token,
space_api_key=api_key,
space_access_token_expires_at=expires_at,
)
),
f'space:{space_account_uuid}',
)
await self._update_space_provider_for_account(existing_user, api_key)
return await self.get_user_by_space_account_uuid(space_account_uuid)
@@ -538,9 +540,38 @@ class UserService:
async def get_first_user(self) -> user.User | None:
"""Get the first user (for single-user mode)"""
result = await self.ap.persistence_mgr.execute_async(sqlalchemy.select(user.User).limit(1))
result_list = result.all()
return result_list[0] if result_list else None
return await self._identity_scalar(
sqlalchemy.select(user.User).limit(1),
f'instance:{self._jwt_identity()[1]}',
)
async def _identity_scalar(
self,
statement: typing.Any,
identity: str,
) -> user.User | None:
"""Execute one exact Account lookup in an explicit discovery transaction."""
digest = hashlib.sha256(identity.encode('utf-8')).hexdigest()
current_session = getattr(self.ap.persistence_mgr, 'current_session', lambda: None)
identity_uow = getattr(self.ap.persistence_mgr, 'identity_discovery_uow', None)
if current_session() is None and callable(identity_uow):
async with identity_uow(digest) as discovery:
return await discovery.session.scalar(statement)
result = await self.ap.persistence_mgr.execute_async(statement)
rows = result.all()
return rows[0] if rows else None
async def _identity_execute(self, statement: typing.Any, identity: str) -> typing.Any:
"""Execute one exact Account mutation in an explicit transaction."""
digest = hashlib.sha256(identity.encode('utf-8')).hexdigest()
current_session = getattr(self.ap.persistence_mgr, 'current_session', lambda: None)
identity_uow = getattr(self.ap.persistence_mgr, 'identity_discovery_uow', None)
if current_session() is None and callable(identity_uow):
async with identity_uow(digest) as discovery:
return await discovery.session.execute(statement)
return await self.ap.persistence_mgr.execute_async(statement)
async def set_password(self, user_email: str, new_password: str, current_password: str | None = None) -> None:
"""Set or change password for a user"""
@@ -557,10 +588,12 @@ class UserService:
await self._verify_password(user_obj.password, current_password)
hashed_password = await self._hash_password(new_password)
await self.ap.persistence_mgr.execute_async(
normalized_email = normalize_email(user_email)
await self._identity_execute(
sqlalchemy.update(user.User)
.where(user.User.normalized_email == normalize_email(user_email))
.values(password=hashed_password)
.where(user.User.normalized_email == normalized_email)
.values(password=hashed_password),
f'email:{normalized_email}',
)
async def bind_space_account(self, user_email: str, code: str) -> user.User:
@@ -596,9 +629,10 @@ class UserService:
raise ValueError('This Space account is already bound to another user')
# Update local account to Space account
await self.ap.persistence_mgr.execute_async(
normalized_email = normalize_email(user_email)
await self._identity_execute(
sqlalchemy.update(user.User)
.where(user.User.normalized_email == normalize_email(user_email))
.where(user.User.normalized_email == normalized_email)
.values(
user=normalize_email(space_email), # Update email to Space email
normalized_email=normalize_email(space_email),
@@ -608,7 +642,8 @@ class UserService:
space_refresh_token=refresh_token,
space_api_key=api_key,
space_access_token_expires_at=expires_at,
)
),
f'email:{normalized_email}',
)
# Update Space model provider API keys
+38 -5
View File
@@ -31,6 +31,9 @@ if typing.TYPE_CHECKING:
# JSON-RPC-ish 401 body returned before the MCP app is reached.
_UNAUTHORIZED_BODY = b'{"error":"unauthorized","message":"A valid LangBot API key is required for MCP access."}'
_ENTITLEMENT_UNAVAILABLE_BODY = (
b'{"error":"entitlement_unavailable","message":"Workspace entitlement is unavailable for MCP access."}'
)
def _extract_api_key(headers: list[tuple[bytes, bytes]]) -> str:
@@ -110,6 +113,29 @@ class MCPMount:
await send({'type': 'http.response.body', 'body': _UNAUTHORIZED_BODY})
return
deployment_admission = getattr(self.ap, 'deployment_admission', None)
try:
if deployment_admission is not None:
deployment_admission.require_active()
entitlement_revision = 0
deployment = getattr(self.ap, 'deployment', None)
if deployment is not None and getattr(deployment, 'multi_workspace_enabled', False):
resolver = getattr(self.ap, 'entitlement_resolver', None)
if resolver is None or identity.instance_uuid != resolver.instance_uuid:
raise RuntimeError('Workspace entitlement resolver is unavailable')
entitlement = await resolver.resolve(identity.workspace_uuid)
entitlement_revision = entitlement.entitlement_revision
except Exception:
await send(
{
'type': 'http.response.start',
'status': 403,
'headers': [(b'content-type', b'application/json')],
}
)
await send({'type': 'http.response.body', 'body': _ENTITLEMENT_UNAVAILABLE_BODY})
return
request_context = RequestContext(
instance_uuid=identity.instance_uuid,
placement_generation=identity.placement_generation,
@@ -125,11 +151,18 @@ class MCPMount:
role=None,
permissions=identity.permissions,
),
entitlement_revision=entitlement_revision,
)
token = bind_request_context(request_context)
try:
await mcp_asgi(scope, receive, send)
finally:
reset_request_context(token)
tenant_scope = getattr(self.ap.persistence_mgr, 'tenant_scope', None)
if not callable(tenant_scope):
raise RuntimeError('MCP request persistence scope is unavailable')
async with tenant_scope(identity.workspace_uuid):
token = bind_request_context(request_context)
try:
await mcp_asgi(scope, receive, send)
if deployment_admission is not None:
deployment_admission.require_active()
finally:
reset_request_context(token)
return dispatcher
+217
View File
@@ -0,0 +1,217 @@
from __future__ import annotations
import asyncio
import datetime as dt
import time
from collections.abc import Callable
from typing import TYPE_CHECKING
from langbot_plugin.box.errors import BoxAdmissionError, BoxRuntimeUnavailableError
from langbot_plugin.box.models import (
SandboxAdmissionGrant,
SandboxAdmissionPolicy,
SandboxAdmissionRevocation,
)
from ..api.http.context import ExecutionContext
from ..cloud.entitlements import EntitlementSnapshot, EntitlementUnavailableError
if TYPE_CHECKING:
from langbot_plugin.box.client import BoxRuntimeClient
from ..core.app import Application
_UTC = dt.timezone.utc
_MANAGED_SANDBOX_FEATURE = 'managed_sandbox'
_MANAGED_SANDBOX_SESSION_LIMIT = 'managed_sandbox_sessions'
_MAX_GRANT_TTL_SEC = 300
class SandboxAdmissionController:
"""Project Cloud entitlements into short-lived Box Runtime grants.
Product and plan names intentionally never cross this boundary. The
closed Control Plane supplies a versioned generic entitlement, while Core
installs only the numeric authority understood by the shared Box Runtime.
No state is allocated for a Workspace until it attempts to use the
managed sandbox. Per-Workspace locks serialize renewal/revocation so a
concurrent first use cannot install conflicting grants.
"""
def __init__(
self,
ap: Application,
client: BoxRuntimeClient,
*,
policy: SandboxAdmissionPolicy,
wall_time: Callable[[], float] = time.time,
) -> None:
self.ap = ap
self.client = client
self.policy = policy
self._wall_time = wall_time
self._locks: dict[str, asyncio.Lock] = {}
self._highest_revisions: dict[str, int] = {}
def _workspace_lock(self, workspace_uuid: str) -> asyncio.Lock:
lock = self._locks.get(workspace_uuid)
if lock is None:
lock = asyncio.Lock()
self._locks[workspace_uuid] = lock
return lock
@staticmethod
def _context_revision(context: ExecutionContext) -> int:
revision = getattr(context, 'entitlement_revision', 0)
if isinstance(revision, bool) or not isinstance(revision, int):
return 0
return max(revision, 0)
def _revocation_revision(self, context: ExecutionContext, candidate_revision: int = 0) -> int:
return max(
1,
self._highest_revisions.get(context.workspace_uuid, 0),
self._context_revision(context),
candidate_revision,
)
async def _revoke_locked(
self,
context: ExecutionContext,
*,
candidate_revision: int = 0,
) -> None:
revision = self._revocation_revision(context, candidate_revision)
revocation = SandboxAdmissionRevocation(
instance_uuid=context.instance_uuid,
workspace_uuid=context.workspace_uuid,
entitlement_revision=revision,
)
try:
result = await self.client.revoke_sandbox_admission_grant(revocation)
if (
not isinstance(result, dict)
or result.get('revoked') is not True
or result.get('workspace_uuid') != context.workspace_uuid
or result.get('entitlement_revision') != revision
):
raise BoxRuntimeUnavailableError('Box Runtime returned an invalid sandbox revocation receipt')
except Exception as exc:
# The caller still fails closed even if the control connection is
# unavailable. A previously installed grant expires independently
# in at most five minutes inside the Runtime.
self.ap.logger.warning(
'Failed to install Box sandbox admission revocation: '
f'workspace_uuid={context.workspace_uuid} revision={revision} error={exc}'
)
self._highest_revisions[context.workspace_uuid] = revision
@staticmethod
def _require_managed_sandbox(snapshot: EntitlementSnapshot) -> None:
snapshot.require_feature(_MANAGED_SANDBOX_FEATURE)
sessions = snapshot.limit(_MANAGED_SANDBOX_SESSION_LIMIT)
if sessions != 1:
raise EntitlementUnavailableError('Workspace entitlement must grant exactly one managed sandbox session')
def _grant_expiry(self, snapshot: EntitlementSnapshot) -> dt.datetime:
now_epoch = int(self._wall_time())
ttl_sec = min(self.policy.max_grant_ttl_sec, _MAX_GRANT_TTL_SEC)
expires_epoch = min(snapshot.expires_at, now_epoch + ttl_sec)
if expires_epoch <= now_epoch:
raise EntitlementUnavailableError('Workspace entitlement expired before sandbox admission')
return dt.datetime.fromtimestamp(expires_epoch, tz=_UTC)
async def require(self, context: ExecutionContext) -> SandboxAdmissionGrant:
"""Validate entitlement freshness and install/renew one Runtime grant."""
resolver = getattr(self.ap, 'entitlement_resolver', None)
if resolver is None:
raise EntitlementUnavailableError('Workspace entitlement resolver is unavailable')
if context.instance_uuid != resolver.instance_uuid:
raise EntitlementUnavailableError('Workspace entitlement targets another LangBot instance')
lock = self._workspace_lock(context.workspace_uuid)
async with lock:
try:
snapshot = await resolver.resolve(
context.workspace_uuid,
minimum_revision=self._context_revision(context),
now=int(self._wall_time()),
)
except EntitlementUnavailableError as exc:
# Only a verified, scoped snapshot can authoritatively revoke
# a revision. Provider timeouts, malformed responses, and
# rollback/equivocation errors fail this request closed but do
# not tombstone a still-valid revision forever.
authoritative_revision = exc.entitlement_revision
if authoritative_revision is not None:
await self._revoke_locked(
context,
candidate_revision=authoritative_revision,
)
raise
try:
self._require_managed_sandbox(snapshot)
except EntitlementUnavailableError:
await self._revoke_locked(
context,
candidate_revision=snapshot.entitlement_revision,
)
raise
grant = SandboxAdmissionGrant(
instance_uuid=context.instance_uuid,
workspace_uuid=context.workspace_uuid,
execution_generation=context.placement_generation,
entitlement_revision=snapshot.entitlement_revision,
expires_at=self._grant_expiry(snapshot),
max_sessions=1,
max_managed_processes=0,
)
result = await self.client.upsert_sandbox_admission_grant(grant)
if (
not isinstance(result, dict)
or result.get('installed') is not True
or result.get('workspace_uuid') != context.workspace_uuid
or result.get('execution_generation') != context.placement_generation
or result.get('entitlement_revision') != snapshot.entitlement_revision
or result.get('max_sessions') != 1
or result.get('max_managed_processes') != 0
):
raise BoxRuntimeUnavailableError('Box Runtime returned an invalid sandbox admission receipt')
self._highest_revisions[context.workspace_uuid] = max(
self._highest_revisions.get(context.workspace_uuid, 0),
snapshot.entitlement_revision,
)
return grant
async def revoke(self, context: ExecutionContext, *, entitlement_revision: int = 0) -> None:
"""Explicitly revoke a Workspace grant using a monotonic tombstone."""
async with self._workspace_lock(context.workspace_uuid):
await self._revoke_locked(context, candidate_revision=entitlement_revision)
def require_cloud_admission_policy(raw_policy: object) -> SandboxAdmissionPolicy:
"""Parse the Cloud Box policy without permitting an OSS downgrade."""
try:
policy = SandboxAdmissionPolicy.model_validate(raw_policy)
except Exception as exc:
raise BoxAdmissionError('Cloud Box sandbox admission policy is invalid') from exc
if not policy.required:
raise BoxAdmissionError('Cloud Box sandbox admission must be required')
if policy.logical_session_id != 'global':
raise BoxAdmissionError('Cloud Box sandbox session ID must be global')
if policy.required_backend != 'nsjail':
raise BoxAdmissionError('Cloud Box sandbox backend must be nsjail')
if policy.max_sessions != 1 or policy.max_managed_processes != 0:
raise BoxAdmissionError('Cloud Box sandbox policy must allow one session and zero managed processes')
if policy.max_grant_ttl_sec > _MAX_GRANT_TTL_SEC:
raise BoxAdmissionError('Cloud Box sandbox admission grant TTL must not exceed 300 seconds')
if policy.workspace_quota_mb <= 0:
raise BoxAdmissionError('Cloud Box sandbox workspace quota must be a positive integer')
return policy
+285
View File
@@ -0,0 +1,285 @@
from __future__ import annotations
import contextlib
import errno
import os
import stat
from collections.abc import Iterable
class UnsafeWorkspacePathError(OSError):
"""A tenant-controlled path could not be opened without following links."""
_DIRECTORY_FLAGS = (
os.O_RDONLY | getattr(os, 'O_DIRECTORY', 0) | getattr(os, 'O_NOFOLLOW', 0) | getattr(os, 'O_CLOEXEC', 0)
)
_FILE_READ_FLAGS = os.O_RDONLY | getattr(os, 'O_NOFOLLOW', 0) | getattr(os, 'O_CLOEXEC', 0)
_FILE_WRITE_FLAGS = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, 'O_NOFOLLOW', 0) | getattr(os, 'O_CLOEXEC', 0)
_MAX_REMOVAL_ENTRIES = 4096
_MAX_REMOVAL_DEPTH = 16
def _component(value: str) -> str:
normalized = str(value or '').strip()
if (
not normalized
or normalized in {'.', '..'}
or '/' in normalized
or '\\' in normalized
or '\x00' in normalized
or len(os.fsencode(normalized)) > 240
):
raise UnsafeWorkspacePathError('Unsafe Workspace path component')
return normalized
def _unsafe(path: str, exc: BaseException | None = None) -> UnsafeWorkspacePathError:
error = UnsafeWorkspacePathError(f'Workspace path is not a link-free directory: {path}')
if exc is not None:
error.__cause__ = exc
return error
@contextlib.contextmanager
def _root_fd(root: str):
try:
fd = os.open(root, _DIRECTORY_FLAGS)
except OSError as exc:
raise _unsafe(root, exc)
try:
if not stat.S_ISDIR(os.fstat(fd).st_mode):
raise _unsafe(root)
yield fd
finally:
os.close(fd)
def _open_dir_at(parent_fd: int, name: str, *, create: bool) -> int:
name = _component(name)
if create:
try:
os.mkdir(name, mode=0o700, dir_fd=parent_fd)
except FileExistsError:
pass
try:
fd = os.open(name, _DIRECTORY_FLAGS, dir_fd=parent_fd)
except OSError as exc:
raise _unsafe(name, exc)
if not stat.S_ISDIR(os.fstat(fd).st_mode):
os.close(fd)
raise _unsafe(name)
return fd
def _remove_entry(
parent_fd: int,
name: str,
*,
budget: list[int] | None = None,
depth: int = 0,
) -> None:
"""Remove an entry recursively without following a symlink at any depth."""
name = _component(name)
budget = budget if budget is not None else [_MAX_REMOVAL_ENTRIES]
if depth > _MAX_REMOVAL_DEPTH or budget[0] <= 0:
raise UnsafeWorkspacePathError('Workspace cleanup exceeded its inode budget')
budget[0] -= 1
for _ in range(4):
try:
child_fd = os.open(name, _DIRECTORY_FLAGS, dir_fd=parent_fd)
except FileNotFoundError:
return
except OSError as exc:
if exc.errno not in {errno.ELOOP, errno.ENOTDIR, errno.EACCES}:
raise
try:
os.unlink(name, dir_fd=parent_fd)
return
except FileNotFoundError:
return
except IsADirectoryError:
continue
else:
try:
_clear_dir(child_fd, budget=budget, depth=depth + 1)
finally:
os.close(child_fd)
try:
os.rmdir(name, dir_fd=parent_fd)
return
except FileNotFoundError:
return
except NotADirectoryError:
continue
raise UnsafeWorkspacePathError('Workspace entry changed while it was being removed')
def _clear_dir(directory_fd: int, *, budget: list[int], depth: int) -> None:
# ``scandir(fd)`` enumerates the already-open directory. Names are then
# resolved relative to the same fd, so a tenant cannot redirect the walk by
# swapping an ancestor symlink between validation and use.
# Do not materialize the whole directory: an attacker-controlled outbox
# may contain an inode bomb even when its byte size is tiny. Removal is
# deliberately budgeted and fails closed once the per-operation cap is
# reached; hard filesystem/inode quota remains a Cloud readiness gate.
with os.scandir(directory_fd) as iterator:
for entry in iterator:
_remove_entry(directory_fd, entry.name, budget=budget, depth=depth)
@contextlib.contextmanager
def _query_fd(root: str, subdir: str, query_key: str, *, create: bool, reset: bool = False):
subdir = _component(subdir)
query_key = _component(query_key)
with _root_fd(root) as root_fd:
subdir_fd = _open_dir_at(root_fd, subdir, create=create)
try:
if reset:
_remove_entry(subdir_fd, query_key)
query_fd = _open_dir_at(subdir_fd, query_key, create=create)
try:
yield query_fd
finally:
os.close(query_fd)
finally:
os.close(subdir_fd)
def write_files(
root: str,
subdir: str,
query_key: str,
files: Iterable[tuple[str, bytes]],
) -> None:
"""Atomically recreate one query directory and write regular files only."""
with _query_fd(root, subdir, query_key, create=True, reset=True) as query_fd:
for raw_name, data in files:
name = _component(raw_name)
try:
file_fd = os.open(name, _FILE_WRITE_FLAGS, 0o600, dir_fd=query_fd)
except OSError as exc:
raise UnsafeWorkspacePathError(f'Could not create a link-free Workspace file: {name}') from exc
with os.fdopen(file_fd, 'wb') as file_obj:
file_obj.write(data)
def _read_directory(
directory_fd: int,
*,
prefix: str,
max_file_bytes: int,
max_files: int,
max_total_bytes: int,
output: list[tuple[str, bytes]],
total: list[int],
remaining_entries: list[int],
remaining_directories: list[int],
depth: int,
) -> None:
if depth > 8:
return
with os.scandir(directory_fd) as iterator:
for entry in iterator:
if len(output) >= max_files or total[0] >= max_total_bytes:
return
if remaining_entries[0] <= 0:
raise UnsafeWorkspacePathError('Sandbox outbox exceeds the directory-entry limit')
remaining_entries[0] -= 1
name = _component(entry.name)
relative = f'{prefix}/{name}' if prefix else name
if entry.is_symlink():
continue
if entry.is_dir(follow_symlinks=False):
if remaining_directories[0] <= 0:
raise UnsafeWorkspacePathError('Sandbox outbox exceeds the directory limit')
remaining_directories[0] -= 1
try:
child_fd = os.open(name, _DIRECTORY_FLAGS, dir_fd=directory_fd)
except OSError:
continue
try:
_read_directory(
child_fd,
prefix=relative,
max_file_bytes=max_file_bytes,
max_files=max_files,
max_total_bytes=max_total_bytes,
output=output,
total=total,
remaining_entries=remaining_entries,
remaining_directories=remaining_directories,
depth=depth + 1,
)
finally:
os.close(child_fd)
continue
try:
file_fd = os.open(name, _FILE_READ_FLAGS, dir_fd=directory_fd)
except OSError:
continue
try:
metadata = os.fstat(file_fd)
if not stat.S_ISREG(metadata.st_mode) or metadata.st_size > max_file_bytes:
continue
remaining = max_total_bytes - total[0]
if metadata.st_size > remaining:
continue
with os.fdopen(file_fd, 'rb', closefd=False) as file_obj:
data = file_obj.read(max_file_bytes + 1)
if len(data) > max_file_bytes or len(data) > remaining:
continue
output.append((relative, data))
total[0] += len(data)
finally:
os.close(file_fd)
def read_regular_files(
root: str,
subdir: str,
query_key: str,
*,
max_file_bytes: int,
max_files: int,
max_total_bytes: int,
max_entries: int = 512,
max_directories: int = 64,
) -> list[tuple[str, bytes]]:
"""Read bounded regular files without following tenant-created links."""
output: list[tuple[str, bytes]] = []
try:
with _query_fd(root, subdir, query_key, create=False) as query_fd:
_read_directory(
query_fd,
prefix='',
max_file_bytes=max_file_bytes,
max_files=max_files,
max_total_bytes=max_total_bytes,
output=output,
total=[0],
remaining_entries=[max_entries],
remaining_directories=[max_directories],
depth=0,
)
except (FileNotFoundError, UnsafeWorkspacePathError):
# A missing directory is an empty outbox. An unsafe existing path is
# deliberately surfaced to the caller rather than followed.
if os.path.lexists(os.path.join(root, subdir, query_key)):
raise
return output
def reset_directory(root: str, subdir: str, query_key: str) -> None:
with _query_fd(root, subdir, query_key, create=True, reset=True):
return
def purge_subdirectory(root: str, subdir: str) -> None:
"""Remove one known subtree without following a hostile replacement link."""
with _root_fd(root) as root_fd:
_remove_entry(root_fd, _component(subdir))
+379 -74
View File
@@ -5,8 +5,10 @@ import collections
import contextlib
import datetime as _dt
import enum
import hashlib
import json
import os
import secrets
from typing import TYPE_CHECKING
import pydantic
@@ -14,11 +16,14 @@ import pydantic
from langbot_plugin.box.client import BoxRuntimeClient
from langbot_plugin.entities.io.context import ActionContext
from langbot_plugin.box.tenancy import box_namespace
from langbot_plugin.box.security import BOX_SHARED_WORKSPACE_PROBE_PREFIX
from .admission import SandboxAdmissionController, require_cloud_admission_policy
from .connector import BoxRuntimeConnector, _get_box_config
from . import secure_fs
from ..telemetry import features as telemetry_features
from ..api.http.context import ExecutionContext
from ..api.http.service.tenant import TenantContext, require_workspace_uuid
from langbot_plugin.box.errors import BoxError, BoxValidationError
from langbot_plugin.box.errors import BoxAdmissionError, BoxError, BoxValidationError
from langbot_plugin.box.models import (
BUILTIN_PROFILES,
BoxExecutionResult,
@@ -52,6 +57,7 @@ class BoxService:
output_limit_chars: int = 4000,
):
self.ap = ap
self._cloud_managed = bool(getattr(getattr(ap, 'deployment', None), 'multi_workspace_enabled', False))
self._enabled = self._load_enabled()
self._runtime_connector: BoxRuntimeConnector | None = None
if client is None:
@@ -68,6 +74,18 @@ class BoxService:
self.profile = self._load_profile()
self.custom_image = self._load_custom_image()
self.workspace_quota_mb = self._load_workspace_quota_mb()
self._admission_policy = (
require_cloud_admission_policy(_get_box_config(ap).get('admission')) if self._cloud_managed else None
)
self._admission = (
SandboxAdmissionController(
ap,
self.client,
policy=self._admission_policy,
)
if self._cloud_managed and self._admission_policy is not None
else None
)
self._recent_errors: collections.deque[dict] = collections.deque(maxlen=_MAX_RECENT_ERRORS)
self._shutdown_task = None
self._reconnect_task: asyncio.Task | None = None
@@ -88,6 +106,10 @@ class BoxService:
``available = False`` to consumers, but distinguished in get_status."""
return self._enabled
@property
def managed_admission_required(self) -> bool:
return self._cloud_managed
async def initialize(self):
if not self._enabled:
# Disabled by config: do NOT connect to a remote runtime, do NOT
@@ -106,19 +128,24 @@ class BoxService:
else:
await self.client.initialize()
self._ensure_default_workspace()
await self._verify_cloud_runtime()
self._available = True
self._connector_error = ''
self.ap.logger.info(
f'LangBot Box runtime initialized: profile={self.profile.name} '
f'default_workspace={self.default_workspace or "(none)"}'
)
await self._purge_attachment_dirs()
# Cloud query directories use globally opaque query UUIDs. Never
# sweep all tenants when a future replica joins the same logical
# instance; that could delete another replica's in-flight files.
if not self._cloud_managed:
await self._purge_attachment_dirs()
except Exception as exc:
self.ap.logger.warning(f'LangBot Box runtime unavailable, sandbox features disabled: {exc}')
self._available = False
self._connector_error = str(exc)
if self._runtime_connector is not None:
await self._on_runtime_disconnect(self._runtime_connector)
if self._cloud_managed:
raise
async def _on_runtime_disconnect(self, connector: BoxRuntimeConnector) -> None:
"""Called by the connector when the Box runtime connection drops.
@@ -146,9 +173,9 @@ class BoxService:
self.ap.logger.info(f'Attempting to reconnect to Box runtime in {delay}s...')
await asyncio.sleep(delay)
try:
await connector.reconnect()
self._ensure_default_workspace()
await self._purge_attachment_dirs()
connector.dispose()
await connector.initialize()
await self._verify_cloud_runtime()
self._available = True
self._connector_error = ''
skill_mgr = getattr(self.ap, 'skill_mgr', None)
@@ -165,6 +192,83 @@ class BoxService:
self._reconnecting = False
self._reconnect_task = None
async def _verify_cloud_runtime(self) -> None:
if not self._cloud_managed:
return
self._ensure_cloud_shared_workspace()
await self._challenge_cloud_shared_workspace()
backend_info = await self.client.get_backend_info()
if (
not isinstance(backend_info, dict)
or backend_info.get('name') != 'nsjail'
or backend_info.get('available') is not True
):
raise BoxValidationError('Cloud Box nsjail isolation readiness failed')
async def _challenge_cloud_shared_workspace(self) -> None:
"""Prove Core and Box Runtime see the same durable filesystem.
Equal configured path strings are not evidence of a shared container
volume. Core creates one high-entropy, no-follow marker under its
canonical root and the authenticated Runtime host-control action reads
that basename only. Any mismatch fails Cloud startup/reconnect closed.
"""
if self.default_workspace is None:
raise BoxValidationError('Cloud Box shared default_workspace is unavailable')
marker_name = f'{BOX_SHARED_WORKSPACE_PROBE_PREFIX}{secrets.token_hex(16)}'
marker_payload = secrets.token_bytes(64)
expected_digest = hashlib.sha256(marker_payload).hexdigest()
directory_flags = os.O_RDONLY | getattr(os, 'O_DIRECTORY', 0)
nofollow = getattr(os, 'O_NOFOLLOW', 0)
root_fd: int | None = None
marker_fd: int | None = None
marker_created = False
try:
root_fd = os.open(self.default_workspace, directory_flags | nofollow)
marker_fd = os.open(
marker_name,
os.O_WRONLY | os.O_CREAT | os.O_EXCL | nofollow,
0o600,
dir_fd=root_fd,
)
marker_created = True
remaining = memoryview(marker_payload)
while remaining:
written = os.write(marker_fd, remaining)
if written <= 0:
raise BoxValidationError('Failed to write Cloud Box shared-volume probe')
remaining = remaining[written:]
os.fsync(marker_fd)
os.close(marker_fd)
marker_fd = None
result = await self.client.verify_shared_workspace(marker_name)
if (
not isinstance(result, dict)
or result.get('marker_name') != marker_name
or result.get('size') != len(marker_payload)
or not secrets.compare_digest(str(result.get('sha256') or ''), expected_digest)
):
raise BoxValidationError(
'Cloud Box Core and Runtime do not share the configured durable Workspace volume'
)
except BoxValidationError:
raise
except Exception as exc:
raise BoxValidationError('Cloud Box shared durable Workspace volume verification failed') from exc
finally:
if marker_fd is not None:
os.close(marker_fd)
if root_fd is not None:
if marker_created:
try:
os.unlink(marker_name, dir_fd=root_fd)
except FileNotFoundError:
pass
os.close(root_fd)
@property
def available(self) -> bool:
return self._available
@@ -211,10 +315,14 @@ class BoxService:
bot_uuid=getattr(context, 'bot_uuid', None),
pipeline_uuid=getattr(context, 'pipeline_uuid', None),
query_uuid=getattr(context, 'query_uuid', None),
entitlement_revision=getattr(context, 'entitlement_revision', 0),
)
@classmethod
def _query_execution_context(cls, query: pipeline_query.Query) -> ExecutionContext:
attached_context = getattr(query, '_execution_context', None)
if isinstance(attached_context, ExecutionContext):
return cls._execution_context(attached_context)
return cls._execution_context(
ExecutionContext(
instance_uuid=str(getattr(query, 'instance_uuid', '') or ''),
@@ -223,6 +331,7 @@ class BoxService:
bot_uuid=getattr(query, 'bot_uuid', None),
pipeline_uuid=getattr(query, 'pipeline_uuid', None),
query_uuid=getattr(query, 'query_uuid', None),
entitlement_revision=getattr(query, 'entitlement_revision', 0),
)
)
@@ -252,6 +361,116 @@ class BoxService:
raise BoxValidationError('Box execution context belongs to a stale Workspace placement')
return execution_context
async def require_workspace_sandbox(self, context: TenantContext) -> ExecutionContext:
"""Fence a Workspace and install its short-lived Cloud admission grant.
OSS keeps the existing singleton behavior and does not require a
Control Plane entitlement. Cloud always resolves a fresh generic
entitlement before a sandbox-visible operation.
"""
execution_context = await self._validated_execution_context(context)
await self._require_validated_workspace_sandbox(execution_context)
return execution_context
async def _require_validated_workspace_sandbox(self, execution_context: ExecutionContext) -> None:
if not self._available:
raise BoxError('Box runtime is not available. Install and start Docker to use sandbox features.')
if self._cloud_managed:
if self._admission is None:
raise BoxAdmissionError('Cloud Box sandbox admission is unavailable')
await self._admission.require(execution_context)
async def is_workspace_sandbox_available(self, context: TenantContext) -> bool:
"""Return tenant-specific availability for UI and tool discovery.
This method deliberately catches entitlement failures so callers can
hide tools without leaking plan details. Direct execution APIs use
:meth:`require_workspace_sandbox` and retain an explicit failure.
"""
if not self._available:
return False
try:
await self.require_workspace_sandbox(context)
return True
except Exception:
return False
def _managed_policy_payload(
self,
context: TenantContext,
spec_payload: dict,
) -> dict:
"""Reject tenant-owned policy fields and apply the Cloud hard policy."""
payload = dict(spec_payload)
if not self._cloud_managed:
return payload
policy = self._admission_policy
if policy is None:
raise BoxAdmissionError('Cloud Box sandbox admission policy is unavailable')
forged_fields = {
'plan',
'subscription',
'managed_sandbox',
'entitlement',
'entitlement_revision',
'max_sessions',
'max_managed_processes',
'backend',
}
submitted_forged_fields = sorted(forged_fields.intersection(payload))
if submitted_forged_fields:
raise BoxAdmissionError(
'Managed sandbox policy fields are host-controlled: ' + ', '.join(submitted_forged_fields)
)
submitted_session_id = str(payload.get('session_id', '') or '').strip()
if submitted_session_id and submitted_session_id != policy.logical_session_id:
raise BoxAdmissionError('Managed sandbox session_id is runtime-owned')
submitted_network = str(getattr(payload.get('network'), 'value', payload.get('network', 'off')) or 'off')
if submitted_network != 'off':
raise BoxAdmissionError('Managed sandbox network access is disabled')
if payload.get('extra_mounts'):
raise BoxAdmissionError('Managed sandbox additional host mounts are disabled')
submitted_mount_path = str(payload.get('mount_path', '/workspace') or '/workspace')
if submitted_mount_path != '/workspace':
raise BoxAdmissionError('Managed sandbox mount_path is runtime-owned')
canonical_host_path = self._tenant_workspace(context)
if canonical_host_path is None:
raise BoxAdmissionError('Managed sandbox Workspace path is unavailable')
submitted_host_path = str(payload.get('host_path', '') or '').strip()
if submitted_host_path and os.path.realpath(submitted_host_path) != os.path.realpath(canonical_host_path):
raise BoxAdmissionError('Managed sandbox host_path is runtime-owned')
timeout = payload.get('timeout_sec', policy.max_timeout_sec)
if isinstance(timeout, bool) or not isinstance(timeout, int):
raise BoxValidationError('timeout_sec must be an integer')
payload.update(
{
'session_id': policy.logical_session_id,
'network': 'off',
'host_path': canonical_host_path,
'mount_path': '/workspace',
'extra_mounts': [],
'persistent': True,
'timeout_sec': min(timeout, policy.max_timeout_sec),
'cpus': policy.cpus,
'memory_mb': policy.memory_mb,
'pids_limit': policy.pids_limit,
'read_only_rootfs': policy.read_only_rootfs,
'workspace_quota_mb': policy.workspace_quota_mb,
}
)
return payload
def _reject_cloud_managed_process(self) -> None:
if self._cloud_managed:
raise BoxAdmissionError('Managed processes are disabled for Cloud sandboxes')
def _tenant_workspace(self, context: TenantContext) -> str | None:
if self.default_workspace is None:
return None
@@ -268,7 +487,8 @@ class BoxService:
if not self._available:
raise BoxError('Box runtime is not available. Install and start Docker to use sandbox features.')
execution_context = await self._validated_execution_context(self._query_execution_context(query))
spec_payload = dict(spec_payload)
spec_payload = self._managed_policy_payload(execution_context, spec_payload)
await self._require_validated_workspace_sandbox(execution_context)
if spec_payload.get('host_path') in (None, ''):
tenant_workspace = self._tenant_workspace(execution_context)
if tenant_workspace is not None:
@@ -295,6 +515,11 @@ class BoxService:
spec,
action_context=self._action_context(execution_context),
)
# A placement may be cut over while a long-running sandbox call is
# in flight. Never accept a result produced by the superseded
# generation. Runtime-side generation fencing prevents new work;
# this second Core check closes the response race.
await self._validated_execution_context(execution_context)
except BoxError as exc:
self._record_error(exc, query)
raise
@@ -322,6 +547,8 @@ class BoxService:
by editing the pipeline config directly through the API (which only
gates the web UI).
"""
if self._cloud_managed:
return 'global'
forced_template = self._forced_box_session_id_template()
if forced_template:
template = forced_template
@@ -373,6 +600,8 @@ class BoxService:
skills it discovered on its own filesystem, so the path is valid there
by construction.
"""
if self._cloud_managed:
return []
skill_mgr = getattr(self.ap, 'skill_mgr', None)
if skill_mgr is None:
return []
@@ -403,13 +632,21 @@ class BoxService:
)
return mounts
async def execute_tool(self, parameters: dict, query: pipeline_query.Query) -> dict:
async def execute_tool(
self,
parameters: dict,
query: pipeline_query.Query,
*,
skill_name: str | None = None,
) -> dict:
"""Execute an agent-facing ``exec`` tool call.
Translates the agent-facing ``command`` field to the internal
``BoxSpec.cmd`` field and injects the session id from the query.
"""
spec_payload: dict = {'cmd': parameters['command']}
if skill_name is not None:
spec_payload['skill_name'] = skill_name
# Pass through allowed agent-facing fields
for key in ('workdir', 'timeout_sec', 'env'):
@@ -435,7 +672,8 @@ class BoxService:
"""Execute trusted internal Box work inside one Workspace namespace."""
execution_context = await self._validated_execution_context(context)
payload = dict(spec_payload)
payload = self._managed_policy_payload(execution_context, spec_payload)
await self._require_validated_workspace_sandbox(execution_context)
if payload.get('host_path') in (None, ''):
tenant_workspace = self._tenant_workspace(execution_context)
if tenant_workspace is not None:
@@ -443,7 +681,9 @@ class BoxService:
if self.shares_filesystem_with_box:
os.makedirs(tenant_workspace, exist_ok=True)
spec = self.build_spec(payload, skip_host_mount_validation=skip_host_mount_validation)
return await self.client.execute(spec, action_context=self._action_context(execution_context))
result = await self.client.execute(spec, action_context=self._action_context(execution_context))
await self._validated_execution_context(execution_context)
return result
# ── Attachment passthrough (inbound / outbound) ──────────────────
#
@@ -473,10 +713,22 @@ class BoxService:
# Hard cap on a single attachment. The HTTP upload endpoints already cap
# uploads at 10MiB; keep parity.
_ATTACHMENT_MAX_BYTES = 10 * _MIB
_ATTACHMENT_MAX_FILES = 20
_ATTACHMENT_MAX_TOTAL_BYTES = 50 * _MIB
# Conservative cap for the exec FALLBACK path only (ARG_MAX / stdout
# truncation). The host-filesystem path has no such limit.
_EXEC_FALLBACK_MAX_BYTES = 256 * 1024
def _attachment_query_key(self, query: pipeline_query.Query) -> str:
query_uuid = str(getattr(query, 'query_uuid', '') or '').strip()
if query_uuid:
if query_uuid in {'.', '..'} or '/' in query_uuid or '\\' in query_uuid or '\x00' in query_uuid:
raise BoxValidationError('Query attachment identity is invalid')
return query_uuid
if self._cloud_managed:
raise BoxValidationError('Cloud attachment transfer requires query_uuid')
return str(query.query_id)
def _host_query_dir(self, subdir: str, query: pipeline_query.Query) -> str | None:
"""Host path for ``/workspace/<subdir>/<query_id>`` when LangBot can
access the bind-mounted workspace directly, else ``None``.
@@ -488,9 +740,9 @@ class BoxService:
E2B and remote runtimes, where we must fall back to the exec channel.
"""
root = self._tenant_workspace(self._query_execution_context(query))
if not root or not os.path.isdir(root):
if not root or not os.path.isdir(root) or os.path.islink(root):
return None
return os.path.join(root, subdir, str(query.query_id))
return os.path.join(root, subdir, self._attachment_query_key(query))
async def _purge_attachment_dirs(self) -> None:
"""Remove leftover inbox/outbox directories on startup.
@@ -509,14 +761,12 @@ class BoxService:
if not root or not os.path.isdir(root):
return
import shutil
host_survivors: list[str] = []
def _host_purge() -> list[str]:
candidates = [
os.path.join(root, self.INBOX_SUBDIR),
os.path.join(root, self.OUTBOX_SUBDIR),
candidates: list[tuple[str, str]] = [
(root, self.INBOX_SUBDIR),
(root, self.OUTBOX_SUBDIR),
]
tenants_root = os.path.join(root, 'tenants')
if os.path.isdir(tenants_root):
@@ -526,17 +776,18 @@ class BoxService:
continue
candidates.extend(
[
os.path.join(tenant_entry.path, self.INBOX_SUBDIR),
os.path.join(tenant_entry.path, self.OUTBOX_SUBDIR),
(tenant_entry.path, self.INBOX_SUBDIR),
(tenant_entry.path, self.OUTBOX_SUBDIR),
]
)
survivors: list[str] = []
for path in candidates:
if not os.path.isdir(path):
continue
shutil.rmtree(path, ignore_errors=True)
if os.path.exists(path):
survivors.append(path)
for candidate_root, subdir in candidates:
path = os.path.join(candidate_root, subdir)
try:
secure_fs.purge_subdirectory(candidate_root, subdir)
except OSError:
if os.path.lexists(path):
survivors.append(path)
return survivors
try:
@@ -651,14 +902,15 @@ class BoxService:
(the webchat session uses small sequential ids) never inherits stale
files from an earlier turn.
"""
import shutil
shutil.rmtree(host_dir, ignore_errors=True)
os.makedirs(host_dir, exist_ok=True)
query_key = os.path.basename(host_dir)
subdir = os.path.basename(os.path.dirname(host_dir))
tenant_root = os.path.dirname(os.path.dirname(host_dir))
try:
secure_fs.write_files(tenant_root, subdir, query_key, files)
except secure_fs.UnsafeWorkspacePathError as exc:
raise BoxValidationError('Sandbox attachment path contains an unsafe symbolic link') from exc
written: list[str] = []
for name, data in files:
with open(os.path.join(host_dir, name), 'wb') as fh:
fh.write(data)
for name, _data in files:
written.append(f'{target_mount_dir}/{name}')
return written
@@ -729,6 +981,8 @@ class BoxService:
"""
if not self._available:
return []
if self._cloud_managed:
await self.require_workspace_sandbox(self._query_execution_context(query))
import langbot_plugin.api.entities.builtin.platform.message as platform_message
@@ -777,7 +1031,8 @@ class BoxService:
if not pending:
return []
target_dir = f'{self.INBOX_MOUNT_DIR}/{query.query_id}'
query_key = self._attachment_query_key(query)
target_dir = f'{self.INBOX_MOUNT_DIR}/{query_key}'
written = await self._write_files_into_sandbox(query, self.INBOX_SUBDIR, target_dir, pending)
written_basenames = {os.path.basename(p) for p in written}
@@ -805,6 +1060,8 @@ class BoxService:
"""
if not self._available:
return []
if self._cloud_managed:
await self.require_workspace_sandbox(self._query_execution_context(query))
host_dir = self._host_query_dir(self.OUTBOX_SUBDIR, query)
if host_dir is not None:
@@ -828,22 +1085,21 @@ class BoxService:
"""Read outbox files straight off the bind-mounted host directory."""
import base64 as _b64
entries: list[dict] = []
if not os.path.isdir(host_dir):
return entries
for root, _dirs, names in os.walk(host_dir):
for name in sorted(names):
path = os.path.join(root, name)
try:
if os.path.getsize(path) > self._ATTACHMENT_MAX_BYTES:
continue
with open(path, 'rb') as fh:
data = fh.read()
except OSError:
continue
rel = os.path.relpath(path, host_dir)
entries.append({'name': rel, 'b64': _b64.b64encode(data).decode('ascii')})
return entries
query_key = os.path.basename(host_dir)
subdir = os.path.basename(os.path.dirname(host_dir))
tenant_root = os.path.dirname(os.path.dirname(host_dir))
try:
files = secure_fs.read_regular_files(
tenant_root,
subdir,
query_key,
max_file_bytes=self._ATTACHMENT_MAX_BYTES,
max_files=self._ATTACHMENT_MAX_FILES,
max_total_bytes=self._ATTACHMENT_MAX_TOTAL_BYTES,
)
except secure_fs.UnsafeWorkspacePathError as exc:
raise BoxValidationError('Sandbox outbox contains an unsafe symbolic link') from exc
return [{'name': name, 'b64': _b64.b64encode(data).decode('ascii')} for name, data in files]
async def _read_outbox_via_exec(self, query: pipeline_query.Query) -> list[dict]:
"""Fallback: read the outbox over the exec channel (E2B / remote).
@@ -853,7 +1109,7 @@ class BoxService:
"""
import json as _json
target_dir = f'{self.OUTBOX_MOUNT_DIR}/{query.query_id}'
target_dir = f'{self.OUTBOX_MOUNT_DIR}/{self._attachment_query_key(query)}'
max_bytes = self._EXEC_FALLBACK_MAX_BYTES
script = (
'import base64, json, os\n'
@@ -898,16 +1154,19 @@ class BoxService:
container's root can remove its own files. Best-effort: never raise
into the pipeline.
"""
target_dir = f'{self.OUTBOX_MOUNT_DIR}/{query.query_id}'
target_dir = f'{self.OUTBOX_MOUNT_DIR}/{self._attachment_query_key(query)}'
if host_dir is not None:
import shutil
def _clear() -> bool:
shutil.rmtree(host_dir, ignore_errors=True)
survived = os.path.exists(host_dir) and bool(os.listdir(host_dir))
os.makedirs(host_dir, exist_ok=True)
return survived
query_key = os.path.basename(host_dir)
subdir = os.path.basename(os.path.dirname(host_dir))
tenant_root = os.path.dirname(os.path.dirname(host_dir))
try:
secure_fs.reset_directory(tenant_root, subdir, query_key)
return False
except OSError:
return True
survived = await asyncio.to_thread(_clear)
if not survived:
@@ -977,9 +1236,9 @@ class BoxService:
self._runtime_connector.dispose()
async def get_sessions(self, context: TenantContext) -> list[dict]:
execution_context = await self._validated_execution_context(context)
if not self._available:
return []
execution_context = await self.require_workspace_sandbox(context)
try:
return await self.client.get_sessions(action_context=self._action_context(execution_context))
except Exception:
@@ -1017,7 +1276,8 @@ class BoxService:
skip_host_mount_validation: bool = False,
) -> dict:
execution_context = await self._validated_execution_context(context)
spec_payload = dict(spec_payload)
spec_payload = self._managed_policy_payload(execution_context, spec_payload)
await self._require_validated_workspace_sandbox(execution_context)
if spec_payload.get('host_path') in (None, ''):
tenant_workspace = self._tenant_workspace(execution_context)
if tenant_workspace is not None:
@@ -1033,6 +1293,7 @@ class BoxService:
session_id: str,
process_payload: dict,
) -> BoxManagedProcessInfo:
self._reject_cloud_managed_process()
execution_context = await self._validated_execution_context(context)
process_spec = BoxManagedProcessSpec.model_validate(process_payload)
return await self.client.start_managed_process(
@@ -1047,6 +1308,7 @@ class BoxService:
session_id: str,
process_id: str = 'default',
) -> BoxManagedProcessInfo:
self._reject_cloud_managed_process()
execution_context = await self._validated_execution_context(context)
return await self.client.get_managed_process(
session_id,
@@ -1060,6 +1322,7 @@ class BoxService:
session_id: str,
process_id: str = 'default',
) -> None:
self._reject_cloud_managed_process()
execution_context = await self._validated_execution_context(context)
return await self.client.stop_managed_process(
session_id,
@@ -1101,6 +1364,7 @@ class BoxService:
execution context used by the action RPC that created the process.
"""
self._reject_cloud_managed_process()
execution_context = await self._validated_execution_context(context)
if self._runtime_connector is None:
raise BoxValidationError(
@@ -1117,23 +1381,32 @@ class BoxService:
)
async def list_skills(self, context: TenantContext) -> list[dict]:
execution_context = await self._validated_execution_context(context)
execution_context = await self._validated_skill_execution_context(context)
return await self.client.list_skills(action_context=self._action_context(execution_context))
async def get_skill(self, context: TenantContext, name: str) -> dict | None:
execution_context = await self._validated_execution_context(context)
execution_context = await self._validated_skill_execution_context(context)
return await self.client.get_skill(name, action_context=self._action_context(execution_context))
async def create_skill(self, context: TenantContext, skill: dict) -> dict:
execution_context = await self._validated_execution_context(context)
execution_context = await self._validated_skill_execution_context(context)
payload = dict(skill)
payload.pop('workspace_uuid', None)
if self._cloud_managed and str(payload.get('package_root', '') or '').strip():
raise BoxAdmissionError('Cloud skill package_root is runtime-owned')
if self._cloud_managed:
payload.pop('package_root', None)
return await self.client.create_skill(payload, action_context=self._action_context(execution_context))
async def update_skill(self, context: TenantContext, name: str, skill: dict) -> dict:
execution_context = await self._validated_execution_context(context)
execution_context = await self._validated_skill_execution_context(context)
payload = dict(skill)
payload.pop('workspace_uuid', None)
if self._cloud_managed:
# The runtime already owns the package path for an existing skill.
# A serialized read response may contain it, but it is never an
# authority-bearing update field in shared Cloud mode.
payload.pop('package_root', None)
return await self.client.update_skill(
name,
payload,
@@ -1141,13 +1414,20 @@ class BoxService:
)
async def delete_skill(self, context: TenantContext, name: str) -> None:
execution_context = await self._validated_execution_context(context)
execution_context = await self._validated_skill_execution_context(context)
await self.client.delete_skill(name, action_context=self._action_context(execution_context))
async def scan_skill_directory(self, context: TenantContext, path: str) -> dict:
execution_context = await self._validated_execution_context(context)
execution_context = await self._validated_skill_execution_context(context)
if self._cloud_managed:
raise BoxAdmissionError('Scanning arbitrary host skill directories is disabled in Cloud')
return await self.client.scan_skill_directory(path, action_context=self._action_context(execution_context))
async def _validated_skill_execution_context(self, context: TenantContext) -> ExecutionContext:
execution_context = await self._validated_execution_context(context)
await self._require_validated_workspace_sandbox(execution_context)
return execution_context
async def list_skill_files(
self,
context: TenantContext,
@@ -1156,7 +1436,7 @@ class BoxService:
include_hidden: bool = False,
max_entries: int = 200,
) -> dict:
execution_context = await self._validated_execution_context(context)
execution_context = await self._validated_skill_execution_context(context)
return await self.client.list_skill_files(
name,
path,
@@ -1166,7 +1446,7 @@ class BoxService:
)
async def read_skill_file(self, context: TenantContext, name: str, path: str) -> dict:
execution_context = await self._validated_execution_context(context)
execution_context = await self._validated_skill_execution_context(context)
return await self.client.read_skill_file(
name,
path,
@@ -1174,7 +1454,7 @@ class BoxService:
)
async def write_skill_file(self, context: TenantContext, name: str, path: str, content: str) -> dict:
execution_context = await self._validated_execution_context(context)
execution_context = await self._validated_skill_execution_context(context)
return await self.client.write_skill_file(
name,
path,
@@ -1190,7 +1470,7 @@ class BoxService:
source_subdir: str = '',
target_suffix: str = 'upload',
) -> list[dict]:
execution_context = await self._validated_execution_context(context)
execution_context = await self._validated_skill_execution_context(context)
return await self.client.preview_skill_zip(
file_bytes,
filename,
@@ -1209,7 +1489,7 @@ class BoxService:
source_subdir: str = '',
target_suffix: str = 'upload',
) -> list[dict]:
execution_context = await self._validated_execution_context(context)
execution_context = await self._validated_skill_execution_context(context)
return await self.client.install_skill_zip(
file_bytes,
filename,
@@ -1430,6 +1710,20 @@ class BoxService:
allowed_roots = ', '.join(self.allowed_mount_roots)
raise BoxValidationError(f'box.local.default_workspace is outside allowed_mount_roots: {allowed_roots}')
def _ensure_cloud_shared_workspace(self) -> None:
"""Require the Core-side view of the Cloud Box durable volume."""
if self.default_workspace is None:
raise BoxValidationError('Cloud Box requires box.local.default_workspace')
if not os.path.isabs(self.default_workspace) or not os.path.isdir(self.default_workspace):
raise BoxValidationError('Cloud Box shared default_workspace must be an existing absolute directory')
if not os.access(self.default_workspace, os.R_OK | os.W_OK | os.X_OK):
raise BoxValidationError('Cloud Box shared default_workspace must be writable by LangBot Core')
if not self.allowed_mount_roots or not any(
_is_path_under(self.default_workspace, allowed_root) for allowed_root in self.allowed_mount_roots
):
raise BoxValidationError('Cloud Box shared default_workspace is outside allowed_mount_roots')
def _validate_host_mount(self, spec: BoxSpec):
if spec.host_path is None:
return
@@ -1572,13 +1866,13 @@ class BoxService:
and error.get('workspace_uuid') == execution_context.workspace_uuid
]
def get_system_guidance(self, query_id=None) -> str:
def get_system_guidance(self, query: pipeline_query.Query | int | str | None = None) -> str:
"""Return LLM system-prompt guidance for the exec tool.
All execution-specific prompt text is kept here so that callers
(e.g. LocalAgentRunner) stay free of box domain knowledge.
``query_id`` is the current turn's pipeline query id. When provided,
``query`` is the current turn's pipeline query. When provided,
the guidance ALWAYS advertises the per-query outbox path so the agent
knows how to deliver generated files back to the user even on turns
where the user sent no inbound attachment (e.g. "generate a QR code"),
@@ -1599,8 +1893,17 @@ class BoxService:
'modify local files in the working directory, use exec with /workspace paths directly; do not ask the '
'user for directory parameters unless they explicitly need a different directory.'
)
if query_id is not None:
outbox_dir = f'{self.OUTBOX_MOUNT_DIR}/{query_id}'
if query is not None:
if not isinstance(query, (int, str)):
query_key = self._attachment_query_key(query)
else:
# Backwards compatibility for OSS callers/tests that passed
# the old process-local integer identity. Cloud callers must
# pass the full Query so an opaque UUID is always advertised.
if self._cloud_managed:
raise BoxValidationError('Cloud outbox guidance requires a pipeline Query')
query_key = str(query)
outbox_dir = f'{self.OUTBOX_MOUNT_DIR}/{query_key}'
guidance += (
f' If you produce any file (image, audio, document, etc.) that should be sent back to the user, '
f'write it into {outbox_dir}/ (create the directory if needed). Every file placed there will be '
@@ -1618,6 +1921,8 @@ class BoxService:
async def get_status(self, context: TenantContext) -> dict:
execution_context = await self._validated_execution_context(context)
if self._cloud_managed and self._available:
await self._require_validated_workspace_sandbox(execution_context)
action_context = self._action_context(execution_context)
recent_error_count = len(self.get_recent_errors(execution_context))
if not self._available:
+15 -9
View File
@@ -126,24 +126,30 @@ def should_prepare_python_env(host_path: str | None) -> bool:
return bool(list_python_manifest_files(normalized_root))
def wrap_python_command_with_env(command: str, *, mount_path: str = '/workspace') -> str:
def wrap_python_command_with_env(
command: str,
*,
mount_path: str = '/workspace',
state_path: str | None = None,
) -> str:
"""Wrap a command with a reusable sandbox-local Python env bootstrap.
This is the generic "workspace is a Python project" path used by mutable
workspaces such as skills. Read-only installation strategies stay in the
higher-level caller because they are application policy, not workspace
semantics.
``mount_path`` is always the source tree used for manifest hashing and
installation. ``state_path`` may point at a separate writable directory
for read-only source mounts; when omitted, legacy mutable-workspace behavior
stores the environment beside the source.
"""
writable_state_path = state_path or mount_path
bootstrap = textwrap.dedent(
f"""
set -e
_LB_VENV_DIR="{mount_path}/.venv"
_LB_META_DIR="{mount_path}/.langbot"
_LB_VENV_DIR="{writable_state_path}/.venv"
_LB_META_DIR="{writable_state_path}/.langbot"
_LB_META_FILE="$_LB_META_DIR/python-env.json"
_LB_LOCK_DIR="$_LB_META_DIR/python-env.lock"
_LB_TMP_DIR="{mount_path}/.tmp"
_LB_PIP_CACHE_DIR="{mount_path}/.cache/pip"
_LB_TMP_DIR="{writable_state_path}/.tmp"
_LB_PIP_CACHE_DIR="{writable_state_path}/.cache/pip"
mkdir -p "$_LB_META_DIR" "$_LB_TMP_DIR" "$_LB_PIP_CACHE_DIR"
_LB_SYSTEM_PYTHON="$(command -v python3 || command -v python || true)"
+161
View File
@@ -3,6 +3,8 @@ from __future__ import annotations
import dataclasses
import importlib.metadata
import inspect
import os
import threading
import time
from collections.abc import Awaitable, Callable
from typing import Any, Protocol
@@ -13,12 +15,17 @@ from .entitlements import EntitlementProvider, OpenSourceEntitlementProvider
CLOUD_BOOTSTRAP_ENTRY_POINT = 'langbot.cloud_bootstrap'
REQUIRED_TENANT_ISOLATION_VERSION = 2
SUPPORTED_PGVECTOR_DIMENSIONS = frozenset({384, 512, 768, 1024, 1536})
class CloudBootstrapError(RuntimeError):
"""Fail-closed Cloud bootstrap validation error."""
class CloudRuntimeUnavailableError(CloudBootstrapError):
"""The verified Cloud receipt no longer admits runtime work."""
@dataclasses.dataclass(frozen=True, slots=True)
class OpenSourceDeployment:
"""Default deployment selected when no closed bootstrap is installed."""
@@ -88,8 +95,70 @@ class VerifiedCloudDeployment:
raise CloudBootstrapError('Cloud runtime requires database.use=postgresql')
if config.get('vdb', {}).get('use') != self.required_vector_backend:
raise CloudBootstrapError('Cloud runtime requires vdb.use=pgvector')
pgvector_config = config.get('vdb', {}).get('pgvector', {})
if pgvector_config.get('use_business_database') is not True:
raise CloudBootstrapError('Cloud runtime requires vdb.pgvector.use_business_database=true')
dimensions = pgvector_config.get('allowed_dimensions')
if (
not isinstance(dimensions, list)
or not dimensions
or any(isinstance(item, bool) or not isinstance(item, int) for item in dimensions)
or not set(dimensions).issubset(SUPPORTED_PGVECTOR_DIMENSIONS)
):
supported = ', '.join(str(item) for item in sorted(SUPPORTED_PGVECTOR_DIMENSIONS))
raise CloudBootstrapError(f'Cloud pgvector allowed_dimensions must be a non-empty subset of: {supported}')
if config.get('mcp', {}).get('stdio', {}).get('enabled', True) is not False:
raise CloudBootstrapError('Cloud runtime requires mcp.stdio.enabled=false')
plugin_worker = config.get('plugin', {}).get('worker', {})
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')
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()
if not runtime_endpoint:
raise CloudBootstrapError('Cloud runtime requires a shared external box.runtime.endpoint')
admission = box_config.get('admission', {})
required_admission = {
'required': True,
'logical_session_id': 'global',
'required_backend': 'nsjail',
'max_sessions': 1,
'max_managed_processes': 0,
}
if any(admission.get(name) != value for name, value in required_admission.items()):
raise CloudBootstrapError(
'Cloud runtime requires grant-enforced Box admission with one global session and zero managed processes'
)
grant_ttl = admission.get('max_grant_ttl_sec')
if isinstance(grant_ttl, bool) or not isinstance(grant_ttl, int) or not 1 <= grant_ttl <= 300:
raise CloudBootstrapError('Cloud Box admission max_grant_ttl_sec must be between 1 and 300')
workspace_quota_mb = admission.get('workspace_quota_mb')
if isinstance(workspace_quota_mb, bool) or not isinstance(workspace_quota_mb, int) or workspace_quota_mb <= 0:
raise CloudBootstrapError('Cloud Box admission workspace_quota_mb must be a positive integer')
local_config = box_config.get('local', {})
host_root = str(local_config.get('host_root', '') or '').strip()
default_workspace = str(local_config.get('default_workspace', '') or '').strip()
allowed_mount_roots = local_config.get('allowed_mount_roots')
if not host_root or not os.path.isabs(host_root):
raise CloudBootstrapError('Cloud Box local.host_root must be an absolute shared-volume path')
if not default_workspace or not os.path.isabs(default_workspace):
raise CloudBootstrapError('Cloud Box local.default_workspace must be an absolute shared-volume path')
if (
not isinstance(allowed_mount_roots, list)
or not allowed_mount_roots
or any(not isinstance(root, str) or not os.path.isabs(root) for root in allowed_mount_roots)
):
raise CloudBootstrapError('Cloud Box local.allowed_mount_roots must contain absolute shared-volume paths')
resolved_workspace = os.path.realpath(default_workspace)
if not any(
resolved_workspace == os.path.realpath(root)
or resolved_workspace.startswith(f'{os.path.realpath(root)}{os.sep}')
for root in allowed_mount_roots
):
raise CloudBootstrapError('Cloud Box local.default_workspace must be under allowed_mount_roots')
class CloudBootstrapProvider(Protocol):
@@ -101,6 +170,98 @@ class CloudBootstrapProvider(Protocol):
) -> VerifiedCloudDeployment | Awaitable[VerifiedCloudDeployment]: ...
class DeploymentAdmissionGuard:
"""Continuously enforce one verified deployment receipt.
Startup verification alone is insufficient because a long-running process
could otherwise keep serving after the signed Manifest expires. The guard
tracks both wall-clock expiry and a monotonic deadline so moving the system
clock backwards cannot extend an already admitted receipt.
A closed bootstrap may atomically replace the receipt with a strictly newer
Manifest generation after performing its own signature verification. The
logical instance and deployment mode cannot change during the process.
"""
def __init__(
self,
instance_uuid: str,
deployment: OpenSourceDeployment | VerifiedCloudDeployment,
*,
wall_time: Callable[[], float] = time.time,
monotonic_time: Callable[[], float] = time.monotonic,
) -> None:
self.instance_uuid = instance_uuid
self._wall_time = wall_time
self._monotonic_time = monotonic_time
self._lock = threading.Lock()
self._deployment = deployment
self._deadline: float | None = None
self._install_initial(deployment)
@property
def deployment(self) -> OpenSourceDeployment | VerifiedCloudDeployment:
with self._lock:
return self._deployment
def _install_initial(self, deployment: OpenSourceDeployment | VerifiedCloudDeployment) -> None:
now = int(self._wall_time())
if isinstance(deployment, VerifiedCloudDeployment):
deployment.validate(self.instance_uuid, now=now)
self._deadline = self._monotonic_time() + (deployment.expires_at - now)
elif not isinstance(deployment, OpenSourceDeployment):
raise TypeError('Deployment admission requires a verified deployment object')
@staticmethod
def _receipt_identity(deployment: VerifiedCloudDeployment) -> tuple[Any, ...]:
return (
deployment.instance_uuid,
deployment.manifest_jti,
deployment.manifest_generation,
deployment.expires_at,
deployment.release,
tuple(sorted(deployment.capabilities)),
deployment.tenant_isolation_version,
deployment.verification_key_id,
)
def replace(self, deployment: VerifiedCloudDeployment) -> None:
"""Atomically install a verified, non-rollback Cloud receipt."""
now = int(self._wall_time())
deployment.validate(self.instance_uuid, now=now)
with self._lock:
current = self._deployment
if not isinstance(current, VerifiedCloudDeployment):
raise CloudRuntimeUnavailableError('Deployment mode cannot change while LangBot is running')
if deployment.manifest_generation < current.manifest_generation:
raise CloudRuntimeUnavailableError('Cloud Manifest generation rolled back')
if deployment.manifest_generation == current.manifest_generation and self._receipt_identity(
deployment
) != self._receipt_identity(current):
raise CloudRuntimeUnavailableError('Cloud Manifest generation has conflicting contents')
self._deployment = deployment
self._deadline = self._monotonic_time() + (deployment.expires_at - now)
def require_active(self) -> OpenSourceDeployment | VerifiedCloudDeployment:
"""Return the active deployment or fail closed after Manifest expiry."""
now = int(self._wall_time())
monotonic_now = self._monotonic_time()
with self._lock:
deployment = self._deployment
deadline = self._deadline
if isinstance(deployment, OpenSourceDeployment):
return deployment
try:
deployment.validate(self.instance_uuid, now=now)
except CloudBootstrapError as exc:
raise CloudRuntimeUnavailableError(str(exc)) from exc
if deadline is None or monotonic_now >= deadline:
raise CloudRuntimeUnavailableError('Verified Cloud Manifest is expired')
return deployment
async def _invoke_provider(
loaded: Any,
*,
+26 -3
View File
@@ -3,6 +3,7 @@ from __future__ import annotations
import asyncio
import json
import time
from collections.abc import Callable
from typing import Protocol, runtime_checkable
import pydantic
@@ -11,6 +12,10 @@ import pydantic
class EntitlementUnavailableError(RuntimeError):
"""Raised when a trusted, currently-active entitlement is unavailable."""
def __init__(self, message: str, *, entitlement_revision: int | None = None) -> None:
super().__init__(message)
self.entitlement_revision = entitlement_revision
class EntitlementSnapshot(pydantic.BaseModel):
"""Plan-agnostic capability projection consumed by open-source Core.
@@ -60,9 +65,15 @@ class EntitlementSnapshot(pydantic.BaseModel):
if self.instance_uuid != instance_uuid or self.workspace_uuid != workspace_uuid:
raise EntitlementUnavailableError('Entitlement scope does not match the Workspace execution context')
if self.status != 'active':
raise EntitlementUnavailableError('Workspace entitlement is not active')
raise EntitlementUnavailableError(
'Workspace entitlement is not active',
entitlement_revision=self.entitlement_revision,
)
if current_time < self.not_before or current_time >= self.expires_at:
raise EntitlementUnavailableError('Workspace entitlement is not currently valid')
raise EntitlementUnavailableError(
'Workspace entitlement is not currently valid',
entitlement_revision=self.entitlement_revision,
)
return self
def require_feature(self, feature: str) -> None:
@@ -95,9 +106,16 @@ class OpenSourceEntitlementProvider:
class EntitlementResolver:
"""Validate scope/freshness and reject revision rollback or equivocation."""
def __init__(self, instance_uuid: str, provider: EntitlementProvider) -> None:
def __init__(
self,
instance_uuid: str,
provider: EntitlementProvider,
*,
deployment_admission: Callable[[], object] | None = None,
) -> None:
self.instance_uuid = instance_uuid
self.provider = provider
self._deployment_admission = deployment_admission
self._lock = asyncio.Lock()
self._snapshots: dict[str, tuple[int, str, EntitlementSnapshot]] = {}
@@ -112,7 +130,12 @@ class EntitlementResolver:
minimum_revision: int = 0,
now: int | None = None,
) -> EntitlementSnapshot:
if self._deployment_admission is not None:
self._deployment_admission()
candidate = await self.provider.get_workspace_entitlement(workspace_uuid)
if self._deployment_admission is not None:
# A provider call may cross the Manifest expiry boundary.
self._deployment_admission()
if not isinstance(candidate, EntitlementSnapshot):
raise EntitlementUnavailableError('Entitlement provider returned an invalid snapshot')
# Deep-copy untrusted provider-owned containers before caching them.
+7
View File
@@ -1,6 +1,7 @@
from __future__ import annotations
import typing
import inspect
from ..core import app
from . import operator
@@ -63,6 +64,12 @@ class CommandManager:
) -> typing.AsyncGenerator[command_context.CommandReturn, None]:
"""执行命令"""
require_context = getattr(self.ap.plugin_connector, 'require_workspace_context', None)
if require_context is not None:
result = require_context(context)
if inspect.isawaitable(result):
await result
command_list = await self.ap.plugin_connector.list_commands(bound_plugins)
for command in command_list:
+12 -24
View File
@@ -4,7 +4,6 @@ import logging
import asyncio
import traceback
import os
import sqlalchemy
from ..platform import botmgr as im_mgr
from ..platform.webhook_pusher import WebhookPusher
@@ -51,7 +50,6 @@ from ..workspace import collaboration as workspace_collaboration_module
from ..cloud import bootstrap as cloud_bootstrap_module
from ..cloud import entitlements as cloud_entitlements_module
from ..api.http.context import ExecutionContext, PrincipalContext, PrincipalType
from ..entity.persistence.workspace import WorkspaceExecutionState, WorkspaceExecutionStatus
class Application:
@@ -132,6 +130,8 @@ class Application:
deployment: cloud_bootstrap_module.OpenSourceDeployment | cloud_bootstrap_module.VerifiedCloudDeployment = None
deployment_admission: cloud_bootstrap_module.DeploymentAdmissionGuard = None
entitlement_resolver: cloud_entitlements_module.EntitlementResolver | None = None
vector_db_mgr: vectordb_mgr.VectorDBManager = None
@@ -251,18 +251,12 @@ class Application:
check_interval_seconds = check_interval_hours * 3600
while True:
try:
execution_states = await self.persistence_mgr.execute_async(
sqlalchemy.select(WorkspaceExecutionState).where(
WorkspaceExecutionState.instance_uuid == self.workspace_service.instance_uuid,
WorkspaceExecutionState.state == WorkspaceExecutionStatus.ACTIVE.value,
WorkspaceExecutionState.write_fenced == sqlalchemy.false(),
)
)
for execution_state in execution_states.all():
bindings = await self.workspace_service.list_active_execution_bindings()
for binding in bindings:
context = ExecutionContext(
instance_uuid=execution_state.instance_uuid,
workspace_uuid=execution_state.workspace_uuid,
placement_generation=execution_state.active_generation,
instance_uuid=binding.instance_uuid,
workspace_uuid=binding.workspace_uuid,
placement_generation=binding.placement_generation,
trigger_principal=PrincipalContext(PrincipalType.SYSTEM),
)
deleted = await self.monitoring_service.cleanup_expired_records(
@@ -299,18 +293,12 @@ class Application:
check_interval_seconds = check_interval_hours * 3600
while True:
try:
execution_states = await self.persistence_mgr.execute_async(
sqlalchemy.select(WorkspaceExecutionState).where(
WorkspaceExecutionState.instance_uuid == self.workspace_service.instance_uuid,
WorkspaceExecutionState.state == WorkspaceExecutionStatus.ACTIVE.value,
WorkspaceExecutionState.write_fenced == sqlalchemy.false(),
)
)
for execution_state in execution_states.all():
bindings = await self.workspace_service.list_active_execution_bindings()
for binding in bindings:
context = ExecutionContext(
instance_uuid=execution_state.instance_uuid,
workspace_uuid=execution_state.workspace_uuid,
placement_generation=execution_state.active_generation,
instance_uuid=binding.instance_uuid,
workspace_uuid=binding.workspace_uuid,
placement_generation=binding.placement_generation,
trigger_principal=PrincipalContext(PrincipalType.SYSTEM),
)
deleted = await self.maintenance_service.cleanup_expired_files(context)
+9 -1
View File
@@ -59,8 +59,16 @@ class BuildAppStage(stage.BootingStage):
instance_config=ap.instance_config.data,
)
ap.deployment = deployment
ap.deployment_admission = cloud_bootstrap.DeploymentAdmissionGuard(
constants.instance_id,
deployment,
)
ap.entitlement_resolver = (
EntitlementResolver(constants.instance_id, deployment.entitlement_provider)
EntitlementResolver(
constants.instance_id,
deployment.entitlement_provider,
deployment_admission=ap.deployment_admission.require_active,
)
if deployment.multi_workspace_enabled
else None
)
+11 -2
View File
@@ -21,6 +21,7 @@ _RUNTIME_POLICY_DEFAULTS = {
'max_pids': 128,
'max_open_files': 256,
'max_file_size_mb': 512,
'require_hard_limits': False,
}
},
'mcp': {'stdio': {'enabled': True}},
@@ -103,11 +104,19 @@ def _apply_env_overrides_to_config(cfg: dict) -> dict:
if '__' not in env_key:
continue
print(f'apply env overrides to config: env_key: {env_key}, env_value: {env_value}')
# Convert environment variable name to config path
# e.g., CONCURRENCY__PIPELINE -> ['concurrency', 'pipeline']
keys = [key.lower() for key in env_key.split('__')]
# macOS and some launchers expose variables such as
# ``__CF_USER_TEXT_ENCODING``. They are not LangBot config paths and
# must not create an empty top-level YAML key when config is dumped.
if any(not key for key in keys):
continue
# Values may contain database passwords, runtime control tokens, or
# provider credentials. Keep the useful audit breadcrumb without ever
# copying the secret into startup logs.
print(f'apply env override to config: env_key: {env_key}')
# Navigate to the target value and validate the path
current = cfg
+74
View File
@@ -0,0 +1,74 @@
from __future__ import annotations
import asyncio
import contextvars
import typing
T = typing.TypeVar('T')
def create_detached_task(
coro: typing.Coroutine[typing.Any, typing.Any, T],
*,
loop: asyncio.AbstractEventLoop | None = None,
name: str | None = None,
after_commit_manager: typing.Any | None = None,
) -> asyncio.Task[T]:
"""Create a task that inherits no request-local ContextVars.
A normal ``asyncio.create_task`` copies the caller's context. That is
unsafe for work which outlives an HTTP request because it can copy the
request's active database transaction or trusted tenant scope into a
different asyncio task. Detached work must receive durable identity such
as ``ExecutionContext`` through explicit arguments and establish its own
tenant scope or unit of work whenever it accesses persistence.
"""
task_loop = loop or asyncio.get_running_loop()
gate: asyncio.Future[None] | None = None
# Inspect the type so dynamic Mock/AsyncMock attributes do not turn into a
# fake gate in lightweight tests or embedders.
gate_factory = getattr(type(after_commit_manager), 'create_after_commit_gate', None)
if callable(gate_factory):
gate = gate_factory(after_commit_manager)
task_coro = _wait_for_commit(coro, gate) if gate is not None else coro
return task_loop.create_task(task_coro, name=name, context=contextvars.Context())
async def _wait_for_commit(
coro: typing.Coroutine[typing.Any, typing.Any, T],
gate: asyncio.Future[None],
) -> T:
try:
await gate
except BaseException:
coro.close()
raise
return await coro
async def run_in_workspace_uow(
ap: typing.Any,
workspace_uuid: str,
operation: typing.Callable[[], typing.Awaitable[T]],
) -> T:
"""Run one short persistence section in a detached Cloud task scope.
This helper deliberately scopes only the supplied operation. Callers
should not wrap long-running network or runtime work in a database
transaction.
"""
persistence_mgr = getattr(ap, 'persistence_mgr', None)
if persistence_mgr is None:
return await operation()
cloud_runtime = getattr(getattr(persistence_mgr, 'mode', None), 'value', None) == 'cloud_runtime'
if not cloud_runtime:
return await operation()
tenant_uow = getattr(persistence_mgr, 'tenant_uow', None)
if not callable(tenant_uow):
raise RuntimeError('Detached Cloud tasks require an explicit tenant unit of work')
async with tenant_uow(workspace_uuid):
return await operation()
+7 -1
View File
@@ -7,6 +7,7 @@ import time
from . import app
from . import entities as core_entities
from .task_boundary import create_detached_task
class TaskContext:
@@ -125,7 +126,12 @@ class TaskWrapper:
TaskWrapper._id_index += 1
self.ap = ap
self.task_context = context or TaskContext()
self.task = self.ap.event_loop.create_task(coro)
self.task = create_detached_task(
coro,
loop=self.ap.event_loop,
name=name or None,
after_commit_manager=getattr(self.ap, 'persistence_mgr', None),
)
self.task_type = task_type
self.kind = kind
self.name = name
+26 -1
View File
@@ -1,3 +1,6 @@
import hashlib
import uuid
import sqlalchemy
from .base import Base
@@ -15,6 +18,17 @@ class PluginSetting(Base):
)
plugin_author = sqlalchemy.Column(sqlalchemy.String(255), primary_key=True)
plugin_name = sqlalchemy.Column(sqlalchemy.String(255), primary_key=True)
installation_uuid = sqlalchemy.Column(
sqlalchemy.String(36),
nullable=False,
default=lambda: str(uuid.uuid4()),
)
artifact_digest = sqlalchemy.Column(
sqlalchemy.String(64),
nullable=False,
default=lambda: hashlib.sha256(f'pending:{uuid.uuid4()}'.encode()).hexdigest(),
)
runtime_revision = sqlalchemy.Column(sqlalchemy.Integer, nullable=False, default=1)
enabled = sqlalchemy.Column(sqlalchemy.Boolean, nullable=False, default=True)
priority = sqlalchemy.Column(sqlalchemy.Integer, nullable=False, default=0)
config = sqlalchemy.Column(sqlalchemy.JSON, nullable=False, default=dict)
@@ -28,4 +42,15 @@ class PluginSetting(Base):
onupdate=sqlalchemy.func.now(),
)
__table_args__ = (sqlalchemy.Index('ix_plugin_settings_workspace_enabled', 'workspace_uuid', 'enabled'),)
__table_args__ = (
sqlalchemy.UniqueConstraint('installation_uuid', name='uq_plugin_settings_installation_uuid'),
sqlalchemy.CheckConstraint('runtime_revision >= 1', name='ck_plugin_settings_runtime_revision_positive'),
sqlalchemy.CheckConstraint('length(artifact_digest) = 64', name='ck_plugin_settings_artifact_digest_length'),
sqlalchemy.Index('ix_plugin_settings_workspace_enabled', 'workspace_uuid', 'enabled'),
sqlalchemy.Index(
'ix_plugin_settings_workspace_installation',
'workspace_uuid',
'installation_uuid',
unique=True,
),
)
@@ -29,6 +29,9 @@ class KnowledgeBase(Base):
)
creation_settings = sqlalchemy.Column(sqlalchemy.JSON, nullable=True, default=None)
retrieval_settings = sqlalchemy.Column(sqlalchemy.JSON, nullable=True, default=None)
# Server-selected pgvector dimension. ``None`` means no embedding has been
# written yet; the first pgvector upsert binds it atomically.
embedding_dimension = sqlalchemy.Column(sqlalchemy.Integer, nullable=True)
# Field sets for different operations
MUTABLE_FIELDS = {'name', 'description', 'retrieval_settings'}
@@ -40,6 +43,7 @@ class KnowledgeBase(Base):
ALL_DB_FIELDS = CREATE_FIELDS | {
'workspace_uuid',
'legacy_vector_collection',
'embedding_dimension',
'emoji',
'created_at',
'updated_at',
@@ -57,6 +61,10 @@ class KnowledgeBase(Base):
sqlite_where=sqlalchemy.text('collection_id IS NOT NULL'),
postgresql_where=sqlalchemy.text('collection_id IS NOT NULL'),
),
sqlalchemy.CheckConstraint(
'embedding_dimension IS NULL OR embedding_dimension > 0',
name='ck_knowledge_bases_embedding_dimension_positive',
),
)
@@ -1,11 +1,13 @@
"""enforce PostgreSQL tenant isolation with row-level security
"""enforce PostgreSQL tenant isolation with exact discovery contracts
Revision ID: 0011_postgres_tenant_rls
Revises: 0010_scope_resources
Create Date: 2026-07-19
The table list is deliberately duplicated from the runtime contract. Alembic
revisions must remain self-contained after application code evolves.
The table and policy lists are deliberately duplicated from the runtime
contract. Alembic revisions must remain self-contained after application code
evolves. Discovery policies are SELECT-only and reveal the minimum rows needed
to turn an authenticated credential into one Workspace transaction.
"""
from __future__ import annotations
@@ -20,8 +22,18 @@ depends_on = None
_POLICY_NAME = 'langbot_workspace_isolation'
_ACCOUNT_POLICY_NAME = 'langbot_account_discovery'
_API_KEY_POLICY_NAME = 'langbot_api_key_discovery'
_INVITATION_POLICY_NAME = 'langbot_invitation_discovery'
_INSTANCE_POLICY_NAME = 'langbot_instance_discovery'
_TENANT_SETTING = 'langbot.workspace_uuid'
_ACCOUNT_SETTING = 'langbot.account_uuid'
_API_KEY_HASH_SETTING = 'langbot.api_key_hash'
_INVITATION_HASH_SETTING = 'langbot.invitation_hash'
_INSTANCE_SETTING = 'langbot.instance_uuid'
_OSS_WORKSPACE_METADATA_KEY = 'oss_workspace_uuid'
_TENANT_TABLE_COLUMNS: dict[str, str] = {
'workspaces': 'uuid',
'workspace_memberships': 'workspace_uuid',
@@ -54,17 +66,42 @@ _TENANT_TABLE_COLUMNS: dict[str, str] = {
}
def _setting(name: str) -> str:
return f"NULLIF(current_setting('{name}', true), '')"
def _tenant_expression(column: str) -> str:
return f'{column}::text = {_setting(_TENANT_SETTING)}'
_DISCOVERY_POLICIES: dict[str, dict[str, str]] = {
'workspace_memberships': {
_ACCOUNT_POLICY_NAME: (f"account_uuid::text = {_setting(_ACCOUNT_SETTING)} AND status = 'active'"),
},
'workspace_execution_states': {
_INSTANCE_POLICY_NAME: (
f"instance_uuid = {_setting(_INSTANCE_SETTING)} AND state = 'active' AND write_fenced = false"
),
},
'api_keys': {
_API_KEY_POLICY_NAME: (
f"key_hash = {_setting(_API_KEY_HASH_SETTING)} AND status = 'active' "
'AND (expires_at IS NULL OR expires_at > CURRENT_TIMESTAMP)'
),
},
'workspace_invitations': {
_INVITATION_POLICY_NAME: f'token_hash = {_setting(_INVITATION_HASH_SETTING)}',
},
}
def _quote_identifier(conn: sa.Connection, identifier: str) -> str:
return conn.dialect.identifier_preparer.quote(identifier)
def _record_oss_workspace_scope(conn: sa.Connection) -> None:
"""Keep PostgreSQL OSS usable after FORCE RLS is enabled.
"""Keep PostgreSQL OSS usable after FORCE RLS is enabled."""
The runtime reads this instance-level metadata and installs the singleton
Workspace as a transaction-local default. Cloud databases have no local
Workspace and therefore do not receive this compatibility value.
"""
local_workspaces = (
conn.execute(sa.text("SELECT uuid FROM workspaces WHERE source = 'local' ORDER BY uuid")).scalars().all()
)
@@ -84,6 +121,43 @@ def _record_oss_workspace_scope(conn: sa.Connection) -> None:
raise RuntimeError('Stored OSS Workspace scope does not match the local Workspace')
def _drop_all_policies(conn: sa.Connection, table_name: str) -> None:
policy_names = conn.execute(
sa.text(
"""
SELECT p.polname
FROM pg_policy p
JOIN pg_class c ON c.oid = p.polrelid
JOIN pg_namespace n ON n.oid = c.relnamespace
WHERE n.nspname = current_schema() AND c.relname = :table_name
"""
),
{'table_name': table_name},
).scalars()
table = _quote_identifier(conn, table_name)
for policy in policy_names:
op.execute(sa.text(f'DROP POLICY {_quote_identifier(conn, policy)} ON {table}'))
def _create_policy(
conn: sa.Connection,
table_name: str,
policy_name: str,
expression: str,
*,
command: str,
) -> None:
table = _quote_identifier(conn, table_name)
policy = _quote_identifier(conn, policy_name)
if command == 'ALL':
sql = f'CREATE POLICY {policy} ON {table} AS PERMISSIVE FOR ALL TO PUBLIC USING ({expression}) WITH CHECK ({expression})'
elif command == 'SELECT':
sql = f'CREATE POLICY {policy} ON {table} AS PERMISSIVE FOR SELECT TO PUBLIC USING ({expression})'
else: # pragma: no cover - migration-local invariant
raise AssertionError(f'Unsupported RLS command: {command}')
op.execute(sa.text(sql))
def upgrade() -> None:
conn = op.get_bind()
if conn.dialect.name != 'postgresql':
@@ -96,22 +170,20 @@ def upgrade() -> None:
_record_oss_workspace_scope(conn)
policy_name = _quote_identifier(conn, _POLICY_NAME)
for table_name, tenant_column in _TENANT_TABLE_COLUMNS.items():
table = _quote_identifier(conn, table_name)
column = _quote_identifier(conn, tenant_column)
tenant_value = f"CAST(NULLIF(current_setting('{_TENANT_SETTING}', true), '') AS VARCHAR(36))"
_drop_all_policies(conn, table_name)
op.execute(sa.text(f'ALTER TABLE {table} ENABLE ROW LEVEL SECURITY'))
op.execute(sa.text(f'ALTER TABLE {table} FORCE ROW LEVEL SECURITY'))
op.execute(sa.text(f'DROP POLICY IF EXISTS {policy_name} ON {table}'))
op.execute(
sa.text(
f'CREATE POLICY {policy_name} ON {table} '
f'USING ({column} = {tenant_value}) '
f'WITH CHECK ({column} = {tenant_value})'
)
_create_policy(
conn,
table_name,
_POLICY_NAME,
_tenant_expression(_quote_identifier(conn, tenant_column)),
command='ALL',
)
for policy_name, expression in _DISCOVERY_POLICIES.get(table_name, {}).items():
_create_policy(conn, table_name, policy_name, expression, command='SELECT')
def downgrade() -> None:
@@ -120,11 +192,10 @@ def downgrade() -> None:
return
existing_tables = set(sa.inspect(conn).get_table_names())
policy_name = _quote_identifier(conn, _POLICY_NAME)
for table_name in _TENANT_TABLE_COLUMNS:
if table_name not in existing_tables:
continue
table = _quote_identifier(conn, table_name)
op.execute(sa.text(f'DROP POLICY IF EXISTS {policy_name} ON {table}'))
_drop_all_policies(conn, table_name)
op.execute(sa.text(f'ALTER TABLE {table} NO FORCE ROW LEVEL SECURITY'))
op.execute(sa.text(f'ALTER TABLE {table} DISABLE ROW LEVEL SECURITY'))
@@ -0,0 +1,179 @@
"""add immutable plugin installation identity
Revision ID: 0012_plugin_identity
Revises: 0011_postgres_tenant_rls
Create Date: 2026-07-19
The migration gives every legacy row a random, stable installation UUID. A
legacy artifact digest is only a valid SHA-256-shaped recovery marker; Core
replaces it with the package digest and increments ``runtime_revision`` before
the next package apply.
"""
from __future__ import annotations
import hashlib
import uuid
import sqlalchemy as sa
from alembic import op
revision = '0012_plugin_identity'
down_revision = '0011_postgres_tenant_rls'
branch_labels = None
depends_on = None
_TABLE = 'plugin_settings'
_INSTALLATION_INDEX = 'ix_plugin_settings_workspace_installation'
_INSTALLATION_UNIQUE = 'uq_plugin_settings_installation_uuid'
_REVISION_CHECK = 'ck_plugin_settings_runtime_revision_positive'
_DIGEST_CHECK = 'ck_plugin_settings_artifact_digest_length'
def _column_names(conn: sa.Connection) -> set[str]:
inspector = sa.inspect(conn)
if _TABLE not in inspector.get_table_names():
return set()
return {column['name'] for column in inspector.get_columns(_TABLE)}
def _legacy_digest(installation_uuid: str) -> str:
return hashlib.sha256(f'legacy-installation:{installation_uuid}'.encode()).hexdigest()
def _suspend_postgres_rls(conn: sa.Connection) -> tuple[bool, bool]:
"""Let the release migration backfill every tenant row after revision 0011."""
if conn.dialect.name != 'postgresql':
return False, False
row = conn.execute(
sa.text('SELECT relrowsecurity, relforcerowsecurity FROM pg_class WHERE oid = to_regclass(:table_name)'),
{'table_name': _TABLE},
).one()
rls_enabled, rls_forced = bool(row.relrowsecurity), bool(row.relforcerowsecurity)
if rls_forced:
op.execute(sa.text(f'ALTER TABLE {_TABLE} NO FORCE ROW LEVEL SECURITY'))
if rls_enabled:
op.execute(sa.text(f'ALTER TABLE {_TABLE} DISABLE ROW LEVEL SECURITY'))
return rls_enabled, rls_forced
def _restore_postgres_rls(conn: sa.Connection, state: tuple[bool, bool]) -> None:
if conn.dialect.name != 'postgresql':
return
rls_enabled, rls_forced = state
if rls_enabled:
op.execute(sa.text(f'ALTER TABLE {_TABLE} ENABLE ROW LEVEL SECURITY'))
if rls_forced:
op.execute(sa.text(f'ALTER TABLE {_TABLE} FORCE ROW LEVEL SECURITY'))
def _backfill(conn: sa.Connection) -> None:
table = sa.table(
_TABLE,
sa.column('workspace_uuid', sa.String(36)),
sa.column('plugin_author', sa.String(255)),
sa.column('plugin_name', sa.String(255)),
sa.column('installation_uuid', sa.String(36)),
sa.column('artifact_digest', sa.String(64)),
sa.column('runtime_revision', sa.Integer()),
)
rows = conn.execute(
sa.select(
table.c.workspace_uuid,
table.c.plugin_author,
table.c.plugin_name,
table.c.installation_uuid,
table.c.artifact_digest,
table.c.runtime_revision,
)
).all()
for row in rows:
installation_uuid = str(row.installation_uuid or uuid.uuid4())
values: dict[str, object] = {}
if not row.installation_uuid:
values['installation_uuid'] = installation_uuid
if not row.artifact_digest or len(str(row.artifact_digest)) != 64:
values['artifact_digest'] = _legacy_digest(installation_uuid)
if row.runtime_revision is None or row.runtime_revision < 1:
values['runtime_revision'] = 1
if values:
conn.execute(
table.update()
.where(table.c.workspace_uuid == row.workspace_uuid)
.where(table.c.plugin_author == row.plugin_author)
.where(table.c.plugin_name == row.plugin_name)
.values(**values)
)
def _constraint_names(conn: sa.Connection, kind: str) -> set[str]:
inspector = sa.inspect(conn)
getter = inspector.get_unique_constraints if kind == 'unique' else inspector.get_check_constraints
return {str(item.get('name')) for item in getter(_TABLE) if item.get('name')}
def upgrade() -> None:
conn = op.get_bind()
columns = _column_names(conn)
if not columns:
return
if 'installation_uuid' not in columns:
op.add_column(_TABLE, sa.Column('installation_uuid', sa.String(36), nullable=True))
if 'artifact_digest' not in columns:
op.add_column(_TABLE, sa.Column('artifact_digest', sa.String(64), nullable=True))
if 'runtime_revision' not in columns:
op.add_column(_TABLE, sa.Column('runtime_revision', sa.Integer(), nullable=True))
rls_state = _suspend_postgres_rls(conn)
try:
_backfill(conn)
finally:
_restore_postgres_rls(conn, rls_state)
# Fresh databases are created from current metadata before Alembic runs;
# guards make the revision safe for that path and for interrupted upgrades.
indexes = {index['name'] for index in sa.inspect(conn).get_indexes(_TABLE)}
unique_constraints = _constraint_names(conn, 'unique')
check_constraints = _constraint_names(conn, 'check')
with op.batch_alter_table(_TABLE) as batch:
batch.alter_column('installation_uuid', existing_type=sa.String(36), nullable=False)
batch.alter_column('artifact_digest', existing_type=sa.String(64), nullable=False)
batch.alter_column('runtime_revision', existing_type=sa.Integer(), nullable=False)
if _REVISION_CHECK not in check_constraints:
batch.create_check_constraint(_REVISION_CHECK, 'runtime_revision >= 1')
if _DIGEST_CHECK not in check_constraints:
batch.create_check_constraint(_DIGEST_CHECK, 'length(artifact_digest) = 64')
if _INSTALLATION_UNIQUE not in unique_constraints:
batch.create_unique_constraint(_INSTALLATION_UNIQUE, ['installation_uuid'])
if _INSTALLATION_INDEX not in indexes and _INSTALLATION_INDEX not in unique_constraints:
batch.create_index(
_INSTALLATION_INDEX,
['workspace_uuid', 'installation_uuid'],
unique=True,
)
def downgrade() -> None:
conn = op.get_bind()
columns = _column_names(conn)
if not columns:
return
indexes = {index['name'] for index in sa.inspect(conn).get_indexes(_TABLE)}
checks = _constraint_names(conn, 'check')
uniques = _constraint_names(conn, 'unique')
with op.batch_alter_table(_TABLE) as batch:
if _INSTALLATION_INDEX in indexes:
batch.drop_index(_INSTALLATION_INDEX)
if _DIGEST_CHECK in checks:
batch.drop_constraint(_DIGEST_CHECK, type_='check')
if _REVISION_CHECK in checks:
batch.drop_constraint(_REVISION_CHECK, type_='check')
if _INSTALLATION_UNIQUE in uniques:
batch.drop_constraint(_INSTALLATION_UNIQUE, type_='unique')
for column_name in ('runtime_revision', 'artifact_digest', 'installation_uuid'):
if column_name in columns:
batch.drop_column(column_name)
@@ -0,0 +1,382 @@
"""create tenant-scoped pgvector storage in the business database
Revision ID: 0013_tenant_pgvector
Revises: 0012_plugin_identity
Create Date: 2026-07-19
The Cloud application role never executes this DDL. A release migration role
installs pgvector once, creates an untyped vector column, and builds a bounded
set of expression/partial ANN indexes. Existing legacy rows are migrated only
when each row maps unambiguously to one knowledge base.
"""
from __future__ import annotations
import contextlib
import typing
import sqlalchemy as sa
from alembic import op
from pgvector.sqlalchemy import Vector
revision = '0013_tenant_pgvector'
down_revision = '0012_plugin_identity'
branch_labels = None
depends_on = None
_VECTOR_TABLE = 'langbot_vectors'
_LEGACY_TABLE = 'langbot_vectors_legacy_0013'
_TENANT_POLICY = 'langbot_workspace_isolation'
_TENANT_SETTING = 'langbot.workspace_uuid'
_KB_DIMENSION_CHECK = 'ck_knowledge_bases_embedding_dimension_positive'
_VECTOR_DIMENSION_CHECK = 'ck_langbot_vectors_embedding_dimension'
_VECTOR_ALLOWED_CHECK = 'ck_langbot_vectors_embedding_dimension_enabled'
_ALLOWED_DIMENSIONS = (384, 512, 768, 1024, 1536)
_LEGACY_SOURCE_TABLES = (
'knowledge_bases',
'knowledge_base_files',
'knowledge_base_chunks',
)
def _quote(conn: sa.Connection, identifier: str) -> str:
return conn.dialect.identifier_preparer.quote(identifier)
def _columns(conn: sa.Connection, table_name: str) -> set[str]:
inspector = sa.inspect(conn)
if table_name not in inspector.get_table_names():
return set()
return {column['name'] for column in inspector.get_columns(table_name)}
def _checks(conn: sa.Connection, table_name: str) -> set[str]:
return {str(item['name']) for item in sa.inspect(conn).get_check_constraints(table_name) if item.get('name')}
def _ensure_knowledge_base_dimension(conn: sa.Connection) -> None:
columns = _columns(conn, 'knowledge_bases')
if 'embedding_dimension' not in columns:
op.add_column('knowledge_bases', sa.Column('embedding_dimension', sa.Integer(), nullable=True))
if _KB_DIMENSION_CHECK not in _checks(conn, 'knowledge_bases'):
with op.batch_alter_table('knowledge_bases') as batch:
batch.create_check_constraint(
_KB_DIMENSION_CHECK,
'embedding_dimension IS NULL OR embedding_dimension > 0',
)
def _create_vector_table() -> None:
enabled = ', '.join(str(item) for item in _ALLOWED_DIMENSIONS)
op.create_table(
_VECTOR_TABLE,
sa.Column('workspace_uuid', sa.String(36), nullable=False),
sa.Column('knowledge_base_uuid', sa.String(255), nullable=False),
sa.Column('vector_id', sa.String(255), nullable=False),
sa.Column('embedding_dimension', sa.Integer(), nullable=False),
sa.Column('embedding', Vector(), nullable=False),
sa.Column('text', sa.Text(), nullable=True),
sa.Column('file_id', sa.String(255), nullable=True),
sa.Column('chunk_uuid', sa.String(255), nullable=True),
sa.PrimaryKeyConstraint(
'workspace_uuid',
'knowledge_base_uuid',
'vector_id',
name='pk_langbot_vectors',
),
sa.ForeignKeyConstraint(
['workspace_uuid', 'knowledge_base_uuid'],
['knowledge_bases.workspace_uuid', 'knowledge_bases.uuid'],
name='fk_langbot_vectors_workspace_kb',
ondelete='CASCADE',
),
sa.CheckConstraint(
'vector_dims(embedding) = embedding_dimension',
name=_VECTOR_DIMENSION_CHECK,
),
sa.CheckConstraint(
f'embedding_dimension IN ({enabled})',
name=_VECTOR_ALLOWED_CHECK,
),
)
op.create_index(
'ix_langbot_vectors_workspace_kb_file',
_VECTOR_TABLE,
['workspace_uuid', 'knowledge_base_uuid', 'file_id'],
)
op.create_index(
'ix_langbot_vectors_workspace_kb_chunk',
_VECTOR_TABLE,
['workspace_uuid', 'knowledge_base_uuid', 'chunk_uuid'],
)
def _legacy_mapping_predicate() -> str:
return """
kb.collection_id = legacy.collection
OR EXISTS (
SELECT 1
FROM knowledge_base_files AS files
LEFT JOIN knowledge_base_chunks AS chunks
ON chunks.workspace_uuid = files.workspace_uuid
AND chunks.file_id = files.uuid
WHERE files.workspace_uuid = kb.workspace_uuid
AND files.kb_id = kb.uuid
AND (files.uuid = legacy.file_id OR chunks.uuid = legacy.chunk_uuid)
)
"""
def _legacy_source_rls_states(conn: sa.Connection) -> dict[str, tuple[bool, bool]]:
rows = (
conn.execute(
sa.text(
"""
SELECT c.relname, c.relrowsecurity, c.relforcerowsecurity
FROM pg_class AS c
JOIN pg_namespace AS n ON n.oid = c.relnamespace
WHERE n.nspname = current_schema()
AND c.relname IN :table_names
AND c.relkind IN ('r', 'p')
"""
).bindparams(sa.bindparam('table_names', expanding=True)),
{'table_names': _LEGACY_SOURCE_TABLES},
)
.mappings()
.all()
)
states = {str(row['relname']): (bool(row['relrowsecurity']), bool(row['relforcerowsecurity'])) for row in rows}
missing = set(_LEGACY_SOURCE_TABLES) - set(states)
if missing:
raise RuntimeError(f'Legacy pgvector source tables are missing: {sorted(missing)!r}')
return states
@contextlib.contextmanager
def _suspend_legacy_source_rls(conn: sa.Connection) -> typing.Iterator[None]:
"""Temporarily let the table-owning migrator map all legacy tenant rows.
Revision 0011 enables and forces RLS on each source table. The release
migrator intentionally has neither superuser nor BYPASSRLS, so even a table
owner cannot read those rows until FORCE RLS is paused. Preserve both flags
independently and restore them in ``finally`` so mixed pre-existing states
survive successful, rejected, and interrupted legacy migrations.
"""
states = _legacy_source_rls_states(conn)
try:
for table_name in _LEGACY_SOURCE_TABLES:
table = _quote(conn, table_name)
conn.execute(sa.text(f'ALTER TABLE {table} NO FORCE ROW LEVEL SECURITY'))
conn.execute(sa.text(f'ALTER TABLE {table} DISABLE ROW LEVEL SECURITY'))
yield
finally:
for table_name in _LEGACY_SOURCE_TABLES:
table = _quote(conn, table_name)
rls_enabled, rls_forced = states[table_name]
enabled_clause = 'ENABLE' if rls_enabled else 'DISABLE'
forced_clause = 'FORCE' if rls_forced else 'NO FORCE'
conn.execute(sa.text(f'ALTER TABLE {table} {enabled_clause} ROW LEVEL SECURITY'))
conn.execute(sa.text(f'ALTER TABLE {table} {forced_clause} ROW LEVEL SECURITY'))
def _migrate_legacy_rows(conn: sa.Connection) -> None:
predicate = _legacy_mapping_predicate()
ambiguous = conn.execute(
sa.text(
f"""
WITH candidates AS (
SELECT legacy.id, kb.workspace_uuid, kb.uuid AS knowledge_base_uuid
FROM {_LEGACY_TABLE} AS legacy
JOIN knowledge_bases AS kb ON ({predicate})
), candidate_counts AS (
SELECT id, COUNT(*) AS count
FROM candidates
GROUP BY id
)
SELECT legacy.id, COALESCE(candidate_counts.count, 0) AS candidate_count
FROM {_LEGACY_TABLE} AS legacy
LEFT JOIN candidate_counts ON candidate_counts.id = legacy.id
WHERE COALESCE(candidate_counts.count, 0) <> 1
LIMIT 1
"""
)
).first()
if ambiguous is not None:
raise RuntimeError(
'Legacy pgvector row cannot be mapped to exactly one Workspace/knowledge base: '
f'{ambiguous.id!r} has {ambiguous.candidate_count} candidates'
)
conn.execute(
sa.text(
f"""
INSERT INTO {_VECTOR_TABLE} (
workspace_uuid,
knowledge_base_uuid,
vector_id,
embedding_dimension,
embedding,
text,
file_id,
chunk_uuid
)
SELECT
kb.workspace_uuid,
kb.uuid,
legacy.id,
vector_dims(legacy.embedding),
legacy.embedding,
legacy.text,
legacy.file_id,
legacy.chunk_uuid
FROM {_LEGACY_TABLE} AS legacy
JOIN knowledge_bases AS kb ON ({predicate})
"""
)
)
mixed_dimension = conn.execute(
sa.text(
f"""
SELECT workspace_uuid, knowledge_base_uuid
FROM {_VECTOR_TABLE}
GROUP BY workspace_uuid, knowledge_base_uuid
HAVING MIN(embedding_dimension) <> MAX(embedding_dimension)
LIMIT 1
"""
)
).first()
if mixed_dimension is not None:
raise RuntimeError('Legacy knowledge base contains mixed embedding dimensions')
conn.execute(
sa.text(
f"""
UPDATE knowledge_bases AS kb
SET embedding_dimension = dimensions.embedding_dimension
FROM (
SELECT workspace_uuid, knowledge_base_uuid, MIN(embedding_dimension) AS embedding_dimension
FROM {_VECTOR_TABLE}
GROUP BY workspace_uuid, knowledge_base_uuid
) AS dimensions
WHERE kb.workspace_uuid = dimensions.workspace_uuid
AND kb.uuid = dimensions.knowledge_base_uuid
AND kb.embedding_dimension IS NULL
"""
)
)
def _drop_all_policies(conn: sa.Connection) -> None:
table = _quote(conn, _VECTOR_TABLE)
policies = conn.execute(
sa.text(
"""
SELECT p.polname
FROM pg_policy p
JOIN pg_class c ON c.oid = p.polrelid
JOIN pg_namespace n ON n.oid = c.relnamespace
WHERE n.nspname = current_schema() AND c.relname = :table_name
"""
),
{'table_name': _VECTOR_TABLE},
).scalars()
for policy_name in policies:
op.execute(sa.text(f'DROP POLICY {_quote(conn, policy_name)} ON {table}'))
def _enable_rls(conn: sa.Connection) -> None:
table = _quote(conn, _VECTOR_TABLE)
policy = _quote(conn, _TENANT_POLICY)
expression = f"workspace_uuid::text = NULLIF(current_setting('{_TENANT_SETTING}', true), '')"
_drop_all_policies(conn)
op.execute(sa.text(f'ALTER TABLE {table} ENABLE ROW LEVEL SECURITY'))
op.execute(sa.text(f'ALTER TABLE {table} FORCE ROW LEVEL SECURITY'))
op.execute(
sa.text(
f'CREATE POLICY {policy} ON {table} AS PERMISSIVE FOR ALL TO PUBLIC '
f'USING ({expression}) WITH CHECK ({expression})'
)
)
def _create_ann_indexes(conn: sa.Connection) -> None:
table = _quote(conn, _VECTOR_TABLE)
for dimension in _ALLOWED_DIMENSIONS:
index = _quote(conn, f'ix_langbot_vectors_hnsw_cosine_{dimension}')
op.execute(
sa.text(
f'CREATE INDEX {index} ON {table} USING hnsw '
f'((embedding::vector({dimension})) vector_cosine_ops) '
f'WHERE embedding_dimension = {dimension}'
)
)
def upgrade() -> None:
conn = op.get_bind()
if 'knowledge_bases' not in sa.inspect(conn).get_table_names():
if conn.dialect.name == 'postgresql':
# The supported PostgreSQL release path creates the business
# tables before stamping 0010 and reaching this migration. Missing
# knowledge_bases therefore means the operator bypassed the
# release bootstrap or the schema is incomplete; stamping head in
# that state would make Cloud runtime validation unrecoverable.
raise RuntimeError('PostgreSQL release migration requires the knowledge_bases table')
# A direct empty SQLite Alembic walk is still used by migration tooling;
# Core creates the complete ORM schema on its following compatibility
# pass, including the portable embedding_dimension field.
return
_ensure_knowledge_base_dimension(conn)
# pgvector storage is PostgreSQL-only, but ``embedding_dimension`` is an
# ORM field used by both deployment modes. Existing OSS SQLite databases
# must receive the column before this revision becomes a no-op.
if conn.dialect.name != 'postgresql':
return
op.execute(sa.text('CREATE EXTENSION IF NOT EXISTS vector'))
columns = _columns(conn, _VECTOR_TABLE)
if columns:
scoped_columns = {
'workspace_uuid',
'knowledge_base_uuid',
'vector_id',
'embedding_dimension',
'embedding',
}
if scoped_columns.issubset(columns):
raise RuntimeError('Tenant pgvector table exists before its owning release migration')
legacy_columns = {'id', 'collection', 'embedding'}
if not legacy_columns.issubset(columns):
raise RuntimeError('Existing pgvector table has an unsupported schema')
if _LEGACY_TABLE in sa.inspect(conn).get_table_names():
raise RuntimeError(f'Interrupted pgvector migration left {_LEGACY_TABLE!r} behind')
op.rename_table(_VECTOR_TABLE, _LEGACY_TABLE)
_create_vector_table()
if _LEGACY_TABLE in sa.inspect(conn).get_table_names():
with _suspend_legacy_source_rls(conn):
_migrate_legacy_rows(conn)
op.drop_table(_LEGACY_TABLE)
_create_ann_indexes(conn)
_enable_rls(conn)
def downgrade() -> None:
conn = op.get_bind()
if conn.dialect.name == 'postgresql' and _VECTOR_TABLE in sa.inspect(conn).get_table_names():
_drop_all_policies(conn)
op.drop_table(_VECTOR_TABLE)
columns = _columns(conn, 'knowledge_bases')
if 'embedding_dimension' in columns:
checks = _checks(conn, 'knowledge_bases')
with op.batch_alter_table('knowledge_bases') as batch:
if _KB_DIMENSION_CHECK in checks:
batch.drop_constraint(_KB_DIMENSION_CHECK, type_='check')
batch.drop_column('embedding_dimension')
+8 -1
View File
@@ -2,6 +2,7 @@ from __future__ import annotations
import abc
import sqlalchemy
import sqlalchemy.ext.asyncio as sqlalchemy_asyncio
from ..core import app
@@ -30,8 +31,14 @@ class BaseDatabaseManager(abc.ABC):
engine: sqlalchemy_asyncio.AsyncEngine
def __init__(self, ap: app.Application) -> None:
def __init__(
self,
ap: app.Application,
*,
url_override: sqlalchemy.engine.URL | None = None,
) -> None:
self.ap = ap
self.url_override = url_override
@abc.abstractmethod
async def initialize(self) -> None:
@@ -1,8 +1,10 @@
from __future__ import annotations
import sqlalchemy
import sqlalchemy.ext.asyncio as sqlalchemy_asyncio
from .. import database
from ..postgresql_url import normalize_asyncpg_url
@database.manager_class('postgresql')
@@ -10,12 +12,29 @@ class PostgreSQLDatabaseManager(database.BaseDatabaseManager):
"""PostgreSQL database manager"""
async def initialize(self) -> None:
postgresql_config = self.ap.instance_config.data.get('database', {}).get('postgresql', {})
host = postgresql_config.get('host', '127.0.0.1')
port = postgresql_config.get('port', 5432)
user = postgresql_config.get('user', 'postgres')
password = postgresql_config.get('password', 'postgres')
database = postgresql_config.get('database', 'postgres')
engine_url = f'postgresql+asyncpg://{user}:{password}@{host}:{port}/{database}'
if self.url_override is not None:
engine_url = self.url_override
else:
postgresql_config = self.ap.instance_config.data.get('database', {}).get('postgresql', {})
explicit_url = postgresql_config.get('url')
if explicit_url:
if not isinstance(explicit_url, str):
raise ValueError('database.postgresql.url must be a string')
try:
engine_url = sqlalchemy.engine.make_url(explicit_url)
except Exception:
raise ValueError('database.postgresql.url is invalid') from None
try:
engine_url = normalize_asyncpg_url(engine_url)
except ValueError:
raise ValueError('database.postgresql.url must use valid PostgreSQL asyncpg options') from None
else:
engine_url = sqlalchemy.URL.create(
'postgresql+asyncpg',
username=postgresql_config.get('user', 'postgres'),
password=postgresql_config.get('password', 'postgres'),
host=postgresql_config.get('host', '127.0.0.1'),
port=postgresql_config.get('port', 5432),
database=postgresql_config.get('database', 'postgres'),
)
self.engine = sqlalchemy_asyncio.create_async_engine(engine_url)
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,25 @@
"""Safe PostgreSQL URL normalization shared by runtime and migration jobs."""
from __future__ import annotations
import sqlalchemy
def normalize_asyncpg_url(url: sqlalchemy.engine.URL) -> sqlalchemy.engine.URL:
"""Select asyncpg and translate the common libpq TLS query spelling."""
if url.drivername == 'postgresql':
url = url.set(drivername='postgresql+asyncpg')
elif url.drivername != 'postgresql+asyncpg':
raise ValueError('PostgreSQL URL must use PostgreSQL with the asyncpg driver')
query = dict(url.query)
sslmode = query.pop('sslmode', None)
if sslmode is not None:
if 'ssl' in query and query['ssl'] != sslmode:
raise ValueError('PostgreSQL URL cannot specify conflicting ssl and sslmode options')
# SQLAlchemy expands URL query keys into asyncpg keyword arguments.
# asyncpg calls this keyword ``ssl`` even though PostgreSQL DSNs
# conventionally spell the same mode ``sslmode``.
query['ssl'] = sslmode
return url.set(query=query)
@@ -0,0 +1,193 @@
"""One-shot, operator-only Cloud PostgreSQL release migration entrypoint."""
from __future__ import annotations
import asyncio
import os
import re
from collections.abc import Mapping
import sqlalchemy
from ..cloud.bootstrap import SUPPORTED_PGVECTOR_DIMENSIONS
from ..core import app as core_app
from ..core.stages.load_config import LoadConfigStage
from ..core.stages.setup_logger import SetupLoggerStage
from .mgr import PersistenceManager, PersistenceMode
from .postgresql_url import normalize_asyncpg_url
DEFAULT_OPERATOR_DSN_ENV = 'LANGBOT_CLOUD_MIGRATION_DSN'
_ENV_NAME = re.compile(r'^[A-Z_][A-Z0-9_]*$')
class CloudReleaseMigrationConfigurationError(RuntimeError):
"""Raised before any database operation when migration input is unsafe."""
def _url_endpoint(url: sqlalchemy.engine.URL, *, label: str) -> tuple[str, int]:
"""Return a comparison-safe PostgreSQL endpoint without leaking its DSN."""
try:
host = (url.host or '').strip().casefold()
port = url.port or 5432
except (TypeError, ValueError):
raise CloudReleaseMigrationConfigurationError(f'{label} PostgreSQL host or port is invalid') from None
if not host or isinstance(port, bool) or not isinstance(port, int) or not 1 <= port <= 65535:
raise CloudReleaseMigrationConfigurationError(f'{label} PostgreSQL host or port is invalid')
return host, port
def _operator_database_url(
instance_config: dict,
*,
environ: Mapping[str, str],
) -> sqlalchemy.engine.URL:
database_config = instance_config.get('database')
if not isinstance(database_config, dict) or database_config.get('use') != 'postgresql':
raise CloudReleaseMigrationConfigurationError(
'Cloud release migration requires explicit database.use=postgresql; SQLite fallback is forbidden'
)
runtime_config = database_config.get('postgresql')
if not isinstance(runtime_config, dict):
raise CloudReleaseMigrationConfigurationError('Cloud runtime PostgreSQL configuration is missing')
migration_config = database_config.get('cloud_migration', {})
if not isinstance(migration_config, dict):
raise CloudReleaseMigrationConfigurationError('database.cloud_migration must be a mapping')
dsn_env = migration_config.get('operator_dsn_env', DEFAULT_OPERATOR_DSN_ENV)
if not isinstance(dsn_env, str) or not _ENV_NAME.fullmatch(dsn_env):
raise CloudReleaseMigrationConfigurationError(
'database.cloud_migration.operator_dsn_env must name an uppercase environment variable'
)
raw_dsn = environ.get(dsn_env, '').strip()
if not raw_dsn:
raise CloudReleaseMigrationConfigurationError(
f'Cloud release migration requires the operator DSN in environment variable {dsn_env}'
)
try:
operator_url = sqlalchemy.engine.make_url(raw_dsn)
except Exception:
# Never echo a malformed DSN because it may contain an unescaped secret.
raise CloudReleaseMigrationConfigurationError('Cloud release migration operator DSN is invalid') from None
try:
operator_url = normalize_asyncpg_url(operator_url)
except ValueError:
raise CloudReleaseMigrationConfigurationError(
'Cloud release migration operator DSN must use valid PostgreSQL asyncpg options'
) from None
operator_user = (operator_url.username or '').strip()
operator_database = (operator_url.database or '').strip()
operator_host, operator_port = _url_endpoint(operator_url, label='Cloud release migration operator')
runtime_url_value = runtime_config.get('url')
if runtime_url_value:
if not isinstance(runtime_url_value, str):
raise CloudReleaseMigrationConfigurationError('Cloud runtime PostgreSQL URL must be a string')
try:
runtime_url = sqlalchemy.engine.make_url(runtime_url_value)
except Exception:
raise CloudReleaseMigrationConfigurationError('Cloud runtime PostgreSQL URL is invalid') from None
if runtime_url.drivername not in {'postgresql', 'postgresql+asyncpg'}:
raise CloudReleaseMigrationConfigurationError('Cloud runtime database URL must use PostgreSQL')
runtime_user = (runtime_url.username or '').strip()
runtime_database = (runtime_url.database or '').strip()
runtime_host, runtime_port = _url_endpoint(runtime_url, label='Cloud runtime')
else:
runtime_user = str(runtime_config.get('user', 'postgres') or '').strip()
runtime_database = str(runtime_config.get('database', 'postgres') or '').strip()
runtime_host = str(runtime_config.get('host', '') or '').strip().casefold()
runtime_port = runtime_config.get('port', 5432)
if not operator_user or not operator_database:
raise CloudReleaseMigrationConfigurationError(
'Cloud release migration operator DSN must include a user, host, and database'
)
if (
not runtime_user
or not runtime_database
or not runtime_host
or isinstance(runtime_port, bool)
or not isinstance(runtime_port, int)
or not 1 <= runtime_port <= 65535
):
raise CloudReleaseMigrationConfigurationError(
'Cloud runtime PostgreSQL user, host, port, and database are required'
)
if operator_user == runtime_user:
raise CloudReleaseMigrationConfigurationError(
'Cloud release migration requires a distinct operator role from the runtime PostgreSQL role'
)
if operator_database != runtime_database:
raise CloudReleaseMigrationConfigurationError(
'Cloud release migration operator DSN must target the configured runtime database'
)
if operator_host != runtime_host or operator_port != runtime_port:
# The first Cloud release intentionally requires the migrator and
# runtime to use the same PostgreSQL endpoint. Supporting a direct
# operator endpoint plus a runtime pooler requires a database-backed
# immutable cluster identity check; accepting aliases here would turn a
# same-named database on another cluster into a silent migration target.
raise CloudReleaseMigrationConfigurationError(
'Cloud release migration operator DSN must target the configured runtime PostgreSQL endpoint'
)
vdb_config = instance_config.get('vdb')
if not isinstance(vdb_config, dict) or vdb_config.get('use') != 'pgvector':
raise CloudReleaseMigrationConfigurationError('Cloud release migration requires vdb.use=pgvector')
pgvector_config = vdb_config.get('pgvector')
if not isinstance(pgvector_config, dict) or pgvector_config.get('use_business_database') is not True:
raise CloudReleaseMigrationConfigurationError(
'Cloud release migration requires vdb.pgvector.use_business_database=true'
)
dimensions = pgvector_config.get('allowed_dimensions')
if (
not isinstance(dimensions, list)
or not dimensions
or any(isinstance(item, bool) or not isinstance(item, int) for item in dimensions)
or not set(dimensions).issubset(SUPPORTED_PGVECTOR_DIMENSIONS)
):
raise CloudReleaseMigrationConfigurationError(
'Cloud release migration pgvector dimensions are outside the release-created index set'
)
return operator_url
async def run_cloud_release_migration(
ap: core_app.Application,
*,
environ: Mapping[str, str] | None = None,
) -> None:
"""Run and validate one release migration with an isolated operator DSN."""
operator_url = _operator_database_url(
ap.instance_config.data,
environ=os.environ if environ is None else environ,
)
manager = PersistenceManager(
ap,
mode=PersistenceMode.RELEASE_MIGRATION,
database_url=operator_url,
)
ap.persistence_mgr = manager
try:
await manager.initialize()
ap.logger.info('Cloud PostgreSQL release migration reached and validated the exact release head.')
finally:
db = getattr(manager, 'db', None)
engine = getattr(db, 'engine', None)
if engine is not None:
await engine.dispose()
async def run_cloud_release_migration_from_config(loop: asyncio.AbstractEventLoop) -> None:
"""Load only process configuration/logging, then run the one-shot job."""
ap = core_app.Application()
ap.event_loop = loop
await LoadConfigStage().run(ap)
await SetupLoggerStage().run(ap)
await run_cloud_release_migration(ap)
+396 -19
View File
@@ -1,5 +1,9 @@
from __future__ import annotations
import asyncio
import contextvars
import dataclasses
import enum
import types
import typing
@@ -8,7 +12,17 @@ import sqlalchemy.ext.asyncio as sqlalchemy_asyncio
TENANT_SETTING = 'langbot.workspace_uuid'
ACCOUNT_SETTING = 'langbot.account_uuid'
API_KEY_HASH_SETTING = 'langbot.api_key_hash'
INVITATION_HASH_SETTING = 'langbot.invitation_hash'
INSTANCE_SETTING = 'langbot.instance_uuid'
IDENTITY_DIGEST_SETTING = 'langbot.identity_digest'
TENANT_POLICY_NAME = 'langbot_workspace_isolation'
ACCOUNT_DISCOVERY_POLICY_NAME = 'langbot_account_discovery'
API_KEY_DISCOVERY_POLICY_NAME = 'langbot_api_key_discovery'
INVITATION_DISCOVERY_POLICY_NAME = 'langbot_invitation_discovery'
INSTANCE_DISCOVERY_POLICY_NAME = 'langbot_instance_discovery'
# Keep this contract explicit. A new tenant-owned table must be added to both
# this runtime list and the corresponding Alembic migration before release.
@@ -41,26 +55,269 @@ TENANT_TABLE_COLUMNS: dict[str, str] = {
'monitoring_errors': 'workspace_uuid',
'monitoring_embedding_calls': 'workspace_uuid',
'monitoring_feedback': 'workspace_uuid',
# Created by 0013 rather than ORM metadata; it is still part of the same
# business-database RLS contract and permits no discovery policies.
'langbot_vectors': 'workspace_uuid',
}
class TenantUnitOfWork:
"""Bind one Workspace scope to one database transaction.
class PersistenceScopeKind(enum.StrEnum):
WORKSPACE = 'workspace'
ACCOUNT_DISCOVERY = 'account_discovery'
API_KEY_DISCOVERY = 'api_key_discovery'
INVITATION_DISCOVERY = 'invitation_discovery'
INSTANCE_DISCOVERY = 'instance_discovery'
IDENTITY_DISCOVERY = 'identity_discovery'
PostgreSQL receives a transaction-local setting used by RLS policies.
SQLite keeps the same transaction boundary without a database setting so
OSS callers can adopt this API without changing database engines.
@dataclasses.dataclass(frozen=True, slots=True)
class PersistenceScope:
"""A trusted, transaction-local database visibility scope."""
kind: PersistenceScopeKind
settings: tuple[tuple[str, str], ...]
@classmethod
def workspace(cls, workspace_uuid: str) -> PersistenceScope:
return cls._one(PersistenceScopeKind.WORKSPACE, TENANT_SETTING, workspace_uuid)
@classmethod
def account(cls, account_uuid: str) -> PersistenceScope:
return cls._one(PersistenceScopeKind.ACCOUNT_DISCOVERY, ACCOUNT_SETTING, account_uuid)
@classmethod
def api_key(cls, key_hash: str) -> PersistenceScope:
return cls._one(PersistenceScopeKind.API_KEY_DISCOVERY, API_KEY_HASH_SETTING, key_hash)
@classmethod
def invitation(cls, token_hash: str) -> PersistenceScope:
return cls._one(PersistenceScopeKind.INVITATION_DISCOVERY, INVITATION_HASH_SETTING, token_hash)
@classmethod
def instance(cls, instance_uuid: str) -> PersistenceScope:
return cls._one(PersistenceScopeKind.INSTANCE_DISCOVERY, INSTANCE_SETTING, instance_uuid)
@classmethod
def identity(cls, identity_digest: str) -> PersistenceScope:
return cls._one(PersistenceScopeKind.IDENTITY_DISCOVERY, IDENTITY_DIGEST_SETTING, identity_digest)
@classmethod
def _one(cls, kind: PersistenceScopeKind, setting: str, value: str) -> PersistenceScope:
return cls(kind, ((setting, cls._normalize(value, kind.value)),))
@staticmethod
def _normalize(value: str, label: str) -> str:
if not isinstance(value, str):
raise TypeError(f'{label} must be a string')
normalized = value.strip()
if not normalized:
raise ValueError(f'{label} must not be empty')
if len(normalized) > 512:
raise ValueError(f'{label} exceeds the database scope limit')
return normalized
class TenantScopeRequiredError(RuntimeError):
"""Raised when Cloud business data is accessed without a trusted scope."""
class CrossScopeTransactionError(RuntimeError):
"""Raised when code tries to change scope inside an active transaction."""
class TransactionRollbackOnlyError(RuntimeError):
"""Raised when a caught failure made the outer transaction unsafe to commit."""
@dataclasses.dataclass(slots=True)
class ActiveScopedTransaction:
scope: PersistenceScope
session: sqlalchemy_asyncio.AsyncSession
engine: sqlalchemy_asyncio.AsyncEngine
owner_task: asyncio.Task[typing.Any] | None
depth: int = 1
rollback_only: bool = False
rollback_only_cause: BaseException | None = None
after_commit_waiters: list[asyncio.Future[None]] = dataclasses.field(default_factory=list)
def mark_rollback_only(self, cause: BaseException | None = None) -> None:
"""Remember that this transaction must never release commit-gated work."""
self.rollback_only = True
if self.rollback_only_cause is None:
self.rollback_only_cause = cause
ActiveTransactionVar = contextvars.ContextVar[ActiveScopedTransaction | None]
# ``AsyncSession`` delegates database I/O to a greenlet-backed synchronous
# Engine, but Python context variables are preserved across that boundary. A
# per-engine ``handle_error`` listener therefore covers DBAPI failures from all
# direct Session APIs (execute/scalar/get/flush) as well as callers which obtain
# the transaction-bound Connection. The owner-task and engine checks prevent
# copied child-task contexts or unrelated database engines from poisoning this
# transaction.
_DATABASE_OPERATION_TRANSACTION: contextvars.ContextVar[ActiveScopedTransaction | None] = contextvars.ContextVar(
'langbot_database_operation_transaction',
default=None,
)
def _mark_current_transaction_rollback_only(exception_context: typing.Any) -> None:
state = _DATABASE_OPERATION_TRANSACTION.get()
if state is None:
return
try:
current_task = asyncio.current_task()
except RuntimeError: # pragma: no cover - async engines always run in an event loop
return
if state.owner_task is not current_task or exception_context.engine is not state.engine.sync_engine:
return
cause = exception_context.sqlalchemy_exception or exception_context.original_exception
state.mark_rollback_only(cause)
def _install_rollback_only_error_listener(engine: sqlalchemy_asyncio.AsyncEngine) -> None:
sync_engine = engine.sync_engine
if not sqlalchemy.event.contains(sync_engine, 'handle_error', _mark_current_transaction_rollback_only):
sqlalchemy.event.listen(sync_engine, 'handle_error', _mark_current_transaction_rollback_only)
@dataclasses.dataclass(slots=True)
class ActivePersistenceScope:
"""A trusted visibility scope without an open database transaction."""
scope: PersistenceScope
owner_task: asyncio.Task[typing.Any] | None
depth: int = 1
ActivePersistenceScopeVar = contextvars.ContextVar[ActivePersistenceScope | None]
class PersistenceScopeBoundary:
"""Carry a trusted persistence scope across long-running async work.
Unlike :class:`TenantUnitOfWork`, this boundary never opens a database
session. ``PersistenceManager.execute_async`` materializes the scope as a
short transaction for each database call. This makes the boundary suitable
for request and pipeline lifetimes which can spend substantial time waiting
on model providers, tools, or streaming clients.
Context variables are copied into child tasks, so implicit use from a child
task is rejected by the manager. A child may establish its own explicit
boundary because the scope value still comes from trusted application code.
"""
def __init__(self, engine: sqlalchemy_asyncio.AsyncEngine, workspace_uuid: str) -> None:
workspace_uuid = workspace_uuid.strip()
if not workspace_uuid:
raise ValueError('TenantUnitOfWork requires a non-empty workspace_uuid')
def __init__(
self,
scope: PersistenceScope,
*,
active_scope: ActivePersistenceScopeVar,
active_transaction: ActiveTransactionVar,
) -> None:
self.scope = scope
self._active_scope = active_scope
self._active_transaction = active_transaction
self._active_state: ActivePersistenceScope | None = None
self._context_token: contextvars.Token[ActivePersistenceScope | None] | None = None
self._used = False
self._owns_scope = False
async def __aenter__(self) -> PersistenceScopeBoundary:
if self._used:
raise RuntimeError('PersistenceScopeBoundary instances cannot be reused')
self._used = True
transaction = self._active_transaction.get()
if transaction is not None:
if transaction.owner_task is not asyncio.current_task():
raise CrossScopeTransactionError(
'Scoped database transactions cannot be inherited by child tasks; open an explicit task scope'
)
if transaction.scope != self.scope:
raise CrossScopeTransactionError(
f'Cannot enter {self.scope.kind.value} scope while {transaction.scope.kind.value} scope is active'
)
active = self._active_scope.get()
if active is not None and active.owner_task is asyncio.current_task():
if active.scope != self.scope:
raise CrossScopeTransactionError(
f'Cannot enter {self.scope.kind.value} scope while {active.scope.kind.value} scope is active'
)
active.depth += 1
self._active_state = active
return self
state = ActivePersistenceScope(
scope=self.scope,
owner_task=asyncio.current_task(),
)
self._active_state = state
self._context_token = self._active_scope.set(state)
self._owns_scope = True
return self
async def __aexit__(
self,
exc_type: type[BaseException] | None,
exc_value: BaseException | None,
traceback: types.TracebackType | None,
) -> bool:
del exc_type, exc_value, traceback
state = self._active_state
if state is None:
raise RuntimeError('PersistenceScopeBoundary has no active state')
if state.owner_task is not asyncio.current_task():
raise CrossScopeTransactionError('Scoped persistence boundaries cannot be exited by a different task')
if not self._owns_scope:
state.depth -= 1
self._active_state = None
return False
if self._context_token is None:
raise RuntimeError('PersistenceScopeBoundary has no context token')
self._active_scope.reset(self._context_token)
state.depth = 0
self._active_state = None
return False
class TenantUnitOfWork:
"""Bind one trusted visibility scope to one database transaction.
PostgreSQL receives transaction-local settings consumed by RLS policies.
SQLite keeps the same transaction boundary so OSS code exercises the same
request/task ownership and rollback semantics. A manager-owned ContextVar
lets all legacy ``execute_async`` helpers reuse this exact session.
"""
def __init__(
self,
engine: sqlalchemy_asyncio.AsyncEngine,
workspace_uuid: str | None = None,
*,
scope: PersistenceScope | None = None,
active_transaction: ActiveTransactionVar | None = None,
active_scope: ActivePersistenceScopeVar | None = None,
) -> None:
if (workspace_uuid is None) == (scope is None):
raise ValueError('TenantUnitOfWork requires exactly one Workspace or persistence scope')
self._engine = engine
self.workspace_uuid = workspace_uuid
self.scope = scope or PersistenceScope.workspace(typing.cast(str, workspace_uuid))
self.workspace_uuid = self.scope.settings[0][1] if self.scope.kind == PersistenceScopeKind.WORKSPACE else None
self._active_transaction = active_transaction
self._active_scope = active_scope
self._session: sqlalchemy_asyncio.AsyncSession | None = None
self._active_state: ActiveScopedTransaction | None = None
self._context_token: contextvars.Token[ActiveScopedTransaction | None] | None = None
self._database_operation_token: contextvars.Token[ActiveScopedTransaction | None] | None = None
self._used = False
self._owns_transaction = False
_install_rollback_only_error_listener(engine)
@property
def session(self) -> sqlalchemy_asyncio.AsyncSession:
@@ -73,16 +330,67 @@ class TenantUnitOfWork:
raise RuntimeError('TenantUnitOfWork instances cannot be reused')
self._used = True
active_scope = self._active_scope.get() if self._active_scope is not None else None
if active_scope is not None and active_scope.owner_task is asyncio.current_task():
if active_scope.scope != self.scope:
# A request/runtime Workspace boundary is transaction-free and
# may temporarily enter a narrower trusted discovery UoW (for
# example, resolving an account record by its verified UUID).
# It must still never switch directly to another Workspace.
workspace_to_discovery = (
active_scope.scope.kind == PersistenceScopeKind.WORKSPACE
and self.scope.kind != PersistenceScopeKind.WORKSPACE
)
if not workspace_to_discovery:
raise CrossScopeTransactionError(
f'Cannot enter {self.scope.kind.value} scope '
f'while {active_scope.scope.kind.value} scope is active'
)
active = self._active_transaction.get() if self._active_transaction is not None else None
if active is not None:
if active.owner_task is asyncio.current_task():
if active.scope != self.scope:
raise CrossScopeTransactionError(
f'Cannot enter {self.scope.kind.value} scope while {active.scope.kind.value} scope is active'
)
active.depth += 1
self._active_state = active
self._session = active.session
return self
# ContextVars are copied into child tasks. Merely calling a helper
# there must fail (PersistenceManager enforces that), but an
# explicit UoW is allowed to replace the inherited pointer with an
# independent transaction owned by the child task.
session = sqlalchemy_asyncio.AsyncSession(self._engine, expire_on_commit=False)
self._session = session
try:
await session.begin()
if self._engine.dialect.name == 'postgresql':
await session.execute(
sqlalchemy.text(f"SELECT set_config('{TENANT_SETTING}', :workspace_uuid, true)"),
{'workspace_uuid': self.workspace_uuid},
)
for setting_name, setting_value in self.scope.settings:
await session.execute(
sqlalchemy.text('SELECT set_config(:setting_name, :setting_value, true)'),
{'setting_name': setting_name, 'setting_value': setting_value},
)
state = ActiveScopedTransaction(
scope=self.scope,
session=session,
engine=self._engine,
owner_task=asyncio.current_task(),
)
self._active_state = state
if self._active_transaction is not None:
self._context_token = self._active_transaction.set(state)
self._database_operation_token = _DATABASE_OPERATION_TRANSACTION.set(state)
self._owns_transaction = True
except BaseException:
if self._database_operation_token is not None:
_DATABASE_OPERATION_TRANSACTION.reset(self._database_operation_token)
self._database_operation_token = None
if self._active_transaction is not None and self._context_token is not None:
self._active_transaction.reset(self._context_token)
self._context_token = None
await session.rollback()
await session.close()
self._session = None
@@ -95,19 +403,88 @@ class TenantUnitOfWork:
exc_value: BaseException | None,
traceback: types.TracebackType | None,
) -> bool:
state = self._active_state
if state is None:
raise RuntimeError('TenantUnitOfWork has no active transaction state')
self._require_owner_task(state)
if not self._owns_transaction:
if exc_type is not None:
state.mark_rollback_only(exc_value)
state.depth -= 1
self._session = None
self._active_state = None
return False
session = self.session
if exc_type is not None:
state.mark_rollback_only(exc_value)
rollback_only = state.rollback_only
committed = False
try:
if exc_type is None:
if exc_type is None and not rollback_only:
await session.commit()
committed = True
else:
await session.rollback()
finally:
await session.close()
self._session = None
try:
if self._active_transaction is not None and self._context_token is not None:
self._active_transaction.reset(self._context_token)
await session.close()
finally:
if self._database_operation_token is not None:
_DATABASE_OPERATION_TRANSACTION.reset(self._database_operation_token)
self._database_operation_token = None
self._session = None
self._active_state = None
state.depth = 0
self._complete_after_commit_waiters(state, committed=committed)
if exc_type is None and rollback_only:
raise TransactionRollbackOnlyError(
'A scoped database operation failed or rolled back; '
'the transaction was rolled back and after-commit work was cancelled'
) from state.rollback_only_cause
return False
async def execute(self, *args: typing.Any, **kwargs: typing.Any) -> sqlalchemy.engine.Result[typing.Any]:
return await self.session.execute(*args, **kwargs)
try:
return await self.session.execute(*args, **kwargs)
except BaseException as exc:
self._mark_rollback_only(exc)
raise
async def flush(self) -> None:
await self.session.flush()
try:
await self.session.flush()
except BaseException as exc:
self._mark_rollback_only(exc)
raise
def _mark_rollback_only(self, cause: BaseException) -> None:
state = self._active_state
if state is not None:
state.mark_rollback_only(cause)
@staticmethod
def _require_owner_task(state: ActiveScopedTransaction) -> None:
if state.owner_task is not asyncio.current_task():
raise CrossScopeTransactionError(
'Scoped database transactions cannot be inherited by child tasks; open an explicit task scope'
)
@staticmethod
def _complete_after_commit_waiters(
state: ActiveScopedTransaction,
*,
committed: bool,
) -> None:
for waiter in state.after_commit_waiters:
if waiter.done():
continue
if committed:
waiter.set_result(None)
else:
waiter.cancel()
state.after_commit_waiters.clear()
+10 -2
View File
@@ -13,6 +13,7 @@ import langbot_plugin.api.entities.builtin.platform.message as platform_message
import langbot_plugin.api.entities.builtin.provider.session as provider_session
from ..api.http.context import ExecutionContext
from ..core.task_boundary import create_detached_task, run_in_workspace_uow
from .pool import ExecutionContextMismatchError
from ..workspace.errors import WorkspaceError, WorkspaceInvariantError
@@ -203,7 +204,10 @@ class MessageAggregator:
if len(buffer.messages) >= MAX_BUFFER_MESSAGES:
force_flush = True
else:
buffer.timer_task = asyncio.create_task(self._delayed_flush(aggregation_key, delay, execution_context))
buffer.timer_task = create_detached_task(
self._delayed_flush(aggregation_key, delay, execution_context),
after_commit_manager=getattr(self.ap, 'persistence_mgr', None),
)
if force_flush:
await self._flush_buffer(aggregation_key, execution_context)
@@ -218,7 +222,11 @@ class MessageAggregator:
try:
await asyncio.sleep(delay)
await self._flush_buffer(aggregation_key, execution_context)
await run_in_workspace_uow(
self.ap,
execution_context.workspace_uuid,
lambda: self._flush_buffer(aggregation_key, execution_context),
)
except asyncio.CancelledError:
pass
except WorkspaceError as exc:
+28 -13
View File
@@ -43,22 +43,37 @@ class Controller:
try:
async with self.semaphore:
execution_context = await self._assert_query_execution_active(selected_query)
pipeline_uuid = selected_query.pipeline_uuid
queued_context = get_query_execution_context(selected_query)
if pipeline_uuid:
pipeline = await self.ap.pipeline_mgr.get_pipeline_by_uuid(
execution_context,
pipeline_uuid,
)
if pipeline:
await pipeline.run(selected_query)
else:
self.ap.logger.warning(
f'Pipeline {pipeline_uuid} not found for query {selected_query.query_id}, query dropped'
async def run_scoped_query() -> None:
execution_context = await self._assert_query_execution_active(selected_query)
pipeline_uuid = selected_query.pipeline_uuid
if pipeline_uuid:
pipeline = await self.ap.pipeline_mgr.get_pipeline_by_uuid(
execution_context,
pipeline_uuid,
)
if pipeline:
await pipeline.run(selected_query)
else:
self.ap.logger.warning(
f'Pipeline {pipeline_uuid} not found for query {selected_query.query_id}, query dropped'
)
else:
self.ap.logger.warning(f'No pipeline_uuid for query {selected_query.query_id}, query dropped')
tenant_scope = getattr(self.ap.persistence_mgr, 'tenant_scope', None)
cloud_runtime = (
getattr(getattr(self.ap.persistence_mgr, 'mode', None), 'value', None) == 'cloud_runtime'
)
if cloud_runtime:
if not callable(tenant_scope):
raise RuntimeError('Cloud query processing requires an explicit tenant scope')
async with tenant_scope(queued_context.workspace_uuid):
await run_scoped_query()
else:
self.ap.logger.warning(f'No pipeline_uuid for query {selected_query.query_id}, query dropped')
await run_scoped_query()
except WorkspaceError as exc:
self.ap.logger.info(
f'Dropped query {selected_query.query_id} because its Workspace execution binding is stale: {exc}'
+28
View File
@@ -506,6 +506,34 @@ class PipelineManager:
async def load_pipelines_from_db(self):
self.ap.logger.info('Loading pipelines from db...')
self.pipelines = []
list_bindings = getattr(self.ap.workspace_service, 'list_active_execution_bindings', None)
tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
cloud_runtime = getattr(getattr(self.ap.persistence_mgr, 'mode', None), 'value', None) == 'cloud_runtime'
if cloud_runtime:
if not callable(list_bindings) or not callable(tenant_uow):
raise RuntimeError('Cloud pipeline loading requires explicit instance discovery and tenant UoWs')
for binding in await list_bindings():
async with tenant_uow(binding.workspace_uuid):
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(persistence_pipeline.LegacyPipeline)
.where(persistence_pipeline.LegacyPipeline.workspace_uuid == binding.workspace_uuid)
.order_by(persistence_pipeline.LegacyPipeline.uuid)
)
for pipeline in result.all():
await self.load_pipeline(
ExecutionContext(
instance_uuid=binding.instance_uuid,
workspace_uuid=binding.workspace_uuid,
placement_generation=binding.placement_generation,
pipeline_uuid=pipeline.uuid,
trigger_principal=PrincipalContext(PrincipalType.SYSTEM),
),
pipeline,
)
return
# Compatibility path for isolated manager tests and older embedders.
result = await self.ap.persistence_mgr.execute_async(sqlalchemy.select(persistence_pipeline.LegacyPipeline))
pipelines = result.all()
+66 -9
View File
@@ -2,6 +2,7 @@ from __future__ import annotations
import asyncio
import dataclasses
import functools
import json
import re
import traceback
@@ -14,6 +15,7 @@ from ..discover import engine
from ..entity.persistence import bot as persistence_bot
from ..entity.persistence import pipeline as persistence_pipeline
from ..entity.persistence import workspace as persistence_workspace
from ..entity.errors import platform as platform_errors
from ..api.http.context import ExecutionContext, PrincipalContext, PrincipalType, RequestContext
@@ -256,6 +258,24 @@ class RuntimeBot:
await self.logger.error(f'Failed to record discarded message: {e}')
async def initialize(self):
def tenant_scoped_listener(listener):
"""Bind adapter callbacks to a Workspace without holding a DB transaction."""
@functools.wraps(listener)
async def wrapped(*args, **kwargs):
tenant_scope = getattr(self.ap.persistence_mgr, 'tenant_scope', None)
cloud_runtime = (
getattr(getattr(self.ap.persistence_mgr, 'mode', None), 'value', None) == 'cloud_runtime'
)
if cloud_runtime:
if not callable(tenant_scope):
raise RuntimeError('Cloud platform callbacks require an explicit tenant scope')
async with tenant_scope(self.workspace_uuid):
return await listener(*args, **kwargs)
return await listener(*args, **kwargs)
return wrapped
async def on_friend_message(
event: platform_events.FriendMessage,
adapter: abstract_platform_adapter.AbstractMessagePlatformAdapter,
@@ -396,8 +416,8 @@ class RuntimeBot:
else:
await self.logger.info('Pipeline skipped for group message due to webhook response')
self.adapter.register_listener(platform_events.FriendMessage, on_friend_message)
self.adapter.register_listener(platform_events.GroupMessage, on_group_message)
self.adapter.register_listener(platform_events.FriendMessage, tenant_scoped_listener(on_friend_message))
self.adapter.register_listener(platform_events.GroupMessage, tenant_scoped_listener(on_group_message))
# Register feedback listener (only effective on adapters that support it)
async def on_feedback(
@@ -444,7 +464,7 @@ class RuntimeBot:
except Exception:
await self.logger.error(f'Failed to record feedback: {traceback.format_exc()}')
self.adapter.register_listener(platform_events.FeedbackEvent, on_feedback)
self.adapter.register_listener(platform_events.FeedbackEvent, tenant_scoped_listener(on_feedback))
async def run(self):
async def exception_wrapper():
@@ -642,14 +662,51 @@ class PlatformManager:
self.bots = []
result = await self.ap.persistence_mgr.execute_async(sqlalchemy.select(persistence_bot.Bot))
instance_uow = getattr(self.ap.persistence_mgr, 'instance_discovery_uow', None)
tenant_scope = getattr(self.ap.persistence_mgr, 'tenant_scope', None)
if callable(instance_uow) and callable(tenant_scope):
async with instance_uow(self.ap.workspace_service.instance_uuid) as discovery:
workspace_uuids = list(
(
await discovery.session.scalars(
sqlalchemy.select(persistence_workspace.WorkspaceExecutionState.workspace_uuid)
.where(
persistence_workspace.WorkspaceExecutionState.instance_uuid
== self.ap.workspace_service.instance_uuid,
persistence_workspace.WorkspaceExecutionState.state
== persistence_workspace.WorkspaceExecutionStatus.ACTIVE.value,
persistence_workspace.WorkspaceExecutionState.write_fenced.is_(False),
)
.order_by(persistence_workspace.WorkspaceExecutionState.workspace_uuid)
)
).all()
)
else:
# Compatibility for lightweight tests and pre-tenancy managers.
result = await self.ap.persistence_mgr.execute_async(sqlalchemy.select(persistence_bot.Bot))
workspace_uuids = sorted({bot.workspace_uuid for bot in result.all()})
bots = result.all()
for bot in bots:
# load all bots here, enable or disable will be handled in runtime
for workspace_uuid in workspace_uuids:
try:
binding = await self.ap.workspace_service.get_execution_binding(bot.workspace_uuid)
if callable(tenant_scope):
async with tenant_scope(workspace_uuid):
await self._load_workspace_bots(workspace_uuid)
else:
await self._load_workspace_bots(workspace_uuid)
except Exception as e:
self.ap.logger.error(
f'Failed to load Workspace bots for {workspace_uuid}: {e}\n{traceback.format_exc()}'
)
async def _load_workspace_bots(self, workspace_uuid: str) -> None:
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(persistence_bot.Bot)
.where(persistence_bot.Bot.workspace_uuid == workspace_uuid)
.order_by(persistence_bot.Bot.uuid)
)
for bot in result.all():
try:
binding = await self.ap.workspace_service.get_execution_binding(workspace_uuid)
execution_context = ExecutionContext(
instance_uuid=binding.instance_uuid,
workspace_uuid=binding.workspace_uuid,
@@ -286,19 +286,30 @@ class OpenClawWeixinAdapter(abstract_platform_adapter.AbstractMessagePlatformAda
if execution_context.bot_uuid != self._bot_uuid:
raise RuntimeError('Weixin Bot UUID does not match its ExecutionContext')
binding = await ap.workspace_service.get_execution_binding(
execution_context.workspace_uuid,
expected_generation=execution_context.placement_generation,
)
if binding.instance_uuid != execution_context.instance_uuid:
raise RuntimeError('Weixin Bot Workspace belongs to another LangBot instance')
async def persist() -> None:
binding = await ap.workspace_service.get_execution_binding(
execution_context.workspace_uuid,
expected_generation=execution_context.placement_generation,
)
if binding.instance_uuid != execution_context.instance_uuid:
raise RuntimeError('Weixin Bot Workspace belongs to another LangBot instance')
await ap.persistence_mgr.execute_async(
sqlalchemy.update(persistence_bot.Bot)
.where(persistence_bot.Bot.workspace_uuid == execution_context.workspace_uuid)
.where(persistence_bot.Bot.uuid == self._bot_uuid)
.values(adapter_config=self.config)
)
await ap.persistence_mgr.execute_async(
sqlalchemy.update(persistence_bot.Bot)
.where(persistence_bot.Bot.workspace_uuid == execution_context.workspace_uuid)
.where(persistence_bot.Bot.uuid == self._bot_uuid)
.values(adapter_config=self.config)
)
cloud_runtime = getattr(getattr(ap.persistence_mgr, 'mode', None), 'value', None) == 'cloud_runtime'
if cloud_runtime:
tenant_uow = getattr(ap.persistence_mgr, 'tenant_uow', None)
if not callable(tenant_uow):
raise RuntimeError('Cloud adapter persistence requires an explicit tenant UoW')
async with tenant_uow(execution_context.workspace_uuid):
await persist()
else:
await persist()
except Exception as e:
await self.logger.warning(f'Failed to persist adapter config: {e}')
File diff suppressed because it is too large Load Diff
+97
View File
@@ -0,0 +1,97 @@
from __future__ import annotations
import re
from typing import Any
from urllib.parse import unquote, urlparse
_GITHUB_OWNER_PATTERN = re.compile(r'^[A-Za-z0-9](?:[A-Za-z0-9-]{0,38})$')
_GITHUB_REPO_PATTERN = re.compile(r'^[A-Za-z0-9._-]{1,100}$')
def _positive_github_id(value: object, field_name: str) -> int:
if isinstance(value, bool):
raise ValueError(f'{field_name} must be a positive GitHub identifier')
try:
identifier = int(value)
except (TypeError, ValueError) as exc:
raise ValueError(f'{field_name} must be a positive GitHub identifier') from exc
if identifier <= 0 or str(value).strip() != str(identifier):
raise ValueError(f'{field_name} must be a positive GitHub identifier')
return identifier
def validate_github_release_asset_url(
asset_url: object,
*,
owner: str,
repo: str,
release_tag: str,
) -> str:
"""Accept only a GitHub browser release URL tied to the requested release."""
normalized_url = str(asset_url or '').strip()
parsed = urlparse(normalized_url)
try:
port = parsed.port
except ValueError as exc:
raise ValueError('asset_url has an invalid port') from exc
if (
parsed.scheme != 'https'
or (parsed.hostname or '').lower() != 'github.com'
or parsed.username is not None
or parsed.password is not None
or port not in {None, 443}
or parsed.fragment
):
raise ValueError('asset_url must be an HTTPS GitHub release asset URL')
decoded_path = unquote(parsed.path)
expected_prefix = f'/{owner}/{repo}/releases/download/{release_tag}/'
if not decoded_path.casefold().startswith(
f'/{owner}/{repo}/releases/download/'.casefold()
) or not decoded_path.startswith(expected_prefix):
raise ValueError('asset_url does not match the requested GitHub release')
if decoded_path == expected_prefix or decoded_path.endswith('/'):
raise ValueError('asset_url must identify a GitHub release asset')
return normalized_url
def validate_github_plugin_install_info(install_info: dict[str, Any]) -> dict[str, Any]:
"""Normalize a GitHub install request without trusting a tenant-provided URL."""
owner = str(install_info.get('owner') or '').strip()
repo = str(install_info.get('repo') or '').strip()
release_tag = str(install_info.get('release_tag') or '').strip()
if _GITHUB_OWNER_PATTERN.fullmatch(owner) is None:
raise ValueError('owner must be a valid GitHub repository owner')
if _GITHUB_REPO_PATTERN.fullmatch(repo) is None:
raise ValueError('repo must be a valid GitHub repository name')
if not release_tag or '\x00' in release_tag or len(release_tag) > 255:
raise ValueError('release_tag must identify a GitHub release')
release_id_value = install_info.get('release_id')
asset_id_value = install_info.get('asset_id')
normalized = dict(install_info)
normalized.update(
{
'owner': owner,
'repo': repo,
'release_tag': release_tag,
'github_url': f'https://github.com/{owner}/{repo}',
}
)
if release_id_value is not None or asset_id_value is not None:
if release_id_value is None or asset_id_value is None:
raise ValueError('release_id and asset_id must be provided together')
normalized['release_id'] = _positive_github_id(release_id_value, 'release_id')
normalized['asset_id'] = _positive_github_id(asset_id_value, 'asset_id')
normalized.pop('asset_url', None)
return normalized
normalized['asset_url'] = validate_github_release_asset_url(
install_info.get('asset_url'),
owner=owner,
repo=repo,
release_tag=release_tag,
)
return normalized
+362 -162
View File
@@ -4,15 +4,26 @@ import inspect
import typing
from typing import Any
import base64
import contextlib
import contextvars
import traceback
import uuid
from dataclasses import dataclass
import sqlalchemy
from langbot_plugin.runtime.io import handler
from langbot_plugin.runtime.io.connection import Connection
from langbot_plugin.entities.io.context import ActionContext
from langbot_plugin.entities.io.context import (
ActionContext,
ApplyPluginInstallationRequest,
InstallationBinding,
PluginInstallationDesiredState,
PluginWorkerPolicy,
ReconcilePluginInstallationsRequest,
RemovePluginInstallationRequest,
RuntimeConfig,
RuntimeIdentity,
)
from langbot_plugin.entities.io.actions.enums import (
CommonAction,
RuntimeToLangBotAction,
@@ -61,6 +72,9 @@ class _PluginInstallationIdentity:
workspace_uuid: str
plugin_author: str
plugin_name: str
installation_uuid: str
runtime_revision: int
artifact_digest: str
_UNTRUSTED_SCOPE_FIELDS = frozenset(
@@ -71,12 +85,13 @@ _UNTRUSTED_SCOPE_FIELDS = frozenset(
'workspace_uuid',
'placement_generation',
'installation_uuid',
'runtime_revision',
'artifact_digest',
}
)
_RUNTIME_SCOPED_ACTIONS = frozenset(
{
CommonAction.PING.value,
CommonAction.FILE_CHUNK.value,
RuntimeToLangBotAction.INITIALIZE_PLUGIN_SETTINGS.value,
RuntimeToLangBotAction.GET_PLUGIN_SETTINGS.value,
@@ -89,113 +104,129 @@ class RuntimeConnectionHandler(handler.Handler):
ap: app.Application
@staticmethod
def derive_installation_uuid(
action_context: ActionContext,
plugin_author: str,
plugin_name: str,
) -> str:
"""Derive the current stable installation capability without trusting plugin data."""
return str(
uuid.uuid5(
uuid.NAMESPACE_URL,
'langbot:plugin-installation:'
f'{action_context.instance_uuid}:{action_context.workspace_uuid}:'
f'{plugin_author}/{plugin_name}',
)
)
def validate_inbound_action_context(
self,
action: str,
action_context: ActionContext | None,
) -> ActionContext | None:
"""Keep a validated child installation capability from the request envelope."""
"""Require a complete installation tuple on every tenant action."""
validated = super().validate_inbound_action_context(action, action_context)
if action_context is not None and action_context.installation_uuid is not None:
# The base handler already checked instance, Workspace and placement
# generation against this connector's immutable binding. The
# installation is resolved against trusted Host state asynchronously
# before dispatch.
if action == CommonAction.PING.value:
if action_context is not None:
raise ValueError('PING does not accept an installation binding')
return None
if isinstance(action_context, InstallationBinding):
return action_context
return validated
if self._allow_legacy_oss_context(action, action_context):
return action_context
raise ValueError(f'{action} requires a complete InstallationBinding context')
def _allow_legacy_oss_context(self, action: str, action_context: ActionContext | None) -> bool:
"""Keep pre-v4 local plugins usable without weakening shared Runtime."""
if not (
getattr(getattr(self.ap, 'deployment', None), 'mode', 'oss') == 'oss'
and isinstance(action_context, ActionContext)
and not isinstance(action_context, InstallationBinding)
):
return False
if action in {
RuntimeToLangBotAction.INITIALIZE_PLUGIN_SETTINGS.value,
RuntimeToLangBotAction.GET_PLUGIN_SETTINGS.value,
CommonAction.FILE_CHUNK.value,
}:
return True
# Pre-v4 OSS workers receive this capability from Core's settings
# response, but their pinned SDK cannot carry revision/digest fields.
# Shared Runtime never enters this compatibility branch.
return bool(action_context.installation_uuid)
def _require_runtime_action_context(self) -> ActionContext:
action_context = self.current_action_context or self.bound_action_context
action_context = self.current_action_context
if action_context is None:
raise ValueError('Plugin Runtime action is missing a trusted Workspace context')
bound = self.require_bound_action_context()
if not bound.same_workspace(action_context):
raise ValueError('Plugin Runtime action context does not match its connector binding')
return action_context
def _remember_installation(
self,
action_context: ActionContext,
plugin_author: str,
plugin_name: str,
) -> str:
installation_uuid = self.derive_installation_uuid(
action_context,
plugin_author,
plugin_name,
)
self._installation_bindings[installation_uuid] = _PluginInstallationIdentity(
workspace_uuid=action_context.workspace_uuid,
plugin_author=plugin_author,
plugin_name=plugin_name,
)
return installation_uuid
async def _resolve_installation_identity(
self,
action_context: ActionContext,
action_context: InstallationBinding,
) -> _PluginInstallationIdentity:
installation_uuid = action_context.installation_uuid
if installation_uuid is None:
raise ValueError('Plugin action is missing installation_uuid')
identity = self._installation_bindings.get(installation_uuid)
if identity is not None:
if identity.workspace_uuid != action_context.workspace_uuid:
raise ValueError('Plugin installation belongs to another Workspace')
cached = self._installation_bindings.get(action_context.installation_uuid)
if cached is not None:
cached_binding, identity = cached
if cached_binding != action_context:
raise ValueError('Plugin installation revision or artifact is stale')
return identity
# A control connection may reconnect while the Runtime keeps plugin
# processes alive. Rebuild the capability map from Workspace-scoped
# settings instead of accepting an installation asserted by the peer.
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(
persistence_plugin.PluginSetting.plugin_author,
persistence_plugin.PluginSetting.plugin_name,
).where(persistence_plugin.PluginSetting.workspace_uuid == action_context.workspace_uuid)
)
for setting in result.all():
candidate_uuid = self.derive_installation_uuid(
action_context,
setting.plugin_author,
setting.plugin_name,
persistence_plugin.PluginSetting.installation_uuid,
persistence_plugin.PluginSetting.runtime_revision,
persistence_plugin.PluginSetting.artifact_digest,
)
if candidate_uuid == installation_uuid:
return self._installation_bindings.setdefault(
installation_uuid,
_PluginInstallationIdentity(
workspace_uuid=action_context.workspace_uuid,
plugin_author=setting.plugin_author,
plugin_name=setting.plugin_name,
),
)
.where(persistence_plugin.PluginSetting.workspace_uuid == action_context.workspace_uuid)
.where(persistence_plugin.PluginSetting.installation_uuid == action_context.installation_uuid)
)
setting = result.first()
if setting is None:
raise ValueError('Plugin installation is not registered in this Workspace')
if (
setting.runtime_revision != action_context.runtime_revision
or setting.artifact_digest != action_context.artifact_digest
):
raise ValueError('Plugin installation revision or artifact is stale')
identity = _PluginInstallationIdentity(
workspace_uuid=action_context.workspace_uuid,
plugin_author=setting.plugin_author,
plugin_name=setting.plugin_name,
installation_uuid=setting.installation_uuid,
runtime_revision=setting.runtime_revision,
artifact_digest=setting.artifact_digest,
)
self._installation_bindings[action_context.installation_uuid] = (action_context, identity)
return identity
raise ValueError('Plugin installation is not registered in this Workspace')
async def _resolve_legacy_oss_installation_identity(
self,
action_context: ActionContext,
) -> _PluginInstallationIdentity:
if not action_context.installation_uuid:
raise ValueError('Legacy OSS plugin action is missing its installation capability')
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(
persistence_plugin.PluginSetting.plugin_author,
persistence_plugin.PluginSetting.plugin_name,
persistence_plugin.PluginSetting.installation_uuid,
persistence_plugin.PluginSetting.runtime_revision,
persistence_plugin.PluginSetting.artifact_digest,
)
.where(persistence_plugin.PluginSetting.workspace_uuid == action_context.workspace_uuid)
.where(persistence_plugin.PluginSetting.installation_uuid == action_context.installation_uuid)
)
setting = result.first()
if setting is None:
raise ValueError('Plugin installation is not registered in this Workspace')
return _PluginInstallationIdentity(
workspace_uuid=action_context.workspace_uuid,
plugin_author=setting.plugin_author,
plugin_name=setting.plugin_name,
installation_uuid=setting.installation_uuid,
runtime_revision=setting.runtime_revision,
artifact_digest=setting.artifact_digest,
)
async def _require_plugin_action_context(
self,
) -> tuple[ActionContext, _PluginInstallationIdentity]:
action_context = self._require_runtime_action_context()
identity = await self._resolve_installation_identity(action_context)
if isinstance(action_context, InstallationBinding):
identity = await self._resolve_installation_identity(action_context)
elif self._allow_legacy_oss_context('plugin_action', action_context):
identity = await self._resolve_legacy_oss_installation_identity(action_context)
else:
raise ValueError('Plugin action requires a complete InstallationBinding context')
return action_context, identity
async def _require_active_action_context(self, action_context: ActionContext) -> None:
@@ -215,25 +246,53 @@ class RuntimeConnectionHandler(handler.Handler):
):
raise ValueError('Plugin Runtime action uses a stale Workspace execution binding')
@contextlib.asynccontextmanager
async def _tenant_action_scope(self, action_context: ActionContext):
"""Bind Runtime-origin work to the trusted Workspace database scope.
Cloud persistence deliberately rejects unscoped access and PostgreSQL
RLS reads the Workspace id from each short database transaction. The
wire envelope is validated before this helper is entered, so plugin
payload fields can never select the scope. A transaction-free boundary
avoids reserving one pooled connection across provider and network waits.
"""
persistence_mgr = getattr(self.ap, 'persistence_mgr', None)
if persistence_mgr is None:
yield
return
tenant_scope_descriptor = getattr(type(persistence_mgr), 'tenant_scope', None)
if not callable(tenant_scope_descriptor):
# Lightweight test doubles and older OSS persistence managers do
# not expose the transaction-free scope API.
yield
return
async with persistence_mgr.tenant_scope(action_context.workspace_uuid):
yield
def _secure_plugin_actions(self) -> None:
"""Wrap plugin-origin actions with installation validation and payload scrubbing."""
for action_name, action_handler in list(self.actions.items()):
if action_name in _RUNTIME_SCOPED_ACTIONS:
if action_name == CommonAction.PING.value:
continue
async def secured_action(
data: dict[str, Any],
*,
_action_handler=action_handler,
_runtime_scoped=action_name in _RUNTIME_SCOPED_ACTIONS,
) -> handler.ActionResponse:
action_context, _ = await self._require_plugin_action_context()
await self._require_active_action_context(action_context)
safe_data = {key: value for key, value in data.items() if key not in _UNTRUSTED_SCOPE_FIELDS}
response = _action_handler(safe_data)
if inspect.isawaitable(response):
response = await response
return response
action_context = self._require_runtime_action_context()
async with self._tenant_action_scope(action_context):
if not _runtime_scoped:
action_context, _ = await self._require_plugin_action_context()
await self._require_active_action_context(action_context)
safe_data = {key: value for key, value in data.items() if key not in _UNTRUSTED_SCOPE_FIELDS}
response = _action_handler(safe_data)
if inspect.isawaitable(response):
response = await response
return response
self.actions[action_name] = secured_action
@@ -253,6 +312,31 @@ class RuntimeConnectionHandler(handler.Handler):
raise ValueError('Plugin installation setting was not found')
return setting
async def _get_plugin_setting_by_name(
self,
action_context: ActionContext,
plugin_author: str,
plugin_name: str,
):
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(persistence_plugin.PluginSetting)
.where(persistence_plugin.PluginSetting.workspace_uuid == action_context.workspace_uuid)
.where(persistence_plugin.PluginSetting.plugin_author == plugin_author)
.where(persistence_plugin.PluginSetting.plugin_name == plugin_name)
)
return result.first()
@staticmethod
def _require_setting_binding(setting, action_context: InstallationBinding) -> None:
if setting is None:
raise ValueError('Plugin installation setting was not found')
if (
setting.installation_uuid != action_context.installation_uuid
or setting.runtime_revision != action_context.runtime_revision
or setting.artifact_digest != action_context.artifact_digest
):
raise ValueError('Plugin installation binding does not match the active setting')
async def _resolve_query(
self,
data: dict[str, Any],
@@ -355,18 +439,24 @@ class RuntimeConnectionHandler(handler.Handler):
connection: Connection,
disconnect_callback: typing.Callable[[], typing.Coroutine[typing.Any, typing.Any, bool]],
ap: app.Application,
action_context: ActionContext,
):
super().__init__(connection, disconnect_callback)
self.ap = ap
self.bind_action_context(action_context.without_installation())
self._installation_bindings: dict[str, _PluginInstallationIdentity] = {}
self._outbound_installation_context: contextvars.ContextVar[InstallationBinding | None] = (
contextvars.ContextVar(
f'{self.__class__.__name__}_{id(self)}_outbound_installation',
default=None,
)
)
self._installation_bindings: dict[
str,
tuple[InstallationBinding, _PluginInstallationIdentity],
] = {}
@self.action(RuntimeToLangBotAction.INITIALIZE_PLUGIN_SETTINGS)
async def initialize_plugin_settings(data: dict[str, Any]) -> handler.ActionResponse:
"""Initialize plugin settings"""
action_context = self._require_runtime_action_context()
await self._require_active_action_context(action_context)
# check if exists plugin setting
plugin_author = data['plugin_author']
plugin_name = data['plugin_name']
@@ -374,39 +464,37 @@ class RuntimeConnectionHandler(handler.Handler):
install_info = data['install_info']
try:
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(persistence_plugin.PluginSetting)
.where(persistence_plugin.PluginSetting.workspace_uuid == action_context.workspace_uuid)
.where(persistence_plugin.PluginSetting.plugin_author == plugin_author)
.where(persistence_plugin.PluginSetting.plugin_name == plugin_name)
setting = await self._get_plugin_setting_by_name(
action_context,
plugin_author,
plugin_name,
)
setting = result.first()
if setting is not None:
# delete plugin setting
if isinstance(action_context, InstallationBinding):
self._require_setting_binding(setting, action_context)
elif setting is None:
# OSS debug and pre-v4 data/plugins are the only callers
# allowed to create settings without a desired-state row.
await self.ap.persistence_mgr.execute_async(
sqlalchemy.delete(persistence_plugin.PluginSetting)
sqlalchemy.insert(persistence_plugin.PluginSetting).values(
workspace_uuid=action_context.workspace_uuid,
plugin_author=plugin_author,
plugin_name=plugin_name,
install_source=install_source,
install_info=install_info,
enabled=True,
priority=0,
config={},
)
)
else:
await self.ap.persistence_mgr.execute_async(
sqlalchemy.update(persistence_plugin.PluginSetting)
.where(persistence_plugin.PluginSetting.workspace_uuid == action_context.workspace_uuid)
.where(persistence_plugin.PluginSetting.plugin_author == plugin_author)
.where(persistence_plugin.PluginSetting.plugin_name == plugin_name)
.values(install_source=install_source, install_info=install_info)
)
# create plugin setting
await self.ap.persistence_mgr.execute_async(
sqlalchemy.insert(persistence_plugin.PluginSetting).values(
workspace_uuid=action_context.workspace_uuid,
plugin_author=plugin_author,
plugin_name=plugin_name,
install_source=install_source,
install_info=install_info,
# inherit from existing setting
enabled=setting.enabled if setting is not None else True,
priority=setting.priority if setting is not None else 0,
config=setting.config if setting is not None else {}, # noqa: F821
)
)
return handler.ActionResponse.success(
data={},
)
@@ -421,16 +509,16 @@ class RuntimeConnectionHandler(handler.Handler):
"""Get plugin settings"""
action_context = self._require_runtime_action_context()
await self._require_active_action_context(action_context)
plugin_author = data['plugin_author']
plugin_name = data['plugin_name']
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(persistence_plugin.PluginSetting)
.where(persistence_plugin.PluginSetting.workspace_uuid == action_context.workspace_uuid)
.where(persistence_plugin.PluginSetting.plugin_author == plugin_author)
.where(persistence_plugin.PluginSetting.plugin_name == plugin_name)
setting = await self._get_plugin_setting_by_name(
action_context,
plugin_author,
plugin_name,
)
if isinstance(action_context, InstallationBinding):
self._require_setting_binding(setting, action_context)
data = {
'enabled': True,
@@ -438,21 +526,20 @@ class RuntimeConnectionHandler(handler.Handler):
'plugin_config': {},
'install_source': 'local',
'install_info': {},
'installation_uuid': self._remember_installation(
action_context,
plugin_author,
plugin_name,
),
'installation_uuid': None,
'runtime_revision': None,
'artifact_digest': None,
}
setting = result.first()
if setting is not None:
data['enabled'] = setting.enabled
data['priority'] = setting.priority
data['plugin_config'] = setting.config
data['install_source'] = setting.install_source
data['install_info'] = setting.install_info
data['installation_uuid'] = setting.installation_uuid
data['runtime_revision'] = setting.runtime_revision
data['artifact_digest'] = setting.artifact_digest
return handler.ActionResponse.success(
data=data,
@@ -1369,24 +1456,144 @@ class RuntimeConnectionHandler(handler.Handler):
self._secure_plugin_actions()
async def ping(self) -> dict[str, Any]:
"""Ping the runtime"""
return await self.call_action(
CommonAction.PING,
{},
timeout=10,
@contextlib.contextmanager
def installation_scope(self, binding: InstallationBinding | None):
"""Attach one immutable installation tuple to nested wire actions."""
token = self._outbound_installation_context.set(binding)
try:
yield
finally:
self._outbound_installation_context.reset(token)
def register_installation_binding(
self,
binding: InstallationBinding,
*,
plugin_author: str,
plugin_name: str,
) -> None:
"""Install Core's current desired-state fence for inbound actions."""
existing = self._installation_bindings.get(binding.installation_uuid)
if existing is not None:
existing_binding = existing[0]
if existing_binding.workspace_uuid != binding.workspace_uuid:
raise ValueError('Plugin installation cannot move between Workspaces')
if binding.placement_generation < existing_binding.placement_generation or (
binding.placement_generation == existing_binding.placement_generation
and binding.runtime_revision < existing_binding.runtime_revision
):
raise ValueError('Cannot register a stale plugin installation binding')
self._installation_bindings[binding.installation_uuid] = (
binding,
_PluginInstallationIdentity(
workspace_uuid=binding.workspace_uuid,
plugin_author=plugin_author,
plugin_name=plugin_name,
installation_uuid=binding.installation_uuid,
runtime_revision=binding.runtime_revision,
artifact_digest=binding.artifact_digest,
),
)
async def set_runtime_config(self, cloud_service_url: str | None) -> dict[str, Any]:
"""Push runtime configuration (e.g. marketplace URL) to the runtime."""
data = {}
if cloud_service_url:
data['cloud_service_url'] = cloud_service_url
return await self.call_action(
LangBotToRuntimeAction.SET_RUNTIME_CONFIG,
data,
timeout=10,
def unregister_installation_binding(self, binding: InstallationBinding) -> None:
existing = self._installation_bindings.get(binding.installation_uuid)
if existing is not None and existing[0] == binding:
self._installation_bindings.pop(binding.installation_uuid, None)
def resolve_outbound_action_context(
self,
action_context: InstallationBinding | ActionContext | dict[str, Any] | None,
) -> InstallationBinding | ActionContext | None:
if action_context is not None:
return super().resolve_outbound_action_context(action_context)
inbound_context = self.current_action_context
if inbound_context is not None:
return inbound_context
return self._outbound_installation_context.get()
def require_outbound_installation_context(self) -> InstallationBinding:
binding = self._outbound_installation_context.get()
if not isinstance(binding, InstallationBinding):
raise ValueError('Host plugin action requires an InstallationBinding scope')
return binding
async def ping(self) -> dict[str, Any]:
"""Ping the runtime"""
with self.installation_scope(None):
return await self.call_action(
CommonAction.PING,
{},
timeout=10,
)
async def set_runtime_config(
self,
*,
runtime_identity: RuntimeIdentity,
worker_policy: PluginWorkerPolicy,
runtime_profile: typing.Literal['oss_dev', 'shared'],
cloud_service_url: str | None,
) -> dict[str, Any]:
"""Push the instance-scoped, immutable Runtime handshake."""
runtime_config = RuntimeConfig(
runtime_identity=runtime_identity,
worker_policy=worker_policy,
runtime_profile=runtime_profile,
cloud_service_url=cloud_service_url,
)
with self.installation_scope(None):
return await self.call_action(
LangBotToRuntimeAction.SET_RUNTIME_CONFIG,
runtime_config.model_dump(exclude_none=True),
timeout=10,
)
async def reconcile_plugin_installations(
self,
installations: tuple[PluginInstallationDesiredState, ...],
) -> dict[str, Any]:
request = ReconcilePluginInstallationsRequest(installations=installations)
with self.installation_scope(None):
return await self.call_action(
LangBotToRuntimeAction.RECONCILE_PLUGIN_INSTALLATIONS,
request.model_dump(),
timeout=120,
)
async def apply_plugin_installation(
self,
binding: InstallationBinding,
*,
artifact_package: bytes | None,
enabled: bool,
) -> dict[str, Any]:
with self.installation_scope(binding):
artifact_file_key = None
if artifact_package is not None:
artifact_file_key = await self.send_file(artifact_package, 'lbpkg')
request = ApplyPluginInstallationRequest(
artifact_file_key=artifact_file_key,
enabled=enabled,
)
return await self.call_action(
LangBotToRuntimeAction.APPLY_PLUGIN_INSTALLATION,
request.model_dump(exclude_none=True),
timeout=120,
)
async def remove_plugin_installation(
self,
binding: InstallationBinding,
) -> dict[str, Any]:
with self.installation_scope(binding):
return await self.call_action(
LangBotToRuntimeAction.REMOVE_PLUGIN_INSTALLATION,
RemovePluginInstallationRequest().model_dump(),
timeout=120,
)
async def install_plugin(
self, install_source: str, install_info: dict[str, Any]
@@ -1455,7 +1662,7 @@ class RuntimeConnectionHandler(handler.Handler):
async def set_plugin_config(self, plugin_author: str, plugin_name: str, config: dict[str, Any]) -> dict[str, Any]:
"""Set plugin config"""
action_context = self.require_bound_action_context()
action_context = self.require_outbound_installation_context()
# update plugin setting
await self.ap.persistence_mgr.execute_async(
sqlalchemy.update(persistence_plugin.PluginSetting)
@@ -1642,12 +1849,7 @@ class RuntimeConnectionHandler(handler.Handler):
async def cleanup_plugin_data(self, plugin_author: str, plugin_name: str) -> None:
"""Cleanup plugin settings and binary storage"""
action_context = self.require_bound_action_context()
installation_uuid = self.derive_installation_uuid(
action_context,
plugin_author,
plugin_name,
)
action_context = self.require_outbound_installation_context()
# Delete plugin settings
await self.ap.persistence_mgr.execute_async(
sqlalchemy.delete(persistence_plugin.PluginSetting)
@@ -1664,9 +1866,6 @@ class RuntimeConnectionHandler(handler.Handler):
.where(persistence_bstorage.BinaryStorage.owner_type == 'plugin')
.where(persistence_bstorage.BinaryStorage.owner == owner)
)
installation_bindings = getattr(self, '_installation_bindings', None)
if installation_bindings is not None:
installation_bindings.pop(installation_uuid, None)
async def call_tool(
self,
@@ -1744,11 +1943,12 @@ class RuntimeConnectionHandler(handler.Handler):
async def get_debug_info(self) -> dict[str, Any]:
"""Get debug information including debug key and WS URL"""
result = await self.call_action(
LangBotToRuntimeAction.GET_DEBUG_INFO,
{},
timeout=10,
)
with self.installation_scope(None):
result = await self.call_action(
LangBotToRuntimeAction.GET_DEBUG_INFO,
{},
timeout=10,
)
return result
# ================= RAG Capability Callers (LangBot -> Runtime) =================
@@ -197,6 +197,23 @@ class ModelManager:
self.llm_model_dict = {}
self.embedding_model_dict = {}
self.rerank_model_dict = {}
list_bindings = getattr(self.ap.workspace_service, 'list_active_execution_bindings', None)
tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
cloud_runtime = getattr(getattr(self.ap.persistence_mgr, 'mode', None), 'value', None) == 'cloud_runtime'
if cloud_runtime:
if not callable(list_bindings) or not callable(tenant_uow):
raise RuntimeError('Cloud model loading requires explicit instance discovery and tenant UoWs')
for binding in await list_bindings():
context = self._context_from_binding(
binding,
trigger_principal=PrincipalContext(principal_type=PrincipalType.SYSTEM),
)
async with tenant_uow(binding.workspace_uuid):
await self._load_workspace_models(context)
return
# Compatibility path for isolated manager tests and older embedders.
contexts: dict[str, ExecutionContext] = {}
async def context_for(workspace_uuid: str | None) -> ExecutionContext:
@@ -263,6 +280,61 @@ class ModelManager:
except Exception as exc:
self.ap.logger.error(f'Failed to load model {model_entity.uuid}: {exc}\n{traceback.format_exc()}')
async def _load_workspace_models(self, context: ExecutionContext) -> None:
"""Load one Workspace while its tenant transaction is active."""
providers_result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(persistence_model.ModelProvider).where(
persistence_model.ModelProvider.workspace_uuid == context.workspace_uuid
)
)
for provider_entity in providers_result.all():
try:
runtime_provider = await self._build_provider(context, provider_entity)
self.provider_dict[self._cache_key(context, provider_entity.uuid)] = runtime_provider
except provider_errors.RequesterNotFoundError as exc:
self.ap.logger.warning(
f'Requester {exc.requester_name} not found, skipping provider {provider_entity.uuid}'
)
except Exception as exc:
self.ap.logger.error(f'Failed to load provider {provider_entity.uuid}: {exc}\n{traceback.format_exc()}')
await self._load_workspace_model_kind(
context,
persistence_model.LLMModel,
self.llm_model_dict,
self._build_llm_model,
)
await self._load_workspace_model_kind(
context,
persistence_model.EmbeddingModel,
self.embedding_model_dict,
self._build_embedding_model,
)
await self._load_workspace_model_kind(
context,
persistence_model.RerankModel,
self.rerank_model_dict,
self._build_rerank_model,
)
async def _load_workspace_model_kind(self, context, entity_type, cache: dict, builder) -> None:
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(entity_type).where(entity_type.workspace_uuid == context.workspace_uuid)
)
for model_entity in result.all():
try:
provider = self.provider_dict.get(self._cache_key(context, model_entity.provider_uuid))
if provider is None:
self.ap.logger.warning(
f'Provider {model_entity.provider_uuid} not found for model {model_entity.uuid}'
)
continue
runtime_model = builder(context, model_entity, provider)
cache[self._cache_key(context, model_entity.uuid)] = runtime_model
except Exception as exc:
self.ap.logger.error(f'Failed to load model {model_entity.uuid}: {exc}\n{traceback.format_exc()}')
async def sync_new_models_from_space(self, context: ExecutionContext) -> None:
"""Sync legacy Space models for the explicitly selected OSS Workspace."""
@@ -211,7 +211,7 @@ class LocalAgentRunner(runner.RequestRunner):
req_messages.append(
provider_message.Message(
role='system',
content=self.ap.box_service.get_system_guidance(query.query_id),
content=self.ap.box_service.get_system_guidance(query),
)
)
+47 -7
View File
@@ -26,6 +26,7 @@ from pydantic import AnyUrl
from .. import loader
from ....core import app
from ....core.task_boundary import create_detached_task, run_in_workspace_uow
from ....api.http.context import ExecutionContext
from ....api.http.service.tenant import TenantContext, require_workspace_uuid
from ....workspace.errors import WorkspaceError, WorkspaceInvariantError
@@ -1497,11 +1498,41 @@ class MCPLoader(loader.ToolLoader):
self.sessions = {}
result = await self.ap.persistence_mgr.execute_async(sqlalchemy.select(persistence_mcp.MCPServer))
servers = result.all()
server_configs: list[tuple[typing.Any, typing.Any, dict]] = []
list_bindings = getattr(self.ap.workspace_service, 'list_active_execution_bindings', None)
tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
cloud_runtime = getattr(getattr(self.ap.persistence_mgr, 'mode', None), 'value', None) == 'cloud_runtime'
if cloud_runtime:
if not callable(list_bindings) or not callable(tenant_uow):
raise RuntimeError('Cloud MCP loading requires explicit instance discovery and tenant UoWs')
for binding in await list_bindings():
async with tenant_uow(binding.workspace_uuid):
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(persistence_mcp.MCPServer)
.where(persistence_mcp.MCPServer.workspace_uuid == binding.workspace_uuid)
.order_by(persistence_mcp.MCPServer.uuid)
)
for server in result.all():
server_configs.append(
(
binding,
server,
self.ap.persistence_mgr.serialize_model(persistence_mcp.MCPServer, server),
)
)
else:
# Compatibility path for isolated loader tests and older embedders.
result = await self.ap.persistence_mgr.execute_async(sqlalchemy.select(persistence_mcp.MCPServer))
for server in result.all():
server_configs.append(
(
None,
server,
self.ap.persistence_mgr.serialize_model(persistence_mcp.MCPServer, server),
)
)
for server in servers:
config = self.ap.persistence_mgr.serialize_model(persistence_mcp.MCPServer, server)
for binding, server, config in server_configs:
if config.get('mode') == 'stdio' and not stdio_mcp_enabled(self.ap):
self.ap.logger.info(
f'Skipping disabled stdio MCP server {server.uuid}; '
@@ -1509,7 +1540,8 @@ class MCPLoader(loader.ToolLoader):
)
continue
try:
binding = await self.ap.workspace_service.get_execution_binding(server.workspace_uuid)
if binding is None:
binding = await self.ap.workspace_service.get_execution_binding(server.workspace_uuid)
execution_context = ExecutionContext(
instance_uuid=binding.instance_uuid,
workspace_uuid=binding.workspace_uuid,
@@ -1521,7 +1553,10 @@ class MCPLoader(loader.ToolLoader):
)
continue
task = asyncio.create_task(self.host_mcp_server(execution_context, config))
task = create_detached_task(
self.host_mcp_server(execution_context, config),
after_commit_manager=getattr(self.ap, 'persistence_mgr', None),
)
self._hosted_mcp_tasks.append(task)
@staticmethod
@@ -1542,7 +1577,12 @@ class MCPLoader(loader.ToolLoader):
return [session for key, session in self.sessions.items() if key[:3] == scope_key]
async def host_mcp_server(self, context: TenantContext, server_config: dict):
execution_context = await self._assert_execution_active(context)
requested_context = _execution_context_from_tenant(context)
execution_context = await run_in_workspace_uow(
self.ap,
requested_context.workspace_uuid,
lambda: self._assert_execution_active(requested_context),
)
configured_workspace = str(server_config.get('workspace_uuid') or '').strip()
if configured_workspace and configured_workspace != execution_context.workspace_uuid:
raise ValueError('MCP server configuration belongs to another Workspace')
+504 -192
View File
@@ -1,8 +1,15 @@
from __future__ import annotations
import base64
import contextlib
import errno
import json
import os
import posixpath
import stat
from collections.abc import Iterator
from dataclasses import dataclass
from pathlib import PurePosixPath
import langbot_plugin.api.entities.builtin.resource.tool as resource_tool
from langbot_plugin.api.entities.events import pipeline_query
@@ -34,6 +41,158 @@ _GREP_MAX_MATCHES = 200
_GREP_MAX_FILES = 5000
_GREP_MAX_LINE_CHARS = 500
_DIRECTORY_OPEN_FLAGS = (
os.O_RDONLY | getattr(os, 'O_DIRECTORY', 0) | getattr(os, 'O_NOFOLLOW', 0) | getattr(os, 'O_CLOEXEC', 0)
)
_FILE_OPEN_FLAGS = getattr(os, 'O_NOFOLLOW', 0) | getattr(os, 'O_CLOEXEC', 0) | getattr(os, 'O_NONBLOCK', 0)
_SECURE_HOST_FILE_OPS_AVAILABLE = bool(
getattr(os, 'O_NOFOLLOW', 0)
and os.open in os.supports_dir_fd
and os.stat in os.supports_dir_fd
and os.mkdir in os.supports_dir_fd
and os.listdir in os.supports_fd
and os.scandir in os.supports_fd
)
@dataclass(frozen=True)
class _HostLocation:
root: str
relative_parts: tuple[str, ...]
selected_skill: dict | None
workspace_anchor: str | None = None
def _unsafe_host_path(path: str, exc: BaseException | None = None) -> ValueError:
error = ValueError(f'Path escapes the workspace boundary or contains a symbolic link: {path}')
if exc is not None:
error.__cause__ = exc
return error
def _relative_workspace_parts(path: str) -> tuple[str, ...]:
normalized = posixpath.normpath(str(path or '/workspace').strip() or '/workspace')
if normalized == '/workspace':
return ()
if not normalized.startswith('/workspace/'):
raise ValueError('Path escapes the workspace boundary.')
parts = tuple(part for part in normalized.removeprefix('/workspace/').split('/') if part)
if any(part in {'.', '..'} or '\x00' in part for part in parts):
raise ValueError('Path escapes the workspace boundary.')
return parts
def _is_symlink_at(parent_fd: int, name: str) -> bool:
try:
return stat.S_ISLNK(os.stat(name, dir_fd=parent_fd, follow_symlinks=False).st_mode)
except FileNotFoundError:
return False
def _open_directory_at(parent_fd: int, name: str, *, create: bool) -> int:
if create:
try:
os.mkdir(name, mode=0o777, dir_fd=parent_fd)
except FileExistsError:
pass
try:
directory_fd = os.open(name, _DIRECTORY_OPEN_FLAGS, dir_fd=parent_fd)
except OSError as exc:
if exc.errno == errno.ELOOP or _is_symlink_at(parent_fd, name):
raise _unsafe_host_path(name, exc)
raise
if not stat.S_ISDIR(os.fstat(directory_fd).st_mode):
os.close(directory_fd)
raise NotADirectoryError(name)
return directory_fd
def _open_directory_parts(root_fd: int, parts: tuple[str, ...], *, create: bool) -> int:
current_fd = os.dup(root_fd)
try:
for part in parts:
next_fd = _open_directory_at(current_fd, part, create=create)
os.close(current_fd)
current_fd = next_fd
return current_fd
except BaseException:
os.close(current_fd)
raise
@contextlib.contextmanager
def _open_host_root(location: _HostLocation, *, create: bool) -> Iterator[int]:
"""Open and pin the tenant root before resolving tenant-controlled names."""
if location.workspace_anchor is None:
root_path = os.path.realpath(location.root)
try:
root_fd = os.open(root_path, _DIRECTORY_OPEN_FLAGS)
except OSError as exc:
if exc.errno == errno.ELOOP:
raise _unsafe_host_path(location.root, exc)
raise
else:
anchor_path = os.path.abspath(location.workspace_anchor)
root_path = os.path.abspath(location.root)
try:
if os.path.commonpath((anchor_path, root_path)) != anchor_path:
raise _unsafe_host_path(location.root)
except ValueError as exc:
raise _unsafe_host_path(location.root, exc)
anchor_real_path = os.path.realpath(anchor_path)
try:
anchor_fd = os.open(anchor_real_path, _DIRECTORY_OPEN_FLAGS)
except OSError as exc:
if exc.errno == errno.ELOOP:
raise _unsafe_host_path(location.workspace_anchor, exc)
raise
try:
root_relative = os.path.relpath(root_path, anchor_path)
root_parts = () if root_relative == '.' else tuple(root_relative.split(os.sep))
root_fd = _open_directory_parts(anchor_fd, root_parts, create=create)
finally:
os.close(anchor_fd)
try:
if not stat.S_ISDIR(os.fstat(root_fd).st_mode):
raise _unsafe_host_path(location.root)
yield root_fd
finally:
os.close(root_fd)
@contextlib.contextmanager
def _open_location_fd(
root_fd: int,
relative_parts: tuple[str, ...],
flags: int,
*,
create_parents: bool = False,
mode: int = 0o666,
) -> Iterator[int]:
if not relative_parts:
target_fd = os.dup(root_fd)
else:
parent_fd = _open_directory_parts(root_fd, relative_parts[:-1], create=create_parents)
try:
try:
target_fd = os.open(relative_parts[-1], flags | _FILE_OPEN_FLAGS, mode, dir_fd=parent_fd)
except OSError as exc:
if exc.errno == errno.ELOOP or _is_symlink_at(parent_fd, relative_parts[-1]):
raise _unsafe_host_path(relative_parts[-1], exc)
raise
finally:
os.close(parent_fd)
try:
yield target_fd
finally:
os.close(target_fd)
class NativeToolLoader(loader.ToolLoader):
def __init__(self, ap):
@@ -59,6 +218,9 @@ class NativeToolLoader(loader.ToolLoader):
@staticmethod
def _execution_context(query: pipeline_query.Query) -> ExecutionContext:
attached_context = getattr(query, '_execution_context', None)
if isinstance(attached_context, ExecutionContext):
return attached_context
return ExecutionContext(
instance_uuid=str(getattr(query, 'instance_uuid', '') or ''),
workspace_uuid=str(getattr(query, 'workspace_uuid', '') or ''),
@@ -66,6 +228,7 @@ class NativeToolLoader(loader.ToolLoader):
bot_uuid=getattr(query, 'bot_uuid', None),
pipeline_uuid=getattr(query, 'pipeline_uuid', None),
query_uuid=getattr(query, 'query_uuid', None),
entitlement_revision=getattr(query, 'entitlement_revision', 0),
)
async def get_tools(self, bound_plugins: list[str] | None = None) -> list[resource_tool.LLMTool]:
@@ -86,6 +249,13 @@ class NativeToolLoader(loader.ToolLoader):
return name in _ALL_TOOL_NAMES and await self._is_sandbox_available()
async def invoke_tool(self, name: str, parameters: dict, query: pipeline_query.Query):
require_sandbox = getattr(
getattr(self.ap, 'box_service', None),
'require_workspace_sandbox',
None,
)
if callable(require_sandbox):
await require_sandbox(self._execution_context(query))
if name == EXEC_TOOL_NAME:
self.ap.logger.info(
'exec tool invoked: '
@@ -111,6 +281,7 @@ class NativeToolLoader(loader.ToolLoader):
async def _invoke_exec(self, parameters: dict, query: pipeline_query.Query) -> dict:
command = str(parameters['command'])
workdir = str(parameters.get('workdir', '/workspace') or '/workspace')
selected_skill_name: str | None = None
# Validate that skill references target activated skills.
selected_skill, _ = skill_loader.resolve_virtual_skill_path(
@@ -140,31 +311,51 @@ class NativeToolLoader(loader.ToolLoader):
if not package_root:
raise ValueError(f'Activated skill "{selected_skill_name}" has no package_root.')
# Pass only the logical name across the authenticated Core→Runtime
# boundary. In Cloud mode the shared Box Runtime resolves the
# Workspace-scoped package root and constructs the read-only mount;
# Core host paths are never accepted as mount authority.
# Wrap command with Python venv bootstrap if the skill has a Python project.
# The venv is created inside the skill's mount path.
skill_mount = f'/workspace/.skills/{selected_skill_name}'
if skill_loader.should_prepare_skill_python_env(package_root):
python_project = selected_skill.get('python_project') is True
if 'python_project' not in selected_skill and bool(
getattr(self.ap.box_service, 'shares_filesystem_with_box', False)
):
# Backward compatibility for a same-process OSS Runtime that
# predates trusted Box metadata. Never probe a path reported by
# an external Runtime from the Core filesystem.
python_project = skill_loader.should_prepare_skill_python_env(package_root)
if python_project:
parameters = dict(parameters)
parameters['command'] = skill_loader.wrap_skill_command_with_python_env(command, mount_path=skill_mount)
parameters['command'] = skill_loader.wrap_skill_command_with_python_env(
command,
mount_path=skill_mount,
state_path=f'/workspace/.skill-envs/{selected_skill_name}',
)
# All exec calls (with or without skills) go through the same container
# via execute_tool. Skills are mounted at /workspace/.skills/{name}/
# via extra_mounts built by BoxService.
result = await self.ap.box_service.execute_tool(parameters, query)
result = await self.ap.box_service.execute_tool(
parameters,
query,
skill_name=selected_skill_name,
)
result = self._normalize_exec_result(result)
if selected_skill is not None:
self._refresh_skill_from_disk(query, selected_skill)
return result
def _resolve_host_path(
def _resolve_host_location(
self,
query: pipeline_query.Query,
sandbox_path: str,
*,
include_visible: bool,
include_activated: bool,
) -> tuple[str, dict | None]:
) -> _HostLocation:
selected_skill, rewritten_path = skill_loader.resolve_virtual_skill_path(
self.ap,
query,
@@ -174,26 +365,26 @@ class NativeToolLoader(loader.ToolLoader):
)
box_service = self.ap.box_service
host_root = (
selected_skill.get('package_root')
if selected_skill is not None
else box_service._tenant_workspace(self._execution_context(query))
)
if selected_skill is not None:
if not self._can_interpret_skill_host_paths():
raise ValueError(
'Skill package paths are owned by the Box Runtime; '
'this operation requires a Runtime skill-file API.'
)
host_root = selected_skill.get('package_root')
workspace_anchor = None
else:
host_root = box_service._tenant_workspace(self._execution_context(query))
workspace_anchor = getattr(box_service, 'default_workspace', None)
if not host_root:
raise ValueError('No host workspace configured for file operations.')
mount_path = '/workspace'
if not rewritten_path.startswith(mount_path):
raise ValueError(f'Path must be under {mount_path}.')
relative = rewritten_path[len(mount_path) :].lstrip('/')
host_path = os.path.realpath(os.path.join(host_root, relative))
host_root = os.path.realpath(host_root)
if not (host_path == host_root or host_path.startswith(host_root + os.sep)):
raise ValueError('Path escapes the workspace boundary.')
return host_path, selected_skill
return _HostLocation(
root=str(host_root),
relative_parts=_relative_workspace_parts(rewritten_path),
selected_skill=selected_skill,
workspace_anchor=str(workspace_anchor) if workspace_anchor else None,
)
def _resolve_skill_relative_path(
self,
@@ -213,21 +404,259 @@ class NativeToolLoader(loader.ToolLoader):
if selected_skill is None:
return None
mount_path = '/workspace'
if not rewritten_path.startswith(mount_path):
raise ValueError(f'Path must be under {mount_path}.')
relative = rewritten_path[len(mount_path) :].lstrip('/') or '.'
relative = '/'.join(_relative_workspace_parts(rewritten_path)) or '.'
return selected_skill, relative
def _can_interpret_skill_host_paths(self) -> bool:
"""Require an explicitly proven shared Core/Runtime filesystem view."""
return _SECURE_HOST_FILE_OPS_AVAILABLE and bool(
getattr(self.ap.box_service, 'shares_filesystem_with_box', False)
)
def _should_use_box_workspace_files(self, selected_skill: dict | None) -> bool:
if selected_skill is not None:
return False
box_service = getattr(self.ap, 'box_service', None)
if box_service is None or not hasattr(box_service, 'execute_tool'):
return False
if not _SECURE_HOST_FILE_OPS_AVAILABLE:
# Preserve the OSS API on platforms without openat/O_NOFOLLOW by
# running inside the tenant-scoped Box mount, never via a racy
# host-path fallback.
return True
default_workspace = getattr(box_service, 'default_workspace', None)
return bool(default_workspace and not os.path.isdir(os.path.realpath(default_workspace)))
def _read_host_location(self, location: _HostLocation, parameters: dict) -> dict:
with _open_host_root(location, create=False) as root_fd:
with _open_location_fd(root_fd, location.relative_parts, os.O_RDONLY) as target_fd:
metadata = os.fstat(target_fd)
if stat.S_ISDIR(metadata.st_mode):
return self._build_directory_result(os.listdir(target_fd))
if not stat.S_ISREG(metadata.st_mode):
raise ValueError('Path must reference a regular file or directory.')
return self._read_text_file_preview(target_fd, parameters, metadata=metadata)
def _write_host_location(self, location: _HostLocation, content: str, parameters: dict) -> None:
if not location.relative_parts:
raise ValueError('Path must reference a file under /workspace.')
encoding, mode = self._write_options(parameters)
if encoding == 'base64':
try:
payload = base64.b64decode(content, validate=True)
except Exception as exc:
raise ValueError(f'invalid base64 content: {exc}') from exc
else:
payload = content.encode('utf-8')
flags = os.O_WRONLY | os.O_CREAT
if mode == 'append':
flags |= os.O_APPEND
with _open_host_root(location, create=True) as root_fd:
with _open_location_fd(
root_fd,
location.relative_parts,
flags,
create_parents=True,
) as target_fd:
if not stat.S_ISREG(os.fstat(target_fd).st_mode):
raise ValueError('Path must reference a regular file.')
if mode != 'append':
os.ftruncate(target_fd, 0)
os.lseek(target_fd, 0, os.SEEK_SET)
self._write_all(target_fd, payload)
def _edit_host_location(
self,
location: _HostLocation,
old_string: str,
new_string: str,
) -> tuple[bool, str | None]:
if not location.relative_parts:
raise ValueError('Path must reference a file under /workspace.')
with _open_host_root(location, create=False) as root_fd:
with _open_location_fd(root_fd, location.relative_parts, os.O_RDWR) as target_fd:
if not stat.S_ISREG(os.fstat(target_fd).st_mode):
return False, 'File not found.'
with os.fdopen(os.dup(target_fd), 'r', encoding='utf-8', errors='replace') as file_obj:
content = file_obj.read()
count = content.count(old_string)
if count == 0:
return False, 'old_string not found in file.'
if count > 1:
return False, f'old_string matches {count} locations; provide a more unique string.'
payload = content.replace(old_string, new_string, 1).encode('utf-8')
os.ftruncate(target_fd, 0)
os.lseek(target_fd, 0, os.SEEK_SET)
self._write_all(target_fd, payload)
return True, None
@staticmethod
def _write_all(file_fd: int, payload: bytes) -> None:
view = memoryview(payload)
while view:
written = os.write(file_fd, view)
if written <= 0:
raise OSError('Could not write the complete workspace file.')
view = view[written:]
@staticmethod
def _rglob_matches(relative_path: str, pattern: str) -> bool:
candidates = {pattern}
pending = [pattern]
while pending:
candidate = pending.pop()
marker = candidate.find('**/')
while marker >= 0:
without_recursive_segment = candidate[:marker] + candidate[marker + 3 :]
if without_recursive_segment not in candidates:
candidates.add(without_recursive_segment)
pending.append(without_recursive_segment)
marker = candidate.find('**/', marker + 3)
return any(candidate and PurePosixPath(relative_path).match(candidate) for candidate in candidates)
def _glob_host_location(self, location: _HostLocation, pattern: str, sandbox_base: str) -> dict:
hits: list[tuple[str, float]] = []
def walk(directory_fd: int, prefix: str) -> None:
with os.scandir(directory_fd) as entries:
for entry in entries:
name = entry.name
if name in _SKIP_DIRS:
continue
try:
child_fd = os.open(name, os.O_RDONLY | _FILE_OPEN_FLAGS, dir_fd=directory_fd)
except OSError:
continue
try:
metadata = os.fstat(child_fd)
relative = f'{prefix}/{name}' if prefix else name
if self._rglob_matches(relative, pattern):
hits.append((relative, metadata.st_mtime))
if stat.S_ISDIR(metadata.st_mode):
walk(child_fd, relative)
finally:
os.close(child_fd)
with _open_host_root(location, create=False) as root_fd:
with _open_location_fd(root_fd, location.relative_parts, os.O_RDONLY) as target_fd:
if not stat.S_ISDIR(os.fstat(target_fd).st_mode):
return {'ok': False, 'error': f'Path is not a directory: {sandbox_base}'}
walk(target_fd, '')
hits.sort(key=lambda item: item[1], reverse=True)
total = len(hits)
sandbox_paths: list[str] = []
output_bytes = 0
truncated_by_bytes = False
for relative, _mtime in hits[:_GLOB_MAX_MATCHES]:
sandbox_path = self._sandbox_child_path(sandbox_base, relative)
entry_bytes = len(sandbox_path.encode('utf-8')) + (1 if sandbox_paths else 0)
if output_bytes + entry_bytes > _DEFAULT_TOOL_RESULT_MAX_BYTES:
truncated_by_bytes = True
break
sandbox_paths.append(sandbox_path)
output_bytes += entry_bytes
return {
'ok': True,
'matches': sandbox_paths,
'preview': '\n'.join(sandbox_paths),
'total': total,
'truncated': total > len(sandbox_paths) or truncated_by_bytes,
'truncated_by': 'bytes' if truncated_by_bytes else ('matches' if total > len(sandbox_paths) else None),
}
def _grep_host_location(
self,
location: _HostLocation,
regex,
include: str | None,
sandbox_base: str,
) -> dict:
matches: list[dict] = []
output_bytes = 0
truncated_by: str | None = None
files_seen = 0
def grep_file(file_fd: int, sandbox_path: str) -> bool:
nonlocal output_bytes, truncated_by
with os.fdopen(os.dup(file_fd), 'r', encoding='utf-8', errors='ignore') as handle:
for lineno, line in enumerate(handle, 1):
if not regex.search(line):
continue
content, line_truncated = self._truncate_grep_line(line.rstrip())
entry = {'file': sandbox_path, 'line': lineno, 'content': content}
entry_bytes = len(json.dumps(entry, ensure_ascii=False).encode('utf-8')) + 1
if output_bytes + entry_bytes > _DEFAULT_TOOL_RESULT_MAX_BYTES:
truncated_by = 'bytes'
return True
if line_truncated and truncated_by is None:
truncated_by = 'line'
matches.append(entry)
output_bytes += entry_bytes
if len(matches) >= _GREP_MAX_MATCHES:
truncated_by = truncated_by or 'matches'
return True
return False
def walk(directory_fd: int, prefix: str) -> bool:
nonlocal files_seen
with os.scandir(directory_fd) as entries:
for entry in entries:
name = entry.name
if name in _SKIP_DIRS:
continue
try:
child_fd = os.open(name, os.O_RDONLY | _FILE_OPEN_FLAGS, dir_fd=directory_fd)
except OSError:
continue
try:
metadata = os.fstat(child_fd)
relative = f'{prefix}/{name}' if prefix else name
if stat.S_ISDIR(metadata.st_mode):
if walk(child_fd, relative):
return True
continue
if not stat.S_ISREG(metadata.st_mode):
continue
if include and not self._rglob_matches(relative, include):
continue
files_seen += 1
if grep_file(child_fd, self._sandbox_child_path(sandbox_base, relative)):
return True
if files_seen >= _GREP_MAX_FILES:
return True
finally:
os.close(child_fd)
return False
with _open_host_root(location, create=False) as root_fd:
with _open_location_fd(root_fd, location.relative_parts, os.O_RDONLY) as target_fd:
metadata = os.fstat(target_fd)
if stat.S_ISREG(metadata.st_mode):
grep_file(target_fd, sandbox_base)
elif stat.S_ISDIR(metadata.st_mode):
walk(target_fd, '')
else:
return {'ok': False, 'error': f'Path not found: {sandbox_base}'}
return {
'ok': True,
'matches': matches,
'total': len(matches),
'truncated': truncated_by is not None,
'truncated_by': truncated_by,
}
@staticmethod
def _sandbox_child_path(base: str, relative: str) -> str:
return f'{str(base).rstrip("/")}/{relative}'
async def _run_workspace_file_script(self, script: str, query: pipeline_query.Query) -> dict:
result = await self.ap.box_service.execute_tool(
{
@@ -531,11 +960,15 @@ else:
)
if skill_request is not None and hasattr(self.ap.box_service, 'read_skill_file'):
selected_skill, relative = skill_request
host_path = self._resolve_skill_host_path(selected_skill, relative)
if host_path and os.path.exists(host_path):
if os.path.isdir(host_path):
return self._build_directory_result(os.listdir(host_path))
return self._read_text_file_preview(host_path, parameters)
if self._can_interpret_skill_host_paths():
host_location = self._resolve_skill_host_location(selected_skill, relative)
else:
host_location = None
if host_location is not None:
try:
return self._read_host_location(host_location, parameters)
except FileNotFoundError:
pass
try:
result = await self.ap.box_service.read_skill_file(
@@ -556,20 +989,18 @@ else:
except Exception as exc:
return {'ok': False, 'error': str(exc)}
host_path, selected_skill = self._resolve_host_path(
host_location = self._resolve_host_location(
query,
path,
include_visible=True,
include_activated=True,
)
if self._should_use_box_workspace_files(selected_skill):
if self._should_use_box_workspace_files(host_location.selected_skill):
return await self._read_workspace_via_box(path, parameters, query)
if not os.path.exists(host_path):
try:
return self._read_host_location(host_location, parameters)
except (FileNotFoundError, NotADirectoryError):
return {'ok': False, 'error': f'File not found: {path}'}
if os.path.isdir(host_path):
entries = os.listdir(host_path)
return self._build_directory_result(entries)
return self._read_text_file_preview(host_path, parameters)
async def _invoke_write(self, parameters: dict, query: pipeline_query.Query) -> dict:
path = parameters['path']
@@ -591,20 +1022,19 @@ else:
await self.ap.skill_mgr.reload_skills(execution_context)
return {'ok': True, 'path': path}
host_path, selected_skill = self._resolve_host_path(
host_location = self._resolve_host_location(
query,
path,
include_visible=False,
include_activated=True,
)
if self._should_use_box_workspace_files(selected_skill):
if self._should_use_box_workspace_files(host_location.selected_skill):
return await self._write_workspace_via_box(path, content, parameters, query)
os.makedirs(os.path.dirname(host_path), exist_ok=True)
try:
self._write_host_file(host_path, content, parameters)
self._write_host_location(host_location, content, parameters)
except ValueError as exc:
return {'ok': False, 'error': str(exc)}
self._refresh_skill_from_disk(query, selected_skill)
self._refresh_skill_from_disk(query, host_location.selected_skill)
return {'ok': True, 'path': path}
async def _invoke_edit(self, parameters: dict, query: pipeline_query.Query) -> dict:
@@ -652,27 +1082,21 @@ else:
await self.ap.skill_mgr.reload_skills(execution_context)
return {'ok': True, 'path': path}
host_path, selected_skill = self._resolve_host_path(
host_location = self._resolve_host_location(
query,
path,
include_visible=False,
include_activated=True,
)
if self._should_use_box_workspace_files(selected_skill):
if self._should_use_box_workspace_files(host_location.selected_skill):
return await self._edit_workspace_via_box(path, old_string, new_string, query)
if not os.path.isfile(host_path):
try:
changed, error = self._edit_host_location(host_location, old_string, new_string)
except (FileNotFoundError, NotADirectoryError):
return {'ok': False, 'error': f'File not found: {path}'}
with open(host_path, 'r', encoding='utf-8', errors='replace') as f:
content = f.read()
count = content.count(old_string)
if count == 0:
return {'ok': False, 'error': 'old_string not found in file.'}
if count > 1:
return {'ok': False, 'error': f'old_string matches {count} locations; provide a more unique string.'}
new_content = content.replace(old_string, new_string, 1)
with open(host_path, 'w', encoding='utf-8') as f:
f.write(new_content)
self._refresh_skill_from_disk(query, selected_skill)
if not changed:
return {'ok': False, 'error': error or f'File not found: {path}'}
self._refresh_skill_from_disk(query, host_location.selected_skill)
return {'ok': True, 'path': path}
def _refresh_skill_from_disk(self, query: pipeline_query.Query, selected_skill: dict | None) -> None:
@@ -931,55 +1355,19 @@ else:
path = str(parameters.get('path', '/workspace') or '/workspace')
self.ap.logger.info(f'glob tool invoked: query_id={query.query_id} pattern={pattern} path={path}')
host_path, selected_skill = self._resolve_host_path(
host_location = self._resolve_host_location(
query,
path,
include_visible=True,
include_activated=True,
)
if self._should_use_box_workspace_files(selected_skill):
if self._should_use_box_workspace_files(host_location.selected_skill):
return await self._glob_workspace_via_box(path, pattern, query)
if not os.path.isdir(host_path):
try:
return self._glob_host_location(host_location, pattern, path)
except (FileNotFoundError, NotADirectoryError):
return {'ok': False, 'error': f'Path is not a directory: {path}'}
from pathlib import Path
base = Path(host_path)
hits = list(base.rglob(pattern))
# Filter out skipped directories
hits = [h for h in hits if not any(skip in h.parts for skip in _SKIP_DIRS)]
# Sort by mtime, newest first
hits.sort(key=lambda p: p.stat().st_mtime if p.exists() else 0, reverse=True)
total = len(hits)
shown = hits[:_GLOB_MAX_MATCHES]
# Convert back to sandbox paths
sandbox_paths = []
output_bytes = 0
truncated_by_bytes = False
for h in shown:
rel = os.path.relpath(str(h), host_path)
sandbox_path = os.path.join(path, rel)
entry_bytes = len(sandbox_path.encode('utf-8')) + (1 if sandbox_paths else 0)
if output_bytes + entry_bytes > _DEFAULT_TOOL_RESULT_MAX_BYTES:
truncated_by_bytes = True
break
sandbox_paths.append(sandbox_path)
output_bytes += entry_bytes
return {
'ok': True,
'matches': sandbox_paths,
'preview': '\n'.join(sandbox_paths),
'total': total,
'truncated': total > len(sandbox_paths) or truncated_by_bytes,
'truncated_by': 'bytes' if truncated_by_bytes else ('matches' if total > len(sandbox_paths) else None),
}
async def _invoke_grep(self, parameters: dict, query: pipeline_query.Query) -> dict:
pattern = parameters['pattern']
path = str(parameters.get('path', '/workspace') or '/workspace')
@@ -987,99 +1375,36 @@ else:
self.ap.logger.info(f'grep tool invoked: query_id={query.query_id} pattern={pattern} path={path}')
import re
from pathlib import Path
try:
regex = re.compile(pattern)
except re.error as e:
return {'ok': False, 'error': f'Invalid regex: {e}'}
host_path, selected_skill = self._resolve_host_path(
host_location = self._resolve_host_location(
query,
path,
include_visible=True,
include_activated=True,
)
if self._should_use_box_workspace_files(selected_skill):
if self._should_use_box_workspace_files(host_location.selected_skill):
return await self._grep_workspace_via_box(path, pattern, include, query)
if not os.path.exists(host_path):
try:
return self._grep_host_location(host_location, regex, include, path)
except (FileNotFoundError, NotADirectoryError):
return {'ok': False, 'error': f'Path not found: {path}'}
base = Path(host_path)
if base.is_file():
files = [base]
else:
files = self._grep_walk(base, include)
matches = []
output_bytes = 0
truncated_by = None
for fp in files:
try:
handle = fp.open('r', encoding='utf-8', errors='ignore')
except OSError:
continue
with handle:
for lineno, line in enumerate(handle, 1):
if regex.search(line):
rel = os.path.relpath(str(fp), host_path)
sandbox_path = os.path.join(path, rel)
content, line_truncated = self._truncate_grep_line(line.rstrip())
entry = {
'file': sandbox_path,
'line': lineno,
'content': content,
}
entry_bytes = len(json.dumps(entry, ensure_ascii=False).encode('utf-8')) + 1
if output_bytes + entry_bytes > _DEFAULT_TOOL_RESULT_MAX_BYTES:
truncated_by = 'bytes'
break
if line_truncated and truncated_by is None:
truncated_by = 'line'
matches.append(entry)
output_bytes += entry_bytes
if len(matches) >= _GREP_MAX_MATCHES:
truncated_by = truncated_by or 'matches'
break
if truncated_by == 'bytes' or len(matches) >= _GREP_MAX_MATCHES:
break
if truncated_by == 'bytes' or len(matches) >= _GREP_MAX_MATCHES:
break
return {
'ok': True,
'matches': matches,
'total': len(matches),
'truncated': truncated_by is not None,
'truncated_by': truncated_by,
}
@staticmethod
def _grep_walk(root, include: str | None) -> list:
"""Walk dir tree for grep, skipping junk dirs."""
results = []
for item in root.rglob(include or '*'):
if any(skip in item.parts for skip in _SKIP_DIRS):
continue
if item.is_file():
results.append(item)
if len(results) >= _GREP_MAX_FILES:
break
return results
@staticmethod
def _resolve_skill_host_path(selected_skill: dict, relative: str) -> str | None:
def _resolve_skill_host_location(selected_skill: dict, relative: str) -> _HostLocation | None:
package_root = str(selected_skill.get('package_root', '') or '').strip()
if not package_root:
return None
host_root = os.path.realpath(package_root)
host_path = os.path.realpath(os.path.join(host_root, relative))
if not (host_path == host_root or host_path.startswith(host_root + os.sep)):
raise ValueError('Path escapes the skill package boundary.')
return host_path
relative_path = '/workspace' if relative in {'', '.'} else f'/workspace/{relative}'
return _HostLocation(
root=package_root,
relative_parts=_relative_workspace_parts(relative_path),
selected_skill=selected_skill,
)
def _normalize_exec_result(self, result: dict) -> dict:
normalized = dict(result)
@@ -1119,9 +1444,9 @@ else:
'truncated_by': 'bytes' if truncated else None,
}
def _read_text_file_preview(self, host_path: str, parameters: dict) -> dict:
def _read_text_file_preview(self, file_fd: int, parameters: dict, *, metadata: os.stat_result) -> dict:
if self._read_encoding(parameters) == 'base64':
return self._read_binary_file_chunk(host_path, parameters)
return self._read_binary_file_chunk(file_fd, parameters, metadata=metadata)
offset = self._positive_int(parameters.get('offset'), default=1)
max_lines = self._positive_int(
@@ -1141,7 +1466,7 @@ else:
truncated_by: str | None = None
next_offset: int | None = None
with open(host_path, 'r', encoding='utf-8', errors='replace') as f:
with os.fdopen(os.dup(file_fd), 'r', encoding='utf-8', errors='replace') as f:
for line_number, line in enumerate(f, 1):
if line_number < offset:
continue
@@ -1182,15 +1507,15 @@ else:
'max_bytes': max_bytes,
}
def _read_binary_file_chunk(self, host_path: str, parameters: dict) -> dict:
def _read_binary_file_chunk(self, file_fd: int, parameters: dict, *, metadata: os.stat_result) -> dict:
byte_offset = self._non_negative_int(parameters.get('byte_offset'), default=0)
max_bytes = self._positive_int(
parameters.get('max_bytes'),
default=_DEFAULT_TOOL_RESULT_MAX_BYTES,
max_value=_DEFAULT_TOOL_RESULT_MAX_BYTES,
)
size_bytes = os.path.getsize(host_path)
with open(host_path, 'rb') as f:
size_bytes = metadata.st_size
with os.fdopen(os.dup(file_fd), 'rb') as f:
f.seek(byte_offset)
data = f.read(max_bytes + 1)
chunk = data[:max_bytes]
@@ -1207,19 +1532,6 @@ else:
'max_bytes': max_bytes,
}
def _write_host_file(self, host_path: str, content: str, parameters: dict) -> None:
encoding, mode = self._write_options(parameters)
if encoding == 'base64':
try:
data = base64.b64decode(content, validate=True)
except Exception as exc:
raise ValueError(f'invalid base64 content: {exc}') from exc
with open(host_path, 'ab' if mode == 'append' else 'wb') as f:
f.write(data)
return
with open(host_path, 'a' if mode == 'append' else 'w', encoding='utf-8') as f:
f.write(content)
@staticmethod
def _read_encoding(parameters: dict) -> str:
return 'base64' if parameters.get('encoding') == 'base64' else 'text'
@@ -201,5 +201,14 @@ def should_prepare_skill_python_env(package_root: str | None) -> bool:
return box_workspace.should_prepare_python_env(package_root)
def wrap_skill_command_with_python_env(command: str, *, mount_path: str = '/workspace') -> str:
return box_workspace.wrap_python_command_with_env(command, mount_path=mount_path).rstrip()
def wrap_skill_command_with_python_env(
command: str,
*,
mount_path: str = '/workspace',
state_path: str | None = None,
) -> str:
return box_workspace.wrap_python_command_with_env(
command,
mount_path=mount_path,
state_path=state_path,
).rstrip()
@@ -73,12 +73,34 @@ class SkillToolLoader(loader.ToolLoader):
return self._sandbox_available
async def invoke_tool(self, name: str, parameters: dict, query) -> typing.Any:
require_sandbox = getattr(
getattr(self.ap, 'box_service', None),
'require_workspace_sandbox',
None,
)
if callable(require_sandbox):
await require_sandbox(self._execution_context(query))
if name == ACTIVATE_SKILL_TOOL_NAME:
return await self._invoke_activate_skill(parameters, query)
if name == REGISTER_SKILL_TOOL_NAME:
return await self._invoke_register_skill(parameters, query)
raise ValueError(f'Unknown skill tool: {name}')
@staticmethod
def _execution_context(query) -> ExecutionContext:
attached_context = getattr(query, '_execution_context', None)
if isinstance(attached_context, ExecutionContext):
return attached_context
return ExecutionContext(
instance_uuid=str(getattr(query, 'instance_uuid', '') or ''),
workspace_uuid=str(getattr(query, 'workspace_uuid', '') or ''),
placement_generation=getattr(query, 'placement_generation', 0) or 0,
bot_uuid=getattr(query, 'bot_uuid', None),
pipeline_uuid=getattr(query, 'pipeline_uuid', None),
query_uuid=getattr(query, 'query_uuid', None),
entitlement_revision=getattr(query, 'entitlement_revision', 0),
)
async def shutdown(self):
pass
@@ -136,7 +158,8 @@ class SkillToolLoader(loader.ToolLoader):
raise ValueError('path is required')
# Resolve sandbox path to host path
host_path = self._resolve_workspace_directory(sandbox_path)
execution_context = self._execution_context(query)
host_path = self._resolve_workspace_directory(sandbox_path, execution_context)
# Get or create skill service
skill_service = getattr(self.ap, 'skill_service', None)
@@ -144,14 +167,6 @@ class SkillToolLoader(loader.ToolLoader):
raise ValueError('Skill service not available')
# Scan and register the skill
execution_context = ExecutionContext(
instance_uuid=str(getattr(query, 'instance_uuid', '') or ''),
workspace_uuid=str(getattr(query, 'workspace_uuid', '') or ''),
placement_generation=getattr(query, 'placement_generation', 0) or 0,
bot_uuid=getattr(query, 'bot_uuid', None),
pipeline_uuid=getattr(query, 'pipeline_uuid', None),
query_uuid=getattr(query, 'query_uuid', None),
)
scanned = await skill_service.scan_directory_async(execution_context, host_path)
# Override name if provided
@@ -178,10 +193,19 @@ class SkillToolLoader(loader.ToolLoader):
'skill': created,
}
def _resolve_workspace_directory(self, sandbox_path: str) -> str:
def _resolve_workspace_directory(
self,
sandbox_path: str,
execution_context: ExecutionContext,
) -> str:
"""Resolve sandbox path to host filesystem path."""
box_service = getattr(self.ap, 'box_service', None)
workspace_root = getattr(box_service, 'default_workspace', None)
tenant_workspace = getattr(box_service, '_tenant_workspace', None)
workspace_root = (
tenant_workspace(execution_context)
if callable(tenant_workspace)
else getattr(box_service, 'default_workspace', None)
)
if not workspace_root:
raise ValueError('No default workspace configured')
+57 -12
View File
@@ -2,6 +2,7 @@ from __future__ import annotations
import typing
import time
import inspect
from typing import TYPE_CHECKING
import langbot_plugin.api.entities.builtin.resource.tool as resource_tool
@@ -35,6 +36,36 @@ class ToolManager:
def __init__(self, ap: app.Application):
self.ap = ap
async def _bind_plugin_workspace(self, context: TenantContext) -> None:
"""Select the tenant before any plugin catalog lookup.
Tool discovery happens before invocation, so relying on ``call_tool``
to bind the Workspace is too late and can expose another task's
catalog in a shared Runtime.
"""
connector = getattr(self.ap, 'plugin_connector', None)
require_context = getattr(connector, 'require_workspace_context', None)
if require_context is None:
return
result = require_context(context)
if inspect.isawaitable(result):
await result
async def _workspace_sandbox_available(self, context: TenantContext) -> bool:
"""Resolve the Workspace capability before exposing sandbox tools."""
box_service = getattr(self.ap, 'box_service', None)
checker = getattr(box_service, 'is_workspace_sandbox_available', None)
if not callable(checker):
# Compatibility for OSS embedders and isolated manager tests. The
# BoxService execution path remains the final authority.
return True
try:
return bool(await checker(context))
except Exception:
return False
async def initialize(self):
from langbot.pkg.utils import importutil
from langbot.pkg.provider.tools import loaders
@@ -65,10 +96,13 @@ class ToolManager:
include_skill_authoring: bool = False,
include_mcp_resource_tools: bool = True,
) -> list[resource_tool.LLMTool]:
await self._bind_plugin_workspace(context)
all_functions: list[resource_tool.LLMTool] = []
all_functions.extend(await self.native_tool_loader.get_tools())
if include_skill_authoring:
sandbox_available = await self._workspace_sandbox_available(context)
if sandbox_available:
all_functions.extend(await self.native_tool_loader.get_tools())
if include_skill_authoring and sandbox_available:
all_functions.extend(await self.skill_tool_loader.get_tools())
all_functions.extend(await self.plugin_tool_loader.get_tools(bound_plugins))
all_functions.extend(
@@ -89,6 +123,7 @@ class ToolManager:
include_skill_authoring: bool = False,
include_mcp_resource_tools: bool = False,
) -> list[dict[str, typing.Any]]:
await self._bind_plugin_workspace(context)
catalog: list[dict[str, typing.Any]] = []
def append_tools(source: str, source_name: str, tools: list[resource_tool.LLMTool]) -> None:
@@ -104,8 +139,10 @@ class ToolManager:
}
)
append_tools('builtin', 'LangBot', await self.native_tool_loader.get_tools())
if include_skill_authoring:
sandbox_available = await self._workspace_sandbox_available(context)
if sandbox_available:
append_tools('builtin', 'LangBot', await self.native_tool_loader.get_tools())
if include_skill_authoring and sandbox_available:
append_tools('skill', 'LangBot', await self.skill_tool_loader.get_tools())
catalog.extend(await self.plugin_tool_loader.get_tool_catalog(bound_plugins))
@@ -121,14 +158,20 @@ class ToolManager:
async def get_tool_by_name(self, context: TenantContext, name: str) -> tool_loader.ToolLookupResult | None:
"""Get tool by name from any active loader."""
for active_loader in (
self.native_tool_loader,
self.plugin_tool_loader,
self.skill_tool_loader,
):
await self._bind_plugin_workspace(context)
sandbox_available = await self._workspace_sandbox_available(context)
if sandbox_available:
tool = await self.native_tool_loader.get_tool(name)
if tool:
return tool
for active_loader in (self.plugin_tool_loader,):
tool = await active_loader.get_tool(name)
if tool:
return tool
if sandbox_available:
tool = await self.skill_tool_loader.get_tool(name)
if tool:
return tool
return await self.mcp_tool_loader.get_tool(context, name)
@@ -237,7 +280,10 @@ class ToolManager:
async def execute_func_call(self, name: str, parameters: dict, query: pipeline_query.Query) -> typing.Any:
from langbot.pkg.telemetry import features as telemetry_features
if await self.native_tool_loader.has_tool(name):
execution_context = get_query_execution_context(query)
await self._bind_plugin_workspace(execution_context)
sandbox_available = await self._workspace_sandbox_available(execution_context)
if sandbox_available and await self.native_tool_loader.has_tool(name):
telemetry_features.increment(query, 'tool_calls', 'native')
return await self._invoke_tool_with_monitoring(
source='native',
@@ -255,7 +301,6 @@ class ToolManager:
query=query,
invoke=lambda: self.plugin_tool_loader.invoke_tool(name, parameters, query),
)
execution_context = get_query_execution_context(query)
if await self.mcp_tool_loader.has_tool(execution_context, name):
telemetry_features.increment(query, 'tool_calls', 'mcp')
return await self._invoke_tool_with_monitoring(
@@ -265,7 +310,7 @@ class ToolManager:
query=query,
invoke=lambda: self.mcp_tool_loader.invoke_tool(name, parameters, query),
)
if await self.skill_tool_loader.has_tool(name):
if sandbox_available and await self.skill_tool_loader.has_tool(name):
telemetry_features.increment(query, 'tool_calls', 'skill')
return await self._invoke_tool_with_monitoring(
source='skill',
+81 -23
View File
@@ -1,4 +1,5 @@
from __future__ import annotations
import asyncio
import io
import mimetypes
import os.path
@@ -14,6 +15,7 @@ from langbot.pkg.api.http.authz import WorkspaceRequiredError
from langbot.pkg.api.http.context import ExecutionContext, RequestContext
from langbot.pkg.api.http.service.tenant import TenantContext, require_workspace_uuid
from langbot.pkg.core import app, taskmgr
from langbot.pkg.core.task_boundary import run_in_workspace_uow
from langbot.pkg.entity.persistence import rag as persistence_rag
from langbot.pkg.workspace.errors import WorkspaceNotFoundError
@@ -90,16 +92,23 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
task_context: taskmgr.TaskContext,
parser_plugin_id: str | None = None,
):
await self._assert_execution_context(execution_context)
await run_in_workspace_uow(
self.ap,
execution_context.workspace_uuid,
lambda: self._assert_execution_context(execution_context),
)
self._require_upload_object_key(execution_context, file.file_name)
try:
# set file status to processing
await self.ap.persistence_mgr.execute_async(
sqlalchemy.update(persistence_rag.File)
.where(persistence_rag.File.workspace_uuid == execution_context.workspace_uuid)
.where(persistence_rag.File.uuid == file.uuid)
.values(status='processing')
)
status_visible = False
for retry_delay in (0.0, 0.01, 0.05, 0.1):
if retry_delay:
await asyncio.sleep(retry_delay)
if await self._set_file_status(execution_context, file.uuid, 'processing'):
status_visible = True
break
if not status_visible:
raise WorkspaceNotFoundError('Knowledge file was not committed before its background task started')
task_context.set_current_action('Processing file')
@@ -152,13 +161,8 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
raise Exception(error_msg)
# set file status to completed
await self._assert_execution_context(execution_context)
await self.ap.persistence_mgr.execute_async(
sqlalchemy.update(persistence_rag.File)
.where(persistence_rag.File.workspace_uuid == execution_context.workspace_uuid)
.where(persistence_rag.File.uuid == file.uuid)
.values(status='completed')
)
if not await self._set_file_status(execution_context, file.uuid, 'completed'):
raise WorkspaceNotFoundError('Knowledge file not found')
except Exception as e:
self.ap.logger.error(f'Error storing file {file.uuid}: {e}')
@@ -166,16 +170,10 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
# A stale placement is fenced from all writes, including failure
# status updates from an old background task.
try:
await self._assert_execution_context(execution_context)
if not await self._set_file_status(execution_context, file.uuid, 'failed'):
raise WorkspaceNotFoundError('Knowledge file not found')
except Exception:
self.ap.logger.warning(f'Skipping stale RAG task status update for file {file.uuid}')
else:
await self.ap.persistence_mgr.execute_async(
sqlalchemy.update(persistence_rag.File)
.where(persistence_rag.File.workspace_uuid == execution_context.workspace_uuid)
.where(persistence_rag.File.uuid == file.uuid)
.values(status='failed')
)
raise
finally:
@@ -191,6 +189,37 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface):
except (WorkspaceRequiredError, WorkspaceNotFoundError):
self.ap.logger.warning(f'Skipping stale RAG upload cleanup for file {file.uuid}')
async def _set_file_status(
self,
execution_context: ExecutionContext,
file_uuid: str,
status: str,
) -> bool:
"""Commit one detached-task status transition in its own tenant UoW."""
async def update() -> bool:
await self._assert_execution_context(execution_context)
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.update(persistence_rag.File)
.where(persistence_rag.File.workspace_uuid == execution_context.workspace_uuid)
.where(persistence_rag.File.uuid == file_uuid)
.values(status=status)
)
return getattr(result, 'rowcount', 0) > 0
persistence_mgr = self.ap.persistence_mgr
managed_mode = getattr(getattr(persistence_mgr, 'mode', None), 'value', None) in {
'cloud_runtime',
'oss_compat',
}
tenant_uow = getattr(persistence_mgr, 'tenant_uow', None)
if managed_mode:
if not callable(tenant_uow):
raise RuntimeError('Knowledge tasks require an explicit tenant UoW')
async with tenant_uow(execution_context.workspace_uuid):
return await update()
return await update()
async def store_file(
self,
execution_context: ExecutionContext,
@@ -735,7 +764,36 @@ class RAGManager:
self.knowledge_bases = {}
# Load knowledge bases
list_bindings = getattr(self.ap.workspace_service, 'list_active_execution_bindings', None)
tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
cloud_runtime = getattr(getattr(self.ap.persistence_mgr, 'mode', None), 'value', None) == 'cloud_runtime'
if cloud_runtime:
if not callable(list_bindings) or not callable(tenant_uow):
raise RuntimeError('Cloud knowledge loading requires explicit instance discovery and tenant UoWs')
for binding in await list_bindings():
async with tenant_uow(binding.workspace_uuid):
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(persistence_rag.KnowledgeBase)
.where(persistence_rag.KnowledgeBase.workspace_uuid == binding.workspace_uuid)
.order_by(persistence_rag.KnowledgeBase.uuid)
)
for knowledge_base in result.all():
try:
await self.load_knowledge_base(
ExecutionContext(
instance_uuid=binding.instance_uuid,
workspace_uuid=binding.workspace_uuid,
placement_generation=binding.placement_generation,
),
knowledge_base,
)
except Exception as e:
self.ap.logger.error(
f'Error loading knowledge base {knowledge_base.uuid}: {e}\n{traceback.format_exc()}'
)
return
# Compatibility path for isolated manager tests and older embedders.
result = await self.ap.persistence_mgr.execute_async(sqlalchemy.select(persistence_rag.KnowledgeBase))
knowledge_bases = result.all()
+63 -32
View File
@@ -3,6 +3,7 @@
from __future__ import annotations
import asyncio
import contextlib
import json
import typing
import httpx
@@ -10,6 +11,7 @@ import sqlalchemy
from ..core import app as core_app
from ..entity.persistence.metadata import Metadata
from ..persistence.tenant_uow import CrossScopeTransactionError
from ..utils import constants
SURVEY_TRIGGERED_KEY = 'survey_triggered_events'
@@ -36,15 +38,41 @@ class SurveyManager:
await self._load_triggered_events()
await self._load_bot_response_count()
@contextlib.asynccontextmanager
async def _instance_transaction(self):
"""Bind instance-global survey metadata to an explicit Cloud transaction."""
persistence_mgr = self.ap.persistence_mgr
try:
active_session = getattr(persistence_mgr, 'current_session', lambda: None)()
except CrossScopeTransactionError:
# A newly-created child task inherited its parent's ContextVar;
# opening an explicit UoW below gives it an independent session.
active_session = None
if active_session is not None:
yield
return
cloud_runtime = getattr(getattr(persistence_mgr, 'mode', None), 'value', None) == 'cloud_runtime'
if cloud_runtime:
instance_uow = getattr(persistence_mgr, 'instance_discovery_uow', None)
if not callable(instance_uow):
raise RuntimeError('Cloud survey metadata requires an explicit instance UoW')
async with instance_uow(self.ap.workspace_service.instance_uuid):
yield
return
yield
async def _load_triggered_events(self):
"""Load previously triggered events from metadata table."""
try:
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(Metadata).where(Metadata.key == SURVEY_TRIGGERED_KEY)
)
row = result.first()
if row:
self._triggered_events = set(json.loads(row[0].value))
async with self._instance_transaction():
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(Metadata.value).where(Metadata.key == SURVEY_TRIGGERED_KEY)
)
value = result.scalar_one_or_none()
if value is not None:
self._triggered_events = set(json.loads(value))
except Exception:
self._triggered_events = set()
@@ -52,17 +80,18 @@ class SurveyManager:
"""Persist triggered events to metadata table."""
try:
value = json.dumps(list(self._triggered_events))
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(Metadata).where(Metadata.key == SURVEY_TRIGGERED_KEY)
)
if result.first():
await self.ap.persistence_mgr.execute_async(
sqlalchemy.update(Metadata).where(Metadata.key == SURVEY_TRIGGERED_KEY).values(value=value)
)
else:
await self.ap.persistence_mgr.execute_async(
sqlalchemy.insert(Metadata).values(key=SURVEY_TRIGGERED_KEY, value=value)
async with self._instance_transaction():
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(Metadata.value).where(Metadata.key == SURVEY_TRIGGERED_KEY)
)
if result.scalar_one_or_none() is not None:
await self.ap.persistence_mgr.execute_async(
sqlalchemy.update(Metadata).where(Metadata.key == SURVEY_TRIGGERED_KEY).values(value=value)
)
else:
await self.ap.persistence_mgr.execute_async(
sqlalchemy.insert(Metadata).values(key=SURVEY_TRIGGERED_KEY, value=value)
)
except Exception as e:
self.ap.logger.debug(f'Failed to save survey triggered events: {e}')
@@ -75,12 +104,13 @@ class SurveyManager:
async def _load_bot_response_count(self):
"""Load the persisted successful bot response count from metadata table."""
try:
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(Metadata).where(Metadata.key == BOT_RESPONSE_COUNT_KEY)
)
row = result.first()
if row:
self._bot_response_count = int(row[0].value)
async with self._instance_transaction():
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(Metadata.value).where(Metadata.key == BOT_RESPONSE_COUNT_KEY)
)
value = result.scalar_one_or_none()
if value is not None:
self._bot_response_count = int(value)
except Exception:
self._bot_response_count = 0
@@ -88,17 +118,18 @@ class SurveyManager:
"""Persist the successful bot response count to metadata table."""
try:
value = str(self._bot_response_count)
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(Metadata).where(Metadata.key == BOT_RESPONSE_COUNT_KEY)
)
if result.first():
await self.ap.persistence_mgr.execute_async(
sqlalchemy.update(Metadata).where(Metadata.key == BOT_RESPONSE_COUNT_KEY).values(value=value)
)
else:
await self.ap.persistence_mgr.execute_async(
sqlalchemy.insert(Metadata).values(key=BOT_RESPONSE_COUNT_KEY, value=value)
async with self._instance_transaction():
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(Metadata.value).where(Metadata.key == BOT_RESPONSE_COUNT_KEY)
)
if result.scalar_one_or_none() is not None:
await self.ap.persistence_mgr.execute_async(
sqlalchemy.update(Metadata).where(Metadata.key == BOT_RESPONSE_COUNT_KEY).values(value=value)
)
else:
await self.ap.persistence_mgr.execute_async(
sqlalchemy.insert(Metadata).values(key=BOT_RESPONSE_COUNT_KEY, value=value)
)
except Exception as e:
self.ap.logger.debug(f'Failed to save survey bot response count: {e}')
+14
View File
@@ -30,6 +30,20 @@ HEARTBEAT_INTERVAL_SECONDS = 24 * 3600
async def _count(ap: core_app.Application, table) -> int:
"""Count rows in a persistence table; -1 when unavailable."""
try:
persistence_mgr = ap.persistence_mgr
cloud_runtime = getattr(getattr(persistence_mgr, 'mode', None), 'value', None) == 'cloud_runtime'
if cloud_runtime:
tenant_uow = getattr(persistence_mgr, 'tenant_uow', None)
if not callable(tenant_uow):
return -1
total = 0
for binding in await ap.workspace_service.list_active_execution_bindings():
async with tenant_uow(binding.workspace_uuid):
result = await persistence_mgr.execute_async(
sqlalchemy.select(sqlalchemy.func.count()).select_from(table)
)
total += int(result.scalar() or 0)
return total
result = await ap.persistence_mgr.execute_async(sqlalchemy.select(sqlalchemy.func.count()).select_from(table))
return int(result.scalar() or 0)
except Exception:
+183 -29
View File
@@ -64,9 +64,26 @@ class VectorDBManager:
# Get pgvector configuration
pgvector_config = kb_config.get('pgvector', {})
use_business_database = pgvector_config.get('use_business_database', False)
allowed_dimensions = pgvector_config.get(
'allowed_dimensions',
[384, 512, 768, 1024, 1536],
)
common_options = {
'use_business_database': use_business_database,
'allowed_dimensions': allowed_dimensions,
}
if use_business_database:
self.vector_db = PgVectorDatabase(self.ap, **common_options)
self.ap.logger.info('Initialized pgvector on the shared business PostgreSQL database.')
return
connection_string = pgvector_config.get('connection_string')
if connection_string:
self.vector_db = PgVectorDatabase(self.ap, connection_string=connection_string)
self.vector_db = PgVectorDatabase(
self.ap,
connection_string=connection_string,
**common_options,
)
else:
# Use individual parameters
host = pgvector_config.get('host', 'localhost')
@@ -75,7 +92,13 @@ class VectorDBManager:
user = pgvector_config.get('user', 'postgres')
password = pgvector_config.get('password', 'postgres')
self.vector_db = PgVectorDatabase(
self.ap, host=host, port=port, database=database, user=user, password=password
self.ap,
host=host,
port=port,
database=database,
user=user,
password=password,
**common_options,
)
self.ap.logger.info('Initialized pgvector database backend.')
@@ -156,23 +179,24 @@ class VectorDBManager:
"""
await self._validate_execution_context(execution_context)
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(
persistence_rag.KnowledgeBase.collection_id,
persistence_rag.KnowledgeBase.legacy_vector_collection,
persistence_workspace.Workspace.source,
async with self.ap.persistence_mgr.tenant_uow(execution_context.workspace_uuid):
result = await self.ap.persistence_mgr.execute_async(
sqlalchemy.select(
persistence_rag.KnowledgeBase.collection_id,
persistence_rag.KnowledgeBase.legacy_vector_collection,
persistence_workspace.Workspace.source,
)
.join(
persistence_workspace.Workspace,
persistence_workspace.Workspace.uuid == persistence_rag.KnowledgeBase.workspace_uuid,
)
.where(
persistence_rag.KnowledgeBase.workspace_uuid == execution_context.workspace_uuid,
persistence_rag.KnowledgeBase.uuid == knowledge_base_uuid,
persistence_workspace.Workspace.instance_uuid == execution_context.instance_uuid,
)
.limit(1)
)
.join(
persistence_workspace.Workspace,
persistence_workspace.Workspace.uuid == persistence_rag.KnowledgeBase.workspace_uuid,
)
.where(
persistence_rag.KnowledgeBase.workspace_uuid == execution_context.workspace_uuid,
persistence_rag.KnowledgeBase.uuid == knowledge_base_uuid,
persistence_workspace.Workspace.instance_uuid == execution_context.instance_uuid,
)
.limit(1)
)
row = result.first()
if row is None:
raise WorkspaceNotFoundError('Knowledge base not found')
@@ -194,6 +218,58 @@ class VectorDBManager:
return self.physical_collection_name(execution_context, knowledge_base_uuid)
def _pgvector_database(self):
from .vdbs.pgvector_db import PgVectorDatabase
return self.vector_db if isinstance(self.vector_db, PgVectorDatabase) else None
async def _resolve_pgvector_scope(
self,
execution_context: ExecutionContext,
knowledge_base_uuid: str,
*,
expected_dimension: int | None,
initialize_dimension: bool,
):
"""Bind and verify the server-owned knowledge-base vector dimension."""
from .vdbs.pgvector_db import PgVectorScope
pgvector = self._pgvector_database()
if pgvector is None: # pragma: no cover - private call invariant
raise RuntimeError('pgvector scope requested for another vector backend')
if expected_dimension is not None and expected_dimension not in pgvector.allowed_dimensions:
raise ValueError(f'Embedding dimension {expected_dimension} is not enabled for this deployment')
async with self.ap.persistence_mgr.tenant_uow(execution_context.workspace_uuid):
query = sqlalchemy.select(persistence_rag.KnowledgeBase.embedding_dimension).where(
persistence_rag.KnowledgeBase.workspace_uuid == execution_context.workspace_uuid,
persistence_rag.KnowledgeBase.uuid == knowledge_base_uuid,
)
current_dimension = (await self.ap.persistence_mgr.execute_async(query)).scalar_one_or_none()
if current_dimension is None and expected_dimension is not None and initialize_dimension:
await self.ap.persistence_mgr.execute_async(
sqlalchemy.update(persistence_rag.KnowledgeBase)
.where(
persistence_rag.KnowledgeBase.workspace_uuid == execution_context.workspace_uuid,
persistence_rag.KnowledgeBase.uuid == knowledge_base_uuid,
persistence_rag.KnowledgeBase.embedding_dimension.is_(None),
)
.values(embedding_dimension=expected_dimension)
)
current_dimension = (await self.ap.persistence_mgr.execute_async(query)).scalar_one_or_none()
if expected_dimension is not None and current_dimension != expected_dimension:
if current_dimension is None:
raise ValueError('Knowledge base has no selected pgvector embedding dimension')
raise ValueError(f'Knowledge base embedding dimension is {current_dimension}, not {expected_dimension}')
return PgVectorScope(
workspace_uuid=execution_context.workspace_uuid,
knowledge_base_uuid=knowledge_base_uuid,
embedding_dimension=current_dimension,
)
async def upsert(
self,
execution_context: ExecutionContext,
@@ -219,6 +295,25 @@ class VectorDBManager:
}
for item in source_metadata
]
pgvector = self._pgvector_database()
if pgvector is not None:
if not vectors:
return
scope = await self._resolve_pgvector_scope(
execution_context,
knowledge_base_uuid,
expected_dimension=len(vectors[0]),
initialize_dimension=True,
)
await pgvector.add_embeddings(
collection=collection_name,
ids=ids,
embeddings_list=vectors,
metadatas=scoped_metadata,
documents=documents,
scope=scope,
)
return
await self.vector_db.add_embeddings(
collection=collection_name,
ids=ids,
@@ -248,15 +343,34 @@ class VectorDBManager:
execution_context,
knowledge_base_uuid,
)
results = await self.vector_db.search(
collection=collection_name,
query_embedding=query_vector,
k=limit,
search_type=search_type,
query_text=query_text,
filter=filter,
vector_weight=vector_weight,
)
pgvector = self._pgvector_database()
if pgvector is not None:
scope = await self._resolve_pgvector_scope(
execution_context,
knowledge_base_uuid,
expected_dimension=len(query_vector),
initialize_dimension=False,
)
results = await pgvector.search(
collection=collection_name,
query_embedding=query_vector,
k=limit,
search_type=search_type,
query_text=query_text,
filter=filter,
vector_weight=vector_weight,
scope=scope,
)
else:
results = await self.vector_db.search(
collection=collection_name,
query_embedding=query_vector,
k=limit,
search_type=search_type,
query_text=query_text,
filter=filter,
vector_weight=vector_weight,
)
if not results or 'ids' not in results or not results['ids']:
return []
@@ -297,8 +411,20 @@ class VectorDBManager:
execution_context,
knowledge_base_uuid,
)
pgvector = self._pgvector_database()
scope = None
if pgvector is not None:
scope = await self._resolve_pgvector_scope(
execution_context,
knowledge_base_uuid,
expected_dimension=None,
initialize_dimension=False,
)
for file_id in file_ids:
await self.vector_db.delete_by_file_id(collection_name, file_id)
if pgvector is not None:
await pgvector.delete_by_file_id(collection_name, file_id, scope=scope)
else:
await self.vector_db.delete_by_file_id(collection_name, file_id)
async def delete_collection(
self,
@@ -311,7 +437,17 @@ class VectorDBManager:
execution_context,
knowledge_base_uuid,
)
await self.vector_db.delete_collection(collection_name)
pgvector = self._pgvector_database()
if pgvector is not None:
scope = await self._resolve_pgvector_scope(
execution_context,
knowledge_base_uuid,
expected_dimension=None,
initialize_dimension=False,
)
await pgvector.delete_collection(collection_name, scope=scope)
else:
await self.vector_db.delete_collection(collection_name)
async def delete_by_filter(
self,
@@ -328,6 +464,15 @@ class VectorDBManager:
execution_context,
knowledge_base_uuid,
)
pgvector = self._pgvector_database()
if pgvector is not None:
scope = await self._resolve_pgvector_scope(
execution_context,
knowledge_base_uuid,
expected_dimension=None,
initialize_dimension=False,
)
return await pgvector.delete_by_filter(collection_name, filter, scope=scope)
return await self.vector_db.delete_by_filter(collection_name, filter)
async def list_by_filter(
@@ -347,4 +492,13 @@ class VectorDBManager:
execution_context,
knowledge_base_uuid,
)
pgvector = self._pgvector_database()
if pgvector is not None:
scope = await self._resolve_pgvector_scope(
execution_context,
knowledge_base_uuid,
expected_dimension=None,
initialize_dimension=False,
)
return await pgvector.list_by_filter(collection_name, filter, limit, offset, scope=scope)
return await self.vector_db.list_by_filter(collection_name, filter, limit, offset)
+322 -276
View File
@@ -1,22 +1,31 @@
from __future__ import annotations
from typing import Any, Dict
from sqlalchemy import create_engine, text, Column, String, Text
from sqlalchemy.orm import declarative_base
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
import contextlib
import dataclasses
from collections.abc import AsyncIterator
from typing import Any
import sqlalchemy
from pgvector.sqlalchemy import Vector
from langbot.pkg.vector.vdb import VectorDatabase
from langbot.pkg.vector.filter_utils import normalize_filter, strip_unsupported_fields
from sqlalchemy.dialects.postgresql import insert as postgresql_insert
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
from sqlalchemy.orm import declarative_base
from langbot.pkg.core import app
from langbot.pkg.vector.filter_utils import normalize_filter, strip_unsupported_fields
from langbot.pkg.vector.vdb import VectorDatabase
Base = declarative_base()
DEFAULT_ALLOWED_DIMENSIONS = (384, 512, 768, 1024, 1536)
# pgvector schema only stores these metadata fields.
_PG_SUPPORTED_FIELDS = {'text', 'file_id', 'chunk_uuid'}
# Callers use canonical metadata key 'uuid' but pgvector stores it as 'chunk_uuid'.
_PG_FIELD_ALIASES = {'uuid': 'chunk_uuid'}
# Map schema field names to SQLAlchemy columns (resolved lazily from PgVectorEntry).
_PG_COLUMN_MAP = {
'text': 'text',
'file_id': 'file_id',
@@ -24,21 +33,50 @@ _PG_COLUMN_MAP = {
}
@dataclasses.dataclass(frozen=True, slots=True)
class PgVectorScope:
"""Trusted relational tenant key for one knowledge-base operation."""
workspace_uuid: str
knowledge_base_uuid: str
embedding_dimension: int | None = None
def __post_init__(self) -> None:
for field_name in ('workspace_uuid', 'knowledge_base_uuid'):
value = getattr(self, field_name)
if not isinstance(value, str) or not value.strip():
raise ValueError(f'{field_name} must not be empty')
object.__setattr__(self, field_name, value.strip())
dimension = self.embedding_dimension
if dimension is not None and (isinstance(dimension, bool) or not isinstance(dimension, int) or dimension <= 0):
raise ValueError('embedding_dimension must be a positive integer')
class PgVectorEntry(Base):
"""SQLAlchemy model for pgvector entries"""
"""Tenant-scoped pgvector row created only by release/OSS migrations."""
__tablename__ = 'langbot_vectors'
id = Column(String, primary_key=True)
collection = Column(String, index=True, nullable=False)
embedding = Column(Vector(1536)) # Default dimension, will be created dynamically
text = Column(Text)
file_id = Column(String, index=True)
chunk_uuid = Column(String)
workspace_uuid = sqlalchemy.Column(sqlalchemy.String(36), primary_key=True)
knowledge_base_uuid = sqlalchemy.Column(sqlalchemy.String(255), primary_key=True)
vector_id = sqlalchemy.Column(sqlalchemy.String(255), primary_key=True)
embedding_dimension = sqlalchemy.Column(sqlalchemy.Integer, nullable=False)
embedding = sqlalchemy.Column(Vector(), nullable=False)
text = sqlalchemy.Column(sqlalchemy.Text)
file_id = sqlalchemy.Column(sqlalchemy.String(255), index=True)
chunk_uuid = sqlalchemy.Column(sqlalchemy.String(255))
__table_args__ = (
sqlalchemy.CheckConstraint(
'vector_dims(embedding) = embedding_dimension',
name='ck_langbot_vectors_embedding_dimension',
),
)
def _build_pg_conditions(filter_dict: dict[str, Any]) -> list:
"""Translate canonical filter dict into a list of SQLAlchemy conditions."""
"""Translate canonical filter dict into SQLAlchemy conditions."""
triples = normalize_filter(filter_dict)
triples = strip_unsupported_fields(triples, _PG_SUPPORTED_FIELDS, _PG_FIELD_ALIASES)
@@ -65,83 +103,139 @@ def _build_pg_conditions(filter_dict: dict[str, Any]) -> list:
class PgVectorDatabase(VectorDatabase):
"""PostgreSQL with pgvector extension database implementation"""
"""PostgreSQL vector adapter with explicit Workspace/RLS scope.
Cloud reuses the business database engine and never performs DDL. OSS can
still opt into a standalone pgvector database; that compatibility mode may
create a fresh schema, but it uses the same explicit tenant keys.
"""
def __init__(
self,
ap: app.Application,
connection_string: str = None,
connection_string: str | None = None,
host: str = 'localhost',
port: int = 5432,
database: str = 'langbot',
user: str = 'postgres',
password: str = 'postgres',
):
"""Initialize pgvector database
Args:
ap: Application instance
connection_string: Full PostgreSQL connection string (overrides other params)
host: PostgreSQL host
port: PostgreSQL port
database: Database name
user: Database user
password: Database password
"""
*,
use_business_database: bool = False,
allowed_dimensions: list[int] | tuple[int, ...] = DEFAULT_ALLOWED_DIMENSIONS,
) -> None:
self.ap = ap
self.use_business_database = use_business_database
self.allowed_dimensions = self._normalize_allowed_dimensions(allowed_dimensions)
self.engine = None
self.async_engine = None
self.AsyncSessionLocal: async_sessionmaker[AsyncSession] | None = None
if use_business_database:
persistence_mgr = getattr(ap, 'persistence_mgr', None)
if persistence_mgr is None:
raise RuntimeError('Shared pgvector requires the initialized business persistence manager')
business_engine = persistence_mgr.get_db_engine()
if business_engine.dialect.name != 'postgresql':
raise RuntimeError('Shared pgvector requires the PostgreSQL business database')
self.async_engine = business_engine
self.ap.logger.info('Connected pgvector adapter to the shared PostgreSQL business database')
return
# Build connection string if not provided
if connection_string:
self.connection_string = connection_string
else:
self.connection_string = f'postgresql+psycopg://{user}:{password}@{host}:{port}/{database}'
self.async_connection_string = self.connection_string.replace('postgresql://', 'postgresql+asyncpg://').replace(
'postgresql+psycopg://', 'postgresql+asyncpg://'
)
self._initialize_standalone_db()
self.engine = None
self.async_engine = None
self.SessionLocal = None
self.AsyncSessionLocal = None
self._collections = set()
self._initialize_db()
@staticmethod
def _normalize_allowed_dimensions(dimensions: list[int] | tuple[int, ...]) -> frozenset[int]:
if not isinstance(dimensions, (list, tuple)) or not dimensions:
raise ValueError('pgvector allowed_dimensions must be a non-empty list')
if any(isinstance(item, bool) or not isinstance(item, int) or item <= 0 for item in dimensions):
raise ValueError('pgvector allowed_dimensions must contain positive integers')
unsupported = set(dimensions) - set(DEFAULT_ALLOWED_DIMENSIONS)
if unsupported:
raise ValueError(f'pgvector dimensions do not have release-created ANN indexes: {sorted(unsupported)}')
return frozenset(dimensions)
def _initialize_db(self):
"""Initialize database connection and create tables"""
try:
# Create async engine for async operations
self.async_engine = create_async_engine(self.async_connection_string, echo=False, pool_pre_ping=True)
self.AsyncSessionLocal = async_sessionmaker(self.async_engine, class_=AsyncSession, expire_on_commit=False)
def _initialize_standalone_db(self) -> None:
"""Initialize the explicit OSS external database compatibility path."""
# Create sync engine for table creation
sync_connection_string = self.connection_string.replace('postgresql+asyncpg://', 'postgresql+psycopg://')
self.engine = create_engine(sync_connection_string, echo=False)
from sqlalchemy import create_engine
# Create pgvector extension and tables
with self.engine.connect() as conn:
# Enable pgvector extension
conn.execute(text('CREATE EXTENSION IF NOT EXISTS vector'))
conn.commit()
self.async_engine = create_async_engine(self.async_connection_string, echo=False, pool_pre_ping=True)
self.AsyncSessionLocal = async_sessionmaker(self.async_engine, class_=AsyncSession, expire_on_commit=False)
sync_connection_string = self.connection_string.replace('postgresql+asyncpg://', 'postgresql+psycopg://')
self.engine = create_engine(sync_connection_string, echo=False)
# Create tables
Base.metadata.create_all(self.engine)
with self.engine.begin() as conn:
conn.execute(sqlalchemy.text('CREATE EXTENSION IF NOT EXISTS vector'))
existing_tables = set(sqlalchemy.inspect(conn).get_table_names())
if PgVectorEntry.__tablename__ in existing_tables:
columns = {
column['name'] for column in sqlalchemy.inspect(conn).get_columns(PgVectorEntry.__tablename__)
}
required = {
'workspace_uuid',
'knowledge_base_uuid',
'vector_id',
'embedding_dimension',
'embedding',
}
if not required.issubset(columns):
raise RuntimeError(
'The external pgvector database uses the legacy unscoped schema; '
'migrate it before enabling multi-tenant vector access'
)
Base.metadata.create_all(conn)
self.ap.logger.info('Connected to PostgreSQL with pgvector')
except Exception as e:
self.ap.logger.error(f'Failed to connect to PostgreSQL: {e}')
raise
self.ap.logger.info('Connected to standalone PostgreSQL pgvector database')
def _require_scope(self, scope: PgVectorScope | None, *, require_dimension: bool) -> PgVectorScope:
if not isinstance(scope, PgVectorScope):
raise ValueError('pgvector operations require a trusted PgVectorScope')
dimension = scope.embedding_dimension
if require_dimension and dimension is None:
raise ValueError('pgvector operation requires an embedding dimension')
if dimension is not None and dimension not in self.allowed_dimensions:
raise ValueError(f'Embedding dimension {dimension} is not enabled for this pgvector deployment')
return scope
@staticmethod
def _scope_conditions(scope: PgVectorScope) -> tuple[Any, Any]:
return (
PgVectorEntry.workspace_uuid == scope.workspace_uuid,
PgVectorEntry.knowledge_base_uuid == scope.knowledge_base_uuid,
)
@contextlib.asynccontextmanager
async def _session(self, scope: PgVectorScope) -> AsyncIterator[AsyncSession]:
admission = getattr(self.ap, 'deployment_admission', None)
if admission is not None:
admission.require_active()
if self.use_business_database:
async with self.ap.persistence_mgr.tenant_uow(scope.workspace_uuid) as uow:
yield uow.session
if admission is not None:
admission.require_active()
return
if self.AsyncSessionLocal is None: # pragma: no cover - constructor invariant
raise RuntimeError('Standalone pgvector session factory is unavailable')
async with self.AsyncSessionLocal() as session, session.begin():
yield session
if admission is not None:
admission.require_active()
async def get_or_create_collection(self, collection: str):
"""Get or create a collection (logical grouping in pgvector)
"""Retain the common adapter API; relational rows need no collection DDL."""
Args:
collection: Collection name (knowledge base UUID)
"""
# In pgvector, collections are logical - we just track them
if collection not in self._collections:
self._collections.add(collection)
self.ap.logger.info(f"Registered pgvector collection '{collection}'")
if not isinstance(collection, str) or not collection.strip():
raise ValueError('collection must not be empty')
return collection
async def add_embeddings(
@@ -151,38 +245,59 @@ class PgVectorDatabase(VectorDatabase):
embeddings_list: list[list[float]],
metadatas: list[dict[str, Any]],
documents: list[str] | None = None,
*,
scope: PgVectorScope | None = None,
) -> None:
"""Add vector embeddings to pgvector
Args:
collection: Collection name
ids: List of unique IDs for each vector
embeddings_list: List of embedding vectors
metadatas: List of metadata dictionaries
"""
scope = self._require_scope(scope, require_dimension=True)
await self.get_or_create_collection(collection)
if not ids:
return
if len(ids) != len(embeddings_list) or len(metadatas) != len(ids):
raise ValueError('pgvector ids, embeddings and metadata lengths must match')
if documents is not None and len(documents) != len(ids):
raise ValueError('pgvector documents length must match ids')
if len(set(ids)) != len(ids) or any(not isinstance(item, str) or not item.strip() for item in ids):
raise ValueError('pgvector vector IDs must be unique non-empty strings per upsert')
expected_dimension = scope.embedding_dimension
if any(len(embedding) != expected_dimension for embedding in embeddings_list):
raise ValueError(f'All embeddings must have the selected dimension {expected_dimension}')
async with self.AsyncSessionLocal() as session:
try:
for i, vector_id in enumerate(ids):
metadata = metadatas[i] if i < len(metadatas) else {}
values = []
for index, vector_id in enumerate(ids):
metadata = metadatas[index]
document = documents[index] if documents is not None else None
values.append(
{
'workspace_uuid': scope.workspace_uuid,
'knowledge_base_uuid': scope.knowledge_base_uuid,
'vector_id': vector_id.strip(),
'embedding_dimension': expected_dimension,
'embedding': embeddings_list[index],
'text': metadata.get('text', document or ''),
'file_id': metadata.get('file_id', ''),
'chunk_uuid': metadata.get('uuid', metadata.get('chunk_uuid', '')),
}
)
entry = PgVectorEntry(
id=vector_id,
collection=collection,
embedding=embeddings_list[i],
text=metadata.get('text', ''),
file_id=metadata.get('file_id', ''),
chunk_uuid=metadata.get('uuid', ''),
)
session.add(entry)
await session.commit()
self.ap.logger.info(f"Added {len(ids)} embeddings to pgvector collection '{collection}'")
except Exception as e:
await session.rollback()
self.ap.logger.error(f'Error adding embeddings to pgvector: {e}')
raise
statement = postgresql_insert(PgVectorEntry).values(values)
excluded = statement.excluded
statement = statement.on_conflict_do_update(
index_elements=[
PgVectorEntry.workspace_uuid,
PgVectorEntry.knowledge_base_uuid,
PgVectorEntry.vector_id,
],
set_={
'embedding_dimension': excluded.embedding_dimension,
'embedding': excluded.embedding,
'text': excluded.text,
'file_id': excluded.file_id,
'chunk_uuid': excluded.chunk_uuid,
},
)
async with self._session(scope) as session:
await session.execute(statement)
self.ap.logger.info(f'Upserted {len(ids)} pgvector embeddings for knowledge base {scope.knowledge_base_uuid}')
async def search(
self,
@@ -193,125 +308,79 @@ class PgVectorDatabase(VectorDatabase):
query_text: str = '',
filter: dict[str, Any] | None = None,
vector_weight: float | None = None,
) -> Dict[str, Any]:
"""Search for similar vectors using cosine distance
Args:
collection: Collection name
query_embedding: Query vector
k: Number of top results to return
Returns:
Dictionary with search results in Chroma-compatible format
"""
*,
scope: PgVectorScope | None = None,
) -> dict[str, Any]:
del query_text, vector_weight
scope = self._require_scope(scope, require_dimension=True)
await self.get_or_create_collection(collection)
if search_type != 'vector':
raise ValueError('pgvector currently supports vector search only')
if k <= 0:
raise ValueError('pgvector search limit must be positive')
if len(query_embedding) != scope.embedding_dimension:
raise ValueError(f'Query embedding must have the selected dimension {scope.embedding_dimension}')
async with self.AsyncSessionLocal() as session:
try:
# Use cosine distance for similarity search
from sqlalchemy import select
typed_embedding = sqlalchemy.cast(PgVectorEntry.embedding, Vector(scope.embedding_dimension))
distance = typed_embedding.cosine_distance(query_embedding)
statement = (
sqlalchemy.select(
PgVectorEntry.vector_id,
PgVectorEntry.text,
PgVectorEntry.file_id,
PgVectorEntry.chunk_uuid,
distance.label('distance'),
)
.where(*self._scope_conditions(scope), PgVectorEntry.embedding_dimension == scope.embedding_dimension)
.order_by(distance)
.limit(k)
)
for condition in _build_pg_conditions(filter or {}):
statement = statement.where(condition)
# Query for similar vectors
stmt = (
select(
PgVectorEntry.id,
PgVectorEntry.text,
PgVectorEntry.file_id,
PgVectorEntry.chunk_uuid,
PgVectorEntry.embedding.cosine_distance(query_embedding).label('distance'),
)
.filter(PgVectorEntry.collection == collection)
.order_by(PgVectorEntry.embedding.cosine_distance(query_embedding))
.limit(k)
)
async with self._session(scope) as session:
rows = (await session.execute(statement)).all()
if filter:
for cond in _build_pg_conditions(filter):
stmt = stmt.filter(cond)
ids = [row.vector_id for row in rows]
distances = [float(row.distance) for row in rows]
metadatas = [
{'text': row.text or '', 'file_id': row.file_id or '', 'uuid': row.chunk_uuid or ''} for row in rows
]
return {'ids': [ids], 'distances': [distances], 'metadatas': [metadatas]}
result = await session.execute(stmt)
rows = result.fetchall()
# Convert to Chroma-compatible format
ids = []
distances = []
metadatas = []
for row in rows:
ids.append(row.id)
distances.append(float(row.distance))
metadatas.append(
{'text': row.text or '', 'file_id': row.file_id or '', 'uuid': row.chunk_uuid or ''}
)
result_dict = {'ids': [ids], 'distances': [distances], 'metadatas': [metadatas]}
self.ap.logger.info(f"pgvector search in '{collection}' returned {len(ids)} results")
return result_dict
except Exception as e:
self.ap.logger.error(f'Error searching pgvector: {e}')
raise
async def delete_by_file_id(self, collection: str, file_id: str) -> None:
"""Delete vectors by file_id
Args:
collection: Collection name
file_id: File ID to filter deletion
"""
async def delete_by_file_id(
self,
collection: str,
file_id: str,
*,
scope: PgVectorScope | None = None,
) -> None:
scope = self._require_scope(scope, require_dimension=False)
await self.get_or_create_collection(collection)
statement = sqlalchemy.delete(PgVectorEntry).where(
*self._scope_conditions(scope),
PgVectorEntry.file_id == file_id,
)
async with self._session(scope) as session:
await session.execute(statement)
async with self.AsyncSessionLocal() as session:
try:
from sqlalchemy import delete
stmt = delete(PgVectorEntry).where(
PgVectorEntry.collection == collection, PgVectorEntry.file_id == file_id
)
await session.execute(stmt)
await session.commit()
self.ap.logger.info(
f"Deleted embeddings from pgvector collection '{collection}' with file_id: {file_id}"
)
except Exception as e:
await session.rollback()
self.ap.logger.error(f'Error deleting from pgvector: {e}')
raise
async def delete_by_filter(self, collection: str, filter: dict[str, Any]) -> int:
"""Delete vectors matching a metadata filter.
Args:
collection: Collection name
filter: Canonical metadata filter dict
"""
async def delete_by_filter(
self,
collection: str,
filter: dict[str, Any],
*,
scope: PgVectorScope | None = None,
) -> int:
scope = self._require_scope(scope, require_dimension=False)
await self.get_or_create_collection(collection)
conditions = _build_pg_conditions(filter)
if not conditions:
self.ap.logger.warning(
f"pgvector delete_by_filter on '{collection}': filter produced no conditions, skipping"
)
self.ap.logger.warning('pgvector delete_by_filter produced no supported conditions; skipping')
return 0
await self.get_or_create_collection(collection)
async with self.AsyncSessionLocal() as session:
try:
from sqlalchemy import delete
stmt = delete(PgVectorEntry).where(PgVectorEntry.collection == collection)
for cond in conditions:
stmt = stmt.where(cond)
result = await session.execute(stmt)
await session.commit()
deleted = result.rowcount
self.ap.logger.info(f"Deleted {deleted} embeddings from pgvector collection '{collection}' by filter")
return deleted
except Exception as e:
await session.rollback()
self.ap.logger.error(f'Error deleting from pgvector by filter: {e}')
raise
statement = sqlalchemy.delete(PgVectorEntry).where(*self._scope_conditions(scope), *conditions)
async with self._session(scope) as session:
result = await session.execute(statement)
return int(result.rowcount or 0)
async def list_by_filter(
self,
@@ -319,85 +388,62 @@ class PgVectorDatabase(VectorDatabase):
filter: dict[str, Any] | None = None,
limit: int = 20,
offset: int = 0,
*,
scope: PgVectorScope | None = None,
) -> tuple[list[dict[str, Any]], int]:
scope = self._require_scope(scope, require_dimension=False)
await self.get_or_create_collection(collection)
if limit <= 0 or offset < 0:
raise ValueError('pgvector pagination requires limit > 0 and offset >= 0')
async with self.AsyncSessionLocal() as session:
try:
from sqlalchemy import select, func
conditions = [*self._scope_conditions(scope), *_build_pg_conditions(filter or {})]
statement = (
sqlalchemy.select(
PgVectorEntry.vector_id,
PgVectorEntry.text,
PgVectorEntry.file_id,
PgVectorEntry.chunk_uuid,
)
.where(*conditions)
.order_by(PgVectorEntry.vector_id)
.offset(offset)
.limit(limit)
)
count_statement = sqlalchemy.select(sqlalchemy.func.count()).select_from(PgVectorEntry).where(*conditions)
async with self._session(scope) as session:
rows = (await session.execute(statement)).all()
total = int((await session.execute(count_statement)).scalar_one())
stmt = (
select(
PgVectorEntry.id,
PgVectorEntry.text,
PgVectorEntry.file_id,
PgVectorEntry.chunk_uuid,
)
.filter(PgVectorEntry.collection == collection)
.offset(offset)
.limit(limit)
)
return (
[
{
'id': row.vector_id,
'document': row.text or '',
'metadata': {
'text': row.text or '',
'file_id': row.file_id or '',
'uuid': row.chunk_uuid or '',
},
}
for row in rows
],
total,
)
count_stmt = (
select(func.count()).select_from(PgVectorEntry).filter(PgVectorEntry.collection == collection)
)
async def delete_collection(
self,
collection: str,
*,
scope: PgVectorScope | None = None,
) -> None:
scope = self._require_scope(scope, require_dimension=False)
await self.get_or_create_collection(collection)
statement = sqlalchemy.delete(PgVectorEntry).where(*self._scope_conditions(scope))
async with self._session(scope) as session:
await session.execute(statement)
if filter:
for cond in _build_pg_conditions(filter):
stmt = stmt.filter(cond)
count_stmt = count_stmt.filter(cond)
result = await session.execute(stmt)
rows = result.fetchall()
count_result = await session.execute(count_stmt)
total = count_result.scalar() or 0
items = []
for row in rows:
items.append(
{
'id': row.id,
'document': row.text or '',
'metadata': {
'text': row.text or '',
'file_id': row.file_id or '',
'uuid': row.chunk_uuid or '',
},
}
)
return items, total
except Exception as e:
self.ap.logger.error(f'Error listing from pgvector: {e}')
raise
async def delete_collection(self, collection: str):
"""Delete all vectors in a collection
Args:
collection: Collection name to delete
"""
if collection in self._collections:
self._collections.remove(collection)
async with self.AsyncSessionLocal() as session:
try:
from sqlalchemy import delete
stmt = delete(PgVectorEntry).where(PgVectorEntry.collection == collection)
await session.execute(stmt)
await session.commit()
self.ap.logger.info(f"Deleted pgvector collection '{collection}'")
except Exception as e:
await session.rollback()
self.ap.logger.error(f'Error deleting pgvector collection: {e}')
raise
async def close(self):
"""Close database connections"""
if self.async_engine:
async def close(self) -> None:
if not self.use_business_database and self.async_engine is not None:
await self.async_engine.dispose()
if self.engine:
if self.engine is not None:
self.engine.dispose()
+79 -1
View File
@@ -155,8 +155,20 @@ class WorkspaceCollaborationService:
) -> ResolvedWorkspaceAccess:
"""Resolve a selector against an active Account membership."""
normalized_workspace_uuid = requested_workspace_uuid.strip() if requested_workspace_uuid else None
if normalized_workspace_uuid is None and self.policy.multi_workspace_enabled:
raise WorkspaceNotFoundError('A Workspace selector is required')
tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
if session is None and normalized_workspace_uuid is not None and callable(tenant_uow):
async with tenant_uow(normalized_workspace_uuid) as uow:
return await self.resolve_account_workspace(
account_uuid,
normalized_workspace_uuid,
session=uow.session,
)
async def operation(active_session: AsyncSession) -> ResolvedWorkspaceAccess:
workspace_uuid = requested_workspace_uuid.strip() if requested_workspace_uuid else None
workspace_uuid = normalized_workspace_uuid
if workspace_uuid is None:
if self.policy.multi_workspace_enabled:
raise WorkspaceNotFoundError('A Workspace selector is required')
@@ -195,6 +207,40 @@ class WorkspaceCollaborationService:
*,
session: AsyncSession | None = None,
) -> list[ResolvedWorkspaceAccess]:
current_session = getattr(self.ap.persistence_mgr, 'current_session', lambda: None)
account_uow = getattr(self.ap.persistence_mgr, 'account_discovery_uow', None)
tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
if session is None and current_session() is None and callable(account_uow) and callable(tenant_uow):
# Discovery exposes only this Account's active Membership rows.
# Each resulting Workspace is then re-read under its own tenant
# transaction; directory discovery never grants business access.
async with account_uow(account_uuid) as discovery:
workspace_uuids = list(
(
await discovery.session.scalars(
sqlalchemy.select(WorkspaceMembership.workspace_uuid)
.where(
WorkspaceMembership.account_uuid == account_uuid,
WorkspaceMembership.status == MembershipStatus.ACTIVE.value,
)
.order_by(WorkspaceMembership.workspace_uuid)
)
).all()
)
accesses: list[ResolvedWorkspaceAccess] = []
for workspace_uuid in workspace_uuids:
async with tenant_uow(workspace_uuid) as workspace_uow:
accesses.append(
await self.resolve_account_workspace(
account_uuid,
workspace_uuid,
session=workspace_uow.session,
)
)
accesses.sort(key=lambda access: (access.workspace.created_at, access.workspace.uuid))
return accesses
async def operation(active_session: AsyncSession) -> list[ResolvedWorkspaceAccess]:
statement = (
sqlalchemy.select(WorkspaceMembership, Workspace)
@@ -346,6 +392,18 @@ class WorkspaceCollaborationService:
*,
session: AsyncSession | None = None,
) -> tuple[WorkspaceInvitation, Workspace]:
if session is None:
scoped_session = getattr(self.ap.persistence_mgr, 'current_session', lambda: None)()
invitation_uow = getattr(self.ap.persistence_mgr, 'invitation_discovery_uow', None)
tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
if scoped_session is None and callable(invitation_uow) and callable(tenant_uow):
token_hash = hash_invitation_token(token)
async with invitation_uow(token_hash) as discovery:
invitation = await self._get_invitation_by_token(discovery.session, token, for_update=False)
workspace_uuid = invitation.workspace_uuid
async with tenant_uow(workspace_uuid) as workspace_uow:
return await self.inspect_invitation(token, session=workspace_uow.session)
async def operation(active_session: AsyncSession) -> tuple[WorkspaceInvitation, Workspace]:
invitation = await self._get_invitation_by_token(active_session, token, for_update=True)
self._validate_invitation_state(invitation)
@@ -365,6 +423,18 @@ class WorkspaceCollaborationService:
*,
session: AsyncSession | None = None,
) -> WorkspaceMembership:
if session is None:
scoped_session = getattr(self.ap.persistence_mgr, 'current_session', lambda: None)()
invitation_uow = getattr(self.ap.persistence_mgr, 'invitation_discovery_uow', None)
tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
if scoped_session is None and callable(invitation_uow) and callable(tenant_uow):
token_digest = hash_invitation_token(token)
async with invitation_uow(token_digest) as discovery:
invitation = await self._get_invitation_by_token(discovery.session, token, for_update=False)
workspace_uuid = invitation.workspace_uuid
async with tenant_uow(workspace_uuid) as workspace_uow:
return await self.accept_invitation(token, account_uuid, session=workspace_uow.session)
async def operation(active_session: AsyncSession) -> WorkspaceMembership:
invitation = await self._get_invitation_by_token(active_session, token, for_update=True)
self._validate_invitation_state(invitation)
@@ -677,6 +747,14 @@ class WorkspaceCollaborationService:
) -> T:
if session is not None:
return await operation(session)
current_session = getattr(self.ap.persistence_mgr, 'current_session', lambda: None)()
if current_session is not None:
return await operation(current_session)
if getattr(getattr(self.ap.persistence_mgr, 'mode', None), 'value', None) == 'cloud_runtime':
require_current_session = getattr(self.ap.persistence_mgr, 'require_current_session', None)
if callable(require_current_session):
require_current_session()
raise RuntimeError('Cloud collaboration services require an explicit persistence unit of work')
async with self._session_factory()() as owned_session:
if read_only:
return await operation(owned_session)
+86 -1
View File
@@ -6,6 +6,7 @@ import typing
from collections.abc import Awaitable, Callable
from typing import TypeVar
import sqlalchemy
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
from ..entity.persistence.workspace import (
@@ -68,6 +69,11 @@ class WorkspaceService:
) -> Workspace:
"""Load one Workspace projected onto this LangBot instance."""
tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
if session is None and callable(tenant_uow):
async with tenant_uow(workspace_uuid) as uow:
return await self.get_workspace(workspace_uuid, session=uow.session)
async def operation(repository: WorkspaceRepository) -> Workspace:
workspace = await repository.get_workspace(workspace_uuid)
if workspace is None or workspace.instance_uuid != self.instance_uuid:
@@ -97,6 +103,11 @@ class WorkspaceService:
) -> WorkspaceExecutionState:
"""Load a Workspace execution state and validate its instance binding."""
tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
if session is None and callable(tenant_uow):
async with tenant_uow(workspace_uuid) as uow:
return await self.get_execution_state(workspace_uuid, session=uow.session)
async def operation(repository: WorkspaceRepository) -> WorkspaceExecutionState:
execution_state = await repository.get_execution_state(workspace_uuid)
if execution_state is None:
@@ -109,6 +120,47 @@ class WorkspaceService:
return await self._run(operation, session=session)
async def list_active_execution_bindings(self) -> list[WorkspaceExecutionBinding]:
"""Discover this instance's active Workspaces, then validate each tenant projection.
The instance discovery transaction may only reveal active, unfenced
execution-state identifiers. It is closed before any Workspace data is
read; every returned binding is revalidated inside its own tenant unit
of work.
"""
instance_discovery_uow = getattr(self.ap.persistence_mgr, 'instance_discovery_uow', None)
statement = (
sqlalchemy.select(WorkspaceExecutionState.workspace_uuid)
.where(
WorkspaceExecutionState.instance_uuid == self.instance_uuid,
WorkspaceExecutionState.state == WorkspaceExecutionStatus.ACTIVE.value,
WorkspaceExecutionState.write_fenced == sqlalchemy.false(),
)
.order_by(WorkspaceExecutionState.workspace_uuid)
)
if callable(instance_discovery_uow):
async with instance_discovery_uow(self.instance_uuid) as uow:
result = await uow.session.execute(statement)
workspace_uuids = list(result.scalars().all())
else:
# Compatibility for lightweight test doubles. Production managers
# always expose the explicit discovery scope.
result = await self.ap.persistence_mgr.execute_async(statement)
workspace_uuids = list(result.scalars().all())
bindings: list[WorkspaceExecutionBinding] = []
for workspace_uuid in workspace_uuids:
try:
bindings.append(await self.get_execution_binding(workspace_uuid))
except (
WorkspaceExecutionUnavailableError,
WorkspaceInvariantError,
WorkspaceNotFoundError,
) as exc:
self.ap.logger.warning(f'Skipping invalid Workspace execution projection {workspace_uuid!r}: {exc}')
return bindings
async def get_execution_binding(
self,
workspace_uuid: str | None = None,
@@ -124,6 +176,18 @@ class WorkspaceService:
a Workspace from source or recency.
"""
self._require_deployment_admission()
tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
if session is None and workspace_uuid is not None and callable(tenant_uow):
async with tenant_uow(workspace_uuid) as uow:
return await self.get_execution_binding(
workspace_uuid,
expected_generation=expected_generation,
session=uow.session,
_require_local=_require_local,
)
async def operation(repository: WorkspaceRepository) -> WorkspaceExecutionBinding:
if workspace_uuid is None:
workspaces = await repository.list_local_workspaces(self.instance_uuid)
@@ -180,7 +244,11 @@ class WorkspaceService:
state=execution_state.state,
)
return await self._run(operation, session=session)
binding = await self._run(operation, session=session)
# The database lookup can cross the Manifest expiry boundary. Never
# return a binding that is already invalid at a side-effect boundary.
self._require_deployment_admission()
return binding
async def get_local_execution_binding(
self,
@@ -408,6 +476,18 @@ class WorkspaceService:
if session is not None:
return await operation(WorkspaceRepository(session))
current_session = getattr(self.ap.persistence_mgr, 'current_session', lambda: None)
active_session = current_session()
if active_session is not None:
return await operation(WorkspaceRepository(active_session))
if getattr(getattr(self.ap.persistence_mgr, 'mode', None), 'value', None) == 'cloud_runtime':
# Do not let service-local session factories silently bypass the
# request/task scope enforced by PersistenceManager.
require_current_session = getattr(self.ap.persistence_mgr, 'require_current_session', None)
if callable(require_current_session):
require_current_session()
raise RuntimeError('Cloud Workspace services require an explicit persistence unit of work')
session_factory = async_sessionmaker(
self.ap.persistence_mgr.get_db_engine(),
expire_on_commit=False,
@@ -415,3 +495,8 @@ class WorkspaceService:
async with session_factory() as owned_session:
async with owned_session.begin():
return await operation(WorkspaceRepository(owned_session))
def _require_deployment_admission(self) -> None:
guard = getattr(self.ap, 'deployment_admission', None)
if guard is not None:
guard.require_active()
+36 -5
View File
@@ -60,11 +60,19 @@ database:
sqlite:
path: 'data/langbot.db'
postgresql:
# Optional SQLAlchemy URL (postgresql[+asyncpg]://...). When set, it
# overrides the structured fields and preserves TLS/query options.
url: ''
host: '127.0.0.1'
port: 5432
user: 'postgres'
password: 'postgres'
database: 'postgres'
cloud_migration:
# `langbot migrate --cloud` reads an operator-only PostgreSQL DSN from
# this environment variable. The operator role must differ from the
# runtime role above; never put its password in this file or CLI args.
operator_dsn_env: 'LANGBOT_CLOUD_MIGRATION_DSN'
vdb:
use: chroma
qdrant:
@@ -88,6 +96,11 @@ vdb:
token: ''
db_name: ''
pgvector:
# SaaS/shared-schema deployments reuse database.postgresql. OSS can
# keep this false when deliberately using an external pgvector DB.
use_business_database: false
# Release migrations create one partial ANN index per enabled value.
allowed_dimensions: [384, 512, 768, 1024, 1536]
host: '127.0.0.1'
port: 5433
database: 'langbot'
@@ -133,6 +146,9 @@ plugin:
max_pids: 128
max_open_files: 256
max_file_size_mb: 512
# Cloud shared Runtime sets this to true and fails closed unless
# delegated cgroup v2 controllers are available.
require_hard_limits: false
binary_storage:
# Max bytes for a single plugin binary storage value
max_value_bytes: 10485760
@@ -164,11 +180,26 @@ box:
# External WebSocket runtimes also require LANGBOT_BOX_CONTROL_TOKEN in
# both LangBot and Box. Keep the shared secret out of this config file.
endpoint: '' # External Box Runtime base URL, e.g. 'ws://127.0.0.1:5410'. Leave empty for local auto-managed runtime.
limits:
max_sessions: 64 # Includes persistent sessions. New sessions fail explicitly when this cap is reached.
max_managed_processes: 64 # Maximum concurrently running stdio MCP / managed processes.
max_completed_processes: 256 # Global cap for retained exited-process diagnostics.
completed_process_retention_sec: 300 # Keep exited-process diagnostics before releasing memory.
# Cloud v2 overrides these values through the instance config/environment.
# OSS keeps admission disabled and preserves the existing multi-session
# local behavior. These limits are Runtime-owned and cannot be relaxed by
# a pipeline, Workspace entitlement, or tool call.
admission:
required: false
logical_session_id: 'global'
required_backend: 'nsjail'
max_sessions: 1
max_managed_processes: 0
max_grant_ttl_sec: 300
max_timeout_sec: 120
cpus: 1.0
memory_mb: 512
pids_limit: 128
read_only_rootfs: true
# OSS admission-disabled mode uses 0 for unlimited compatibility.
# Cloud bootstrap requires a positive hard quota.
workspace_quota_mb: 0
readiness_cache_sec: 15
local:
profile: 'default'
image: '' # Custom local sandbox image. Leave empty to use the profile default.
+3
View File
@@ -82,6 +82,9 @@ class FakeApp:
def _create_mock_persistence_manager(self):
persistence_mgr = AsyncMock()
persistence_mgr.execute_async = AsyncMock()
# AsyncMock invents arbitrary callable attributes on access. Keep the
# optional production UoW hook explicitly absent unless a test opts in.
persistence_mgr.tenant_uow = None
return persistence_mgr
def _create_mock_query_pool(self):
@@ -9,6 +9,7 @@ import pytest
import quart
from langbot.pkg.api.http.controller.groups.box import BoxRouterGroup
from langbot.pkg.cloud.entitlements import EntitlementUnavailableError
pytestmark = pytest.mark.integration
@@ -34,6 +35,8 @@ async def box_security_api():
'owner-token': SimpleNamespace(uuid='owner-account', user='owner@example.com'),
}
application = Mock()
application.deployment = SimpleNamespace(multi_workspace_enabled=False)
application.persistence_mgr = SimpleNamespace(tenant_uow=None)
application.user_service.get_authenticated_account = AsyncMock(side_effect=lambda token: accounts[token])
application.workspace_collaboration_service.resolve_account_workspace = AsyncMock(
side_effect=lambda account_uuid, _workspace_uuid: _access(account_uuid)
@@ -41,6 +44,7 @@ async def box_security_api():
application.box_service.get_status = AsyncMock(return_value={'enabled': True})
application.box_service.get_sessions = AsyncMock(return_value=[{'session_id': 'private-session'}])
application.box_service.get_recent_errors = Mock(return_value=[{'error': 'private error'}])
application.box_service.managed_admission_required = False
quart_app = quart.Quart(__name__)
router = BoxRouterGroup(application, quart_app)
@@ -81,3 +85,17 @@ async def test_owner_can_audit_box_sessions_and_errors(box_security_api):
assert errors.status_code == 200
application.box_service.get_sessions.assert_awaited_once()
application.box_service.get_recent_errors.assert_called_once()
@pytest.mark.asyncio
async def test_box_status_returns_explicit_403_when_workspace_has_no_managed_sandbox(box_security_api):
application, client = box_security_api
application.box_service.get_status.side_effect = EntitlementUnavailableError(
'Workspace entitlement does not grant managed_sandbox'
)
response = await client.get('/api/v1/box/status', headers=_headers('viewer-token'))
assert response.status_code == 403
payload = await response.get_json()
assert payload['code'] == 'managed_sandbox_unavailable'
@@ -0,0 +1,214 @@
from __future__ import annotations
import json
import logging
from types import SimpleNamespace
import pytest
import sqlalchemy as sa
from quart import Quart
from langbot.pkg.api.http.controller.groups.system import SystemRouterGroup
from langbot.pkg.api.http.controller.groups.user import UserRouterGroup
from langbot.pkg.api.http.controller.groups.workspaces import WorkspacesRouterGroup
from langbot.pkg.api.http.service.user import UserService
from langbot.pkg.entity.persistence.metadata import WorkspaceMetadata
from langbot.pkg.persistence.mgr import PersistenceManager
from langbot.pkg.utils import constants
from langbot.pkg.workspace.collaboration import WorkspaceCollaborationService
from langbot.pkg.workspace.service import WorkspaceService
pytestmark = [pytest.mark.integration, pytest.mark.asyncio]
def _authorization(token: str, workspace_uuid: str | None = None) -> dict[str, str]:
headers = {'Authorization': f'Bearer {token}'}
if workspace_uuid is not None:
headers['X-Workspace-Id'] = workspace_uuid
return headers
async def test_fresh_oss_workspace_http_journey_uses_real_sqlite_persistence(
tmp_path,
monkeypatch,
):
"""Exercise the first-run Workspace journey through real HTTP handlers."""
instance_uuid = 'fresh-oss-workspace-journey'
monkeypatch.setattr(constants, 'instance_id', instance_uuid)
application = SimpleNamespace(
logger=logging.getLogger('fresh-oss-workspace-journey-test'),
instance_config=SimpleNamespace(
data={
'database': {
'use': 'sqlite',
'sqlite': {'path': str(tmp_path / 'langbot.db')},
},
'system': {
'jwt': {'secret': 'fresh-oss-workspace-secret', 'expire': 3600},
'allow_modify_login_info': True,
'limitation': {},
'outbound_ips': [],
},
'api': {'global_api_key': ''},
'plugin': {'enable_marketplace': True},
'space': {
'url': 'https://space.langbot.app',
'models_gateway_api_url': 'https://api.langbot.cloud/v1',
'disable_models_service': False,
},
'mcp': {'stdio': {'enabled': True}},
}
),
)
persistence = PersistenceManager(application)
application.persistence_mgr = persistence
await persistence.initialize()
try:
application.workspace_service = WorkspaceService(
application,
instance_uuid=instance_uuid,
)
application.workspace_collaboration_service = WorkspaceCollaborationService(
application,
application.workspace_service,
)
application.user_service = UserService(application)
quart_app = Quart(__name__)
await UserRouterGroup(application, quart_app).initialize()
await WorkspacesRouterGroup(application, quart_app).initialize()
await SystemRouterGroup(application, quart_app).initialize()
client = quart_app.test_client()
initialization = await client.get('/api/v1/user/init')
assert initialization.status_code == 200
assert (await initialization.get_json())['data'] == {'initialized': False}
initialized = await client.post(
'/api/v1/user/init',
json={'user': 'owner@example.com', 'password': 'owner-password'},
)
assert initialized.status_code == 200
authenticated = await client.post(
'/api/v1/user/auth',
json={'user': 'owner@example.com', 'password': 'owner-password'},
)
assert authenticated.status_code == 200
token = (await authenticated.get_json())['data']['token']
bootstrap = await client.get(
'/api/v1/workspaces/bootstrap',
headers=_authorization(token),
)
assert bootstrap.status_code == 200
bootstrap_workspaces = (await bootstrap.get_json())['data']['workspaces']
assert len(bootstrap_workspaces) == 1
bootstrap_access = bootstrap_workspaces[0]
workspace_uuid = bootstrap_access['workspace']['uuid']
assert bootstrap_access['workspace'] == {
'uuid': workspace_uuid,
'instance_uuid': instance_uuid,
'name': 'Default Workspace',
'slug': 'default',
'type': 'team',
'status': 'active',
'source': 'local',
}
assert bootstrap_access['membership']['email'] == 'owner@example.com'
assert bootstrap_access['membership']['role'] == 'owner'
assert 'workspace.update' in bootstrap_access['permissions']
current = await client.get(
'/api/v1/workspaces/current',
headers=_authorization(token, workspace_uuid),
)
assert current.status_code == 200
current_data = (await current.get_json())['data']
assert current_data['workspace']['uuid'] == workspace_uuid
assert current_data['membership']['account_uuid'] == bootstrap_access['membership']['account_uuid']
assert current_data['membership']['role'] == 'owner'
user_info = await client.get(
'/api/v1/user/info',
headers=_authorization(token, workspace_uuid),
)
assert user_info.status_code == 200
assert (await user_info.get_json())['data'] == {
'account_uuid': bootstrap_access['membership']['account_uuid'],
'user': 'owner@example.com',
'account_type': 'local',
'has_password': True,
}
initial_system_info = await client.get(
'/api/v1/system/info',
headers=_authorization(token, workspace_uuid),
)
assert initial_system_info.status_code == 200
assert (await initial_system_info.get_json())['data']['wizard_status'] == 'none'
assert (await initial_system_info.get_json())['data']['wizard_progress'] is None
progress = {'step': 2, 'selected_adapter': 'telegram', 'bot_saved': False}
updated_progress = await client.put(
'/api/v1/system/wizard/progress',
headers=_authorization(token, workspace_uuid),
json=progress,
)
assert updated_progress.status_code == 200
persisted_progress = await persistence.execute_async(
sa.select(WorkspaceMetadata.value).where(
WorkspaceMetadata.workspace_uuid == workspace_uuid,
WorkspaceMetadata.key == 'wizard_progress',
)
)
assert json.loads(persisted_progress.scalar_one()) == progress
system_info_with_progress = await client.get(
'/api/v1/system/info',
headers=_authorization(token, workspace_uuid),
)
assert system_info_with_progress.status_code == 200
progress_data = (await system_info_with_progress.get_json())['data']
assert progress_data['wizard_status'] == 'none'
assert progress_data['wizard_progress'] == progress
completed = await client.post(
'/api/v1/system/wizard/completed',
headers=_authorization(token, workspace_uuid),
json={'status': 'completed'},
)
assert completed.status_code == 200
completed_system_info = await client.get(
'/api/v1/system/info',
headers=_authorization(token, workspace_uuid),
)
assert completed_system_info.status_code == 200
completed_data = (await completed_system_info.get_json())['data']
assert completed_data['wizard_status'] == 'completed'
assert completed_data['wizard_progress'] is None
rejected_workspace = await client.post(
'/api/v1/workspaces',
headers=_authorization(token, workspace_uuid),
json={'name': 'Second Workspace'},
)
assert rejected_workspace.status_code == 403
rejected_data = await rejected_workspace.get_json()
assert rejected_data['code'] == 'edition_limit'
persisted_wizard_status = await persistence.execute_async(
sa.select(WorkspaceMetadata.value).where(
WorkspaceMetadata.workspace_uuid == workspace_uuid,
WorkspaceMetadata.key == 'wizard_status',
)
)
assert persisted_wizard_status.scalar_one() == 'completed'
finally:
await persistence.get_db_engine().dispose()
@@ -86,6 +86,7 @@ async def plugin_security_api(plugin_module):
}
application = Mock()
application.deployment = SimpleNamespace(multi_workspace_enabled=False)
application.user_service.get_authenticated_account = AsyncMock(side_effect=lambda token: accounts[token])
application.workspace_collaboration_service.resolve_account_workspace = AsyncMock(
side_effect=lambda account_uuid, _workspace_uuid: _access(account_uuid)
@@ -112,6 +113,7 @@ async def plugin_security_api(plugin_module):
persistence_result = Mock()
persistence_result.scalar_one_or_none.return_value = RAW_CONFIG
application.persistence_mgr.execute_async = AsyncMock(return_value=persistence_result)
application.persistence_mgr.tenant_uow = None
quart_app = quart.Quart(__name__)
router = plugin_module.PluginsRouterGroup(application, quart_app)
@@ -246,3 +248,25 @@ async def test_viewer_cannot_read_plugin_runtime_logs(plugin_security_api):
assert response.status_code == 403
assert (await response.get_json())['code'] == 'permission_denied'
application.plugin_connector.get_plugin_logs.assert_not_awaited()
@pytest.mark.asyncio
async def test_github_install_rejects_internal_asset_url_before_task_creation(
plugin_security_api,
):
application, client, _ = plugin_security_api
response = await client.post(
'/api/v1/plugins/install/github',
headers=_headers('manager-token'),
json={
'asset_url': 'http://169.254.169.254/latest/meta-data',
'owner': 'langbot-app',
'repo': 'demo-plugin',
'release_tag': 'v1.0.0',
},
)
assert response.status_code == 400
assert 'HTTPS GitHub release asset URL' in (await response.get_json())['msg']
application.task_mgr.create_user_task.assert_not_called()
@@ -25,6 +25,8 @@ async def space_oauth_api():
execution=SimpleNamespace(instance_uuid='instance-a', placement_generation=1),
)
application = Mock()
application.deployment = SimpleNamespace(multi_workspace_enabled=False)
application.persistence_mgr = None
application.user_service.get_authenticated_account = AsyncMock(return_value=account)
application.user_service.issue_space_oauth_state = AsyncMock(
side_effect=lambda purpose, **_: f'opaque-{purpose}-state'
+71 -30
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
import json
import logging
from types import SimpleNamespace
@@ -13,11 +14,13 @@ from langbot.pkg.api.http.controller.groups.workspaces import (
InvitationsRouterGroup,
WorkspacesRouterGroup,
)
from langbot.pkg.api.http.controller.groups.system import SystemRouterGroup
from langbot.pkg.api.http.controller.groups.apikeys import ApiKeysRouterGroup
from langbot.pkg.api.http.controller.groups.user import UserRouterGroup
from langbot.pkg.api.http.service.apikey import ApiKeyService
from langbot.pkg.api.http.service.user import ControlPlaneDirectoryRequiredError, UserService
from langbot.pkg.entity.persistence.base import Base
from langbot.pkg.entity.persistence.metadata import WorkspaceMetadata
from langbot.pkg.entity.persistence.user import User
from langbot.pkg.entity.persistence.workspace import (
Workspace,
@@ -25,6 +28,7 @@ from langbot.pkg.entity.persistence.workspace import (
WorkspaceInvitation,
WorkspaceMembership,
)
from langbot.pkg.persistence.mgr import PersistenceManager
from langbot.pkg.workspace.collaboration import WorkspaceCollaborationService
from langbot.pkg.workspace.service import WorkspaceService
from langbot.pkg.workspace.policy import CloudWorkspacePolicy
@@ -33,28 +37,6 @@ from langbot.pkg.workspace.policy import CloudWorkspacePolicy
pytestmark = [pytest.mark.integration, pytest.mark.asyncio]
class _PersistenceManager:
def __init__(self, engine):
self.engine = engine
def get_db_engine(self):
return self.engine
async def execute_async(self, *args, **kwargs):
async with self.engine.connect() as connection:
result = await connection.execute(*args, **kwargs)
await connection.commit()
return result
@staticmethod
def serialize_model(model, row, masked_columns=()):
return {
column.name: getattr(row, column.name)
for column in model.__table__.columns
if column.name not in masked_columns
}
@pytest.fixture
async def workspace_api(tmp_path):
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "workspace-api.db"}')
@@ -62,7 +44,8 @@ async def workspace_api(tmp_path):
await connection.run_sync(Base.metadata.create_all)
application = SimpleNamespace()
application.persistence_mgr = _PersistenceManager(engine)
application.persistence_mgr = PersistenceManager(application)
application.persistence_mgr.db = SimpleNamespace(get_engine=lambda: engine)
application.instance_config = SimpleNamespace(
data={
'system': {
@@ -85,19 +68,27 @@ async def workspace_api(tmp_path):
application.user_service = UserService(application)
application.apikey_service = ApiKeyService(application)
owner = await application.user_service.create_initial_account(
'owner@example.com',
'owner-password',
)
owner_token = await application.user_service.generate_jwt_token(owner)
quart_app = Quart(__name__)
await WorkspacesRouterGroup(application, quart_app).initialize()
await InvitationsRouterGroup(application, quart_app).initialize()
await ApiKeysRouterGroup(application, quart_app).initialize()
await UserRouterGroup(application, quart_app).initialize()
await SystemRouterGroup(application, quart_app).initialize()
yield application, quart_app.test_client(), engine, owner_token
client = quart_app.test_client()
init_response = await client.post(
'/api/v1/user/init',
json={'user': 'owner@example.com', 'password': 'owner-password'},
)
assert init_response.status_code == 200
auth_response = await client.post(
'/api/v1/user/auth',
json={'user': 'owner@example.com', 'password': 'owner-password'},
)
assert auth_response.status_code == 200
owner_token = (await auth_response.get_json())['data']['token']
yield application, client, engine, owner_token
await engine.dispose()
@@ -108,6 +99,56 @@ def _auth(token: str, workspace_uuid: str | None = None) -> dict[str, str]:
return headers
async def test_fresh_sqlite_login_returns_current_workspace_and_user_info(workspace_api):
_, client, _, owner_token = workspace_api
current_response = await client.get('/api/v1/workspaces/current', headers=_auth(owner_token))
assert current_response.status_code == 200
current = (await current_response.get_json())['data']
assert current['workspace']['uuid']
assert current['membership']['email'] == 'owner@example.com'
assert current['membership']['role'] == 'owner'
info_response = await client.get('/api/v1/user/info', headers=_auth(owner_token))
assert info_response.status_code == 200
info = (await info_response.get_json())['data']
assert info['account_uuid'] == current['membership']['account_uuid']
assert info['user'] == 'owner@example.com'
async def test_authenticated_system_info_reads_workspace_wizard_metadata(workspace_api):
application, client, _, owner_token = workspace_api
current_response = await client.get('/api/v1/workspaces/current', headers=_auth(owner_token))
workspace_uuid = (await current_response.get_json())['data']['workspace']['uuid']
progress = {'step': 3, 'selected_adapter': 'telegram'}
await application.persistence_mgr.execute_async(
sqlalchemy.insert(WorkspaceMetadata),
[
{
'workspace_uuid': workspace_uuid,
'key': 'wizard_status',
'value': 'completed',
},
{
'workspace_uuid': workspace_uuid,
'key': 'wizard_progress',
'value': json.dumps(progress),
},
],
)
response = await client.get(
'/api/v1/system/info',
headers=_auth(owner_token, workspace_uuid),
)
assert response.status_code == 200
data = (await response.get_json())['data']
assert data['wizard_status'] == 'completed'
assert data['wizard_progress'] == progress
async def test_owner_invites_second_account_and_secret_is_not_persisted(workspace_api):
application, client, engine, owner_token = workspace_api
@@ -10,6 +10,7 @@ Run: uv run pytest tests/integration/persistence/test_migrations.py -q
from __future__ import annotations
import pytest
import sqlalchemy
from sqlalchemy.ext.asyncio import create_async_engine
from langbot.pkg.entity.persistence.base import Base
@@ -149,6 +150,27 @@ class TestSQLiteMigrationUpgrade:
rev2 = await get_alembic_current(sqlite_engine)
assert rev2 == rev1, f'Expected {rev1}, got {rev2}'
@pytest.mark.asyncio
async def test_upgrade_from_0012_adds_knowledge_base_embedding_dimension(self, sqlite_engine):
"""The PostgreSQL pgvector revision also evolves the OSS ORM schema."""
async with sqlite_engine.begin() as conn:
await conn.exec_driver_sql(
'CREATE TABLE knowledge_bases ('
'uuid VARCHAR(255) PRIMARY KEY, workspace_uuid VARCHAR(36) NOT NULL, name VARCHAR(255) NOT NULL)'
)
await run_alembic_stamp(sqlite_engine, '0012_plugin_identity')
await run_alembic_upgrade(sqlite_engine, 'head')
async with sqlite_engine.connect() as conn:
columns = await conn.run_sync(
lambda sync_conn: {
item['name'] for item in sqlalchemy.inspect(sync_conn).get_columns('knowledge_bases')
}
)
assert 'embedding_dimension' in columns
class TestSQLiteMigrationFreshDatabase:
"""Tests for fresh database workflow."""
@@ -16,19 +16,30 @@ from __future__ import annotations
import logging
import os
import uuid
import asyncio
import contextlib
import datetime
import hashlib
import typing
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
import sqlalchemy as sa
from quart import Quart
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine, create_async_engine
from sqlalchemy import text
from langbot.pkg.entity.persistence.base import Base
from langbot.pkg.entity.persistence.user import User
from langbot.pkg.persistence.mgr import PersistenceManager, PersistenceMode
from langbot.pkg.persistence.tenant_uow import TENANT_POLICY_NAME, TENANT_TABLE_COLUMNS, TenantUnitOfWork
from langbot.pkg.persistence.tenant_uow import (
TENANT_POLICY_NAME,
TENANT_TABLE_COLUMNS,
TenantUnitOfWork,
TransactionRollbackOnlyError,
)
from langbot.pkg.persistence.alembic_runner import (
run_alembic_upgrade,
run_alembic_stamp,
@@ -39,6 +50,32 @@ from alembic.config import Config
from alembic.script import ScriptDirectory
from langbot.pkg.utils import constants
from langbot.pkg.workspace.collaboration import normalize_email
from langbot.pkg.workspace.collaboration import WorkspaceCollaborationService
from langbot.pkg.workspace.policy import CloudWorkspacePolicy
from langbot.pkg.workspace.service import WorkspaceService
from langbot.pkg.api.http.authz import Permission
from langbot.pkg.api.http.controller import group as http_group
from langbot.pkg.api.http.controller.groups.system import SystemRouterGroup
from langbot.pkg.api.http.controller.groups.webhooks import WebhookRouterGroup
from langbot.pkg.api.http.context import ExecutionContext, RequestContext
from langbot.pkg.api.http.service.apikey import ApiKeyService
from langbot.pkg.api.http.service.monitoring import MonitoringService
from langbot.pkg.api.http.service.user import UserService
from langbot.pkg.api.mcp.context import get_request_context as get_mcp_request_context
from langbot.pkg.api.mcp.mount import MCPMount
from langbot.pkg.entity.persistence.apikey import ApiKey
from langbot.pkg.entity.persistence.metadata import WorkspaceMetadata
from langbot.pkg.entity.persistence.monitoring import MonitoringFeedback
from langbot.pkg.entity.persistence.workspace import (
Workspace,
WorkspaceExecutionState,
WorkspaceMembership,
)
from langbot.pkg.platform.botmgr import PlatformManager
from langbot.pkg.pipeline.pipelinemgr import PipelineManager
from langbot.pkg.provider.modelmgr.modelmgr import ModelManager
from langbot.pkg.provider.tools.loaders.mcp import MCPLoader
from langbot.pkg.rag.knowledge.kbmgr import RAGManager
from .resource_migration_support import TENANT_TABLES, create_legacy_resource_schema
@@ -54,6 +91,42 @@ def _get_script_head() -> str:
return ScriptDirectory.from_config(cfg).get_current_head()
async def _grant_runtime_role_business_objects(
conn: AsyncConnection,
role_name: str,
quote: typing.Callable[[str], str],
) -> None:
"""Mirror the release job's object ACLs without overgranting Alembic."""
business_tables = tuple(sorted({table.name for table in Base.metadata.tables.values()} | {'langbot_vectors'}))
quoted_tables = ', '.join(f'public.{quote(table_name)}' for table_name in business_tables)
await conn.execute(text(f'GRANT SELECT, INSERT, UPDATE, DELETE ON TABLE {quoted_tables} TO {quote(role_name)}'))
await conn.execute(text(f'GRANT SELECT ON TABLE public.alembic_version TO {quote(role_name)}'))
sequence_query = text(
"""
SELECT DISTINCT sequence.relname
FROM pg_class sequence
JOIN pg_namespace sequence_namespace ON sequence_namespace.oid = sequence.relnamespace
JOIN pg_depend dependency
ON dependency.classid = 'pg_class'::regclass
AND dependency.objid = sequence.oid
AND dependency.refclassid = 'pg_class'::regclass
AND dependency.deptype IN ('a', 'i')
JOIN pg_class business_table ON business_table.oid = dependency.refobjid
JOIN pg_namespace table_namespace ON table_namespace.oid = business_table.relnamespace
WHERE sequence.relkind = 'S'
AND sequence_namespace.nspname = 'public'
AND table_namespace.nspname = 'public'
AND business_table.relname IN :table_names
ORDER BY sequence.relname
"""
).bindparams(sa.bindparam('table_names', expanding=True))
sequence_names = tuple((await conn.execute(sequence_query, {'table_names': business_tables})).scalars().all())
if sequence_names:
quoted_sequences = ', '.join(f'public.{quote(sequence_name)}' for sequence_name in sequence_names)
await conn.execute(text(f'GRANT USAGE, SELECT ON SEQUENCE {quoted_sequences} TO {quote(role_name)}'))
def _application_for_postgres_url(postgres_url: str, logger_name: str) -> SimpleNamespace:
url = sa.engine.make_url(postgres_url)
return SimpleNamespace(
@@ -115,15 +188,20 @@ async def postgres_engine(postgres_url):
@pytest.fixture
async def clean_tables(postgres_engine):
"""Drop all tables before and after each test for isolation."""
# Drop all tables before test
async with postgres_engine.begin() as conn:
await conn.run_sync(Base.metadata.drop_all)
async def drop_all_tables() -> None:
# Alembic can create tables (notably langbot_vectors) outside the ORM
# metadata, and legacy migration tests intentionally alter constraints.
# Reflect the dedicated test schema instead of relying on stale ORM DDL.
async with postgres_engine.begin() as conn:
table_names = await conn.run_sync(lambda sync_conn: sa.inspect(sync_conn).get_table_names())
quote = postgres_engine.dialect.identifier_preparer.quote
for table_name in table_names:
await conn.execute(text(f'DROP TABLE {quote(table_name)} CASCADE'))
await drop_all_tables()
yield
# Drop all tables after test
async with postgres_engine.begin() as conn:
await conn.run_sync(Base.metadata.drop_all)
await drop_all_tables()
@pytest.fixture
@@ -442,10 +520,16 @@ class TestPostgreSQLResourceTenancyMigration:
await conn.execute(
text(
'INSERT INTO plugin_settings '
'(workspace_uuid, plugin_author, plugin_name, enabled) '
"VALUES (:workspace_uuid, 'author', 'plugin', true)"
'(workspace_uuid, plugin_author, plugin_name, enabled, '
'installation_uuid, artifact_digest, runtime_revision) '
"VALUES (:workspace_uuid, 'author', 'plugin', true, "
':installation_uuid, :artifact_digest, 1)'
),
{'workspace_uuid': second_workspace_uuid},
{
'workspace_uuid': second_workspace_uuid,
'installation_uuid': str(uuid.uuid4()),
'artifact_digest': hashlib.sha256(b'test-plugin-artifact').hexdigest(),
},
)
with pytest.raises(IntegrityError):
@@ -556,10 +640,7 @@ class TestPostgreSQLTenantRuntime:
)
)
await conn.execute(text(f'GRANT USAGE ON SCHEMA public TO {quote(role_name)}'))
await conn.execute(
text(f'GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA public TO {quote(role_name)}')
)
await conn.execute(text(f'GRANT USAGE, SELECT ON ALL SEQUENCES IN SCHEMA public TO {quote(role_name)}'))
await _grant_runtime_role_business_objects(conn, role_name, quote)
created_roles.append(role_name)
try:
@@ -703,6 +784,31 @@ class TestPostgreSQLTenantRuntime:
cloud_manager.create_tables.assert_not_awaited()
cloud_manager._run_alembic_migrations.assert_not_awaited()
with pytest.raises(TransactionRollbackOnlyError, match='after-commit work was cancelled'):
async with cloud_manager.tenant_uow(workspace_a):
duplicate_statement = sa.insert(WorkspaceMetadata).values(
workspace_uuid=workspace_a,
key='rollback-only-unique',
value='must-not-commit',
)
await cloud_manager.execute_async(duplicate_statement)
after_commit_gate = cloud_manager.create_after_commit_gate()
assert after_commit_gate is not None
try:
await cloud_manager.execute_async(duplicate_statement)
except IntegrityError:
pass
assert after_commit_gate.cancelled()
async with cloud_manager.tenant_uow(workspace_a):
assert (
await cloud_manager.execute_async(
sa.select(sa.func.count())
.select_from(WorkspaceMetadata)
.where(WorkspaceMetadata.key == 'rollback-only-unique')
)
).scalar_one() == 0
superuser_manager = PersistenceManager(
_application_for_postgres_url(postgres_url, 'postgres-superuser-runtime-test'),
mode=PersistenceMode.CLOUD_RUNTIME,
@@ -743,3 +849,644 @@ class TestPostgreSQLTenantRuntime:
for role_name in reversed(created_roles):
await conn.execute(text(f'DROP OWNED BY {quote(role_name)}'))
await conn.execute(text(f'DROP ROLE {quote(role_name)}'))
@pytest.mark.asyncio
async def test_cloud_discovery_http_and_transaction_contract(
self,
postgres_url,
postgres_engine,
clean_tables,
clean_alembic_version,
monkeypatch,
):
"""Exercise the complete request/discovery path as a non-owner role."""
instance_uuid = 'cloud-request-rls-test'
other_instance_uuid = 'other-cloud-instance'
workspace_a = '31000000-0000-0000-0000-000000000001'
workspace_b = '32000000-0000-0000-0000-000000000002'
workspace_other = '33000000-0000-0000-0000-000000000003'
workspace_fenced = '34000000-0000-0000-0000-000000000004'
account_a_uuid = '41000000-0000-0000-0000-000000000001'
shared_account_uuid = '42000000-0000-0000-0000-000000000002'
active_secret = 'lbk_active-request-key'
revoked_secret = 'lbk_revoked-request-key'
expired_secret = 'lbk_expired-request-key'
role_name = f'lb_request_{uuid.uuid4().hex[:12]}'
role_password = f'Lb{uuid.uuid4().hex}'
quote = postgres_engine.dialect.identifier_preparer.quote
managers: list[PersistenceManager] = []
role_created = False
_restore_postgres_manager_registry(monkeypatch)
monkeypatch.setattr(constants, 'instance_id', instance_uuid)
release_manager = PersistenceManager(
_application_for_postgres_url(postgres_url, 'postgres-request-release-test'),
mode=PersistenceMode.RELEASE_MIGRATION,
)
managers.append(release_manager)
def role_url() -> str:
return (
sa.engine.make_url(postgres_url)
.set(username=role_name, password=role_password)
.render_as_string(hide_password=False)
)
async def seed_workspace(
workspace_uuid: str,
*,
target_instance: str,
state: str = 'active',
write_fenced: bool = False,
) -> None:
async with release_manager.tenant_uow(workspace_uuid) as uow:
uow.session.add(
Workspace(
uuid=workspace_uuid,
instance_uuid=target_instance,
name=workspace_uuid[-4:],
slug=f'workspace-{workspace_uuid[-4:]}',
type='team',
status='active',
source='cloud_projection',
projection_revision=1,
)
)
await uow.session.flush()
uow.session.add(
WorkspaceExecutionState(
workspace_uuid=workspace_uuid,
instance_uuid=target_instance,
active_generation=1,
state=state,
write_fenced=write_fenced,
source='cloud',
desired_state_revision=1,
)
)
uow.session.add(
WorkspaceMetadata(
workspace_uuid=workspace_uuid,
key='tenant-marker',
value=workspace_uuid,
)
)
try:
await release_manager.initialize()
async with release_manager.get_db_engine().begin() as conn:
await conn.execute(
sa.insert(User),
[
{
'uuid': account_a_uuid,
'user': 'account-a@example.com',
'normalized_email': 'account-a@example.com',
'password': 'closed-directory',
'account_type': 'local',
'status': 'active',
'source': 'cloud_projection',
'projection_revision': 1,
},
{
'uuid': shared_account_uuid,
'user': 'shared@example.com',
'normalized_email': 'shared@example.com',
'password': 'closed-directory',
'account_type': 'local',
'status': 'active',
'source': 'cloud_projection',
'projection_revision': 1,
},
],
)
await seed_workspace(workspace_a, target_instance=instance_uuid)
await seed_workspace(workspace_b, target_instance=instance_uuid)
await seed_workspace(workspace_other, target_instance=other_instance_uuid)
await seed_workspace(
workspace_fenced,
target_instance=instance_uuid,
write_fenced=True,
)
async with release_manager.tenant_uow(workspace_a) as uow:
uow.session.add_all(
[
WorkspaceMembership(
uuid=str(uuid.uuid4()),
workspace_uuid=workspace_a,
account_uuid=account_a_uuid,
role='owner',
status='active',
projection_revision=1,
),
WorkspaceMembership(
uuid=str(uuid.uuid4()),
workspace_uuid=workspace_a,
account_uuid=shared_account_uuid,
role='viewer',
status='active',
projection_revision=1,
),
ApiKey(
uuid=str(uuid.uuid4()),
workspace_uuid=workspace_a,
name='active-key',
key_hash=hashlib.sha256(active_secret.encode()).hexdigest(),
scopes=[Permission.WORKSPACE_VIEW.value],
status='active',
),
]
)
async with release_manager.tenant_uow(workspace_b) as uow:
uow.session.add_all(
[
WorkspaceMembership(
uuid=str(uuid.uuid4()),
workspace_uuid=workspace_b,
account_uuid=shared_account_uuid,
role='viewer',
status='active',
projection_revision=1,
),
ApiKey(
uuid=str(uuid.uuid4()),
workspace_uuid=workspace_b,
name='revoked-key',
key_hash=hashlib.sha256(revoked_secret.encode()).hexdigest(),
scopes=[Permission.WORKSPACE_VIEW.value],
status='revoked',
),
ApiKey(
uuid=str(uuid.uuid4()),
workspace_uuid=workspace_b,
name='expired-key',
key_hash=hashlib.sha256(expired_secret.encode()).hexdigest(),
scopes=[Permission.WORKSPACE_VIEW.value],
status='active',
expires_at=datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
- datetime.timedelta(minutes=1),
),
]
)
async with postgres_engine.connect() as conn:
await conn.execute(text(f"CREATE ROLE {quote(role_name)} LOGIN PASSWORD '{role_password}'"))
role_created = True
await conn.execute(
text(
f'GRANT CONNECT ON DATABASE {quote(sa.engine.make_url(postgres_url).database)} '
f'TO {quote(role_name)}'
)
)
await conn.execute(text(f'GRANT USAGE ON SCHEMA public TO {quote(role_name)}'))
await _grant_runtime_role_business_objects(conn, role_name, quote)
runtime_application = _application_for_postgres_url(role_url(), 'postgres-request-runtime-test')
runtime_application.instance_config.data.update(
{
'system': {
'jwt': {'secret': 'postgres-request-jwt', 'expire': 3600},
},
'api': {'global_api_key': ''},
}
)
runtime_application.logger = logging.getLogger('postgres-request-runtime-test')
cloud_manager = PersistenceManager(runtime_application, mode=PersistenceMode.CLOUD_RUNTIME)
managers.append(cloud_manager)
await cloud_manager.initialize()
runtime_application.persistence_mgr = cloud_manager
cloud_manager.ap = runtime_application
runtime_application.workspace_service = WorkspaceService(
runtime_application,
policy=CloudWorkspacePolicy(),
instance_uuid=instance_uuid,
)
runtime_application.workspace_collaboration_service = WorkspaceCollaborationService(
runtime_application,
runtime_application.workspace_service,
policy=CloudWorkspacePolicy(),
)
runtime_application.user_service = UserService(runtime_application)
runtime_application.apikey_service = ApiKeyService(runtime_application)
runtime_application.monitoring_service = MonitoringService(runtime_application)
bindings = await runtime_application.workspace_service.list_active_execution_bindings()
assert {binding.workspace_uuid for binding in bindings} == {workspace_a, workspace_b}
# No scope is an application error before SQL reaches PostgreSQL.
with pytest.raises(RuntimeError, match='explicit Workspace or discovery'):
await cloud_manager.execute_async(sa.select(WorkspaceMetadata))
# Discovery exposes only its index rows and cannot write.
with pytest.raises(TransactionRollbackOnlyError, match='transaction was rolled back'):
async with cloud_manager.account_discovery_uow(shared_account_uuid) as discovery:
assert set(
(
await discovery.session.scalars(
sa.select(WorkspaceMembership.workspace_uuid).order_by(
WorkspaceMembership.workspace_uuid
)
)
).all()
) == {workspace_a, workspace_b}
with pytest.raises(sa.exc.DBAPIError):
await discovery.session.execute(
sa.insert(WorkspaceMembership).values(
uuid=str(uuid.uuid4()),
workspace_uuid=workspace_a,
account_uuid=shared_account_uuid,
role='viewer',
status='active',
projection_revision=1,
)
)
async with cloud_manager.api_key_discovery_uow(
hashlib.sha256(active_secret.encode()).hexdigest()
) as discovery:
assert await discovery.session.scalar(sa.select(ApiKey.workspace_uuid)) == workspace_a
update_result = await discovery.session.execute(
sa.update(ApiKey)
.where(ApiKey.key_hash == hashlib.sha256(active_secret.encode()).hexdigest())
.values(name='discovery-must-not-write')
)
assert update_result.rowcount == 0
async with cloud_manager.api_key_discovery_uow(
hashlib.sha256(revoked_secret.encode()).hexdigest()
) as discovery:
assert await discovery.session.scalar(sa.select(ApiKey.workspace_uuid)) is None
async with cloud_manager.api_key_discovery_uow(
hashlib.sha256(expired_secret.encode()).hexdigest()
) as discovery:
assert await discovery.session.scalar(sa.select(ApiKey.workspace_uuid)) is None
async with cloud_manager.instance_discovery_uow(instance_uuid) as discovery:
assert set(
(await discovery.session.scalars(sa.select(WorkspaceExecutionState.workspace_uuid))).all()
) == {workspace_a, workspace_b}
update_result = await discovery.session.execute(
sa.update(WorkspaceExecutionState)
.where(WorkspaceExecutionState.workspace_uuid == workspace_a)
.values(write_fenced=True)
)
assert update_result.rowcount == 0
# Instance discovery deliberately cannot see business rows.
assert (await discovery.session.execute(sa.select(WorkspaceMetadata))).all() == []
platform_manager = PlatformManager(runtime_application)
platform_manager._load_workspace_bots = AsyncMock()
await platform_manager.load_bots_from_db()
assert {call.args[0] for call in platform_manager._load_workspace_bots.await_args_list} == {
workspace_a,
workspace_b,
}
# Every startup cache loader traverses the instance index first,
# then reads business resources in one tenant transaction at a time.
await ModelManager(runtime_application).load_models_from_db()
await PipelineManager(runtime_application).load_pipelines_from_db()
await MCPLoader(runtime_application).load_mcp_servers_from_db()
await RAGManager(runtime_application).load_knowledge_bases_from_db()
# An omitted Workspace predicate remains isolated by RLS.
async with cloud_manager.tenant_uow(workspace_a):
values = (await cloud_manager.execute_async(sa.select(WorkspaceMetadata.value))).scalars().all()
assert values == [workspace_a]
async def read_tenant_repeatedly(workspace_uuid: str) -> list[str]:
observed: list[str] = []
for _ in range(5):
async with cloud_manager.tenant_uow(workspace_uuid):
observed.extend(
(await cloud_manager.execute_async(sa.select(WorkspaceMetadata.value))).scalars().all()
)
await asyncio.sleep(0)
return observed
observed_a, observed_b = await asyncio.gather(
read_tenant_repeatedly(workspace_a),
read_tenant_repeatedly(workspace_b),
)
assert observed_a == [workspace_a] * 5
assert observed_b == [workspace_b] * 5
accesses = await runtime_application.workspace_collaboration_service.list_account_workspaces(
shared_account_uuid
)
assert {access.workspace.uuid for access in accesses} == {workspace_a, workspace_b}
assert await runtime_application.apikey_service.authenticate_api_key(revoked_secret) is None
assert await runtime_application.apikey_service.authenticate_api_key(expired_secret) is None
active_identity = await runtime_application.apikey_service.authenticate_api_key(active_secret)
assert active_identity is not None
assert active_identity.workspace_uuid == workspace_a
async def record_feedback(feedback_type: int) -> str | None:
async with cloud_manager.tenant_scope(workspace_a):
return await runtime_application.monitoring_service.record_feedback(
ExecutionContext(
instance_uuid=instance_uuid,
workspace_uuid=workspace_a,
placement_generation=1,
),
feedback_id='concurrent-feedback',
feedback_type=feedback_type,
)
feedback_ids = await asyncio.gather(record_feedback(1), record_feedback(2))
assert feedback_ids[0] == feedback_ids[1]
async with cloud_manager.tenant_uow(workspace_a) as uow:
feedback_rows = (
(
await uow.execute(
sa.select(MonitoringFeedback).where(MonitoringFeedback.feedback_id == 'concurrent-feedback')
)
)
.scalars()
.all()
)
assert len(feedback_rows) == 1
assert feedback_rows[0].feedback_type in {1, 2}
await uow.execute(
sa.delete(MonitoringFeedback).where(MonitoringFeedback.feedback_id == 'concurrent-feedback')
)
class TenantRuntimeRouter(http_group.RouterGroup):
name = 'postgres-tenant-runtime'
path = '/tenant-runtime'
async def initialize(self) -> None:
@self.route('/account', permission=Permission.WORKSPACE_VIEW)
async def account_route(request_context: RequestContext):
assert self.ap.persistence_mgr.current_session() is None
if self.quart_app.config.get('FORCE_HANDLER_FAILURE'):
async with self.ap.persistence_mgr.tenant_uow(request_context.workspace_uuid):
await self.ap.persistence_mgr.execute_async(
sa.insert(WorkspaceMetadata).values(
workspace_uuid=request_context.workspace_uuid,
key='rolled-back-handler',
value='must-not-commit',
)
)
raise RuntimeError('forced handler failure')
values = (
(await self.ap.persistence_mgr.execute_async(sa.select(WorkspaceMetadata.value)))
.scalars()
.all()
)
assert self.ap.persistence_mgr.current_session() is None
return self.success(data={'workspace_uuid': request_context.workspace_uuid, 'values': values})
@self.route('/key', auth_type=http_group.AuthType.API_KEY)
async def key_route(request_context: RequestContext):
values = (
(await self.ap.persistence_mgr.execute_async(sa.select(WorkspaceMetadata.value)))
.scalars()
.all()
)
return self.success(data={'workspace_uuid': request_context.workspace_uuid, 'values': values})
@self.route('/bootstrap', auth_type=http_group.AuthType.ACCOUNT_TOKEN)
async def bootstrap_route(user_email: str):
account = await self.ap.user_service.get_user_by_email(user_email)
accesses = await self.ap.workspace_collaboration_service.list_account_workspaces(account.uuid)
return self.success(data=sorted(access.workspace.uuid for access in accesses))
quart_app = Quart(__name__)
await TenantRuntimeRouter(runtime_application, quart_app).initialize()
webhook_bot_uuid = str(uuid.uuid4())
class TenantAwareWebhookAdapter:
async def handle_unified_webhook(self, **_kwargs):
assert cloud_manager.current_session() is None
await asyncio.sleep(0)
assert cloud_manager.current_session() is None
values = (
(
await cloud_manager.execute_async(
sa.select(WorkspaceMetadata.value).where(WorkspaceMetadata.key == 'tenant-marker')
)
)
.scalars()
.all()
)
assert cloud_manager.current_session() is None
return {'values': values}
runtime_application.platform_mgr = SimpleNamespace(
resolve_public_bot=AsyncMock(
return_value=SimpleNamespace(
workspace_uuid=workspace_a,
placement_generation=1,
enable=True,
adapter=TenantAwareWebhookAdapter(),
)
)
)
await WebhookRouterGroup(runtime_application, quart_app).initialize()
await SystemRouterGroup(runtime_application, quart_app).initialize()
client = quart_app.test_client()
account_a = await runtime_application.user_service.get_user_by_uuid(account_a_uuid)
shared_account = await runtime_application.user_service.get_user_by_uuid(shared_account_uuid)
assert account_a is not None and shared_account is not None
account_token = await runtime_application.user_service.generate_jwt_token(account_a)
shared_token = await runtime_application.user_service.generate_jwt_token(shared_account)
async with cloud_manager.tenant_uow(workspace_a):
await cloud_manager.execute_async(
sa.insert(WorkspaceMetadata).values(
workspace_uuid=workspace_a,
key='wizard_status',
value='completed',
)
)
response = await client.get(
'/api/v1/system/info',
headers={
'Authorization': f'Bearer {account_token}',
'X-Workspace-Id': workspace_a,
},
)
assert response.status_code == 200
assert (await response.get_json())['data']['wizard_status'] == 'completed'
async with cloud_manager.tenant_uow(workspace_a):
await cloud_manager.execute_async(
sa.delete(WorkspaceMetadata).where(
WorkspaceMetadata.workspace_uuid == workspace_a,
WorkspaceMetadata.key == 'wizard_status',
)
)
response = await client.post(f'/bots/{webhook_bot_uuid}')
assert response.status_code == 200
assert await response.get_json() == {'values': [workspace_a]}
response = await client.get(
'/tenant-runtime/account',
headers={
'Authorization': f'Bearer {account_token}',
'X-Workspace-Id': workspace_a,
},
)
assert response.status_code == 200
assert (await response.get_json())['data'] == {
'workspace_uuid': workspace_a,
'values': [workspace_a],
}
response = await client.get(
'/tenant-runtime/account',
headers={
'Authorization': f'Bearer {account_token}',
'X-Workspace-Id': workspace_b,
},
)
assert response.status_code == 404
response = await client.get(
'/tenant-runtime/account',
headers={'Authorization': f'Bearer {shared_token}'},
)
assert response.status_code == 404
response = await client.get(
'/tenant-runtime/bootstrap',
headers={'Authorization': f'Bearer {shared_token}'},
)
assert response.status_code == 200
assert (await response.get_json())['data'] == [workspace_a, workspace_b]
response = await client.get(
'/tenant-runtime/key',
headers={'X-API-Key': active_secret, 'X-Workspace-Id': workspace_b},
)
assert response.status_code == 200
assert (await response.get_json())['data'] == {
'workspace_uuid': workspace_a,
'values': [workspace_a],
}
# The parallel MCP ASGI entrypoint authenticates the same key and
# retains only a trusted Workspace scope for the full tool request.
# Each DB call gets its own short RLS transaction.
mcp_observation: dict[str, typing.Any] = {}
async def fake_mcp_asgi(scope, receive, send):
del scope, receive
context = get_mcp_request_context()
assert cloud_manager.current_session() is None
values = (await cloud_manager.execute_async(sa.select(WorkspaceMetadata.value))).scalars().all()
assert cloud_manager.current_session() is None
await asyncio.sleep(0)
assert cloud_manager.current_session() is None
repeated_values = (
(await cloud_manager.execute_async(sa.select(WorkspaceMetadata.value))).scalars().all()
)
assert repeated_values == values
assert cloud_manager.current_session() is None
mcp_observation.update(
workspace_uuid=context.workspace_uuid,
values=values,
)
await send({'type': 'http.response.start', 'status': 200, 'headers': []})
await send({'type': 'http.response.body', 'body': b'{}'})
async def unused_quart_asgi(scope, receive, send): # pragma: no cover - routing assertion
del scope, receive, send
raise AssertionError('MCP request was routed to Quart')
mount = MCPMount.__new__(MCPMount)
mount.ap = runtime_application
mount._mcp_asgi = fake_mcp_asgi
sent_messages: list[dict[str, typing.Any]] = []
async def receive():
return {'type': 'http.request', 'body': b'', 'more_body': False}
async def send(message):
sent_messages.append(message)
await mount.wrap(unused_quart_asgi)(
{
'type': 'http',
'path': '/mcp',
'headers': [
(b'x-api-key', active_secret.encode()),
(b'x-workspace-id', workspace_b.encode()),
],
},
receive,
send,
)
assert sent_messages[0]['status'] == 200
assert mcp_observation == {'workspace_uuid': workspace_a, 'values': [workspace_a]}
original_api_key_discovery = cloud_manager.api_key_discovery_uow
@contextlib.asynccontextmanager
async def revoke_after_discovery(key_hash: str):
async with original_api_key_discovery(key_hash) as discovery:
yield discovery
async with cloud_manager.tenant_uow(workspace_a):
await cloud_manager.execute_async(
sa.update(ApiKey).where(ApiKey.key_hash == key_hash).values(status='revoked')
)
monkeypatch.setattr(cloud_manager, 'api_key_discovery_uow', revoke_after_discovery)
assert await runtime_application.apikey_service.authenticate_api_key(active_secret) is None
monkeypatch.setattr(cloud_manager, 'api_key_discovery_uow', original_api_key_discovery)
quart_app.config['FORCE_HANDLER_FAILURE'] = True
response = await client.get(
'/tenant-runtime/account',
headers={
'Authorization': f'Bearer {account_token}',
'X-Workspace-Id': workspace_a,
},
)
assert response.status_code == 500
quart_app.config['FORCE_HANDLER_FAILURE'] = False
async with cloud_manager.tenant_uow(workspace_a):
assert (
await cloud_manager.execute_async(
sa.select(sa.func.count())
.select_from(WorkspaceMetadata)
.where(WorkspaceMetadata.key == 'rolled-back-handler')
)
).scalar_one() == 0
# Transaction-local settings are gone when pooled connections are reused.
async with cloud_manager.get_db_engine().connect() as conn:
assert await conn.scalar(text("SELECT current_setting('langbot.workspace_uuid', true)")) in (
None,
'',
)
assert await conn.scalar(text('SELECT COUNT(*) FROM workspace_metadata')) == 0
# Runtime validation rejects both extra permissive policies and a
# modified expression even if the expected policy name remains.
async with postgres_engine.connect() as conn:
await conn.execute(text('CREATE POLICY injected_policy ON bots FOR SELECT USING (true)'))
with pytest.raises(RuntimeError, match='policy set does not match'):
await cloud_manager._validate_postgres_tenant_schema(validate_runtime_role=True)
async with postgres_engine.connect() as conn:
await conn.execute(text('DROP POLICY injected_policy ON bots'))
await conn.execute(text('DROP POLICY langbot_workspace_isolation ON workspace_metadata'))
await conn.execute(
text(
'CREATE POLICY langbot_workspace_isolation ON workspace_metadata '
'FOR ALL TO PUBLIC USING (true) WITH CHECK (true)'
)
)
with pytest.raises(RuntimeError, match='policy definitions are invalid'):
await cloud_manager._validate_postgres_tenant_schema(validate_runtime_role=True)
finally:
for manager in reversed(managers):
await _dispose_manager(manager)
if role_created:
async with postgres_engine.connect() as conn:
await conn.execute(text(f'DROP OWNED BY {quote(role_name)}'))
await conn.execute(text(f'DROP ROLE IF EXISTS {quote(role_name)}'))
@@ -0,0 +1,546 @@
"""Real PostgreSQL/pgvector tenant isolation and CRUD verification."""
from __future__ import annotations
import logging
import os
import uuid
from types import SimpleNamespace
import pytest
import sqlalchemy as sa
from sqlalchemy import text
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
from langbot.pkg.entity.persistence.base import Base
from langbot.pkg.persistence.alembic_runner import get_alembic_current, run_alembic_stamp, run_alembic_upgrade
from langbot.pkg.persistence.mgr import PersistenceManager, PersistenceMode
from langbot.pkg.utils import constants
from langbot.pkg.vector.vdbs.pgvector_db import PgVectorDatabase, PgVectorScope
pytestmark = [pytest.mark.integration, pytest.mark.slow, pytest.mark.asyncio]
def _application(postgres_url: str, logger_name: str) -> SimpleNamespace:
url = sa.engine.make_url(postgres_url)
return SimpleNamespace(
instance_config=SimpleNamespace(
data={
'database': {
'use': 'postgresql',
'postgresql': {
'host': url.host,
'port': url.port,
'user': url.username,
'password': url.password,
'database': url.database,
},
}
}
),
logger=logging.getLogger(logger_name),
)
def _restore_postgres_registry(monkeypatch) -> None:
from langbot.pkg.persistence import mgr as persistence_mgr_module
from langbot.pkg.persistence.databases.postgresql import PostgreSQLDatabaseManager
monkeypatch.setattr(
persistence_mgr_module.database,
'preregistered_managers',
[PostgreSQLDatabaseManager],
)
@pytest.fixture
def postgres_url() -> str:
url = os.environ.get('TEST_POSTGRES_URL')
if not url:
pytest.skip('TEST_POSTGRES_URL not set')
return url
@pytest.fixture
async def postgres_engine(postgres_url: str):
engine = create_async_engine(postgres_url, isolation_level='AUTOCOMMIT')
yield engine
await engine.dispose()
@pytest.fixture
async def clean_database(postgres_engine: AsyncEngine):
async def clean() -> None:
async with postgres_engine.begin() as conn:
await conn.execute(text('DROP TABLE IF EXISTS langbot_vectors_legacy_0013 CASCADE'))
await conn.execute(text('DROP TABLE IF EXISTS langbot_vectors CASCADE'))
await conn.run_sync(Base.metadata.drop_all)
await conn.execute(text('DROP TABLE IF EXISTS alembic_version'))
await clean()
yield
await clean()
async def test_legacy_upgrade_temporarily_suspends_and_restores_source_rls_for_unprivileged_owner(
postgres_url: str,
postgres_engine: AsyncEngine,
clean_database,
) -> None:
workspace_uuid = '30000000-0000-0000-0000-000000000303'
knowledge_base_uuid = 'legacy-knowledge-base'
migrator_role = f'lb_vector_migrator_{uuid.uuid4().hex[:12]}'
migrator_password = f'Lb{uuid.uuid4().hex}'
quote = postgres_engine.dialect.identifier_preparer.quote
database_name = sa.engine.make_url(postgres_url).database
migrator_url = (
sa.engine.make_url(postgres_url)
.set(username=migrator_role, password=migrator_password)
.render_as_string(hide_password=False)
)
migrator_engine: AsyncEngine | None = None
role_created = False
admin_role: str | None = None
source_tables = ('knowledge_bases', 'knowledge_base_files', 'knowledge_base_chunks')
async def read_rls_states(conn, table_names: tuple[str, ...]) -> dict[str, tuple[bool, bool]]:
rows = (
(
await conn.execute(
text(
"""
SELECT c.relname, c.relrowsecurity, c.relforcerowsecurity
FROM pg_class AS c
JOIN pg_namespace AS n ON n.oid = c.relnamespace
WHERE n.nspname = current_schema()
AND c.relname IN :table_names
"""
).bindparams(sa.bindparam('table_names', expanding=True)),
{'table_names': table_names},
)
)
.mappings()
.all()
)
return {str(row['relname']): (bool(row['relrowsecurity']), bool(row['relforcerowsecurity'])) for row in rows}
try:
async with postgres_engine.begin() as conn:
admin_role = await conn.scalar(text('SELECT current_user'))
await conn.execute(text('CREATE EXTENSION IF NOT EXISTS vector'))
await conn.run_sync(Base.metadata.create_all)
await conn.execute(
text(
'ALTER TABLE knowledge_bases DROP CONSTRAINT IF EXISTS '
'ck_knowledge_bases_embedding_dimension_positive'
)
)
await conn.execute(text('ALTER TABLE knowledge_bases DROP COLUMN IF EXISTS embedding_dimension'))
await run_alembic_stamp(postgres_engine, '0010_scope_resources')
await run_alembic_upgrade(postgres_engine, '0012_plugin_identity')
assert await get_alembic_current(postgres_engine) == '0012_plugin_identity'
embedding_a = '[' + ','.join(['0.125'] * 384) + ']'
embedding_b = '[' + ','.join(['0.25'] * 384) + ']'
async with postgres_engine.begin() as conn:
await conn.execute(
text(
"""
INSERT INTO workspaces
(uuid, instance_uuid, name, slug, type, status, source, projection_revision)
VALUES
(:uuid, 'legacy-vector-instance', 'legacy', 'legacy-vector',
'team', 'active', 'cloud_projection', 0)
"""
),
{'uuid': workspace_uuid},
)
await conn.execute(
text(
"""
INSERT INTO knowledge_bases
(uuid, workspace_uuid, name, collection_id, legacy_vector_collection)
VALUES
(:uuid, :workspace_uuid, 'legacy', 'legacy-collection', true)
"""
),
{'uuid': knowledge_base_uuid, 'workspace_uuid': workspace_uuid},
)
await conn.execute(
text(
"""
INSERT INTO knowledge_base_files
(uuid, workspace_uuid, kb_id, file_name, extension, status)
VALUES
('legacy-file', :workspace_uuid, :kb_uuid, 'legacy.txt', 'txt', 'completed')
"""
),
{'workspace_uuid': workspace_uuid, 'kb_uuid': knowledge_base_uuid},
)
await conn.execute(
text(
"""
INSERT INTO knowledge_base_chunks (uuid, workspace_uuid, file_id, text)
VALUES ('legacy-chunk', :workspace_uuid, 'legacy-file', 'chunk text')
"""
),
{'workspace_uuid': workspace_uuid},
)
await conn.execute(
text(
"""
CREATE TABLE langbot_vectors (
id VARCHAR(255) PRIMARY KEY,
collection VARCHAR(255),
embedding vector NOT NULL,
text TEXT,
file_id VARCHAR(255),
chunk_uuid VARCHAR(255)
)
"""
)
)
await conn.execute(
text(
"""
INSERT INTO langbot_vectors
(id, collection, embedding, text, file_id, chunk_uuid)
VALUES
('legacy-by-collection', 'legacy-collection', CAST(:embedding_a AS vector),
'collection row', NULL, NULL),
('legacy-by-chunk', 'unmatched-collection', CAST(:embedding_b AS vector),
'chunk row', NULL, 'legacy-chunk')
"""
),
{'embedding_a': embedding_a, 'embedding_b': embedding_b},
)
# Exercise exact restoration rather than assuming all source tables
# arrived with identical flags.
await conn.execute(text('ALTER TABLE knowledge_base_files NO FORCE ROW LEVEL SECURITY'))
await conn.execute(text('ALTER TABLE knowledge_base_chunks NO FORCE ROW LEVEL SECURITY'))
await conn.execute(text('ALTER TABLE knowledge_base_chunks DISABLE ROW LEVEL SECURITY'))
expected_source_rls = await read_rls_states(conn, source_tables)
assert expected_source_rls == {
'knowledge_bases': (True, True),
'knowledge_base_files': (True, False),
'knowledge_base_chunks': (False, False),
}
await conn.execute(
text(f"CREATE ROLE {quote(migrator_role)} LOGIN PASSWORD '{migrator_password}' NOSUPERUSER NOBYPASSRLS")
)
role_created = True
await conn.execute(text(f'GRANT CONNECT ON DATABASE {quote(database_name)} TO {quote(migrator_role)}'))
await conn.execute(text(f'GRANT USAGE, CREATE ON SCHEMA public TO {quote(migrator_role)}'))
await conn.execute(text(f'GRANT SELECT, UPDATE ON alembic_version TO {quote(migrator_role)}'))
for table_name in (*source_tables, 'langbot_vectors'):
await conn.execute(text(f'ALTER TABLE {quote(table_name)} OWNER TO {quote(migrator_role)}'))
role = (
await conn.execute(
text('SELECT rolsuper, rolbypassrls FROM pg_roles WHERE rolname = :role'),
{'role': migrator_role},
)
).one()
assert role == (False, False)
owned_tables = set(
(
await conn.execute(
text(
"""
SELECT c.relname
FROM pg_class AS c
JOIN pg_namespace AS n ON n.oid = c.relnamespace
WHERE n.nspname = current_schema()
AND c.relname IN :table_names
AND pg_get_userbyid(c.relowner) = :role
"""
).bindparams(sa.bindparam('table_names', expanding=True)),
{'table_names': (*source_tables, 'langbot_vectors'), 'role': migrator_role},
)
).scalars()
)
assert owned_tables == {*source_tables, 'langbot_vectors'}
migrator_engine = create_async_engine(migrator_url)
await run_alembic_upgrade(migrator_engine, '0013_tenant_pgvector')
assert await get_alembic_current(migrator_engine) == '0013_tenant_pgvector'
async with postgres_engine.connect() as conn:
migrated_rows = (
(
await conn.execute(
text(
"""
SELECT workspace_uuid, knowledge_base_uuid, vector_id,
embedding_dimension, text, file_id, chunk_uuid
FROM langbot_vectors
ORDER BY vector_id
"""
)
)
)
.mappings()
.all()
)
assert [dict(row) for row in migrated_rows] == [
{
'workspace_uuid': workspace_uuid,
'knowledge_base_uuid': knowledge_base_uuid,
'vector_id': 'legacy-by-chunk',
'embedding_dimension': 384,
'text': 'chunk row',
'file_id': None,
'chunk_uuid': 'legacy-chunk',
},
{
'workspace_uuid': workspace_uuid,
'knowledge_base_uuid': knowledge_base_uuid,
'vector_id': 'legacy-by-collection',
'embedding_dimension': 384,
'text': 'collection row',
'file_id': None,
'chunk_uuid': None,
},
]
assert (
await conn.scalar(
text('SELECT embedding_dimension FROM knowledge_bases WHERE uuid = :uuid'),
{'uuid': knowledge_base_uuid},
)
== 384
)
assert await conn.scalar(text("SELECT to_regclass('langbot_vectors_legacy_0013') IS NULL")) is True
assert await read_rls_states(conn, source_tables) == expected_source_rls
assert await read_rls_states(conn, ('langbot_vectors',)) == {'langbot_vectors': (True, True)}
assert (
await conn.scalar(
text(
"""
SELECT COUNT(*)
FROM pg_policy AS p
JOIN pg_class AS c ON c.oid = p.polrelid
JOIN pg_namespace AS n ON n.oid = c.relnamespace
WHERE n.nspname = current_schema()
AND c.relname = 'langbot_vectors'
AND p.polname = 'langbot_workspace_isolation'
"""
)
)
== 1
)
async with migrator_engine.connect() as conn:
assert await conn.scalar(text('SELECT COUNT(*) FROM langbot_vectors')) == 0
async with migrator_engine.begin() as conn:
await conn.execute(
text("SELECT set_config('langbot.workspace_uuid', :workspace_uuid, true)"),
{'workspace_uuid': workspace_uuid},
)
assert await conn.scalar(text('SELECT COUNT(*) FROM langbot_vectors')) == 2
finally:
if migrator_engine is not None:
await migrator_engine.dispose()
if role_created:
async with postgres_engine.connect() as conn:
if admin_role is None: # pragma: no cover - setup cannot create the role without an admin
admin_role = await conn.scalar(text('SELECT current_user'))
await conn.execute(text(f'REASSIGN OWNED BY {quote(migrator_role)} TO {quote(admin_role)}'))
await conn.execute(text(f'DROP OWNED BY {quote(migrator_role)}'))
await conn.execute(text(f'DROP ROLE IF EXISTS {quote(migrator_role)}'))
async def test_pgvector_shared_database_is_scoped_indexed_and_ddl_free_at_runtime(
postgres_url: str,
postgres_engine: AsyncEngine,
clean_database,
monkeypatch,
) -> None:
instance_uuid = 'pgvector-tenant-integration'
workspace_a = '10000000-0000-0000-0000-000000000101'
workspace_b = '20000000-0000-0000-0000-000000000202'
kb_a = 'knowledge-base-a'
kb_b = 'knowledge-base-b'
role_suffix = uuid.uuid4().hex[:12]
runtime_role = f'lb_vector_runtime_{role_suffix}'
role_password = f'Lb{uuid.uuid4().hex}'
release_manager: PersistenceManager | None = None
runtime_manager: PersistenceManager | None = None
role_created = False
_restore_postgres_registry(monkeypatch)
monkeypatch.setattr(constants, 'instance_id', instance_uuid)
quote = postgres_engine.dialect.identifier_preparer.quote
database_name = sa.engine.make_url(postgres_url).database
runtime_url = (
sa.engine.make_url(postgres_url)
.set(username=runtime_role, password=role_password)
.render_as_string(hide_password=False)
)
try:
release_app = _application(postgres_url, 'pgvector-release-migration-test')
release_manager = PersistenceManager(release_app, mode=PersistenceMode.RELEASE_MIGRATION)
release_app.persistence_mgr = release_manager
await release_manager.initialize()
for workspace_uuid, kb_uuid, slug in (
(workspace_a, kb_a, 'vector-a'),
(workspace_b, kb_b, 'vector-b'),
):
async with release_manager.tenant_uow(workspace_uuid) as uow:
await uow.execute(
text(
"""
INSERT INTO workspaces
(uuid, instance_uuid, name, slug, type, status, source, projection_revision)
VALUES
(:uuid, :instance_uuid, :name, :slug, 'team', 'active', 'cloud_projection', 0)
"""
),
{
'uuid': workspace_uuid,
'instance_uuid': instance_uuid,
'name': slug,
'slug': slug,
},
)
await uow.execute(
text(
"""
INSERT INTO knowledge_bases
(uuid, workspace_uuid, name, embedding_dimension)
VALUES (:uuid, :workspace_uuid, :name, 384)
"""
),
{'uuid': kb_uuid, 'workspace_uuid': workspace_uuid, 'name': slug},
)
async with postgres_engine.connect() as conn:
await conn.execute(text(f"CREATE ROLE {quote(runtime_role)} LOGIN PASSWORD '{role_password}'"))
await conn.execute(text(f'GRANT CONNECT ON DATABASE {quote(database_name)} TO {quote(runtime_role)}'))
await conn.execute(text(f'GRANT USAGE ON SCHEMA public TO {quote(runtime_role)}'))
business_tables = release_manager._runtime_business_table_names()
quoted_tables = ', '.join(f'public.{quote(table_name)}' for table_name in business_tables)
await conn.execute(
text(f'GRANT SELECT, INSERT, UPDATE, DELETE ON TABLE {quoted_tables} TO {quote(runtime_role)}')
)
await conn.execute(text(f'GRANT SELECT ON TABLE public.alembic_version TO {quote(runtime_role)}'))
sequence_names = await release_manager._runtime_business_sequence_names(conn, business_tables)
if sequence_names:
quoted_sequences = ', '.join(f'public.{quote(sequence_name)}' for sequence_name in sequence_names)
await conn.execute(text(f'GRANT USAGE, SELECT ON SEQUENCE {quoted_sequences} TO {quote(runtime_role)}'))
role_created = True
runtime_app = _application(runtime_url, 'pgvector-runtime-test')
runtime_manager = PersistenceManager(runtime_app, mode=PersistenceMode.CLOUD_RUNTIME)
runtime_app.persistence_mgr = runtime_manager
await runtime_manager.initialize()
adapter = PgVectorDatabase(
runtime_app,
use_business_database=True,
allowed_dimensions=[384],
)
scope_a = PgVectorScope(workspace_a, kb_a, 384)
scope_b = PgVectorScope(workspace_b, kb_b, 384)
# The same vector ID is valid in two Workspaces because the relational
# primary key includes Workspace and knowledge base.
await adapter.add_embeddings(
'opaque-a',
['same-vector'],
[[0.1] * 384],
[{'text': 'workspace-a', 'file_id': 'file-a', 'uuid': 'chunk-a'}],
scope=scope_a,
)
await adapter.add_embeddings(
'opaque-b',
['same-vector'],
[[0.2] * 384],
[{'text': 'workspace-b', 'file_id': 'file-b', 'uuid': 'chunk-b'}],
scope=scope_b,
)
result_a = await adapter.search('opaque-a', [0.1] * 384, scope=scope_a)
result_b = await adapter.search('opaque-b', [0.2] * 384, scope=scope_b)
assert result_a['metadatas'][0][0]['text'] == 'workspace-a'
assert result_b['metadatas'][0][0]['text'] == 'workspace-b'
# Guessing another knowledge-base UUID while retaining A's Workspace
# cannot escape either the explicit conditions or PostgreSQL RLS.
guessed = await adapter.search(
'attacker-controlled-name',
[0.2] * 384,
scope=PgVectorScope(workspace_a, kb_b, 384),
)
assert guessed['ids'] == [[]]
with pytest.raises(ValueError, match='trusted PgVectorScope'):
await adapter.search('opaque-a', [0.1] * 384)
with pytest.raises(ValueError, match='selected dimension'):
await adapter.add_embeddings(
'opaque-a',
['bad-dimension'],
[[0.1] * 383],
[{}],
scope=scope_a,
)
# Deliberately omit the application Workspace predicate. FORCE RLS is
# still the second isolation boundary and returns only A.
async with runtime_manager.tenant_uow(workspace_a) as uow:
rows = (
await uow.execute(
text("SELECT workspace_uuid, vector_id FROM langbot_vectors WHERE vector_id = 'same-vector'")
)
).all()
assert rows == [(workspace_a, 'same-vector')]
await uow.execute(text('SET LOCAL enable_seqscan = off'))
plan = '\n'.join(
(
await uow.execute(
text(
"""
EXPLAIN SELECT vector_id
FROM langbot_vectors
WHERE workspace_uuid = :workspace_uuid
AND knowledge_base_uuid = :knowledge_base_uuid
AND embedding_dimension = 384
ORDER BY (embedding::vector(384)) <=> CAST(:query AS vector(384))
LIMIT 5
"""
),
{
'workspace_uuid': workspace_a,
'knowledge_base_uuid': kb_a,
'query': '[' + ','.join(['0.1'] * 384) + ']',
},
)
).scalars()
)
assert 'ix_langbot_vectors_hnsw_cosine_384' in plan
runtime_engine = runtime_manager.get_db_engine()
async with runtime_engine.connect() as conn:
assert await conn.scalar(text('SELECT COUNT(*) FROM langbot_vectors')) == 0
assert await conn.scalar(text("SELECT current_setting('langbot.workspace_uuid', true)")) in (None, '')
items_a, total_a = await adapter.list_by_filter('opaque-a', scope=scope_a)
assert total_a == 1
assert items_a[0]['metadata']['file_id'] == 'file-a'
await adapter.delete_by_file_id('opaque-a', 'file-a', scope=scope_a)
assert (await adapter.list_by_filter('opaque-a', scope=scope_a))[1] == 0
assert (await adapter.list_by_filter('opaque-b', scope=scope_b))[1] == 1
finally:
if runtime_manager is not None and getattr(runtime_manager, 'db', None) is not None:
await runtime_manager.get_db_engine().dispose()
if release_manager is not None and getattr(release_manager, 'db', None) is not None:
await release_manager.get_db_engine().dispose()
if role_created:
async with postgres_engine.connect() as conn:
await conn.execute(text(f'DROP OWNED BY {quote(runtime_role)}'))
await conn.execute(text(f'DROP ROLE IF EXISTS {quote(runtime_role)}'))
@@ -0,0 +1,94 @@
from __future__ import annotations
import hashlib
import uuid
import pytest
import sqlalchemy as sa
from sqlalchemy.ext.asyncio import create_async_engine
from langbot.pkg.persistence.alembic_runner import (
get_alembic_current,
run_alembic_stamp,
run_alembic_upgrade,
)
pytestmark = [pytest.mark.integration, pytest.mark.asyncio]
async def test_legacy_plugin_settings_receive_stable_random_installation_identities(tmp_path):
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "plugin-identity.db"}')
metadata = sa.MetaData()
plugin_settings = sa.Table(
'plugin_settings',
metadata,
sa.Column('workspace_uuid', sa.String(36), primary_key=True),
sa.Column('plugin_author', sa.String(255), primary_key=True),
sa.Column('plugin_name', sa.String(255), primary_key=True),
sa.Column('enabled', sa.Boolean, nullable=False, server_default=sa.true()),
sa.Column('priority', sa.Integer, nullable=False, server_default='0'),
sa.Column('config', sa.JSON, nullable=False, server_default='{}'),
sa.Column('install_source', sa.String(255), nullable=False, server_default='local'),
sa.Column('install_info', sa.JSON, nullable=False, server_default='{}'),
)
try:
async with engine.begin() as connection:
await connection.run_sync(metadata.create_all)
await connection.execute(
plugin_settings.insert(),
[
{
'workspace_uuid': '11111111-1111-4111-8111-111111111111',
'plugin_author': 'author',
'plugin_name': 'one',
},
{
'workspace_uuid': '22222222-2222-4222-8222-222222222222',
'plugin_author': 'author',
'plugin_name': 'two',
},
],
)
await run_alembic_stamp(engine, '0011_postgres_tenant_rls')
await run_alembic_upgrade(engine, '0012_plugin_identity')
async with engine.connect() as connection:
rows = (
(
await connection.execute(
sa.text(
'SELECT installation_uuid, artifact_digest, runtime_revision '
'FROM plugin_settings ORDER BY workspace_uuid'
)
)
)
.mappings()
.all()
)
columns = await connection.run_sync(
lambda sync_connection: {
column['name']: column for column in sa.inspect(sync_connection).get_columns('plugin_settings')
}
)
indexes = await connection.run_sync(
lambda sync_connection: {
index['name']: index for index in sa.inspect(sync_connection).get_indexes('plugin_settings')
}
)
assert await get_alembic_current(engine) == '0012_plugin_identity'
assert columns['installation_uuid']['nullable'] is False
assert columns['artifact_digest']['nullable'] is False
assert columns['runtime_revision']['nullable'] is False
assert indexes['ix_plugin_settings_workspace_installation']['unique'] == 1
assert len({row['installation_uuid'] for row in rows}) == 2
for row in rows:
uuid.UUID(row['installation_uuid'])
assert row['runtime_revision'] == 1
assert (
row['artifact_digest']
== hashlib.sha256(f'legacy-installation:{row["installation_uuid"]}'.encode()).hexdigest()
)
finally:
await engine.dispose()
@@ -0,0 +1,723 @@
"""Real PostgreSQL coverage for the one-shot Cloud release migration job."""
from __future__ import annotations
import logging
import os
import uuid
from types import SimpleNamespace
import pytest
import sqlalchemy as sa
from sqlalchemy import text
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
from langbot.pkg.persistence import release_migration
from langbot.pkg.persistence.alembic_runner import (
get_alembic_current,
get_alembic_head,
run_alembic_stamp,
run_alembic_upgrade,
)
from langbot.pkg.persistence.mgr import (
PersistenceManager,
PersistenceMode,
_RELEASE_MIGRATION_ADVISORY_LOCK_ID,
)
from langbot.pkg.utils import constants
pytestmark = [pytest.mark.integration, pytest.mark.slow, pytest.mark.asyncio]
_RUNTIME_PASSWORD = 'runtime-secret-not-used-by-migration'
@pytest.fixture
def postgres_url() -> str:
url = os.environ.get('TEST_POSTGRES_URL')
if not url:
pytest.skip('TEST_POSTGRES_URL not set')
return url
@pytest.fixture
async def postgres_engine(postgres_url: str):
engine = create_async_engine(postgres_url, isolation_level='AUTOCOMMIT')
yield engine
await engine.dispose()
@pytest.fixture
async def clean_database(postgres_engine: AsyncEngine):
async def clean() -> None:
async with postgres_engine.begin() as conn:
table_names = await conn.run_sync(lambda sync_conn: sa.inspect(sync_conn).get_table_names())
quote = postgres_engine.dialect.identifier_preparer.quote
for table_name in table_names:
await conn.execute(text(f'DROP TABLE {quote(table_name)} CASCADE'))
await clean()
yield
await clean()
def _restore_postgres_registry(monkeypatch) -> None:
from langbot.pkg.persistence import mgr as persistence_mgr_module
from langbot.pkg.persistence.databases.postgresql import PostgreSQLDatabaseManager
monkeypatch.setattr(
persistence_mgr_module.database,
'preregistered_managers',
[PostgreSQLDatabaseManager],
)
def _application(postgres_url: str, *, runtime_role: str = 'langbot_runtime_not_used_by_migration') -> SimpleNamespace:
url = sa.engine.make_url(postgres_url)
return SimpleNamespace(
instance_config=SimpleNamespace(
data={
'database': {
'use': 'postgresql',
'postgresql': {
'host': url.host,
'port': url.port,
# This is deliberately not the operator role in the DSN.
'user': runtime_role,
'password': _RUNTIME_PASSWORD,
'database': url.database,
},
'cloud_migration': {'operator_dsn_env': 'TEST_RELEASE_OPERATOR_DSN'},
},
'vdb': {
'use': 'pgvector',
'pgvector': {
'use_business_database': True,
'allowed_dimensions': [384, 512, 768, 1024, 1536],
},
},
}
),
logger=logging.getLogger('cloud-release-migration-entrypoint-test'),
persistence_mgr=None,
)
async def test_release_entrypoint_holds_lock_migrates_validates_and_disposes(
postgres_url: str,
postgres_engine: AsyncEngine,
clean_database,
monkeypatch,
) -> None:
_restore_postgres_registry(monkeypatch)
monkeypatch.setattr(constants, 'instance_id', 'release-migration-entrypoint-test')
original_validate = PersistenceManager._validate_release_schema
validation_observed_lock = False
runtime_role = f'lb_release_runtime_{uuid.uuid4().hex[:12]}'
quote = postgres_engine.dialect.identifier_preparer.quote
async def validate_while_asserting_lock(self: PersistenceManager) -> None:
nonlocal validation_observed_lock
async with postgres_engine.connect() as conn:
acquired = await conn.scalar(
text('SELECT pg_try_advisory_lock(:lock_id)'),
{'lock_id': _RELEASE_MIGRATION_ADVISORY_LOCK_ID},
)
if acquired:
await conn.scalar(
text('SELECT pg_advisory_unlock(:lock_id)'),
{'lock_id': _RELEASE_MIGRATION_ADVISORY_LOCK_ID},
)
assert acquired is False
validation_observed_lock = True
await original_validate(self)
monkeypatch.setattr(PersistenceManager, '_validate_release_schema', validate_while_asserting_lock)
async with postgres_engine.connect() as conn:
await conn.execute(text(f"CREATE ROLE {quote(runtime_role)} LOGIN PASSWORD '{_RUNTIME_PASSWORD}'"))
# The release job must make this otherwise bare LOGIN usable without
# relying on pre-provisioned object ACLs.
assert (
await conn.scalar(
text(
"""
SELECT NOT EXISTS (
SELECT 1
FROM pg_class c
JOIN pg_namespace n ON n.oid = c.relnamespace
CROSS JOIN LATERAL aclexplode(c.relacl) acl
JOIN pg_roles grantee ON grantee.oid = acl.grantee
WHERE n.nspname = current_schema()
AND grantee.rolname = :runtime_role
)
"""
),
{'runtime_role': runtime_role},
)
is True
)
ap = _application(postgres_url, runtime_role=runtime_role)
try:
await release_migration.run_cloud_release_migration(
ap,
environ={'TEST_RELEASE_OPERATOR_DSN': postgres_url},
)
assert validation_observed_lock is True
assert await get_alembic_current(postgres_engine) == get_alembic_head()
manager = ap.persistence_mgr
assert isinstance(manager, PersistenceManager)
business_tables = set(manager._runtime_business_table_names())
async with postgres_engine.connect() as conn:
assert await conn.scalar(text("SELECT to_regclass('langbot_vectors') IS NOT NULL")) is True
assert (
await conn.scalar(text("SELECT EXISTS (SELECT 1 FROM pg_extension WHERE extname = 'vector')")) is True
)
runtime_role_state = (
(
await conn.execute(
text(
"""
SELECT
rolcanlogin,
rolsuper,
rolbypassrls,
rolcreatedb,
rolcreaterole,
rolreplication
FROM pg_roles
WHERE rolname = :runtime_role
"""
),
{'runtime_role': runtime_role},
)
)
.mappings()
.one()
)
assert dict(runtime_role_state) == {
'rolcanlogin': True,
'rolsuper': False,
'rolbypassrls': False,
'rolcreatedb': False,
'rolcreaterole': False,
'rolreplication': False,
}
database_privileges = set(
(
await conn.execute(
text(
"""
SELECT acl.privilege_type
FROM pg_database database
CROSS JOIN LATERAL aclexplode(database.datacl) acl
JOIN pg_roles grantee ON grantee.oid = acl.grantee
WHERE database.datname = current_database()
AND grantee.rolname = :runtime_role
AND acl.is_grantable IS FALSE
"""
),
{'runtime_role': runtime_role},
)
)
.scalars()
.all()
)
schema_privileges = set(
(
await conn.execute(
text(
"""
SELECT acl.privilege_type
FROM pg_namespace namespace
CROSS JOIN LATERAL aclexplode(namespace.nspacl) acl
JOIN pg_roles grantee ON grantee.oid = acl.grantee
WHERE namespace.nspname = current_schema()
AND grantee.rolname = :runtime_role
AND acl.is_grantable IS FALSE
"""
),
{'runtime_role': runtime_role},
)
)
.scalars()
.all()
)
assert database_privileges == {'CONNECT'}
assert schema_privileges == {'USAGE'}
assert (
await conn.scalar(
text("SELECT has_database_privilege(:runtime_role, current_database(), 'CREATE')"),
{'runtime_role': runtime_role},
)
is False
)
assert (
await conn.scalar(
text("SELECT has_schema_privilege(:runtime_role, current_schema(), 'CREATE')"),
{'runtime_role': runtime_role},
)
is False
)
# PostgreSQL grants TEMP to PUBLIC by default. The first Cloud
# release deliberately tolerates that inherited compatibility
# privilege while granting no direct TEMP ACL to the runtime role.
assert (
await conn.scalar(
text("SELECT has_database_privilege(:runtime_role, current_database(), 'TEMP')"),
{'runtime_role': runtime_role},
)
is True
)
object_grants = (
(
await conn.execute(
text(
"""
SELECT
c.relname,
c.relkind::text AS relkind,
acl.privilege_type,
acl.is_grantable
FROM pg_class c
JOIN pg_namespace n ON n.oid = c.relnamespace
CROSS JOIN LATERAL aclexplode(c.relacl) acl
JOIN pg_roles grantee ON grantee.oid = acl.grantee
WHERE n.nspname = current_schema()
AND c.relkind IN ('r', 'p', 'S')
AND grantee.rolname = :runtime_role
ORDER BY c.relname, acl.privilege_type
"""
),
{'runtime_role': runtime_role},
)
)
.mappings()
.all()
)
direct_table_grants: dict[str, set[str]] = {}
direct_sequence_grants: dict[str, set[str]] = {}
for grant in object_grants:
assert grant['is_grantable'] is False
target = direct_sequence_grants if grant['relkind'] == 'S' else direct_table_grants
target.setdefault(grant['relname'], set()).add(grant['privilege_type'])
assert set(direct_table_grants) == business_tables | {'alembic_version'}
assert all(
privileges == {'SELECT', 'INSERT', 'UPDATE', 'DELETE'}
for table_name, privileges in direct_table_grants.items()
if table_name != 'alembic_version'
)
assert direct_table_grants['alembic_version'] == {'SELECT'}
assert direct_sequence_grants
assert all(privileges == {'USAGE', 'SELECT'} for privileges in direct_sequence_grants.values())
# The session-level lock must be released before the one-shot job exits.
assert (
await conn.scalar(
text('SELECT pg_try_advisory_lock(:lock_id)'),
{'lock_id': _RELEASE_MIGRATION_ADVISORY_LOCK_ID},
)
is True
)
assert (
await conn.scalar(
text('SELECT pg_advisory_unlock(:lock_id)'),
{'lock_id': _RELEASE_MIGRATION_ADVISORY_LOCK_ID},
)
is True
)
runtime_url = (
sa.engine.make_url(postgres_url)
.set(username=runtime_role, password=_RUNTIME_PASSWORD)
.render_as_string(hide_password=False)
)
runtime_engine = create_async_engine(runtime_url)
try:
async with runtime_engine.begin() as conn:
assert await conn.scalar(text('SELECT current_user')) == runtime_role
await conn.execute(text("INSERT INTO metadata (key, value) VALUES ('runtime-grant-smoke', 'created')"))
await conn.execute(text("UPDATE metadata SET value = 'updated' WHERE key = 'runtime-grant-smoke'"))
assert (
await conn.scalar(text("SELECT value FROM metadata WHERE key = 'runtime-grant-smoke'")) == 'updated'
)
await conn.execute(text("DELETE FROM metadata WHERE key = 'runtime-grant-smoke'"))
await conn.scalar(
text('SELECT nextval(CAST(:sequence_name AS regclass))'),
{'sequence_name': f'public.{sorted(direct_sequence_grants)[0]}'},
)
finally:
await runtime_engine.dispose()
runtime_application = _application(postgres_url, runtime_role=runtime_role)
runtime_manager = PersistenceManager(runtime_application, mode=PersistenceMode.CLOUD_RUNTIME)
runtime_application.persistence_mgr = runtime_manager
try:
# This reads alembic_version as the actual runtime role and then
# reruns the complete grant/catalog validator before startup.
await runtime_manager.initialize()
finally:
await runtime_manager.get_db_engine().dispose()
# The catalog validator must fail closed if the role is no longer
# deployable, even though the remaining grants still look plausible.
async with postgres_engine.connect() as conn:
await conn.execute(text(f'REVOKE DELETE ON TABLE public.metadata FROM {quote(runtime_role)}'))
with pytest.raises(RuntimeError, match="table 'metadata' grants are incomplete"):
await manager._validate_configured_runtime_postgres_role(require_grants=True)
finally:
async with postgres_engine.connect() as conn:
await conn.execute(text(f'DROP OWNED BY {quote(runtime_role)}'))
await conn.execute(text(f'DROP ROLE IF EXISTS {quote(runtime_role)}'))
async def test_release_entrypoint_rejects_privileged_or_table_owning_runtime_role(
postgres_url: str,
postgres_engine: AsyncEngine,
clean_database,
monkeypatch,
) -> None:
_restore_postgres_registry(monkeypatch)
monkeypatch.setattr(constants, 'instance_id', 'release-runtime-role-validation-test')
runtime_role = f'lb_release_runtime_{uuid.uuid4().hex[:12]}'
quote = postgres_engine.dialect.identifier_preparer.quote
async with postgres_engine.connect() as conn:
operator_role = await conn.scalar(text('SELECT current_user'))
await conn.execute(text(f'CREATE ROLE {quote(runtime_role)} LOGIN'))
ap = _application(postgres_url, runtime_role=runtime_role)
try:
await release_migration.run_cloud_release_migration(
ap,
environ={'TEST_RELEASE_OPERATOR_DSN': postgres_url},
)
async with postgres_engine.connect() as conn:
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} SUPERUSER'))
with pytest.raises(RuntimeError, match='must not be superuser or BYPASSRLS'):
await release_migration.run_cloud_release_migration(
ap,
environ={'TEST_RELEASE_OPERATOR_DSN': postgres_url},
)
async with postgres_engine.connect() as conn:
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} NOSUPERUSER BYPASSRLS'))
with pytest.raises(RuntimeError, match='must not be superuser or BYPASSRLS'):
await release_migration.run_cloud_release_migration(
ap,
environ={'TEST_RELEASE_OPERATOR_DSN': postgres_url},
)
async with postgres_engine.connect() as conn:
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} NOBYPASSRLS'))
await conn.execute(text(f'ALTER TABLE bots OWNER TO {quote(runtime_role)}'))
with pytest.raises(RuntimeError, match='owns tenant tables'):
await release_migration.run_cloud_release_migration(
ap,
environ={'TEST_RELEASE_OPERATOR_DSN': postgres_url},
)
finally:
async with postgres_engine.connect() as conn:
table_exists = await conn.scalar(text("SELECT to_regclass('bots') IS NOT NULL"))
if table_exists:
await conn.execute(text(f'ALTER TABLE bots OWNER TO {quote(operator_role)}'))
await conn.execute(text(f'DROP OWNED BY {quote(runtime_role)}'))
await conn.execute(text(f'DROP ROLE IF EXISTS {quote(runtime_role)}'))
async def test_runtime_role_catalog_validator_rejects_delegation_and_escape_hatches(
postgres_url: str,
postgres_engine: AsyncEngine,
clean_database,
monkeypatch,
) -> None:
_restore_postgres_registry(monkeypatch)
monkeypatch.setattr(constants, 'instance_id', 'release-runtime-role-catalog-test')
suffix = uuid.uuid4().hex[:12]
runtime_role = f'lb_release_runtime_{suffix}'
delegated_role = f'lb_release_delegate_{suffix}'
extra_schema = f'lb_release_schema_{suffix}'
extra_view = f'lb_release_view_{suffix}'
owned_routine = f'lb_release_owned_routine_{suffix}'
security_definer = f'lb_release_definer_{suffix}'
foreign_wrapper = f'lb_release_fdw_{suffix}'
foreign_server = f'lb_release_server_{suffix}'
persistent_setting_secret = f'lb_release_secret_{suffix}'
database_name = sa.engine.make_url(postgres_url).database
assert database_name
quote = postgres_engine.dialect.identifier_preparer.quote
async with postgres_engine.connect() as conn:
await conn.execute(text(f"CREATE ROLE {quote(runtime_role)} LOGIN PASSWORD '{_RUNTIME_PASSWORD}'"))
await conn.execute(text(f'CREATE ROLE {quote(delegated_role)}'))
ap = _application(postgres_url, runtime_role=runtime_role)
try:
await release_migration.run_cloud_release_migration(
ap,
environ={'TEST_RELEASE_OPERATOR_DSN': postgres_url},
)
manager = ap.persistence_mgr
assert isinstance(manager, PersistenceManager)
async with postgres_engine.connect() as conn:
await conn.execute(text(f'GRANT pg_read_all_data TO {quote(runtime_role)}'))
with pytest.raises(RuntimeError, match='must not participate in role memberships'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'REVOKE pg_read_all_data FROM {quote(runtime_role)}'))
# Reject delegation in the other direction too: no role may be
# allowed to SET ROLE to the runtime identity or administer it.
await conn.execute(text(f'GRANT {quote(runtime_role)} TO {quote(delegated_role)} WITH ADMIN OPTION'))
with pytest.raises(RuntimeError, match='must not participate in role memberships'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'REVOKE {quote(runtime_role)} FROM {quote(delegated_role)}'))
await conn.execute(
text(f'GRANT SELECT ON TABLE public.metadata TO {quote(runtime_role)} WITH GRANT OPTION')
)
with pytest.raises(RuntimeError, match='GRANT OPTION'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(
text(f'REVOKE GRANT OPTION FOR SELECT ON TABLE public.metadata FROM {quote(runtime_role)}')
)
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} SET search_path TO public'))
with pytest.raises(RuntimeError, match='persistent session overrides'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} RESET search_path'))
await conn.execute(text(f'ALTER DATABASE {quote(database_name)} SET search_path TO public'))
with pytest.raises(RuntimeError, match='persistent session overrides'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'ALTER DATABASE {quote(database_name)} RESET search_path'))
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} SET session_replication_role TO replica'))
with pytest.raises(RuntimeError, match='persistent session overrides'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} RESET session_replication_role'))
await conn.execute(
text(f"ALTER ROLE {quote(runtime_role)} SET application_name TO '{persistent_setting_secret}'")
)
with pytest.raises(RuntimeError, match='persistent session overrides') as error:
await manager._validate_configured_runtime_postgres_role()
assert persistent_setting_secret not in str(error.value)
async with postgres_engine.connect() as conn:
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} RESET application_name'))
await conn.execute(text('CREATE EXTENSION dblink'))
with pytest.raises(RuntimeError, match='extensions must include vector and be limited'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text('DROP EXTENSION dblink'))
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'GRANT CREATE ON DATABASE {quote(database_name)} TO {quote(runtime_role)}'))
await conn.execute(text(f'GRANT CREATE ON SCHEMA public TO {quote(runtime_role)}'))
runtime_url = (
sa.engine.make_url(postgres_url)
.set(username=runtime_role, password=_RUNTIME_PASSWORD)
.render_as_string(hide_password=False)
)
runtime_engine = create_async_engine(runtime_url)
try:
async with runtime_engine.begin() as conn:
# hstore is a trusted extension in the production PG16 image,
# so this creates a real runtime-owned extension catalog row.
await conn.execute(text('CREATE EXTENSION hstore'))
finally:
await runtime_engine.dispose()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'REVOKE CREATE ON SCHEMA public FROM {quote(runtime_role)}'))
await conn.execute(text(f'REVOKE CREATE ON DATABASE {quote(database_name)} FROM {quote(runtime_role)}'))
with pytest.raises(RuntimeError, match='must not own extensions'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text('DROP EXTENSION hstore CASCADE'))
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'CREATE FOREIGN DATA WRAPPER {quote(foreign_wrapper)}'))
await conn.execute(
text(f'CREATE SERVER {quote(foreign_server)} FOREIGN DATA WRAPPER {quote(foreign_wrapper)}')
)
await conn.execute(text(f'CREATE USER MAPPING FOR {quote(runtime_role)} SERVER {quote(foreign_server)}'))
with pytest.raises(RuntimeError, match='foreign data wrappers, servers, or user mappings'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'DROP FOREIGN DATA WRAPPER {quote(foreign_wrapper)} CASCADE'))
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'CREATE SCHEMA {quote(extra_schema)} AUTHORIZATION {quote(runtime_role)}'))
with pytest.raises(RuntimeError, match='non-business schemas'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'DROP SCHEMA {quote(extra_schema)} CASCADE'))
await conn.execute(text(f'CREATE VIEW public.{quote(extra_view)} AS SELECT key FROM public.metadata'))
await conn.execute(text(f'GRANT SELECT ON public.{quote(extra_view)} TO {quote(runtime_role)}'))
with pytest.raises(RuntimeError, match='non-business objects|table privileges are unsafe'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'DROP VIEW public.{quote(extra_view)}'))
await conn.execute(text(f'GRANT SELECT (key) ON public.metadata TO {quote(runtime_role)}'))
with pytest.raises(RuntimeError, match='column-level ACLs'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'REVOKE SELECT (key) ON public.metadata FROM {quote(runtime_role)}'))
# System file access functions are not SECURITY DEFINER, so the
# validator must reject their explicit EXECUTE ACL independently.
await conn.execute(
text(f'GRANT EXECUTE ON FUNCTION pg_catalog.pg_read_file(text) TO {quote(runtime_role)}')
)
runtime_application = _application(postgres_url, runtime_role=runtime_role)
runtime_manager = PersistenceManager(runtime_application, mode=PersistenceMode.CLOUD_RUNTIME)
runtime_application.persistence_mgr = runtime_manager
try:
with pytest.raises(RuntimeError, match='explicit EXECUTE privileges on routines'):
await runtime_manager.initialize()
finally:
await runtime_manager.get_db_engine().dispose()
async with postgres_engine.connect() as conn:
await conn.execute(
text(f'REVOKE EXECUTE ON FUNCTION pg_catalog.pg_read_file(text) FROM {quote(runtime_role)}')
)
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
# replica disables ordinary triggers/rules and foreign-key
# enforcement; it must never reach the runtime identity.
await conn.execute(text(f'GRANT SET ON PARAMETER session_replication_role TO {quote(runtime_role)}'))
with pytest.raises(RuntimeError, match='explicit SET or ALTER SYSTEM parameter privileges'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'REVOKE SET ON PARAMETER session_replication_role FROM {quote(runtime_role)}'))
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(
text(f"CREATE FUNCTION public.{quote(owned_routine)}() RETURNS integer LANGUAGE sql AS 'SELECT 1'")
)
await conn.execute(text(f'ALTER FUNCTION public.{quote(owned_routine)}() OWNER TO {quote(runtime_role)}'))
with pytest.raises(RuntimeError, match='must not own routines'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'DROP FUNCTION public.{quote(owned_routine)}()'))
async with postgres_engine.connect() as conn:
await conn.execute(
text(
f'CREATE FUNCTION public.{quote(security_definer)}() RETURNS integer '
"LANGUAGE sql SECURITY DEFINER AS 'SELECT 1'"
)
)
# Extension membership must not exempt an executable definer from
# the runtime audit, even for an allowlisted extension.
await conn.execute(text(f'ALTER EXTENSION vector ADD FUNCTION public.{quote(security_definer)}()'))
with pytest.raises(RuntimeError, match='SECURITY DEFINER'):
await manager._validate_configured_runtime_postgres_role()
async with postgres_engine.connect() as conn:
await conn.execute(text(f'ALTER EXTENSION vector DROP FUNCTION public.{quote(security_definer)}()'))
finally:
async with postgres_engine.connect() as conn:
await conn.execute(text(f'ALTER DATABASE {quote(database_name)} RESET search_path'))
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} RESET search_path'))
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} RESET session_replication_role'))
await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} RESET application_name'))
await conn.execute(text(f'REVOKE pg_read_all_data FROM {quote(runtime_role)}'))
await conn.execute(text(f'REVOKE {quote(runtime_role)} FROM {quote(delegated_role)}'))
await conn.execute(text(f'REVOKE CREATE ON SCHEMA public FROM {quote(runtime_role)}'))
await conn.execute(text(f'REVOKE CREATE ON DATABASE {quote(database_name)} FROM {quote(runtime_role)}'))
await conn.execute(
text(f'REVOKE EXECUTE ON FUNCTION pg_catalog.pg_read_file(text) FROM {quote(runtime_role)}')
)
await conn.execute(text(f'REVOKE SET ON PARAMETER session_replication_role FROM {quote(runtime_role)}'))
await conn.execute(text('DROP EXTENSION IF EXISTS dblink CASCADE'))
await conn.execute(text('DROP EXTENSION IF EXISTS hstore CASCADE'))
await conn.execute(text(f'DROP FOREIGN DATA WRAPPER IF EXISTS {quote(foreign_wrapper)} CASCADE'))
await conn.execute(text(f'DROP FUNCTION IF EXISTS public.{quote(owned_routine)}()'))
security_definer_is_extension_member = await conn.scalar(
text(
"""
SELECT EXISTS (
SELECT 1
FROM pg_depend dependency
JOIN pg_extension extension ON extension.oid = dependency.refobjid
WHERE dependency.classid = 'pg_proc'::regclass
AND dependency.objid = to_regprocedure(:routine)
AND dependency.refclassid = 'pg_extension'::regclass
AND dependency.deptype = 'e'
AND extension.extname = 'vector'
)
"""
),
{'routine': f'public.{security_definer}()'},
)
if security_definer_is_extension_member:
await conn.execute(text(f'ALTER EXTENSION vector DROP FUNCTION public.{quote(security_definer)}()'))
await conn.execute(text(f'DROP FUNCTION IF EXISTS public.{quote(security_definer)}()'))
await conn.execute(text(f'DROP VIEW IF EXISTS public.{quote(extra_view)}'))
await conn.execute(text(f'DROP SCHEMA IF EXISTS {quote(extra_schema)} CASCADE'))
await conn.execute(text(f'DROP OWNED BY {quote(runtime_role)}'))
await conn.execute(text(f'DROP ROLE IF EXISTS {quote(delegated_role)}'))
await conn.execute(text(f'DROP ROLE IF EXISTS {quote(runtime_role)}'))
async def test_direct_postgres_head_stamp_fails_without_business_schema(
postgres_engine: AsyncEngine,
clean_database,
) -> None:
await run_alembic_stamp(postgres_engine, '0012_plugin_identity')
with pytest.raises(RuntimeError, match='requires the knowledge_bases table'):
await run_alembic_upgrade(postgres_engine)
assert await get_alembic_current(postgres_engine) == '0012_plugin_identity'
async def test_release_entrypoint_fails_immediately_when_another_job_holds_lock(
postgres_url: str,
postgres_engine: AsyncEngine,
clean_database,
monkeypatch,
) -> None:
_restore_postgres_registry(monkeypatch)
ap = _application(postgres_url)
async with postgres_engine.connect() as lock_connection:
assert (
await lock_connection.scalar(
text('SELECT pg_try_advisory_lock(:lock_id)'),
{'lock_id': _RELEASE_MIGRATION_ADVISORY_LOCK_ID},
)
is True
)
try:
with pytest.raises(RuntimeError, match='already holds the advisory lock'):
await release_migration.run_cloud_release_migration(
ap,
environ={'TEST_RELEASE_OPERATOR_DSN': postgres_url},
)
finally:
assert (
await lock_connection.scalar(
text('SELECT pg_advisory_unlock(:lock_id)'),
{'lock_id': _RELEASE_MIGRATION_ADVISORY_LOCK_ID},
)
is True
)
async with postgres_engine.connect() as conn:
assert await conn.run_sync(lambda sync_conn: sa.inspect(sync_conn).get_table_names()) == []
@@ -195,10 +195,16 @@ async def test_sqlite_scoped_keys_allow_cross_workspace_but_reject_same_workspac
await conn.execute(
sa.text(
'INSERT INTO plugin_settings '
'(workspace_uuid, plugin_author, plugin_name, enabled) '
"VALUES (:workspace_uuid, 'author', 'plugin', 1)"
'(workspace_uuid, plugin_author, plugin_name, enabled, '
'installation_uuid, artifact_digest, runtime_revision) '
"VALUES (:workspace_uuid, 'author', 'plugin', 1, "
':installation_uuid, :artifact_digest, 1)'
),
{'workspace_uuid': second_workspace_uuid},
{
'workspace_uuid': second_workspace_uuid,
'installation_uuid': str(uuid.uuid4()),
'artifact_digest': hashlib.sha256(b'test-plugin-artifact').hexdigest(),
},
)
await conn.execute(
sa.text(
@@ -0,0 +1,415 @@
from __future__ import annotations
import asyncio
import datetime as dt
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
import pytest
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
from langbot_plugin.box.backend import BaseSandboxBackend
from langbot_plugin.box.client import ActionRPCBoxClient
from langbot_plugin.box.errors import BoxAdmissionError
from langbot_plugin.box.models import (
BoxExecutionResult,
BoxExecutionStatus,
BoxNetworkMode,
BoxSessionInfo,
BoxSpec,
)
from langbot_plugin.box.runtime import BoxRuntime
from langbot_plugin.box.server import BoxServerHandler
from langbot_plugin.runtime.io.handler import Handler
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.box.service import BoxService
from langbot.pkg.cloud.entitlements import (
EntitlementResolver,
EntitlementSnapshot,
EntitlementUnavailableError,
)
pytestmark = pytest.mark.integration
_UTC = dt.timezone.utc
class _AdmissionBackend(BaseSandboxBackend):
name = 'nsjail'
def __init__(self, logger):
super().__init__(logger)
self.started_specs: list[BoxSpec] = []
self.stopped_sessions: list[str] = []
async def is_available(self) -> bool:
return True
async def get_readiness(self, *, workspace_path=None, strict=False) -> dict:
return {
'available': True,
'cgroup_v2': True,
'namespace_isolation': True,
'mount_isolation': True,
'network_isolation': True,
'hard_workspace_quota': True,
'hard_skill_storage_quota': True,
'bounded_ephemeral_storage': True,
'inode_quota': True,
}
async def start_session(self, spec: BoxSpec) -> BoxSessionInfo:
self.started_specs.append(spec)
now = dt.datetime.now(_UTC)
return BoxSessionInfo(
session_id=spec.session_id,
backend_name=self.name,
backend_session_id=f'jail-{len(self.started_specs)}',
image=spec.image,
network=spec.network,
host_path=spec.host_path,
host_path_mode=spec.host_path_mode,
mount_path=spec.mount_path,
persistent=spec.persistent,
cpus=spec.cpus,
memory_mb=spec.memory_mb,
pids_limit=spec.pids_limit,
read_only_rootfs=spec.read_only_rootfs,
workspace_quota_mb=spec.workspace_quota_mb,
created_at=now,
last_used_at=now,
)
async def exec(self, session: BoxSessionInfo, spec: BoxSpec) -> BoxExecutionResult:
await asyncio.sleep(0)
return BoxExecutionResult(
session_id=session.session_id,
backend_name=self.name,
status=BoxExecutionStatus.COMPLETED,
exit_code=0,
stdout=spec.cmd,
stderr='',
duration_ms=1,
)
async def stop_session(self, session: BoxSessionInfo):
self.stopped_sessions.append(session.session_id)
class _QueueConnection:
def __init__(self, rx: asyncio.Queue[str], tx: asyncio.Queue[str]):
self._rx = rx
self._tx = tx
async def send(self, message: str) -> None:
await self._tx.put(message)
async def receive(self) -> str:
return await self._rx.get()
async def close(self) -> None:
return None
async def _rpc_client(runtime: BoxRuntime):
client_to_server: asyncio.Queue[str] = asyncio.Queue()
server_to_client: asyncio.Queue[str] = asyncio.Queue()
client_connection = _QueueConnection(server_to_client, client_to_server)
server_connection = _QueueConnection(client_to_server, server_to_client)
server_handler = BoxServerHandler(
server_connection,
runtime,
host_control_authenticated=True,
trusted_instance_uuid='instance-a',
)
server_task = asyncio.create_task(server_handler.run())
client_handler = Handler(client_connection)
client_task = asyncio.create_task(client_handler.run())
client = ActionRPCBoxClient(logger=Mock())
client.set_handler(client_handler)
return client, server_task, client_task
class _Entitlements:
def __init__(self):
self.snapshots: dict[str, EntitlementSnapshot] = {}
async def get_workspace_entitlement(self, workspace_uuid: str) -> EntitlementSnapshot:
return self.snapshots[workspace_uuid]
def _snapshot(workspace_uuid: str, *, revision: int = 1, managed: bool = True) -> EntitlementSnapshot:
return EntitlementSnapshot(
instance_uuid='instance-a',
workspace_uuid=workspace_uuid,
entitlement_revision=revision,
status='active',
not_before=1,
expires_at=4_000_000_000,
features={'managed_sandbox': managed},
limits={'managed_sandbox_sessions': 1 if managed else 0},
)
def _context(workspace_uuid: str, *, revision: int = 1) -> ExecutionContext:
return ExecutionContext(
instance_uuid='instance-a',
workspace_uuid=workspace_uuid,
placement_generation=1,
entitlement_revision=revision,
)
def _query(context: ExecutionContext, query_id: int):
query = pipeline_query.Query.model_construct(
query_id=query_id,
bot_uuid='bot-a',
pipeline_uuid='pipeline-a',
launcher_type='person',
launcher_id=f'user-{query_id}',
variables={},
)
object.__setattr__(query, 'instance_uuid', context.instance_uuid)
object.__setattr__(query, 'workspace_uuid', context.workspace_uuid)
object.__setattr__(query, 'placement_generation', context.placement_generation)
object.__setattr__(query, '_execution_context', context)
return query
async def _stack(tmp_path):
shared_root = tmp_path / 'shared-box'
workspace_root = shared_root / 'workspaces'
workspace_root.mkdir(parents=True)
box_config = {
'enabled': True,
'backend': 'nsjail',
'runtime': {'endpoint': 'ws://langbot-box:5410'},
'local': {
'profile': 'default',
'host_root': str(shared_root),
'default_workspace': str(workspace_root),
'allowed_mount_roots': [str(shared_root)],
},
'admission': {
'required': True,
'logical_session_id': 'global',
'required_backend': 'nsjail',
'max_sessions': 1,
'max_managed_processes': 0,
'max_grant_ttl_sec': 300,
'max_timeout_sec': 60,
'cpus': 0.5,
'memory_mb': 256,
'pids_limit': 64,
'read_only_rootfs': True,
'workspace_quota_mb': 32,
'readiness_cache_sec': 0,
},
}
logger = Mock()
backend = _AdmissionBackend(logger)
runtime = BoxRuntime(logger, backends=[backend])
runtime.init(box_config)
await runtime.initialize()
client, server_task, client_task = await _rpc_client(runtime)
entitlements = _Entitlements()
workspace_service = SimpleNamespace(
instance_uuid='instance-a',
get_execution_binding=AsyncMock(
side_effect=lambda workspace_uuid, expected_generation: SimpleNamespace(
instance_uuid='instance-a',
workspace_uuid=workspace_uuid,
placement_generation=expected_generation,
)
),
)
app = SimpleNamespace(
logger=logger,
deployment=SimpleNamespace(multi_workspace_enabled=True),
entitlement_resolver=EntitlementResolver('instance-a', entitlements),
workspace_service=workspace_service,
instance_config=SimpleNamespace(data={'box': box_config, 'system': {'limitation': {}}}),
)
service = BoxService(app, client=client)
await service.initialize()
return service, runtime, backend, entitlements, server_task, client_task
@pytest.mark.asyncio
async def test_concurrent_first_use_creates_one_persistent_global_session(tmp_path):
service, runtime, backend, entitlements, server_task, client_task = await _stack(tmp_path)
context = _context('workspace-a')
entitlements.snapshots[context.workspace_uuid] = _snapshot(context.workspace_uuid)
try:
first, second = await asyncio.gather(
service.execute_tool({'command': 'echo first'}, _query(context, 1)),
service.execute_tool({'command': 'echo second'}, _query(context, 2)),
)
assert first['session_id'] == 'global'
assert second['session_id'] == 'global'
assert len(backend.started_specs) == 1
spec = backend.started_specs[0]
assert spec.persistent is True
assert spec.network == BoxNetworkMode.OFF
assert spec.cpus == 0.5
assert spec.memory_mb == 256
assert spec.pids_limit == 64
assert spec.workspace_quota_mb == 32
finally:
server_task.cancel()
client_task.cancel()
await runtime.shutdown()
@pytest.mark.asyncio
async def test_entitlement_loss_revokes_and_closes_existing_global_session(tmp_path):
service, runtime, backend, entitlements, server_task, client_task = await _stack(tmp_path)
context = _context('workspace-a')
entitlements.snapshots[context.workspace_uuid] = _snapshot(context.workspace_uuid, revision=1)
try:
await service.execute_tool({'command': 'true'}, _query(context, 1))
assert len(runtime.get_sessions()) == 1
entitlements.snapshots[context.workspace_uuid] = _snapshot(
context.workspace_uuid,
revision=2,
managed=False,
)
with pytest.raises(EntitlementUnavailableError):
await service.execute_tool({'command': 'true'}, _query(context, 2))
assert runtime.get_sessions() == []
assert len(backend.stopped_sessions) == 1
finally:
server_task.cancel()
client_task.cancel()
await runtime.shutdown()
@pytest.mark.asyncio
async def test_two_workspaces_get_isolated_physical_sessions_and_paths(tmp_path):
service, runtime, backend, entitlements, server_task, client_task = await _stack(tmp_path)
first = _context('workspace-a')
second = _context('workspace-b')
entitlements.snapshots[first.workspace_uuid] = _snapshot(first.workspace_uuid)
entitlements.snapshots[second.workspace_uuid] = _snapshot(second.workspace_uuid)
try:
result_a = await service.execute_tool({'command': 'tenant-a'}, _query(first, 1))
result_b = await service.execute_tool({'command': 'tenant-b'}, _query(second, 2))
assert result_a['session_id'] == result_b['session_id'] == 'global'
assert len(backend.started_specs) == 2
assert backend.started_specs[0].session_id != backend.started_specs[1].session_id
assert backend.started_specs[0].host_path != backend.started_specs[1].host_path
assert len(runtime.get_sessions()) == 2
finally:
server_task.cancel()
client_task.cancel()
await runtime.shutdown()
@pytest.mark.asyncio
async def test_cloud_skills_reject_host_paths_and_require_managed_entitlement(tmp_path):
service, runtime, backend, entitlements, server_task, client_task = await _stack(tmp_path)
first = _context('workspace-a')
second = _context('workspace-b')
ineligible = _context('workspace-free')
entitlements.snapshots[first.workspace_uuid] = _snapshot(first.workspace_uuid)
entitlements.snapshots[second.workspace_uuid] = _snapshot(second.workspace_uuid)
entitlements.snapshots[ineligible.workspace_uuid] = _snapshot(
ineligible.workspace_uuid,
managed=False,
)
try:
private = await service.create_skill(
second,
{
'name': 'private',
'instructions': 'workspace-b secret',
},
)
own_skill = await service.create_skill(
first,
{
'name': 'runner',
'instructions': 'Run scripts/main.py',
},
)
await service.write_skill_file(first, 'runner', 'scripts/main.py', "print('ok')")
await service.write_skill_file(first, 'runner', 'requirements.txt', 'requests==2.32.0\n')
refreshed_skill = await service.get_skill(first, 'runner')
assert refreshed_skill is not None
assert refreshed_skill['python_project'] is True
await service.execute_tool(
{
'command': 'python /workspace/.skills/runner/scripts/main.py',
'workdir': '/workspace/.skills/runner',
},
_query(first, 91),
skill_name='runner',
)
mounted_spec = backend.started_specs[-1]
assert len(mounted_spec.extra_mounts) == 1
assert mounted_spec.extra_mounts[0].host_path == own_skill['package_root']
assert mounted_spec.extra_mounts[0].mount_path == '/workspace/.skills/runner'
assert mounted_spec.extra_mounts[0].mode.value == 'ro'
with pytest.raises(BoxAdmissionError, match='Scanning arbitrary host'):
await service.scan_skill_directory(first, private['package_root'])
with pytest.raises(BoxAdmissionError, match='package_root is runtime-owned'):
await service.create_skill(
first,
{
'name': 'stolen',
'package_root': private['package_root'],
},
)
assert await service.get_skill(first, 'private') is None
with pytest.raises(EntitlementUnavailableError):
await service.list_skills(ineligible)
finally:
server_task.cancel()
client_task.cancel()
await runtime.shutdown()
@pytest.mark.asyncio
async def test_forged_plan_network_session_and_managed_process_never_reach_runtime(tmp_path):
service, runtime, backend, entitlements, server_task, client_task = await _stack(tmp_path)
context = _context('workspace-a')
entitlements.snapshots[context.workspace_uuid] = _snapshot(context.workspace_uuid)
query = _query(context, 1)
try:
with pytest.raises(BoxAdmissionError, match='host-controlled'):
await service.execute_spec_payload(
{'cmd': 'true', 'session_id': 'global', 'plan': 'pro'},
query,
)
with pytest.raises(BoxAdmissionError, match='network access is disabled'):
await service.execute_spec_payload(
{'cmd': 'true', 'session_id': 'global', 'network': 'on'},
query,
)
with pytest.raises(BoxAdmissionError, match='session_id is runtime-owned'):
await service.execute_spec_payload(
{'cmd': 'true', 'session_id': 'attacker'},
query,
)
with pytest.raises(BoxAdmissionError, match='Managed processes are disabled'):
await service.start_managed_process(
context,
'global',
{'command': 'sleep', 'args': ['60']},
)
assert backend.started_specs == []
assert runtime.get_sessions() == []
finally:
server_task.cancel()
client_task.cancel()
await runtime.shutdown()
+7
View File
@@ -46,6 +46,13 @@ def build_ap() -> SimpleNamespace:
)
ap.apikey_service = SimpleNamespace(authenticate_api_key=authenticate_api_key)
@contextlib.asynccontextmanager
async def tenant_scope(workspace_uuid: str):
assert workspace_uuid == 'workspace-1'
yield
ap.persistence_mgr = SimpleNamespace(tenant_scope=tenant_scope)
ap.bot_service = SimpleNamespace(
get_bots=AsyncMock(return_value=[{'uuid': 'bot-1', 'name': 'Demo Bot', 'adapter': 'telegram'}])
)
@@ -1,5 +1,6 @@
from __future__ import annotations
import contextlib
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
@@ -101,6 +102,60 @@ async def test_public_webhook_error_uses_same_generic_error_contract():
assert 'do-not-return' not in (await response.get_data(as_text=True))
async def test_public_webhook_carries_scope_without_holding_database_session():
class ScopeOnlyPersistenceManager:
mode = SimpleNamespace(value='cloud_runtime')
def __init__(self):
self.active_workspace = None
@contextlib.asynccontextmanager
async def tenant_scope(self, workspace_uuid):
self.active_workspace = workspace_uuid
try:
yield
finally:
self.active_workspace = None
def current_session(self):
return None
persistence_mgr = ScopeOnlyPersistenceManager()
workspace_uuid = '00000000-0000-0000-0000-00000000000a'
bot_uuid = '11111111-1111-4111-8111-111111111111'
class Adapter:
async def handle_unified_webhook(self, **_kwargs):
assert persistence_mgr.active_workspace == workspace_uuid
assert persistence_mgr.current_session() is None
return {'ok': True}
async def get_execution_binding(resolved_workspace_uuid, expected_generation=None):
assert resolved_workspace_uuid == workspace_uuid
assert expected_generation == 4
runtime_bot = SimpleNamespace(
workspace_uuid=workspace_uuid,
placement_generation=4,
enable=True,
adapter=Adapter(),
)
application = SimpleNamespace(
logger=Mock(),
persistence_mgr=persistence_mgr,
platform_mgr=SimpleNamespace(resolve_public_bot=AsyncMock(return_value=runtime_bot)),
workspace_service=SimpleNamespace(get_execution_binding=get_execution_binding),
)
quart_app = quart.Quart(__name__)
await WebhookRouterGroup(application, quart_app).initialize()
response = await quart_app.test_client().post(f'/bots/{bot_uuid}')
assert response.status_code == 200
assert await response.get_json() == {'ok': True}
assert persistence_mgr.active_workspace is None
async def test_authentication_failure_does_not_return_internal_exception_text():
logger = Mock()
application = SimpleNamespace(
@@ -0,0 +1,73 @@
from __future__ import annotations
import asyncio
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
import pytest
import quart
from sqlalchemy.ext.asyncio import create_async_engine
from langbot.pkg.api.http.controller import group
from langbot.pkg.persistence.mgr import PersistenceManager, PersistenceMode
pytestmark = pytest.mark.asyncio
async def test_authenticated_route_does_not_hold_database_session_during_external_wait():
entered = asyncio.Event()
release = asyncio.Event()
observations: list[bool] = []
engine = create_async_engine('sqlite+aiosqlite:///:memory:')
persistence = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
persistence.db = SimpleNamespace(get_engine=lambda: engine)
class BlockingRouter(group.RouterGroup):
name = 'blocking-route-test'
path = '/blocking-route-test'
async def initialize(self) -> None:
@self.route('', methods=['GET'], auth_type=group.AuthType.USER_TOKEN)
async def _():
observations.append(persistence.current_session() is None)
entered.set()
await release.wait()
observations.append(persistence.current_session() is None)
return self.success(data={})
account = SimpleNamespace(uuid='account-a', user='owner@example.com')
access = SimpleNamespace(
execution=SimpleNamespace(instance_uuid='instance-a', placement_generation=1),
workspace=SimpleNamespace(uuid='workspace-a'),
membership=SimpleNamespace(uuid='membership-a', role='owner', projection_revision=1),
)
application = SimpleNamespace(
persistence_mgr=persistence,
deployment=SimpleNamespace(multi_workspace_enabled=False),
user_service=SimpleNamespace(get_authenticated_account=AsyncMock(return_value=account)),
workspace_collaboration_service=SimpleNamespace(resolve_account_workspace=AsyncMock(return_value=access)),
logger=Mock(),
)
quart_app = quart.Quart(__name__)
await BlockingRouter(application, quart_app).initialize()
client = quart_app.test_client()
request = asyncio.create_task(
client.get(
'/blocking-route-test',
headers={'Authorization': 'Bearer token', 'X-Workspace-Id': 'workspace-a'},
)
)
try:
await entered.wait()
assert observations == [True]
release.set()
response = await request
assert response.status_code == 200
assert observations == [True, True]
finally:
release.set()
if not request.done():
await request
await engine.dispose()
@@ -166,6 +166,29 @@ async def test_revoked_expired_and_unknown_keys_fail_closed(api_key_context):
assert await service.authenticate_api_key(expired_secret) is None
async def test_revoke_winning_last_used_update_race_fails_authentication(api_key_context):
application, service, context, _engine = api_key_context
created = await service.create_api_key(context, 'Racing revoke')
original_execute = application.persistence_mgr.execute_async
injected_revoke = False
async def execute_with_revoke(statement, *args, **kwargs):
nonlocal injected_revoke
if (
not injected_revoke
and isinstance(statement, sqlalchemy.sql.dml.Update)
and statement.table.name == ApiKey.__tablename__
):
injected_revoke = True
await original_execute(sqlalchemy.update(ApiKey).where(ApiKey.id == created['id']).values(status='revoked'))
return await original_execute(statement, *args, **kwargs)
application.persistence_mgr.execute_async = execute_with_revoke
assert await service.authenticate_api_key(created['key']) is None
assert injected_revoke is True
async def test_cross_workspace_crud_and_secret_guessing_are_isolated(api_key_context):
application, service, first_context, engine = api_key_context
second_workspace_uuid = str(uuid.uuid4())
@@ -13,6 +13,7 @@ Source: src/langbot/pkg/api/http/service/maintenance.py
from __future__ import annotations
import contextlib
import pytest
from unittest.mock import AsyncMock, Mock, patch, MagicMock
from types import SimpleNamespace
@@ -196,6 +197,51 @@ class TestMaintenanceServiceCleanupExpiredFiles:
assert ap.logger.warning.called
assert 'uploaded_files' in result
async def test_cloud_cleanup_carries_scope_without_holding_database_session(self):
class ScopeOnlyPersistenceManager:
mode = SimpleNamespace(value='cloud_runtime')
def __init__(self):
self.active_workspace = None
@contextlib.asynccontextmanager
async def tenant_scope(self, workspace_uuid):
self.active_workspace = workspace_uuid
try:
yield
finally:
self.active_workspace = None
def current_session(self):
return None
persistence_mgr = ScopeOnlyPersistenceManager()
application = SimpleNamespace(
persistence_mgr=persistence_mgr,
instance_config=SimpleNamespace(data={}),
logger=SimpleNamespace(warning=Mock()),
)
service = MaintenanceService(application)
async def cleanup_uploads(_context, _retention_days):
assert persistence_mgr.active_workspace == TEST_CONTEXT.workspace_uuid
assert persistence_mgr.current_session() is None
return 2
def cleanup_logs(_retention_days):
assert persistence_mgr.active_workspace == TEST_CONTEXT.workspace_uuid
assert persistence_mgr.current_session() is None
return 1
service._cleanup_expired_uploaded_files = cleanup_uploads
service._cleanup_expired_log_files = cleanup_logs
assert await service.cleanup_expired_files(TEST_CONTEXT) == {
'uploaded_files': 2,
'log_files': 1,
}
assert persistence_mgr.active_workspace is None
class TestMaintenanceServiceGetStorageAnalysis:
"""Tests for get_storage_analysis method."""
@@ -13,6 +13,7 @@ Source: src/langbot/pkg/api/http/service/mcp.py
from __future__ import annotations
import asyncio
import copy
import pytest
from unittest.mock import AsyncMock, Mock, MagicMock
@@ -485,6 +486,54 @@ class TestMCPServiceCreateMCPServer:
# Verify - host_mcp_server was called
ap.tool_mgr.mcp_tool_loader.host_mcp_server.assert_called_once()
async def test_create_mcp_server_does_not_start_host_until_transaction_commits(self):
"""The Runtime must not observe a server row that can still roll back."""
gate = asyncio.get_running_loop().create_future()
class PersistenceManagerStub:
def create_after_commit_gate(self):
return gate
ap = SimpleNamespace()
ap.persistence_mgr = PersistenceManagerStub()
ap.instance_config = SimpleNamespace(data={'system': {'limitation': {'max_extensions': -1}}})
observed = []
async def host_mcp_server(context, config):
observed.append((context, config))
ap.tool_mgr = SimpleNamespace(
mcp_tool_loader=SimpleNamespace(
host_mcp_server=host_mcp_server,
_hosted_mcp_tasks=[],
)
)
server_entity = _create_mock_mcp_server(server_uuid='new-uuid', enable=True)
results = [
_create_mock_result([]),
Mock(),
_create_mock_result(first_item=server_entity),
]
ap.persistence_mgr.execute_async = AsyncMock(side_effect=results)
ap.persistence_mgr.serialize_model = Mock(
return_value={'uuid': 'new-uuid', 'name': 'New Server', 'enable': True}
)
service = _service(ap)
await service.create_mcp_server(_CONTEXT, {'name': 'New Server', 'enable': True})
await asyncio.sleep(0)
assert observed == []
gate.set_result(None)
await ap.tool_mgr.mcp_tool_loader._hosted_mcp_tasks[0]
assert observed == [
(
_CONTEXT,
{'uuid': 'new-uuid', 'name': 'New Server', 'enable': True},
)
]
async def test_create_mcp_server_disabled_no_load(self):
"""Does not load server when disabled."""
# Setup
@@ -11,7 +11,9 @@ from langbot.pkg.api.http.authz import WorkspaceRequiredError
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.api.http.service.monitoring import MonitoringService
from langbot.pkg.entity.persistence.base import Base
from langbot.pkg.entity.persistence.monitoring import MonitoringMessage
from langbot.pkg.entity.persistence.workspace import Workspace
from langbot.pkg.persistence.mgr import PersistenceManager
pytestmark = pytest.mark.asyncio
@@ -136,6 +138,26 @@ async def test_same_session_and_resource_ids_do_not_collide(service):
assert (await service.get_message_details(context_a, message_b))['found'] is False
async def test_tool_call_inherits_context_from_connection_message_row(service):
context = _context(WORKSPACE_A)
message_id = await _record_message(service, context, 'tool context')
await service.record_tool_call(
context,
tool_name='search',
tool_source='native',
duration=12,
message_id=message_id,
)
tool_calls, total = await service.get_tool_calls(context)
assert total == 1
assert tool_calls[0]['bot_id'] == 'same-bot'
assert tool_calls[0]['pipeline_id'] == 'same-pipeline'
assert tool_calls[0]['session_id'] == 'same-session'
assert tool_calls[0]['message_id'] == message_id
async def test_feedback_upsert_and_cancel_are_workspace_scoped(service):
context_a = _context(WORKSPACE_A)
context_b = _context(WORKSPACE_B)
@@ -152,3 +174,58 @@ async def test_feedback_upsert_and_cancel_are_workspace_scoped(service):
await service.record_feedback(context_a, feedback_id='same-feedback', feedback_type=3)
assert (await service.get_feedback_stats(context_a))['total_feedback'] == 0
assert (await service.get_feedback_stats(context_b))['total_feedback'] == 1
async def test_cleanup_commits_sqlite_delete_before_vacuum(tmp_path):
engine = create_async_engine(
f'sqlite+aiosqlite:///{tmp_path / "monitoring-cleanup.db"}',
connect_args={'timeout': 0.1},
)
application = SimpleNamespace(
instance_config=SimpleNamespace(data={'database': {'use': 'sqlite'}}),
)
manager = PersistenceManager(application)
manager.db = SimpleNamespace(get_engine=lambda: engine)
application.persistence_mgr = manager
try:
async with engine.begin() as connection:
await connection.run_sync(Base.metadata.create_all)
await connection.execute(
sqlalchemy.insert(Workspace).values(
uuid=WORKSPACE_A,
instance_uuid='instance',
name='A',
slug='a',
source='cloud_projection',
)
)
await connection.execute(
sqlalchemy.insert(MonitoringMessage).values(
id='expired-message',
workspace_uuid=WORKSPACE_A,
timestamp=datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)
- datetime.timedelta(days=30),
bot_id='bot',
bot_name='Bot',
pipeline_id='pipeline',
pipeline_name='Pipeline',
message_content='expired',
session_id='session',
status='success',
level='info',
)
)
deleted = await MonitoringService(application).cleanup_expired_records(
_context(WORKSPACE_A),
retention_days=1,
)
assert deleted['monitoring_messages'] == 1
async with engine.connect() as connection:
remaining = await connection.scalar(
sqlalchemy.select(sqlalchemy.func.count()).select_from(MonitoringMessage)
)
assert remaining == 0
finally:
await engine.dispose()
+11 -4
View File
@@ -45,9 +45,16 @@ async def test_verify_api_key_keeps_db_validation_for_lbk_keys(key_exists):
if key_exists
else None
)
query_result = Mock()
query_result.first.return_value = key
persistence_mgr = SimpleNamespace(execute_async=AsyncMock(side_effect=[query_result, Mock(rowcount=1)]))
discovery_result = Mock()
discovery_result.first.return_value = key
query_results = [discovery_result]
if key_exists:
scoped_result = Mock()
scoped_result.first.return_value = key
update_result = Mock()
update_result.scalar_one_or_none.return_value = key.id
query_results.extend([scoped_result, update_result])
persistence_mgr = SimpleNamespace(execute_async=AsyncMock(side_effect=query_results))
instance_config = SimpleNamespace(data={'api': {'global_api_key': ''}})
workspace_service = SimpleNamespace(
get_execution_binding=AsyncMock(
@@ -69,7 +76,7 @@ async def test_verify_api_key_keeps_db_validation_for_lbk_keys(key_exists):
result = await service.verify_api_key('lbk_valid_format')
assert result is key_exists
assert persistence_mgr.execute_async.await_count == (2 if key_exists else 1)
assert persistence_mgr.execute_async.await_count == (3 if key_exists else 1)
if key_exists:
workspace_service.get_execution_binding.assert_awaited_once_with('workspace-a')
else:
@@ -0,0 +1,121 @@
from __future__ import annotations
import asyncio
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
import sqlalchemy as sa
from sqlalchemy.ext.asyncio import create_async_engine
from langbot.pkg.api.http.service.apikey import ApiKeyIdentity
from langbot.pkg.api.mcp.context import get_request_context
from langbot.pkg.api.mcp.mount import MCPMount
from langbot.pkg.persistence.mgr import PersistenceManager, PersistenceMode
from langbot.pkg.persistence.tenant_uow import PersistenceScopeKind
@pytest.mark.asyncio
async def test_mcp_mount_keeps_request_context_but_no_session_during_stream_wait(tmp_path) -> None:
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "mcp-short-scope.db"}')
table = sa.Table('mcp_scope_probe', sa.MetaData(), sa.Column('id', sa.Integer, primary_key=True))
manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
manager.db = SimpleNamespace(get_engine=lambda: engine)
checked_out = 0
def on_checkout(*_args):
nonlocal checked_out
checked_out += 1
def on_checkin(*_args):
nonlocal checked_out
checked_out -= 1
sa.event.listen(engine.sync_engine, 'checkout', on_checkout)
sa.event.listen(engine.sync_engine, 'checkin', on_checkin)
try:
async with engine.begin() as conn:
await conn.run_sync(table.metadata.create_all)
identity = ApiKeyIdentity(
instance_uuid='instance-1',
workspace_uuid='workspace-1',
placement_generation=7,
api_key_uuid='key-1',
permissions=frozenset({'pipelines:read'}),
)
app = SimpleNamespace(
apikey_service=SimpleNamespace(authenticate_api_key=AsyncMock(return_value=identity)),
persistence_mgr=manager,
deployment_admission=None,
deployment=None,
)
stream_waiting = asyncio.Event()
release_stream = asyncio.Event()
observations: list[tuple[str, str, bool]] = []
async def fake_mcp_asgi(scope, receive, send):
del scope, receive
context = get_request_context()
assert manager.current_scope().kind is PersistenceScopeKind.WORKSPACE
assert manager.current_session() is None
await manager.execute_async(sa.select(table.c.id))
assert manager.current_session() is None
observations.append((context.request_id, context.workspace_uuid, manager.current_session() is None))
stream_waiting.set()
await release_stream.wait()
preserved_context = get_request_context()
observations.append(
(
preserved_context.request_id,
preserved_context.workspace_uuid,
manager.current_session() is None,
)
)
await manager.execute_async(sa.select(table.c.id))
assert manager.current_session() is None
await send({'type': 'http.response.start', 'status': 200, 'headers': []})
await send({'type': 'http.response.body', 'body': b'{}'})
async def unused_quart_asgi(scope, receive, send):
del scope, receive, send
raise AssertionError('MCP request was routed to Quart')
mount = MCPMount.__new__(MCPMount)
mount.ap = app
mount._mcp_asgi = fake_mcp_asgi
sent_messages: list[dict] = []
async def receive():
return {'type': 'http.request', 'body': b'', 'more_body': False}
async def send(message):
sent_messages.append(message)
async def release_after_observation() -> None:
await asyncio.wait_for(stream_waiting.wait(), timeout=2)
assert checked_out == 0
release_stream.set()
release_task = asyncio.create_task(release_after_observation())
await mount.wrap(unused_quart_asgi)(
{
'type': 'http',
'path': '/mcp',
'headers': [(b'x-api-key', b'secret')],
},
receive,
send,
)
await release_task
assert sent_messages[0]['status'] == 200
assert len(observations) == 2
assert observations[0] == observations[1]
assert observations[0][1:] == ('workspace-1', True)
assert checked_out == 0
assert manager.current_scope() is None
with pytest.raises(RuntimeError, match='context is unavailable'):
get_request_context()
finally:
await engine.dispose()
@@ -1,5 +1,6 @@
from __future__ import annotations
from contextlib import asynccontextmanager
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
@@ -106,3 +107,33 @@ async def test_background_plugin_operation_refences_captured_generation(plugin_r
await router._run_fenced_plugin_operation(CONTEXT, operation)
operation.assert_not_awaited()
@pytest.mark.asyncio
async def test_background_plugin_operation_revalidates_inside_short_tenant_uow(plugin_router_cls):
scopes = []
@asynccontextmanager
async def tenant_uow(workspace_uuid):
scopes.append(workspace_uuid)
yield
connector = SimpleNamespace(
require_workspace_context=AsyncMock(side_effect=lambda context: context),
)
operation = AsyncMock(return_value='done')
router = object.__new__(plugin_router_cls)
router.ap = SimpleNamespace(
plugin_connector=connector,
persistence_mgr=SimpleNamespace(
mode=SimpleNamespace(value='cloud_runtime'),
tenant_uow=tenant_uow,
),
)
result = await router._run_fenced_plugin_operation(CONTEXT, operation)
assert result == 'done'
assert scopes == [CONTEXT.workspace_uuid]
connector.require_workspace_context.assert_awaited_once_with(CONTEXT)
operation.assert_awaited_once()
@@ -0,0 +1,69 @@
from __future__ import annotations
from contextlib import asynccontextmanager
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
import pytest
from langbot.pkg.api.http.context import (
PrincipalContext,
PrincipalType,
RequestContext,
WorkspaceContext,
)
from langbot.pkg.api.http.controller.groups.pipelines.websocket_chat import WebSocketChatRouterGroup
@pytest.mark.asyncio
async def test_websocket_pipeline_lookup_opens_workspace_uow_after_auth_scope_closed() -> None:
workspace_uuid = 'workspace-a'
scopes: list[str] = []
in_scope = False
@asynccontextmanager
async def tenant_uow(selected_workspace_uuid: str):
nonlocal in_scope
assert not in_scope
in_scope = True
scopes.append(selected_workspace_uuid)
try:
yield
finally:
in_scope = False
async def get_pipeline(_context, _pipeline_uuid):
assert in_scope
return {'uuid': 'pipeline-a'}
adapter = Mock()
router = object.__new__(WebSocketChatRouterGroup)
router.ap = SimpleNamespace(
persistence_mgr=SimpleNamespace(
mode=SimpleNamespace(value='cloud_runtime'),
tenant_uow=tenant_uow,
),
pipeline_service=SimpleNamespace(get_pipeline=AsyncMock(side_effect=get_pipeline)),
platform_mgr=SimpleNamespace(get_websocket_proxy_bot=AsyncMock(return_value=SimpleNamespace(adapter=adapter))),
)
request_context = RequestContext(
instance_uuid='instance-a',
placement_generation=1,
request_id='request-a',
auth_type='user_token',
principal=PrincipalContext(
principal_type=PrincipalType.ACCOUNT,
account_uuid='account-a',
),
workspace=WorkspaceContext(
workspace_uuid=workspace_uuid,
membership_uuid='membership-a',
role='owner',
permissions=frozenset(),
),
)
result = await router._get_scoped_adapter(request_context, 'pipeline-a')
assert result is adapter
assert scopes == [workspace_uuid]
+230
View File
@@ -0,0 +1,230 @@
from __future__ import annotations
import datetime as dt
import hashlib
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
import pytest
from langbot_plugin.box.models import SandboxAdmissionPolicy
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.box.admission import (
BoxAdmissionError,
SandboxAdmissionController,
require_cloud_admission_policy,
)
from langbot.pkg.box.service import BoxService
from langbot.pkg.cloud.entitlements import (
EntitlementResolver,
EntitlementSnapshot,
EntitlementUnavailableError,
)
_UTC = dt.timezone.utc
_CONTEXT = ExecutionContext(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
placement_generation=3,
entitlement_revision=7,
)
def _snapshot(
*,
revision: int = 7,
managed: bool = True,
sessions: int = 1,
expires_at: int = 2_000,
) -> EntitlementSnapshot:
return EntitlementSnapshot(
instance_uuid='instance-a',
workspace_uuid='workspace-a',
entitlement_revision=revision,
status='active',
not_before=1,
expires_at=expires_at,
features={'managed_sandbox': managed},
limits={'managed_sandbox_sessions': sessions},
)
def _controller(snapshot: EntitlementSnapshot, *, now: float = 1_000.25):
provider = SimpleNamespace(get_workspace_entitlement=AsyncMock(return_value=snapshot))
resolver = EntitlementResolver('instance-a', provider)
client = SimpleNamespace(
upsert_sandbox_admission_grant=AsyncMock(
side_effect=lambda grant: {
'installed': True,
'workspace_uuid': grant.workspace_uuid,
'execution_generation': grant.execution_generation,
'entitlement_revision': grant.entitlement_revision,
'max_sessions': grant.max_sessions,
'max_managed_processes': grant.max_managed_processes,
}
),
revoke_sandbox_admission_grant=AsyncMock(
side_effect=lambda revocation: {
'revoked': True,
'workspace_uuid': revocation.workspace_uuid,
'entitlement_revision': revocation.entitlement_revision,
}
),
)
app = SimpleNamespace(entitlement_resolver=resolver, logger=Mock())
controller = SandboxAdmissionController(
app,
client,
policy=SandboxAdmissionPolicy(required=True, max_grant_ttl_sec=300),
wall_time=lambda: now,
)
return controller, client, provider
def test_cloud_admission_policy_requires_positive_workspace_quota():
with pytest.raises(BoxAdmissionError, match='workspace quota must be a positive integer'):
require_cloud_admission_policy(
{
'required': True,
'workspace_quota_mb': 0,
}
)
policy = require_cloud_admission_policy(
{
'required': True,
'workspace_quota_mb': 32,
}
)
assert policy.workspace_quota_mb == 32
@pytest.mark.asyncio
async def test_active_generic_entitlement_installs_short_lived_numeric_grant():
controller, client, provider = _controller(_snapshot())
grant = await controller.require(_CONTEXT)
assert grant.instance_uuid == _CONTEXT.instance_uuid
assert grant.workspace_uuid == _CONTEXT.workspace_uuid
assert grant.execution_generation == _CONTEXT.placement_generation
assert grant.entitlement_revision == 7
assert grant.max_sessions == 1
assert grant.max_managed_processes == 0
assert grant.expires_at == dt.datetime.fromtimestamp(1_300, tz=_UTC)
assert (grant.expires_at - dt.datetime.fromtimestamp(1_000.25, tz=_UTC)).total_seconds() < 300
provider.get_workspace_entitlement.assert_awaited_once_with('workspace-a')
client.upsert_sandbox_admission_grant.assert_awaited_once_with(grant)
client.revoke_sandbox_admission_grant.assert_not_awaited()
@pytest.mark.asyncio
@pytest.mark.parametrize(
'snapshot',
[
_snapshot(managed=False),
_snapshot(sessions=0),
_snapshot(sessions=2),
],
)
async def test_non_eligible_entitlement_revokes_and_fails_closed(snapshot):
controller, client, _provider = _controller(snapshot)
with pytest.raises(EntitlementUnavailableError):
await controller.require(_CONTEXT)
client.upsert_sandbox_admission_grant.assert_not_awaited()
revocation = client.revoke_sandbox_admission_grant.await_args.args[0]
assert revocation.entitlement_revision == snapshot.entitlement_revision
@pytest.mark.asyncio
async def test_transient_entitlement_failure_does_not_tombstone_valid_revision():
controller, client, provider = _controller(_snapshot())
await controller.require(_CONTEXT)
provider.get_workspace_entitlement.side_effect = RuntimeError('control plane unavailable')
with pytest.raises(RuntimeError, match='control plane unavailable'):
await controller.require(_CONTEXT)
client.revoke_sandbox_admission_grant.assert_not_awaited()
provider.get_workspace_entitlement.side_effect = None
provider.get_workspace_entitlement.return_value = _snapshot()
recovered = await controller.require(_CONTEXT)
assert recovered.entitlement_revision == 7
@pytest.mark.asyncio
async def test_runtime_receipt_mismatch_is_revoked_and_never_admitted():
controller, client, _provider = _controller(_snapshot())
client.upsert_sandbox_admission_grant.return_value = {'installed': True, 'workspace_uuid': 'other'}
client.upsert_sandbox_admission_grant.side_effect = None
with pytest.raises(Exception, match='invalid sandbox admission receipt'):
await controller.require(_CONTEXT)
client.revoke_sandbox_admission_grant.assert_not_awaited()
@pytest.mark.asyncio
async def test_authoritative_cancelled_revision_is_revoked():
cancelled = _snapshot(revision=8).model_copy(update={'status': 'cancelled'})
controller, client, _provider = _controller(cancelled)
with pytest.raises(EntitlementUnavailableError, match='not active'):
await controller.require(_CONTEXT)
revocation = client.revoke_sandbox_admission_grant.await_args.args[0]
assert revocation.entitlement_revision == 8
@pytest.mark.asyncio
async def test_cloud_box_readiness_failure_aborts_service_initialization(tmp_path):
workspace_root = tmp_path / 'box' / 'workspaces'
workspace_root.mkdir(parents=True)
box_config = {
'enabled': True,
'backend': 'nsjail',
'runtime': {'endpoint': 'ws://box:5410'},
'local': {
'host_root': str(tmp_path / 'box'),
'default_workspace': str(workspace_root),
'allowed_mount_roots': [str(tmp_path / 'box')],
},
'admission': {
'required': True,
'logical_session_id': 'global',
'required_backend': 'nsjail',
'max_sessions': 1,
'max_managed_processes': 0,
'max_grant_ttl_sec': 300,
'workspace_quota_mb': 32,
},
}
client = SimpleNamespace(
initialize=AsyncMock(),
verify_shared_workspace=AsyncMock(
side_effect=lambda marker_name: {
'marker_name': marker_name,
'size': (workspace_root / marker_name).stat().st_size,
'sha256': hashlib.sha256((workspace_root / marker_name).read_bytes()).hexdigest(),
}
),
get_backend_info=AsyncMock(return_value={'name': 'docker', 'available': True}),
)
app = SimpleNamespace(
logger=Mock(),
deployment=SimpleNamespace(multi_workspace_enabled=True),
entitlement_resolver=Mock(),
workspace_service=SimpleNamespace(instance_uuid='instance-a'),
instance_config=SimpleNamespace(data={'box': box_config}),
)
service = BoxService(app, client=client)
with pytest.raises(Exception, match='nsjail isolation readiness failed'):
await service.initialize()
assert service.available is False
+138 -13
View File
@@ -117,6 +117,9 @@ class _InProcessBoxRuntimeClient(BoxRuntimeClient):
async def init(self, config: dict) -> None:
self._runtime.init(config)
async def verify_shared_workspace(self, marker_name: str) -> dict:
return self._runtime.verify_shared_workspace(marker_name)
class FakeBackend(BaseSandboxBackend):
def __init__(self, logger: Mock, available: bool = True):
@@ -322,6 +325,43 @@ def test_separated_box_runtime_does_not_create_default_workspace_in_langbot(tmp_
assert not (host_root / 'default').exists()
@pytest.mark.asyncio
async def test_cloud_initialize_fails_when_core_and_runtime_volumes_are_separated(tmp_path):
logger = Mock()
core_root = tmp_path / 'core-box'
runtime_root = tmp_path / 'runtime-box'
(core_root / 'default').mkdir(parents=True)
runtime = BoxRuntime(logger=logger, backends=[FakeBackend(logger)], session_ttl_sec=300)
runtime.init(
{
'local': {
'host_root': str(runtime_root),
'default_workspace': 'default',
'allowed_mount_roots': [str(runtime_root)],
}
}
)
app = make_app(logger, host_root=str(core_root))
app.deployment = SimpleNamespace(multi_workspace_enabled=True)
app.instance_config.data['box'].update(
{
'backend': 'nsjail',
'admission': {'required': True, 'workspace_quota_mb': 32},
}
)
service = BoxService(
app,
client=_InProcessBoxRuntimeClient(logger, runtime),
)
with pytest.raises(BoxValidationError, match='shared durable Workspace volume'):
await service.initialize()
assert service.available is False
assert list((core_root / 'default').glob('.langbot-box-volume-probe-*')) == []
await runtime.shutdown()
def test_separated_box_runtime_allows_box_owned_missing_host_path(tmp_path):
logger = Mock()
runtime = BoxRuntime(logger=logger, backends=[FakeBackend(logger)], session_ttl_sec=300)
@@ -1928,7 +1968,7 @@ class TestInboundOutboundRoundTrip:
assert '/workspace/inbox/' in parameters['command']
return {
'ok': True,
'stdout': '["/workspace/inbox/42/image_1.png"]',
'stdout': '["/workspace/inbox/query-42/image_1.png"]',
'stderr': '',
}
@@ -1938,7 +1978,7 @@ class TestInboundOutboundRoundTrip:
assert len(descriptors) == 1
d = descriptors[0]
assert d['type'] == 'Image'
assert d['path'] == '/workspace/inbox/42/image_1.png'
assert d['path'] == '/workspace/inbox/query-42/image_1.png'
assert d['size'] == len(img_bytes)
@pytest.mark.asyncio
@@ -2051,7 +2091,7 @@ class TestAttachmentHostPath:
assert d['type'] == 'Image'
assert d['size'] == len(big)
# File actually landed on the host workspace.
host_file = os.path.join(ws, 'inbox', str(query.query_id), d['name'])
host_file = os.path.join(ws, 'inbox', str(query.query_uuid), d['name'])
assert os.path.isfile(host_file)
assert open(host_file, 'rb').read() == big
@@ -2063,7 +2103,7 @@ class TestAttachmentHostPath:
service, ws = self._service_with_workspace(tmp_path)
# Seed a stale file under the same query_id (simulates webchat id reuse).
stale_dir = os.path.join(ws, 'inbox', '42')
stale_dir = os.path.join(ws, 'inbox', 'query-42')
os.makedirs(stale_dir, exist_ok=True)
open(os.path.join(stale_dir, 'image_1.png'), 'wb').write(b'STALE-OLD-IMAGE')
@@ -2079,11 +2119,40 @@ class TestAttachmentHostPath:
# No leftover content from the stale image.
assert b'STALE-OLD-IMAGE' not in open(host_file, 'rb').read()
@pytest.mark.asyncio
async def test_inbound_host_replaces_query_symlink_without_touching_other_workspace(self, tmp_path):
import base64
import langbot_plugin.api.entities.builtin.platform.message as platform_message
service, ws = self._service_with_workspace(tmp_path)
other_workspace = tmp_path / 'other-workspace'
other_workspace.mkdir()
protected = other_workspace / 'protected.txt'
protected.write_bytes(b'workspace-b-secret')
inbox = os.path.join(ws, 'inbox')
os.makedirs(inbox, exist_ok=True)
os.symlink(other_workspace, os.path.join(inbox, 'query-42'))
query = make_query()
payload = b'workspace-a-input'
query.message_chain = platform_message.MessageChain(
[platform_message.File(name='input.bin', base64=base64.b64encode(payload).decode())]
)
service.execute_tool = AsyncMock(side_effect=AssertionError('exec must not be used on host path'))
descriptors = await service.materialize_inbound_attachments(query)
assert descriptors[0]['path'] == '/workspace/inbox/query-42/input.bin'
assert protected.read_bytes() == b'workspace-b-secret'
assert not os.path.islink(os.path.join(inbox, 'query-42'))
assert open(os.path.join(inbox, 'query-42', 'input.bin'), 'rb').read() == payload
@pytest.mark.asyncio
async def test_outbound_reads_host_and_clears(self, tmp_path):
service, ws = self._service_with_workspace(tmp_path)
query = make_query()
outbox = os.path.join(ws, 'outbox', str(query.query_id))
outbox = os.path.join(ws, 'outbox', str(query.query_uuid))
os.makedirs(outbox, exist_ok=True)
# A large file that would be truncated on the exec/stdout path:
big_png = b'\x89PNG\r\n\x1a\n' + b'y' * (400 * 1024)
@@ -2103,13 +2172,69 @@ class TestAttachmentHostPath:
# Outbox cleared after collection.
assert os.listdir(outbox) == []
@pytest.mark.asyncio
async def test_outbound_host_never_follows_query_or_file_symlinks(self, tmp_path):
service, ws = self._service_with_workspace(tmp_path)
query = make_query()
other_workspace = tmp_path / 'other-workspace'
other_workspace.mkdir()
secret = other_workspace / 'secret.txt'
secret.write_bytes(b'workspace-b-secret')
outbox_root = os.path.join(ws, 'outbox')
os.makedirs(outbox_root, exist_ok=True)
# A hostile query-directory replacement is rejected rather than read.
query_dir = os.path.join(outbox_root, str(query.query_uuid))
os.symlink(other_workspace, query_dir)
with pytest.raises(BoxValidationError, match='symbolic link'):
await service.collect_outbound_attachments(query)
assert secret.read_bytes() == b'workspace-b-secret'
os.unlink(query_dir)
os.makedirs(query_dir)
os.symlink(secret, os.path.join(query_dir, 'leak.txt'))
service.execute_tool = AsyncMock(side_effect=AssertionError('exec must not be used on host path'))
assert await service.collect_outbound_attachments(query) == []
assert secret.read_bytes() == b'workspace-b-secret'
@pytest.mark.asyncio
async def test_outbound_host_fails_closed_on_inode_bomb(self, tmp_path):
service, ws = self._service_with_workspace(tmp_path)
query = make_query()
outbox = os.path.join(ws, 'outbox', str(query.query_uuid))
os.makedirs(outbox, exist_ok=True)
harmless_target = tmp_path / 'harmless-target'
harmless_target.write_bytes(b'x')
# Symlinks do not count toward the 20 returned files, so this proves
# traversal itself has a bounded entry budget.
for index in range(513):
os.symlink(harmless_target, os.path.join(outbox, f'entry-{index}'))
with pytest.raises(BoxValidationError, match='symbolic link'):
await service.collect_outbound_attachments(query)
assert harmless_target.read_bytes() == b'x'
def test_host_attachment_directories_use_query_uuid_not_process_local_id(self, tmp_path):
service, _ws = self._service_with_workspace(tmp_path)
first = make_query(query_id=7)
second = make_query(query_id=7)
object.__setattr__(first, 'query_uuid', 'replica-a-query')
object.__setattr__(second, 'query_uuid', 'replica-b-query')
first_path = service._host_query_dir(service.OUTBOX_SUBDIR, first)
second_path = service._host_query_dir(service.OUTBOX_SUBDIR, second)
assert first_path is not None and first_path.endswith('/outbox/replica-a-query')
assert second_path is not None and second_path.endswith('/outbox/replica-b-query')
assert first_path != second_path
@pytest.mark.asyncio
async def test_outbound_empty_clears_stale_host_dir(self, tmp_path):
# Reusing a query_id (counter resets on restart) must not re-send files
# a previous run left in the outbox: an empty collection still clears it.
service, ws = self._service_with_workspace(tmp_path)
query = make_query()
outbox = os.path.join(ws, 'outbox', str(query.query_id))
outbox = os.path.join(ws, 'outbox', str(query.query_uuid))
os.makedirs(outbox, exist_ok=True)
# Stale file from a prior turn; the agent produced nothing this turn —
# but _read_outbox_host would still pick it up, so collection must drop
@@ -2158,16 +2283,16 @@ class TestAttachmentHostPath:
os.makedirs(os.path.join(outbox, '0'), exist_ok=True)
# Simulate a host delete that cannot remove the root-owned outbox.
import shutil as _shutil
from langbot.pkg.box import secure_fs
real_rmtree = _shutil.rmtree
real_purge = secure_fs.purge_subdirectory
def fake_rmtree(path, *a, **k):
if os.path.abspath(path) == os.path.abspath(outbox):
return # "permission denied" — silently leaves the dir
return real_rmtree(path, *a, **k)
def fake_purge(root, subdir):
if os.path.abspath(os.path.join(root, subdir)) == os.path.abspath(outbox):
raise PermissionError('root-owned')
return real_purge(root, subdir)
monkeypatch.setattr(_shutil, 'rmtree', fake_rmtree)
monkeypatch.setattr(secure_fs, 'purge_subdirectory', fake_purge)
service.build_spec = Mock()
service.client.execute = AsyncMock()
+156 -1
View File
@@ -1,11 +1,14 @@
from __future__ import annotations
import dataclasses
from types import SimpleNamespace
import pytest
from langbot.pkg.cloud.bootstrap import (
CloudBootstrapError,
CloudRuntimeUnavailableError,
DeploymentAdmissionGuard,
OpenSourceDeployment,
VerifiedCloudDeployment,
resolve_deployment,
@@ -62,8 +65,34 @@ class _EntryPoints(list):
def _cloud_config() -> dict:
return {
'database': {'use': 'postgresql'},
'vdb': {'use': 'pgvector'},
'vdb': {
'use': 'pgvector',
'pgvector': {
'use_business_database': True,
'allowed_dimensions': [384, 768, 1536],
},
},
'mcp': {'stdio': {'enabled': False}},
'plugin': {'worker': {'require_hard_limits': True}},
'box': {
'enabled': True,
'backend': 'nsjail',
'runtime': {'endpoint': 'ws://langbot-box:5410'},
'admission': {
'required': True,
'logical_session_id': 'global',
'required_backend': 'nsjail',
'max_sessions': 1,
'max_managed_processes': 0,
'max_grant_ttl_sec': 300,
'workspace_quota_mb': 32,
},
'local': {
'host_root': '/var/lib/langbot/box',
'default_workspace': '/var/lib/langbot/box/workspaces',
'allowed_mount_roots': ['/var/lib/langbot/box'],
},
},
# Proves mutable product metadata does not participate in selection.
'system': {'edition': 'community'},
}
@@ -99,6 +128,7 @@ async def test_verified_closed_entry_point_activates_cloud_policy():
('database', {'use': 'sqlite'}, 'database.use=postgresql'),
('vdb', {'use': 'chroma'}, 'vdb.use=pgvector'),
('mcp', {'stdio': {'enabled': True}}, 'mcp.stdio.enabled=false'),
('plugin', {'worker': {'require_hard_limits': False}}, 'plugin.worker.require_hard_limits=true'),
],
)
async def test_cloud_runtime_config_is_fail_closed(field, value, message):
@@ -114,6 +144,79 @@ async def test_cloud_runtime_config_is_fail_closed(field, value, message):
)
@pytest.mark.parametrize(
('pgvector_config', 'message'),
[
({'use_business_database': False, 'allowed_dimensions': [1536]}, 'use_business_database=true'),
({'use_business_database': True, 'allowed_dimensions': []}, 'allowed_dimensions'),
({'use_business_database': True, 'allowed_dimensions': [3072]}, 'allowed_dimensions'),
({'use_business_database': True, 'allowed_dimensions': [True]}, 'allowed_dimensions'),
],
)
async def test_cloud_pgvector_contract_is_fail_closed(pgvector_config, message):
config = _cloud_config()
config['vdb']['pgvector'] = pgvector_config
with pytest.raises(CloudBootstrapError, match=message):
await resolve_deployment(
instance_uuid='instance-a',
instance_config=config,
entry_points=lambda: _EntryPoints([_EntryPoint(_Provider())]),
now=1_000,
)
@pytest.mark.parametrize(
('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'),
(
lambda config: config['box']['admission'].update(max_sessions=2),
'grant-enforced Box admission',
),
(
lambda config: config['box']['admission'].update(max_managed_processes=1),
'zero managed processes',
),
(
lambda config: config['box']['admission'].update(max_grant_ttl_sec=301),
'max_grant_ttl_sec',
),
(
lambda config: config['box']['admission'].update(workspace_quota_mb=0),
'workspace_quota_mb must be a positive integer',
),
(
lambda config: config['box']['admission'].update(workspace_quota_mb=True),
'workspace_quota_mb must be a positive integer',
),
(
lambda config: config['box']['local'].update(default_workspace='relative/workspaces'),
'default_workspace must be an absolute',
),
(
lambda config: config['box']['local'].update(
default_workspace='/other/workspaces',
),
'under allowed_mount_roots',
),
],
)
async def test_cloud_box_contract_is_fail_closed(mutate, message):
config = _cloud_config()
mutate(config)
with pytest.raises(CloudBootstrapError, match=message):
await resolve_deployment(
instance_uuid='instance-a',
instance_config=config,
entry_points=lambda: _EntryPoints([_EntryPoint(_Provider())]),
now=1_000,
)
async def test_invalid_provider_never_falls_back_to_oss():
provider = SimpleNamespace(bootstrap=lambda **_: object())
@@ -134,3 +237,55 @@ async def test_duplicate_closed_providers_fail_closed():
entry_points=lambda: _EntryPoints([_EntryPoint(_Provider()), _EntryPoint(_Provider())]),
now=1_000,
)
async def test_deployment_admission_expires_even_after_wall_clock_rollback():
wall = [1_000.0]
monotonic = [50.0]
deployment = dataclasses.replace(
_Provider().bootstrap(instance_uuid='instance-a', instance_config={}),
expires_at=1_010,
)
guard = DeploymentAdmissionGuard(
'instance-a',
deployment,
wall_time=lambda: wall[0],
monotonic_time=lambda: monotonic[0],
)
assert guard.require_active() is deployment
wall[0] = 900.0
monotonic[0] = 60.0
with pytest.raises(CloudRuntimeUnavailableError, match='expired'):
guard.require_active()
async def test_deployment_admission_accepts_only_monotonic_non_conflicting_renewal():
wall = [1_000.0]
monotonic = [50.0]
current = dataclasses.replace(
_Provider().bootstrap(instance_uuid='instance-a', instance_config={}),
expires_at=1_010,
)
guard = DeploymentAdmissionGuard(
'instance-a',
current,
wall_time=lambda: wall[0],
monotonic_time=lambda: monotonic[0],
)
renewed = dataclasses.replace(
current,
manifest_jti='manifest-b',
manifest_generation=4,
expires_at=2_000,
)
guard.replace(renewed)
assert guard.require_active() is renewed
rollback = dataclasses.replace(current, manifest_generation=2)
with pytest.raises(CloudRuntimeUnavailableError, match='rolled back'):
guard.replace(rollback)
conflicting = dataclasses.replace(renewed, manifest_jti='different')
with pytest.raises(CloudRuntimeUnavailableError, match='conflicting'):
guard.replace(conflicting)
@@ -84,3 +84,26 @@ async def test_resolver_rejects_same_revision_with_different_contents():
await resolver.resolve('workspace-a', now=150)
with pytest.raises(EntitlementUnavailableError, match='conflicting contents'):
await resolver.resolve('workspace-a', now=150)
@pytest.mark.asyncio
async def test_resolver_checks_deployment_admission_before_and_after_provider_call():
checks = 0
def require_admission() -> None:
nonlocal checks
checks += 1
if checks == 2:
raise RuntimeError('manifest expired during provider call')
provider = AsyncMock()
provider.get_workspace_entitlement = AsyncMock(return_value=_snapshot())
resolver = EntitlementResolver(
'instance-a',
provider,
deployment_admission=require_admission,
)
with pytest.raises(RuntimeError, match='expired during provider call'):
await resolver.resolve('workspace-a', now=150)
assert checks == 2
+29
View File
@@ -35,6 +35,22 @@ class TestApplyEnvOverridesToConfig:
assert result['system']['name'] == 'custom_name'
def test_override_log_never_prints_secret_value(self, capsys):
"""Environment-backed credentials must not be copied into logs."""
load_config = get_load_config_module()
secret = 'database-password-that-must-not-leak'
cfg = {'database': {'postgresql': {'password': ''}}}
env = {'DATABASE__POSTGRESQL__PASSWORD': secret}
with patch.dict(os.environ, env, clear=True):
result = load_config._apply_env_overrides_to_config(cfg)
captured = capsys.readouterr().out
assert result['database']['postgresql']['password'] == secret
assert 'DATABASE__POSTGRESQL__PASSWORD' in captured
assert secret not in captured
def test_override_int_value(self):
"""Test overriding an int value with proper conversion."""
load_config = get_load_config_module()
@@ -196,6 +212,19 @@ class TestApplyEnvOverridesToConfig:
assert result['system']['name'] == 'default'
def test_skip_env_vars_with_empty_path_segments(self, capsys):
"""Platform variables such as __CF_USER_TEXT_ENCODING are not config."""
load_config = get_load_config_module()
cfg = {'system': {'name': 'default'}}
env = {'__CF_USER_TEXT_ENCODING': '0x1F5:0x0:0x64'}
with patch.dict(os.environ, env, clear=True):
result = load_config._apply_env_overrides_to_config(cfg)
assert result == cfg
assert capsys.readouterr().out == ''
def test_nested_config_path(self):
"""Test overriding deeply nested config."""
load_config = get_load_config_module()
+48
View File
@@ -12,6 +12,7 @@ from __future__ import annotations
import pytest
import asyncio
import contextvars
import sys
from unittest.mock import Mock, MagicMock
from contextlib import contextmanager
@@ -360,6 +361,53 @@ class TestAsyncTaskManager:
wrapper.cancel()
@pytest.mark.asyncio
async def test_create_task_does_not_inherit_request_context(self):
"""Long-lived tasks must receive identity through explicit arguments."""
_, _, AsyncTaskManager = get_taskmgr_classes()
mock_app = create_mock_app()
manager = AsyncTaskManager(mock_app)
request_value = contextvars.ContextVar('request_value', default=None)
token = request_value.set('request-scoped-transaction')
observed = []
async def detached_task(captured_workspace: str) -> None:
observed.append((request_value.get(), captured_workspace))
try:
wrapper = manager.create_task(detached_task('workspace-a'))
await wrapper.task
finally:
request_value.reset(token)
assert observed == [(None, 'workspace-a')]
@pytest.mark.asyncio
async def test_create_task_waits_for_registered_transaction_commit(self):
_, _, AsyncTaskManager = get_taskmgr_classes()
mock_app = create_mock_app()
gate = asyncio.get_running_loop().create_future()
class PersistenceManagerStub:
def create_after_commit_gate(self):
return gate
mock_app.persistence_mgr = PersistenceManagerStub()
manager = AsyncTaskManager(mock_app)
observed = []
async def background_work() -> None:
observed.append('started')
wrapper = manager.create_task(background_work())
await asyncio.sleep(0)
assert observed == []
gate.set_result(None)
await wrapper.task
assert observed == ['started']
@pytest.mark.asyncio
async def test_get_stats_counts_correctly(self):
"""Test get_stats returns correct counts."""
@@ -11,9 +11,19 @@ Note: Uses import isolation to break circular import chains.
from __future__ import annotations
import sys
from unittest.mock import Mock, MagicMock
from contextlib import contextmanager
from typing import Generator
from unittest.mock import MagicMock, Mock
import pytest
@pytest.fixture(autouse=True)
def isolate_database_manager_registry(monkeypatch: pytest.MonkeyPatch) -> None:
"""Keep decorator tests from mutating the process-wide manager registry."""
from langbot.pkg.persistence import database
monkeypatch.setattr(database, 'preregistered_managers', list(database.preregistered_managers))
@contextmanager
@@ -0,0 +1,90 @@
from __future__ import annotations
from types import SimpleNamespace
import pytest
# Persistence manager performs the package's database-manager registration;
# importing a concrete manager first would enter the historical app/mgr cycle.
from langbot.pkg.persistence import mgr as _persistence_mgr # noqa: F401
from langbot.pkg.persistence.databases import postgresql
@pytest.mark.asyncio
async def test_postgresql_manager_parses_explicit_url_without_string_reassembly(monkeypatch) -> None:
captured = None
sentinel_engine = object()
def create_engine(url):
nonlocal captured
captured = url
return sentinel_engine
monkeypatch.setattr(postgresql.sqlalchemy_asyncio, 'create_async_engine', create_engine)
ap = SimpleNamespace(
instance_config=SimpleNamespace(
data={
'database': {
'postgresql': {
'url': 'postgresql://runtime:p%40ss@db.internal:5432/langbot?sslmode=require',
}
}
}
)
)
manager = postgresql.PostgreSQLDatabaseManager(ap)
await manager.initialize()
assert captured.drivername == 'postgresql+asyncpg'
assert captured.password == 'p@ss'
assert captured.query['ssl'] == 'require'
assert 'sslmode' not in captured.query
assert manager.engine is sentinel_engine
@pytest.mark.asyncio
async def test_postgresql_manager_builds_structured_url_with_special_password(monkeypatch) -> None:
captured = None
def create_engine(url):
nonlocal captured
captured = url
return object()
monkeypatch.setattr(postgresql.sqlalchemy_asyncio, 'create_async_engine', create_engine)
ap = SimpleNamespace(
instance_config=SimpleNamespace(
data={
'database': {
'postgresql': {
'host': 'db.internal',
'port': 5432,
'user': 'runtime',
'password': 'p@ss:/?#word',
'database': 'langbot',
}
}
}
)
)
await postgresql.PostgreSQLDatabaseManager(ap).initialize()
assert captured.password == 'p@ss:/?#word'
assert captured.host == 'db.internal'
assert captured.database == 'langbot'
@pytest.mark.asyncio
async def test_postgresql_manager_rejects_non_postgresql_url_without_echoing_secret() -> None:
ap = SimpleNamespace(
instance_config=SimpleNamespace(
data={'database': {'postgresql': {'url': 'sqlite:///operator-super-secret.db'}}}
)
)
manager = postgresql.PostgreSQLDatabaseManager(ap)
with pytest.raises(ValueError, match='valid PostgreSQL') as exc_info:
await manager.initialize()
assert 'operator-super-secret' not in str(exc_info.value)
@@ -0,0 +1,153 @@
from __future__ import annotations
import logging
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
import sqlalchemy
from langbot.__main__ import _build_parser
from langbot.pkg.persistence import release_migration
from langbot.pkg.persistence.mgr import PersistenceManager, PersistenceMode
def _cloud_config(*, database_use: str = 'postgresql', runtime_user: str = 'langbot_runtime') -> dict:
return {
'database': {
'use': database_use,
'postgresql': {
'host': 'runtime-db',
'port': 5432,
'user': runtime_user,
'password': 'runtime-secret',
'database': 'langbot',
},
'cloud_migration': {
'operator_dsn_env': 'TEST_LANGBOT_OPERATOR_DSN',
},
},
'vdb': {
'use': 'pgvector',
'pgvector': {
'use_business_database': True,
'allowed_dimensions': [384, 1536],
},
},
}
def _operator_environ(
*,
user: str = 'langbot_migrator',
database: str = 'langbot',
host: str = 'runtime-db',
port: int = 5432,
) -> dict[str, str]:
return {
'TEST_LANGBOT_OPERATOR_DSN': (f'postgresql://{user}:operator%40secret@{host}:{port}/{database}?sslmode=require')
}
def test_cloud_migration_cli_is_explicit() -> None:
args = _build_parser().parse_args(['migrate', '--cloud'])
assert args.command == 'migrate'
assert args.cloud is True
with pytest.raises(SystemExit) as exc_info:
_build_parser().parse_args(['migrate'])
assert exc_info.value.code == 2
def test_operator_url_is_separate_and_preserves_escaped_secret() -> None:
url = release_migration._operator_database_url(
_cloud_config(),
environ=_operator_environ(),
)
assert url.drivername == 'postgresql+asyncpg'
assert url.username == 'langbot_migrator'
assert url.password == 'operator@secret'
assert url.host == 'runtime-db'
assert url.port == 5432
assert url.database == 'langbot'
assert url.query['ssl'] == 'require'
assert 'sslmode' not in url.query
@pytest.mark.parametrize(
('config', 'environ', 'message'),
[
(_cloud_config(database_use='sqlite'), _operator_environ(), 'SQLite fallback is forbidden'),
(_cloud_config(), {}, 'requires the operator DSN'),
(_cloud_config(), {'TEST_LANGBOT_OPERATOR_DSN': 'not a secret://operator-password'}, 'DSN is invalid'),
(
_cloud_config(),
{'TEST_LANGBOT_OPERATOR_DSN': 'postgresql://operator:secret@runtime-db:not-a-port/langbot'},
'DSN is invalid',
),
(_cloud_config(), _operator_environ(user='langbot_runtime'), 'distinct operator role'),
(_cloud_config(), _operator_environ(database='another_database'), 'configured runtime database'),
(_cloud_config(), _operator_environ(host='other-cluster'), 'runtime PostgreSQL endpoint'),
(_cloud_config(), _operator_environ(port=6432), 'runtime PostgreSQL endpoint'),
],
)
def test_operator_url_rejects_unsafe_configuration(config: dict, environ: dict[str, str], message: str) -> None:
with pytest.raises(release_migration.CloudReleaseMigrationConfigurationError, match=message) as exc_info:
release_migration._operator_database_url(config, environ=environ)
assert 'operator-password' not in str(exc_info.value)
@pytest.mark.asyncio
async def test_release_migration_disposes_operator_engine_on_failure(monkeypatch) -> None:
engine = SimpleNamespace(dispose=AsyncMock())
manager = SimpleNamespace(
db=SimpleNamespace(engine=engine),
initialize=AsyncMock(side_effect=RuntimeError('migration failed')),
)
def manager_factory(*args, **kwargs):
del args, kwargs
return manager
monkeypatch.setattr(release_migration, 'PersistenceManager', manager_factory)
ap = SimpleNamespace(
instance_config=SimpleNamespace(data=_cloud_config()),
logger=logging.getLogger('release-migration-disposal-test'),
persistence_mgr=None,
)
with pytest.raises(RuntimeError, match='migration failed'):
await release_migration.run_cloud_release_migration(ap, environ=_operator_environ())
assert ap.persistence_mgr is manager
engine.dispose.assert_awaited_once()
@pytest.mark.asyncio
async def test_release_mode_rejects_sqlite_before_schema_changes(tmp_path, monkeypatch) -> None:
from langbot.pkg.persistence import mgr as persistence_mgr_module
from langbot.pkg.persistence.databases.sqlite import SQLiteDatabaseManager
monkeypatch.setattr(persistence_mgr_module.database, 'preregistered_managers', [SQLiteDatabaseManager])
ap = SimpleNamespace(
instance_config=SimpleNamespace(
data={
'database': {
'use': 'sqlite',
'sqlite': {'path': str(tmp_path / 'must-not-migrate.db')},
}
}
),
logger=logging.getLogger('release-migration-sqlite-rejection-test'),
)
manager = PersistenceManager(ap, mode=PersistenceMode.RELEASE_MIGRATION)
with pytest.raises(RuntimeError, match='requires PostgreSQL'):
await manager.initialize()
await manager.get_db_engine().dispose()
engine = sqlalchemy.create_engine(f'sqlite:///{tmp_path / "must-not-migrate.db"}')
try:
assert sqlalchemy.inspect(engine).get_table_names() == []
finally:
engine.dispose()
+350 -1
View File
@@ -1,11 +1,23 @@
from __future__ import annotations
import asyncio
import contextvars
from types import SimpleNamespace
import pytest
import sqlalchemy as sa
from sqlalchemy.ext.asyncio import create_async_engine
from langbot.pkg.core.task_boundary import create_detached_task
from langbot.pkg.entity.persistence.user import User
from langbot.pkg.persistence.mgr import PersistenceManager, PersistenceMode
from langbot.pkg.persistence.tenant_uow import TenantUnitOfWork
from langbot.pkg.persistence.tenant_uow import (
CrossScopeTransactionError,
PersistenceScopeKind,
TenantScopeRequiredError,
TenantUnitOfWork,
TransactionRollbackOnlyError,
)
pytestmark = pytest.mark.asyncio
@@ -61,3 +73,340 @@ def test_persistence_mode_must_be_a_trusted_enum() -> None:
manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
assert manager.mode is PersistenceMode.CLOUD_RUNTIME
async def test_manager_reuses_one_session_and_rejects_cross_workspace(tmp_path) -> None:
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "manager-uow.db"}')
table = sa.Table(
'manager_rows',
sa.MetaData(),
sa.Column('id', sa.Integer, primary_key=True),
sa.Column('workspace_uuid', sa.String(36), nullable=False),
)
manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
manager.db = SimpleNamespace(get_engine=lambda: engine)
try:
async with engine.begin() as conn:
await conn.run_sync(table.metadata.create_all)
with pytest.raises(TenantScopeRequiredError, match='explicit Workspace or discovery'):
await manager.execute_async(sa.select(table))
async with manager.tenant_uow('workspace-a') as outer:
await manager.execute_async(sa.insert(table).values(id=1, workspace_uuid='workspace-a'))
async with manager.tenant_uow('workspace-a') as inner:
assert inner.session is outer.session
assert manager.current_session() is outer.session
assert (await manager.execute_async(sa.select(table.c.id))).scalar_one() == 1
with pytest.raises(CrossScopeTransactionError, match='while workspace scope is active'):
async with manager.tenant_uow('workspace-b'):
pass
async with engine.connect() as conn:
assert (await conn.execute(sa.select(table.c.id))).scalars().all() == [1]
finally:
await engine.dispose()
async def test_manager_scoped_execute_preserves_core_row_and_scalar_contract(tmp_path) -> None:
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "manager-result-contract.db"}')
manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
manager.db = SimpleNamespace(get_engine=lambda: engine)
try:
async with engine.begin() as conn:
await conn.run_sync(User.__table__.create)
await conn.execute(
sa.insert(User).values(
uuid='account-a',
user='owner@example.com',
normalized_email='owner@example.com',
password='hashed-password',
)
)
async with manager.tenant_uow('workspace-a') as uow:
result = await manager.execute_async(sa.select(User))
row = result.first()
assert row is not None
assert row.uuid == 'account-a'
assert row.user == 'owner@example.com'
scalar_result = await manager.execute_async(sa.select(User.uuid))
assert scalar_result.scalar_one() == 'account-a'
list_result = await manager.execute_async(sa.select(User.user))
assert list_result.scalars().all() == ['owner@example.com']
# Direct UoW execution remains an ORM API for code that opts into it.
orm_result = await uow.execute(sa.select(User))
assert orm_result.scalars().one().uuid == 'account-a'
finally:
await engine.dispose()
async def test_transaction_free_tenant_scope_opens_one_short_uow_per_database_call(tmp_path) -> None:
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "short-scope.db"}')
table = sa.Table(
'short_scope_rows',
sa.MetaData(),
sa.Column('id', sa.Integer, primary_key=True),
sa.Column('workspace_uuid', sa.String(36), nullable=False),
)
manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
manager.db = SimpleNamespace(get_engine=lambda: engine)
try:
async with engine.begin() as conn:
await conn.run_sync(table.metadata.create_all)
async with manager.tenant_scope('workspace-a'):
assert manager.current_scope() is not None
assert manager.current_scope().kind is PersistenceScopeKind.WORKSPACE
assert manager.current_scope().settings == (('langbot.workspace_uuid', 'workspace-a'),)
assert manager.current_session() is None
await manager.execute_async(sa.insert(table).values(id=1, workspace_uuid='workspace-a'))
assert manager.current_session() is None
# The first statement has already committed. Long external waits
# inside this boundary retain only identity, never a DB session.
await asyncio.sleep(0)
assert manager.current_session() is None
assert (await manager.execute_async(sa.select(table.c.id))).scalar_one() == 1
assert manager.current_session() is None
async with manager.tenant_scope('workspace-a'):
assert manager.current_session() is None
with pytest.raises(CrossScopeTransactionError, match='while workspace scope is active'):
async with manager.tenant_scope('workspace-b'):
pass
with pytest.raises(CrossScopeTransactionError, match='while workspace scope is active'):
async with manager.tenant_uow('workspace-b'):
pass
assert manager.current_scope() is None
with pytest.raises(TenantScopeRequiredError, match='explicit Workspace'):
await manager.execute_async(sa.select(table))
finally:
await engine.dispose()
async def test_transaction_free_scope_requires_explicit_child_task_scope(tmp_path) -> None:
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "short-scope-child.db"}')
table = sa.Table('child_scope_rows', sa.MetaData(), sa.Column('id', sa.Integer, primary_key=True))
manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
manager.db = SimpleNamespace(get_engine=lambda: engine)
try:
async with engine.begin() as conn:
await conn.run_sync(table.metadata.create_all)
async with manager.tenant_scope('workspace-a'):
async def inherited_access() -> None:
await manager.execute_async(sa.select(table))
with pytest.raises(CrossScopeTransactionError, match='cannot be inherited by child tasks'):
await asyncio.create_task(inherited_access())
async def explicitly_scoped_access() -> None:
async with manager.tenant_scope('workspace-a'):
await manager.execute_async(sa.insert(table).values(id=1))
assert manager.current_session() is None
await asyncio.create_task(explicitly_scoped_access())
assert manager.current_session() is None
async with engine.connect() as conn:
assert (await conn.execute(sa.select(table.c.id))).scalars().all() == [1]
finally:
await engine.dispose()
async def test_caught_nested_failure_marks_outer_transaction_rollback_only(tmp_path) -> None:
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "rollback-only.db"}')
table = sa.Table('rollback_rows', sa.MetaData(), sa.Column('id', sa.Integer, primary_key=True))
manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
manager.db = SimpleNamespace(get_engine=lambda: engine)
try:
async with engine.begin() as conn:
await conn.run_sync(table.metadata.create_all)
with pytest.raises(TransactionRollbackOnlyError, match='transaction was rolled back'):
async with manager.tenant_uow('workspace-a'):
await manager.execute_async(sa.insert(table).values(id=1))
try:
async with manager.tenant_uow('workspace-a'):
raise ValueError('caught nested failure')
except ValueError:
pass
async with engine.connect() as conn:
assert (await conn.execute(sa.select(table))).all() == []
finally:
await engine.dispose()
@pytest.mark.parametrize('executor_kind', ['manager', 'uow', 'session'])
async def test_caught_database_error_rolls_back_and_cancels_after_commit_gate(tmp_path, executor_kind: str) -> None:
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / f"db-error-{executor_kind}.db"}')
table = sa.Table('unique_rows', sa.MetaData(), sa.Column('id', sa.Integer, primary_key=True))
manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
manager.db = SimpleNamespace(get_engine=lambda: engine)
try:
async with engine.begin() as conn:
await conn.run_sync(table.metadata.create_all)
with pytest.raises(TransactionRollbackOnlyError, match='after-commit work was cancelled'):
async with manager.tenant_uow('workspace-a') as uow:
statement = sa.insert(table).values(id=1)
if executor_kind == 'manager':
await manager.execute_async(statement)
elif executor_kind == 'uow':
await uow.execute(statement)
else:
await uow.session.execute(statement)
gate = manager.create_after_commit_gate()
assert gate is not None
try:
if executor_kind == 'manager':
await manager.execute_async(statement)
elif executor_kind == 'uow':
await uow.execute(statement)
else:
await uow.session.execute(statement)
except sa.exc.IntegrityError:
pass
assert gate.cancelled()
async with engine.connect() as conn:
assert (await conn.execute(sa.select(table))).all() == []
finally:
await engine.dispose()
async def test_child_task_must_open_its_own_explicit_uow(tmp_path) -> None:
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "child-task.db"}')
table = sa.Table(
'child_rows',
sa.MetaData(),
sa.Column('id', sa.Integer, primary_key=True),
sa.Column('workspace_uuid', sa.String(36), nullable=False),
)
manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
manager.db = SimpleNamespace(get_engine=lambda: engine)
try:
async with engine.begin() as conn:
await conn.run_sync(table.metadata.create_all)
async with manager.tenant_uow('workspace-a'):
async def inherited_access() -> None:
await manager.execute_async(sa.select(table))
with pytest.raises(CrossScopeTransactionError, match='cannot be inherited by child tasks'):
await asyncio.create_task(inherited_access())
async def explicitly_scoped_access() -> None:
async with manager.tenant_uow('workspace-a'):
await manager.execute_async(sa.insert(table).values(id=2, workspace_uuid='workspace-a'))
await asyncio.create_task(explicitly_scoped_access())
async with engine.connect() as conn:
assert (await conn.execute(sa.select(table.c.id))).scalars().all() == [2]
finally:
await engine.dispose()
async def test_detached_task_starts_without_parent_scope_and_rolls_back_its_uow(tmp_path) -> None:
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "detached-task.db"}')
table = sa.Table(
'detached_rows',
sa.MetaData(),
sa.Column('id', sa.Integer, primary_key=True),
sa.Column('workspace_uuid', sa.String(36), nullable=False),
)
manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
manager.db = SimpleNamespace(get_engine=lambda: engine)
try:
async with engine.begin() as conn:
await conn.run_sync(table.metadata.create_all)
async with manager.tenant_uow('workspace-a'):
async def detached_write() -> None:
# Before the detached boundary this access raises
# CrossScopeTransactionError because asyncio copies the
# parent's ActiveScopedTransaction into this child task.
assert manager.current_scope() is None
async with manager.tenant_uow('workspace-a'):
await manager.execute_async(sa.insert(table).values(id=2, workspace_uuid='workspace-a'))
raise RuntimeError('roll back detached write')
task = create_detached_task(detached_write())
with pytest.raises(RuntimeError, match='roll back detached write'):
await task
await manager.execute_async(sa.insert(table).values(id=1, workspace_uuid='workspace-a'))
async with engine.connect() as conn:
assert (await conn.execute(sa.select(table.c.id))).scalars().all() == [1]
finally:
await engine.dispose()
async def test_after_commit_task_waits_for_commit_and_starts_with_empty_context(tmp_path) -> None:
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "after-commit.db"}')
manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
manager.db = SimpleNamespace(get_engine=lambda: engine)
request_value = contextvars.ContextVar('after_commit_request_value', default=None)
observed = []
try:
async with manager.tenant_uow('workspace-a'):
token = request_value.set('request-scope')
async def after_commit_work() -> None:
observed.append((request_value.get(), manager.current_scope()))
try:
task = create_detached_task(
after_commit_work(),
after_commit_manager=manager,
)
await asyncio.sleep(0)
assert observed == []
finally:
request_value.reset(token)
await task
assert observed == [(None, None)]
finally:
await engine.dispose()
async def test_after_commit_task_is_cancelled_and_coroutine_closed_on_rollback(tmp_path) -> None:
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "after-rollback.db"}')
manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
manager.db = SimpleNamespace(get_engine=lambda: engine)
started = False
async def should_not_start() -> None:
nonlocal started
started = True
coro = should_not_start()
try:
with pytest.raises(RuntimeError, match='rollback request'):
async with manager.tenant_uow('workspace-a'):
task = create_detached_task(coro, after_commit_manager=manager)
await asyncio.sleep(0)
assert not task.done()
raise RuntimeError('rollback request')
await asyncio.sleep(0)
assert task.cancelled()
assert not started
assert coro.cr_frame is None
finally:
await engine.dispose()
+39 -5
View File
@@ -13,8 +13,11 @@ from __future__ import annotations
import pytest
import asyncio
import contextvars
from contextlib import asynccontextmanager
from unittest.mock import Mock, AsyncMock
from importlib import import_module
from types import SimpleNamespace
from tests.factories import (
FakeApp,
@@ -894,17 +897,48 @@ class TestMessageAggregatorWorkspaceIsolation:
app = make_aggregator_app()
enable_aggregation(app)
agg = get_aggregator_module().MessageAggregator(app)
delayed_flush = AsyncMock()
request_value = contextvars.ContextVar('aggregator_request_value', default=None)
token = request_value.set('request-scope')
observed = []
async def delayed_flush(*args):
observed.append((request_value.get(), args[2]))
monkeypatch.setattr(agg, '_delayed_flush', delayed_flush)
context = execution_context(pipeline_uuid='test-pipeline')
await agg.add_message(**scoped_message_kwargs(context))
await asyncio.sleep(0)
try:
await agg.add_message(**scoped_message_kwargs(context))
await asyncio.sleep(0)
finally:
request_value.reset(token)
delayed_flush.assert_awaited_once()
assert delayed_flush.await_args.args[2] is context
assert observed == [(None, context)]
await agg.flush_all()
@pytest.mark.asyncio
async def test_delayed_flush_opens_explicit_workspace_uow(self, monkeypatch):
app = make_aggregator_app()
app.persistence_mgr.mode = SimpleNamespace(value='cloud_runtime')
scopes = []
@asynccontextmanager
async def tenant_uow(workspace_uuid):
scopes.append(workspace_uuid)
yield
app.persistence_mgr.tenant_uow = tenant_uow
agg = get_aggregator_module().MessageAggregator(app)
flush = AsyncMock()
monkeypatch.setattr(agg, '_flush_buffer', flush)
context = execution_context('workspace-a', pipeline_uuid='test-pipeline')
key = aggregation_key(context, pipeline_uuid='test-pipeline')
await agg._delayed_flush(key, 0, context)
assert scopes == ['workspace-a']
flush.assert_awaited_once_with(key, context)
@pytest.mark.asyncio
async def test_flush_rejects_context_from_another_workspace(self):
app = make_aggregator_app()

Some files were not shown because too many files have changed in this diff Show More