mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-07-25 13:56:08 +00:00
feat(tenancy): harden shared cloud runtime boundaries
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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 Runtime;Cloud 固定使用 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 Runtime;Cloud 固定使用 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-001:Plugin Runtime 多租户控制面
|
||||
|
||||
状态:`DECIDED — MVP target confirmed, implementation pending`
|
||||
状态:`IMPLEMENTED — Cloud deployment verification pending`
|
||||
|
||||
### 3.1 当前实现事实
|
||||
### 3.1 已实现的基础
|
||||
|
||||
- SDK `RuntimeContext` 和 Core `PluginRuntimeConnector` 当前都绑定一个 Workspace,Runtime 只有一个全局 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 storage;Runtime 本地缓存丢失后可由 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 进入 nsjail;installation 的 home/tmp/data 使用独立可写 mount。
|
||||
cache 可接受的签名/来源、撤销和 GC 规范属于后续发布规则,不改变本轮基于已验证 digest 的只读共享边界。
|
||||
- artifact 与按环境摘要构建的共享依赖环境以只读 mount 进入 nsjail;installation 的 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-002:Box 多租户控制面和套餐边界
|
||||
|
||||
状态:`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 和 generation,Runtime namespace 由可信 context 派生。
|
||||
- 短期 `SandboxAdmissionGrant`、revision tombstone 和原子 session admission 强制每个合资格 Workspace 最多一个 `global` persistent session,managed 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_NOFOLLOW;Cloud 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 session;Runtime 重启后进程和临时目录失效,但 `/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-003:SaaS PostgreSQL 与 pgvector
|
||||
|
||||
状态:`DECIDED — MVP target confirmed, implementation pending`
|
||||
状态:`PARTIALLY IMPLEMENTED — shared schema/pgvector complete; SaaS transaction and deployment gates remain`
|
||||
|
||||
当前实现仍有明确差距:Cloud 模板默认 Chroma;pgvector 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` ACL;runtime role 不得有效执行任何
|
||||
`SECURITY DEFINER` routine,包括被 allowlisted extension 收编的 routine。普通非 `SECURITY DEFINER` 内建函数的隐式执行权限不在此禁令内。
|
||||
- PostgreSQL 新 database 默认向 `PUBLIC` 提供的 `TEMP` 是首版在专用业务 database 上明确接受的兼容性决定,
|
||||
不是 migrator 对 runtime role 的直接 grant;首版不得据此把业务 database 与不受信任工作负载混用。
|
||||
migrator 在释放 advisory lock 前建立并校验上述精确 allowlist;Cloud 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,并拒绝相同 role;migrator/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-004:stdio 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
@@ -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
@@ -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
|
||||
|
||||
|
||||
@@ -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')
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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')
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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
@@ -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:
|
||||
|
||||
@@ -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)"
|
||||
|
||||
@@ -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,
|
||||
*,
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
@@ -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)
|
||||
|
||||
@@ -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
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,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
|
||||
|
||||
@@ -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')
|
||||
@@ -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)
|
||||
|
||||
+1255
-20
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)
|
||||
@@ -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()
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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}'
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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}')
|
||||
|
||||
|
||||
+1347
-385
File diff suppressed because it is too large
Load Diff
@@ -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
@@ -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),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@@ -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')
|
||||
|
||||
@@ -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')
|
||||
|
||||
|
||||
@@ -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',
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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}')
|
||||
|
||||
|
||||
@@ -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
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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'
|
||||
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
|
||||
@@ -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]
|
||||
@@ -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
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
|
||||
@@ -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
Reference in New Issue
Block a user