xinyin025/LangBot
1
0
Fork 0
mirror of https://github.com/langbot-app/LangBot.git synced 2026-07-25 13:56:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(tenancy): connect cloud workspace control plane

Browse Source
This commit is contained in:
Junyan Qin
2026-07-24 19:11:33 +08:00
parent d7cdd206c2
commit 98f45aa88e
40 changed files with 4159 additions and 452 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/multi-tenant/implementation-checklist.md
+11 -7
View File
@@ -8,7 +8,7 @@ verification gates. Exact commands and observed results are recorded in the
- [x] LangBot uses branch feat/multi-tenants.
- [x] langbot-plugin-sdk uses branch feat/multi-tenants.
- [x] langbot-space has no changes made by this implementation; Cloud v2 does not extend the legacy Space deployment scheme.
- [x] langbot-space implements the greenfield Cloud v2 modular-monolith control plane without extending the legacy per-account Pod topology; the old Pod UI remains available when Cloud v2 is disabled and only retained Pods appear in the v2 view.
- [x] Unrelated untracked files in either repository remain untouched.
- [x] Open-source startup cannot enable SaaS multi-workspace through edition flags or unsigned configuration.
@@ -20,18 +20,22 @@ 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.
- [x] The closed Control Plane owns the global Account, Workspace, Membership, and Invitation directory.
- [ ] The closed Control Plane execution-ownership module issues monotonic generations and owner leases for projected Workspaces.
- [ ] Core verifies a signed `InstanceManifest` before the closed bootstrap can inject `CloudWorkspacePolicy`.
- [x] Core verifies a signed `InstanceManifest` before the closed bootstrap can inject `CloudWorkspacePolicy`.
- [ ] Tenant database writes hold a generation-aware shared transaction fence through commit, while execution-owner 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 an execution-generation change through stable published keys or an explicitly atomic key/reference migration.
- [ ] 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.
- [ ] Entitlement checks, usage aggregation, and subscription lifecycle are implemented in the closed Control Plane; production activation still requires provider callback amount/currency/session/expiry binding inside the locked fulfillment transaction.
- [ ] Account registration persists a `new_api.provision_account` outbox item with the Account and personal Workspace, and an in-process reconciler provisions New API idempotently after commit.
- [ ] EPay and Stripe callbacks bind provider identity, amount, currency, channel/session, payment status, and expiry to the locked order before entitlement fulfillment.
- [ ] 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 Plugin Runtime Supervisor automatically restores an unexpectedly exited enabled worker with bounded backoff and cannot create a cross-tenant restart storm.
- [ ] The Plugin Runtime shared profile refuses to run without delegated cgroup v2 CPU, memory-plus-swap, and PID limits, all verified in a real Linux container; production tenant-safe egress remains incomplete.
- [x] The Plugin Runtime Supervisor automatically restores an unexpectedly exited enabled worker with bounded per-installation backoff.
- [ ] Jitter, a global restart concurrency limit, and a Runtime-level circuit breaker prevent a systemic failure from creating a cross-tenant restart storm.
- [ ] Until authenticated Runtime takeover or an owner lease/fence exists, the M0 deployment rolls Core and Plugin Runtime together and forbids an independent Core-only rollout.
- [ ] Plugin installation data has an operator-owned hard disk quota provider that atomically rejects writes over the limit; directory scans are not accepted as enforcement.
- [ ] 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.
@@ -281,7 +285,7 @@ Each row type must have a non-null Workspace UUID, scoped indexes, scoped unique
## 10. Completion evidence
- [x] LangBot and SDK branch refs are recorded in the verification report.
- [x] Space git diff is empty relative to the pre-work snapshot.
- [x] Space contains the closed adapter package and Cloud v2 control plane, billing, migration, and Workspace UI changes; unrelated pre-existing files remain unstaged.
- [x] Migration output is captured for SQLite and PostgreSQL.
- [x] Test commands and results are recorded.
- [x] Browser E2E actions and observed results are recorded.
docs/multi-tenant/implementation-decisions.md
+33 -5
View File
@@ -13,7 +13,7 @@ This log records implementation choices made while delivering the Workspace arch
- Decision: Community builds create exactly one Workspace per LangBot instance and allow multiple Accounts through invitations.
- Reason: This preserves a simple self-hosted deployment while making authorization and ownership explicit. Creating a second Workspace is an edition error, not a hidden fallback.
- SaaS boundary: Multi-Workspace directory, execution ownership, entitlement, and billing are the responsibility of a separate closed SaaS Control Plane. Core consumes a validated projection and remains the final isolation and authorization enforcement point; it does not become the SaaS system of record or billing engine.
- Deployment boundary: Cloud v2 is a greenfield deployment design. The previous per-account instance/pod scheme is not migrated, preserved for compatibility, or extended by this implementation. Existing OAuth, marketplace, and payment concepts may be reused only through explicit adapters where they still fit; `langbot-space` itself is not changed.
- Deployment boundary: Cloud v2 is a greenfield deployment design. The previous per-account instance/pod scheme is not migrated or extended and remains available only for existing subscriptions. Existing OAuth, marketplace, and payment rails are reused through explicit adapters where they still fit; new Workspace, subscription, entitlement, directory, and usage modules live alongside the legacy Pod flow in `langbot-space`.
### Workspace selection is trusted only after authentication
@@ -41,7 +41,7 @@ This log records implementation choices made while delivering the Workspace arch
- Decision: One instance-scoped Plugin Runtime control plane serves all Workspaces in the logical LangBot instance. Each running plugin installation has its own nsjail worker with an immutable binding containing `instance_uuid`, `workspace_uuid`, `execution_generation` (stored as the compatibility field `placement_generation` until the schema rename), `installation_uuid`, `runtime_revision`, and verified artifact digest; enabled-resident is the desired semantic. 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. Enabled-resident is the desired semantic, but the current Supervisor restores an unexpectedly exited worker only on the next Core apply/reconcile; a completion callback, bounded backoff, and cross-tenant restart-storm isolation remain Cloud activation gates.
- 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. The installation Supervisor now restores an unexpectedly exited enabled worker through a completion callback with bounded exponential backoff. Jitter, global restart concurrency limits, and a Runtime-wide circuit breaker are still required to prove that an infrastructure-wide failure cannot create a cross-tenant restart storm.
- 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.
@@ -83,8 +83,8 @@ This log records implementation choices made while delivering the Workspace arch
### Unreleased SDK protocol is pinned reproducibly without publishing
- Decision: The SDK tenancy protocol is versioned as 0.4.15. This task does not create a GitHub release or publish PyPI because the user authorized pushing code, not a package release. After the SDK feature branch is final, LangBot's feature branch temporarily pins the exact pushed SDK Git commit. Before merging to master, the release gate is to publish `langbot-plugin==0.4.15` and replace the Git pin with the registry pin.
- Reason: PyPI 0.4.14 does not contain `ActionContext`; pinning it makes clean installs fail, while pinning an unpublished 0.4.15 makes dependency sync impossible. An exact Git commit is reproducible and keeps the feature branch testable without expanding release authority.
- Decision: The SDK tenancy protocol is versioned as 0.4.18. This task does not create a GitHub release or publish PyPI because the user authorized pushing code, not a package release. After the SDK feature branch is final, LangBot's feature branch temporarily pins the exact pushed SDK Git commit. Before merging to master, the release gate is to publish `langbot-plugin==0.4.18` and replace the Git pin with the registry pin.
- Reason: The current registry release does not contain the complete tenant action context and shared Runtime hardening. An exact Git commit is reproducible and keeps the feature branch testable without expanding release authority.
### Cloud directory writes stay outside Core
@@ -240,5 +240,33 @@ This log records implementation choices made while delivering the Workspace arch
### Shared Plugin Runtime starts only from verified desired state
- Decision: SDK shared mode waits for immutable runtime configuration before inspecting plugin state, never scans or launches legacy `data/plugins`, and rejects legacy install/restart/delete/upgrade control actions. Worker RPC files use installation-private directories with aggregate size enforcement, and resident nsjail workers explicitly disable the default 600-second wall-time limit.
- Remaining gate: Unexpected worker exit is still recovered only by the next Core apply/reconcile. Completion callbacks, bounded backoff, hard installation disk quota, and production Linux/cgroup/egress evidence remain Cloud activation requirements.
- Remaining gate: completion-callback recovery and bounded per-installation backoff are implemented. Cross-tenant restart-storm suppression, hard installation disk quota, and production egress policy remain Cloud activation requirements; Linux nsjail/cgroup CPU, memory-plus-swap, PID, namespace, and cgroup-reaping behavior have real-container evidence.
- Reason: A shared supervisor reduces per-Workspace services only if legacy global paths, writable transfer state, and lifecycle defaults cannot bypass installation isolation.
## 2026-07-24
### Cloud v2 remains a modular monolith with one closed adapter
- Decision: Workspace directory, plans, subscriptions, payment fulfillment, entitlements, usage, and signed runtime feeds are implemented as modules in the existing Space service and PostgreSQL database. The only separately packaged closed component is the thin Core bootstrap/control-plane adapter; it verifies signed data and owns no business state.
- Compatibility: The legacy Pod card and fulfillment path remain visible only to accounts with old subscriptions. New purchases create Workspace subscriptions and never provision a per-user LangBot Pod.
- Reason: A separate tenancy or billing service would add deployment, queue, network, and consistency cost without providing an isolation boundary. The signed adapter keeps SaaS logic closed while Core retains the ORM, RLS, authorization, and runtime enforcement boundary.
### Registration creates data, not infrastructure
- Decision: Account registration and personal Workspace creation share one transaction and include an active owner membership, Free subscription, entitlement snapshot, and outbox notifications. Repair is idempotent for older accounts. Empty Workspace creation starts no Plugin worker, Box sandbox, database, queue, bucket, or tenant service.
- Activation: This automatic Cloud v2 transaction is enabled only when Space has an explicit valid `CLOUD_V2_INSTANCE_UUID`. A legacy marketplace-only Space deployment with no Cloud v2 instance configured keeps its existing Account registration path; Cloud v2 internal endpoints still fail closed instead of inventing an instance identity.
- Billing boundary: Core and runtimes consume only generic capability and numeric-limit entitlements. They never branch on `free` or `pro`; plan names, prices, payment providers, and fulfillment stay in Space.
- Reason: This keeps the marginal cost of a new user close to a few PostgreSQL rows while preserving an immediate, usable Workspace.
### Directory bootstrap is full; steady-state projection is per Workspace
- Decision: Space generates the initial signed directory snapshot and its outbox high-water mark in one read-only PostgreSQL `REPEATABLE READ` transaction. After bootstrap, each event page names affected Workspaces and Core fetches one signed `directory.delta` for that set. Missing requested Workspaces are tombstones; unrelated Workspaces are untouched.
- Cursor boundary: A delta carries no event cursor. Each signed event page carries the transaction-consistent current high-water mark, while Core advances only to the final event it actually consumed, even when the authoritative delta already contains a later revision. Event payload Workspace and revision fields must exactly match the signed event envelope; readiness is renewed only after the replica-local cursor reaches the signed high-water and the shared projection is not ahead.
- Replica boundary: Every Core replica owns a process-local consumer cursor because its verified entitlement cache is process-local. Replicas share the PostgreSQL projection high-water mark and inbox. The state separately records which cursor range was atomically subsumed by a full snapshot, so a lagging replica can add missing receipts within that coverage and refresh its local caches without repeating tenant mutations; a missing receipt beyond snapshot coverage fails closed.
- Reason: A shared consumer cursor would let one replica starve another replica's local cache, while a full snapshot on every change would make steady-state cost grow with every registered Workspace.
### Cloud OAuth can authenticate only an existing projected Account
- Decision: In multi-Workspace Cloud mode, Space OAuth may refresh tokens only for an active `cloud_projection` Account whose Space subject UUID and normalized email exactly match. It cannot create an Account, relink by email, mutate directory identity, or choose a Workspace.
- Redirect boundary: When Cloud v2 configures its Core public URL, Space issues authorization codes only to that exact callback origin or the explicitly retained legacy managed-Pod domain. Community installations remain dynamic OAuth clients because they cannot pre-register with the public Space service: the consent screen displays their hostname, and only HTTPS or loopback HTTP with the fixed callback path is accepted. Remote HTTP, userinfo, fragments, arbitrary callback paths, and unrecognized query parameters always fail closed.
- Reason: Space is the SaaS identity and directory authority, but Core remains the authentication enforcement point. Projected-only matching avoids split-brain identities; redirect allowlisting prevents bearer authorization-code exfiltration.
docs/multi-tenant/pending-architecture-decisions.md
+16 -9
View File
@@ -2,7 +2,7 @@
状态：`DECIDED — Core isolation kernel implemented; SaaS activation gates remain`
创建日期：2026-07-19
最近更新：2026-07-20
最近更新：2026-07-24
本文记录 Cloud v2 多租户架构中已经确认的首期决策、明确淘汰的方案和仍需在后续阶段决定的扩展项。
本文同时记录实现状态。“实现完成”仅指开源 Core/SDK 的隔离内核和 fail-closed 门禁，
@@ -28,10 +28,11 @@
| 编号 | 结论 | 首期状态 |
| ----- | --------------------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------- |
| D-001 | 一个共享 Plugin Runtime 控制面；每个运行中的 plugin installation 独占一个 nsjail 子进程；只有 digest 相同且已验证的代码 artifact 可以只读共享 | `FOUNDATION IMPLEMENTED — Linux/egress, crash recovery and disk-quota gates pending` |
| D-001 | 一个共享 Plugin Runtime 控制面；每个运行中的 plugin installation 独占一个 nsjail 子进程；只有 digest 相同且已验证的代码 artifact 可以只读共享 | `FOUNDATION VERIFIED — egress, total disk-quota and restart-storm gates 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` |
| D-005 | 目录启动使用事务一致的全量快照，运行时按事件涉及的 Workspace 拉取增量；每个 Core replica 独立消费事件，共享 PostgreSQL 投影和 inbox | `IMPLEMENTED — production fault injection pending` |
Workspace 的具体创建、释放、数据导出和单 Workspace 恢复机制不在本轮决定；本文只保证这些后续能力不会改变稳定的
`workspace_uuid`，也不会要求重建租户专属部署。
@@ -76,13 +77,15 @@ flowchart LR
M1 是 M0 的透明扩容，M2 是相同架构下的资源等级；两者都不是新的 LangBot 实例、Cell 或 CloudInstance。
外部 API 只认识稳定的 `instance_uuid` 和 `workspace_uuid`，不认识 replica、worker、pool 或 shard。
Plugin Runtime 与 Core 在 M0 是否共用 Pod 仍可按发布和故障域决定；即使共用 Pod，也必须使用独立容器和 security context，
Core 不能继承 Plugin Runtime 所需的 nsjail/cgroup 权限。Box Runtime 同样使用独立进程身份和安全配置，
Plugin Runtime 与 Core 在 M0 使用独立容器和 security context，Core 不能继承 Plugin Runtime 所需的
nsjail/cgroup 权限。当前 Runtime 在进程生命周期内绑定首次认证的 `runtime_id`，因此 M0 必须把 Core 与
Plugin Runtime 放在同一 rollout/restart unit 中协调重启；在实现受认证 takeover 或 owner lease/fencing 前，
不能单独滚动 Core 并让它接管仍存活的 Runtime。Box Runtime 同样使用独立进程身份和安全配置，
不与 Plugin Runtime 合并成一个高权限进程。
## 3. D-001：Plugin Runtime 多租户控制面
状态：`FOUNDATION IMPLEMENTED — Cloud deployment and worker-recovery gates pending`
状态：`FOUNDATION VERIFIED — Cloud egress, total disk-quota and restart-storm gates pending`
### 3.1 已实现的基础
@@ -116,8 +119,8 @@ Core 不能继承 Plugin Runtime 所需的 nsjail/cgroup 权限。Box Runtime
- 安装并启用插件后，Supervisor 在自己的 Runtime 容器内直接启动一个 nsjail 子进程；
不再为每个插件创建 nested container、Pod、sidecar 或租户级 Runtime service。
- desired semantics 要求 enabled installation 保持 resident，不做 idle eviction；停用、删除、revision/generation 变化或 entitlement 撤销时停止并按需重建。
当前实现只会在 Runtime 重连或 Core apply/reconcile 时恢复意外退出的 worker，尚缺 completion callback、有界 backoff 和跨租户重启风暴抑制；
这属于 Cloud 激活门禁，不能把 desired semantics 描述成已经具备即时自愈。
当前 Supervisor 已通过 completion callback 和有界指数 backoff 恢复意外退出的 worker；尚缺 jitter、全局重启并发上限和 Runtime 级 circuit breaker，
因而系统性依赖或宿主故障下的跨租户重启风暴抑制仍是 Cloud 激活门禁。
- 子进程使用一次性 registration capability 向 Supervisor 注册；capability 由可信 desired state 派生并绑定完整 installation tuple，
不是插件直接建立 Core Host connection，也不能只绑定 author/name/path。Supervisor/Core 据此注入 tenant context，
丢弃插件 payload 中自带的 scope 字段。
@@ -473,9 +476,9 @@ mcp:
- 直接调用 API、重放旧配置和启动 bootstrap 都失败关闭，且不会产生 `mcp-shared` session、nsjail 进程或额外配额占用。
- OSS 默认行为保持兼容；HTTP/SSE MCP 正常工作。
## 7. 四项决策之间的关系
## 7. 五项决策之间的关系
四项决策共同遵循：
五项决策共同遵循：
> 多租户共享可信控制面、连接池、只读 artifact 和基础容量；租户独占不可信执行进程、sandbox、secret、可写文件和数据作用域。
@@ -485,6 +488,7 @@ mcp:
- PostgreSQL 和 pgvector 共享数据库组件，但使用显式 tenant key、应用层 scope 和 RLS 防止共享存储变成共享权限。
- 订阅管理只在闭源 Control Plane 维护套餐与计费规则，并向开源 Core 投影签名/版本化 entitlement；
Core/Runtime 执行通用 capability 和数值限额，不复制套餐名称或计费逻辑。
- 目录同步只在启动或恢复时读取全量快照；常态变更按 Workspace 聚合为签名增量，新增租户不会使每次目录事件退化为全实例重投影。
## 8. 当前不做分布式时仍保留的能力
@@ -497,6 +501,9 @@ mcp:
6. schema migration、任务扫描、监控聚合和运维接口不能假设永远只有一个 Core 进程。
7. 外部 API 不暴露 replica、worker 或 shard 标识；未来扩容不改变 Workspace URL、UUID 或客户端协议。
8. 只有出现容量、可用性、地域或合规需求时才增加 replica/shard；预留协议不等于现在部署额外组件。
9. 每个 Core replica 保存自己的事件消费 cursor，因为 entitlement cache 是进程本地状态；PostgreSQL 中的 projection high-water mark
、snapshot coverage 和 inbox 仍由所有 replica 共享，用于幂等投影和冲突检测。事件页携带签名 high-water，副本追平前不能续期
ready；不能让一个 replica 的共享 cursor 使其他 replica 跳过本地 cache 刷新。
## 9. 本轮明确不做的事情
docs/multi-tenant/verification-report.md
+199 -204
View File
@@ -1,257 +1,252 @@
# Multi-tenant isolation kernel verification report
# Cloud v2 multi-tenant verification report
Date: 2026-07-20
Date: 2026-07-24
Status: `CORE ISOLATION FOUNDATION VERIFIED — SAAS ACTIVATION REMAINS DISABLED`
Status: `FOUNDATION AND CONTROL PLANE VERIFIED — PRODUCTION ACTIVATION REMAINS GATED`
This report records the final implementation and verification evidence for the
shared Workspace isolation foundation. It does not claim that the closed SaaS
Control Plane, billing system, or greenfield Cloud v2 deployment is
production-ready.
This report records the implementation and verification evidence for one
logical LangBot instance serving multiple Workspace tenants. It covers the
open-source Core, the shared Plugin/Box runtimes, the closed Space adapter, and
the Space Cloud v2 modular-monolith control plane. It does not claim that the
production Cloud deployment may enable `CLOUD_V2_ENABLED` yet.
## Repository refs and scope
- LangBot branch: `feat/multi-tenants`
- LangBot implementation commit:
`90a977488212d3f3f50fb92d94e3fcae60f27eab`
- LangBot base: `origin/master` at
`6baeb032a7f76c65337b51c7b58593de4687a61c`
- Plugin SDK branch: `feat/multi-tenants`
- Plugin SDK commit and LangBot dependency pin:
`95a1805af2038c745de2c018a00db1305089a32e`
- No tracked `langbot-space` change was made. Cloud v2 does not extend or
preserve the legacy per-account instance/pod deployment topology.
- LangBot Core: branch `feat/multi-tenants`
- langbot-plugin-sdk: branch `feat/multi-tenants`, commit
`e7d946af4a6b1494fbe74627c1815ace19ac8991`
- langbot-space: branch `feat/cloud-v2-control-plane`
- Closed Core adapter: `langbot-space/cloud-adapter`
- SDK protocol/package version: `0.4.18`
The SDK protocol is versioned as 0.4.15 and is intentionally consumed from the
exact pushed Git commit above. Publishing the release and replacing the Git
pin with a registry pin remain merge-to-master release work; no registry
release was performed here.
Core pins the exact pushed SDK commit above. Cloud v2 is a greenfield
multi-Workspace deployment and does not provision one Pod, database, queue, or
Runtime per tenant. Legacy Space Pods remain a compatibility surface only.
## Implemented boundaries
## Implemented product and security boundaries
### Workspace, identity, and authorization
### Workspace and identity
- Community bootstraps exactly one local Workspace per LangBot instance while
allowing multiple Accounts, memberships, invitations, and role-based
permissions inside it. A mutable edition setting cannot activate SaaS
multi-Workspace routing.
- Each Account registration resolves or creates its local singleton Workspace.
Invitation acceptance is one-time, email-bound, and survives the signed-out
login transition without exposing the invitation secret in URLs after the
initial fragment capture.
- The Account-token Workspace discovery route is intentionally separate from
Workspace-authorized routes. `X-Workspace-Id` becomes authority only after
membership resolution; API keys and public resources derive Workspace from
their trusted owner.
- Ordinary monitoring uses `resource.view`. Export remains `data.export`,
privileged runtime/system audit remains `audit.view`, and frontend controls
follow the same permission split.
- SaaS has one stable `instance_uuid`; `workspace_uuid` is the tenant key.
- Registering an Account creates its personal Workspace and owner Membership in
the same PostgreSQL transaction.
- A Workspace may contain multiple users with owner, admin, and member roles.
Invitation acceptance is token-hash based, email-bound, single-use,
concurrency-safe, and member-limit checked while locked.
- Community remains exactly one local Workspace while supporting multiple
Accounts and fixed RBAC roles. Installing the closed adapter is required to
inject the SaaS policy; configuration alone cannot activate it.
- Disabled or deleted Accounts cannot use password/code/OAuth login, sessions,
refresh/access tokens, or personal access tokens.
### Persistence and PostgreSQL
### Closed Space control plane
- PostgreSQL uses one shared business schema with application scope plus exact
`ENABLE` and `FORCE ROW LEVEL SECURITY` policies. Tenant UoWs keep `SET LOCAL`
and business SQL on one owned root transaction and one task.
- The public tenant Session accepts only structured SQLAlchemy query/DML trees.
Raw/textual SQL, textual modifiers, custom AST/compiler nodes, unapproved
functions/operators/types/casts, hidden dialect values, foreign binds,
loader/execution options, Session event callbacks, ORM SQL-expression
attributes, transaction control, and live-result escape fail closed and make
the root transaction rollback-only.
- The SQL guard is a trusted-Core misuse boundary, not an in-process Python
sandbox. Mapped metadata and registered compilers are trusted boot-time code;
plugins remain out of process. SQLAlchemy dialect-private container fields
are pinned and covered by regression tests before dependency upgrades.
- Legacy local RAG restore remains structured and tenant-fenced while accepting
SQLite string-valued timestamps and both historical PostgreSQL `TEXT` and
fresh-schema `JSON` settings columns.
- The one-shot release migrator owns DDL and grants a distinct runtime role only
the required business-table DML, `alembic_version` read access, sequence
access, database `CONNECT`, and schema `USAGE`. Runtime startup reruns the
role, ACL, schema, session, extension, routine, parameter, and RLS audit.
- pgvector is stored in the same business database with Workspace-scoped keys,
checked dimensions, release-created partial ANN indexes, and RLS.
- `CLOUD_V2_ENABLED` defaults to false. Invalid or incomplete configuration
fails startup, and disabled endpoints return 503.
- Space owns Workspace/Membership/Invitation, versioned plans, subscriptions,
entitlements, usage events, directory outbox, and signed instance manifests.
- Free and Pro are Workspace plans. Pro projects one managed global sandbox;
Free projects none. Both force stdio MCP off in the first release.
- Subscription changes use Workspace advisory locking, expected entitlement
revision, one-live-order uniqueness, idempotent usage ingestion, period-end
downgrade, renewal, and lazy expiry settlement.
- Directory snapshot and per-Workspace delta reads use repeatable-read,
read-only transactions with a high-water cursor. Core stores a replica-local
consumer cursor and shared PostgreSQL projection state, snapshot coverage,
and inbox rows.
- The `/cloud` page selects a Workspace and shows its independent subscription,
entitlement, limits, and usage. When Cloud v2 is disabled or the backend does
not expose the feature field, the complete legacy Welcome/Pod UI is used.
### Shared Plugin Runtime and Box
### PostgreSQL and pgvector
- One instance-scoped Plugin Runtime supervisor serves multiple Workspaces.
Every enabled installation still owns one nsjail worker with an immutable
instance/Workspace/generation/installation/revision/artifact binding.
- Verified same-digest code and dependency environments may be mounted
read-only and reused. Plugin processes and writable installation state are
never merged, including for the same plugin and version.
- Runtime limits come only from instance configuration with native environment
overrides. Plugin manifests cannot raise limits. Shared mode rejects legacy
global plugin directories and legacy lifecycle control paths.
- One shared Box control plane admits at most one persistent logical `global`
sandbox for an entitled Workspace. Admission is entitlement-gated,
generation-fenced, shared-volume-challenged, quota-aware, and fixed to
nsjail. Plain nsjail correctly fails Cloud readiness without a hard quota
provider.
- stdio MCP has an independent gate and is forced off by Cloud bootstrap even
when Box is available.
- SaaS business data uses one PostgreSQL shared schema with application scope
plus forced RLS. Cloud directory writes are separated from local tenant
writes.
- Projected Account and Membership revisions are monotonic. Tombstones remove
memberships, and stale revisions cannot resurrect them.
- pgvector shares the business PostgreSQL database and remains
Workspace-scoped.
- Space runs versioned SQL migrations before application seeds. A fresh
database, a partial pre-migration database, and an existing-baseline path
converge without startup-time `AutoMigrate`.
### Shared Plugin Runtime
- One instance-scoped trusted Supervisor serves multiple Workspaces.
- Every enabled installation has its own nsjail process bound to
`(instance, workspace, execution generation, installation, runtime revision,
artifact digest)`.
- Verified same-digest code and dependency files are mounted read-only and may
be shared; home, tmp, data, process namespace, registration capability, and
cgroup are private.
- Instance configuration owns CPU, memory, PID, open-file, and file-size
limits. Plugin manifests cannot increase them.
- Memory includes swap: nsjail receives `memory.max` and `swap.max=0`.
- Unexpected worker exit is recovered by a completion callback with bounded
per-installation exponential backoff. Remove, reconcile, Runtime shutdown,
and container SIGTERM perform a graceful-to-SIGKILL bounded reap.
### Shared Box Runtime and MCP
- One shared Box control plane serves Workspace-bound logical sessions.
- Cloud grants allow at most one persistent `global` session for an entitled
Workspace and no managed processes in the first release.
- Cloud is fixed to nsjail and network-off. Core and Box prove the shared
durable Workspace mount with an authenticated marker challenge.
- stdio MCP is independently gated and forced off for Cloud v2.
- The current nsjail Box backend does not provide hard byte/inode quotas, so
Cloud readiness correctly fails closed instead of silently using a soft
directory scan.
## Automated verification
### LangBot backend and static checks
### LangBot Core
```text
uv run pytest -q --tb=short
2527 passed, 32 skipped, 177 warnings in 73.37s
uv run --no-sync pytest -q
2590 passed, 32 skipped, 177 warnings
focused tenant UoW and legacy RAG compatibility suite
120 passed, 10 warnings in 1.30s
uv run ruff check .
passed
changed Python format check
15 files already formatted
real PostgreSQL migration, pgvector, and release-migrator suites
21 passed, 11 warnings
uv run --no-sync ruff check .
uv lock --check
git diff --check
passed
```
The full suite includes startup E2E, the fresh SQLite registration/Workspace
journey, API authorization, storage, RAG, Plugin, MCP, and 18 Docker-backed Box
integration cases.
The full suite ran without the closed adapter installed, proving the open-source
single-Workspace/multi-user path remains standalone. Focused closed-adapter,
directory projection, runtime connector, Box cleanup, and configuration suites
also passed with the adapter installed.
### Real PostgreSQL 16 and pgvector
### Plugin SDK and real Linux runtime
```text
test_migrations_postgres.py
test_pgvector_postgres.py
test_release_migration_postgres.py
21 passed, 11 warnings in 11.61s
SDK full suite
1226 passed, 22 existing warnings
Ruff check and format check
git diff --check
passed
```
The real-database suite covers fresh and upgraded RLS schemas, two-Workspace
isolation, pgvector CRUD and ANN indexes, least-privilege runtime bootstrap,
release locking, historical/fresh RAG settings column compatibility, and
transaction-scope escape rejection. Negative tests inject and reject role
membership, grant options, persistent GUCs, extra schemas and relations,
column ACLs, explicit routine and parameter ACLs, runtime-owned and
`SECURITY DEFINER` routines, extra/runtime-owned extensions, and foreign data
wrappers, servers, and user mappings.
A privileged Linux test container with host cgroup namespace ran one shared
Runtime and two Workspace installations:
Alembic has exactly one head: `0013_tenant_pgvector`.
- both workers referenced the same artifact inode;
- home, tmp, and data inodes were distinct;
- each plugin saw only PID 1 in its private PID namespace;
- a tampered binding and an unknown installation were rejected;
- the control token was absent from worker environments;
- cgroups were distinct with `memory.max=134217728`, `memory.swap.max=0`,
`pids.max=32`, and `cpu.max=500000 1000000`;
- touching 256 MiB exited with code 137 without swap growth;
- the 32nd fork failed with `EAGAIN`;
- reconcile and container SIGTERM removed the worker cgroups.
### Plugin SDK and Box Runtime
The same run started the Runtime from a non-root working directory, covering
absolute nsjail mount-source normalization.
### Space backend, adapter, and frontend
```text
Plugin SDK full suite
1160 passed, 18 warnings
Plugin SDK Ruff and action consistency checks
go test ./...
passed
Docker-backed Box integration suite
18 passed, 33 warnings
fresh PostgreSQL app startup, partial-baseline migration,
Cloud v2 migration rerun, and control-plane integration
passed
closed adapter pytest and Ruff
passed
pnpm exec tsc --noEmit
targeted ESLint for changed Cloud files
pnpm build
passed; 47 routes generated
```
The LangBot environment resolves `langbot-plugin==0.4.15` from the same exact
SDK commit recorded above. The SDK local ref and remote
`refs/heads/feat/multi-tenants` were independently verified equal.
The PostgreSQL checks started from an empty database and verified all 32
registered migrations in order, Cloud v2 Free/Pro seeds, legacy plan seeds,
Cloud columns/indexes, and repeatable reruns.
### Frontend
## Cross-service and browser E2E
```text
pnpm lint
0 errors, 34 existing warnings
### Signed Space-to-Core directory projection
VITE_API_BASE_URL=/ pnpm build
passed; 3168 modules transformed
Using an isolated Space PostgreSQL database and a migrated Core PostgreSQL
database:
pnpm exec playwright test --project=chromium
43 passed in 22.2s
```
1. Space issued a signed manifest for the fixed LangBot instance.
2. Space returned a directory snapshot at cursor 5 containing two Workspaces,
two projected Accounts, and owner/admin memberships.
3. Core verified the signature, instance, release, validity window, and
capability before injecting the Cloud Workspace policy.
4. Core stored both Workspaces, all active Memberships, both projected
Accounts, snapshot coverage, inbox entries, and cursor 5.
5. Account-field-only projection revisions and Workspace directory revisions
remained independent, and cross-Workspace Account conflicts failed closed.
The frontend suite includes Viewer monitoring visibility, privileged control
hiding, stable terminal invitation states, email-mismatch login handoff, and a
temporary invitation-accept failure that retains the authenticated session and
succeeds on retry.
### Real browser Cloud v2 flow
### Packaging and deployment configuration
A real local browser operated the Space frontend and backend:
- `uv lock --check` resolved 277 packages and passed.
- Changed TypeScript/TSX files passed Prettier verification.
- `docker compose -f docker/docker-compose.yaml config --quiet` passed with the
existing warning that the top-level `version` field is obsolete.
- The staged high-signal credential scan found no private key or provider
token.
1. An owner logged in and saw the automatically created personal Workspace.
2. A second Account invited the owner as admin to another Workspace.
3. The selector showed both Workspaces and switched between them without stale
subscription or usage state crossing the selection.
4. Each Workspace showed its own Free entitlement and limits.
5. Pro checkout opened for the selected Workspace. With no configured payment
channels, confirmation was visibly disabled and no order was created.
6. No browser errors occurred; only the pre-existing i18next warning for the
short `en` locale was observed.
## Real browser E2E
The legacy feature-flag branch was then covered by API/static checks and the
production build: false or missing `cloud_v2_enabled` renders the original
Welcome/Pod client; a failed web-config request renders an explicit retry
instead of guessing a deployment mode.
A production frontend and backend were run against an isolated temporary
SQLite data directory and operated through a real local Chrome session:
### Deliberate Core startup failure
1. The first Account registered and automatically received the default
singleton Workspace as owner.
2. Skipping the setup wizard reached `/home/monitoring`; reload and a new tab
recovered the authenticated Account, membership, and Workspace selection.
3. The owner invited a Viewer. After signing out, the second Account registered
through the invitation and joined the same Workspace.
4. Owner and Admin could read overview/export/audit/API-key surfaces and reach
resource-write validation. Runtime debug mutation remained disabled.
5. Developer retained ordinary resource management and monitoring, while
export, audit, and API-key management returned 403.
6. Operator and Viewer retained ordinary monitoring reads while resource
writes, export, audit, API-key management, and runtime mutation returned 403.
7. The UI matched the server matrix: Add/API Key/Export controls appeared only
with their permissions, while ordinary monitoring refresh remained visible
to read-only members.
8. A second Workspace creation attempt returned 403 `edition_limit`, proving
Community remains single-Workspace and multi-user.
9. Used, revoked, expired, and email-mismatched invitation paths produced
stable visible states. The mismatch survived login without a contradictory
generic success toast.
Core completed signed manifest verification and directory projection, then
stopped at the Box readiness gate because the current nsjail backend cannot
prove hard Workspace byte/inode quota enforcement. Connector shutdown and
reconnect tasks were cleanly reaped; no event-loop or never-awaited coroutine
warning remained.
The server was stopped, port 15321 was verified free, and temporary artifacts
were moved to the system Trash for recoverable cleanup. No invitation secret,
password, token, or provider credential is recorded in this report.
This is a successful fail-closed acceptance result, not a passing production
Cloud boot.
## Deliberately incomplete SaaS activation gates
## Remaining production activation gates
Multi-Workspace activation remains unavailable in the open-source bootstrap.
Cloud v2 must stay disabled until all of the following are implemented and
independently verified:
Cloud v2 must remain disabled until these gates are closed:
- the closed Control Plane owns the Account/Workspace directory, memberships,
invitations, monotonic execution generations, entitlements, usage,
subscriptions, and billing;
- ordinary tenant writes hold a generation-aware fence through commit, and a
generation-stamped outbox or equivalent atomic publish fence protects
external side effects;
- durable object references survive generation cutovers without stranding
files, images, plugin artifacts, or knowledge-base content;
- every tenant-configurable outbound path has complete SSRF/egress policy, and
Plugin Runtime has production Linux namespace, cgroup v2, and network
isolation evidence;
- an unexpectedly exited enabled plugin worker is restored by a completion
callback with bounded backoff and cross-tenant restart-storm isolation;
- plugin installation data has an operator-owned hard disk quota, not a
directory-scan approximation;
- Box supplies and proves hard byte/inode quotas for Workspace, Skill, root,
home, and temporary storage, and the production shared volume passes the
authenticated marker challenge;
- production PostgreSQL uses a dedicated endpoint/cluster, or a tested
HBA/proxy policy proves the runtime credential can connect only to the target
business database;
- the release migrator is deployed as a one-shot Job with credential issuance,
backup, rollback, failure/retry, and orchestration procedures;
- OAuth exchange, directory projection, leases, snapshots, event replay, and
outbox state use durable shared stores suitable for horizontally scaled
replicas;
- Workspace release, export, deletion, restore, and disaster-recovery semantics
are decided and implemented; and
- production Runtime restart, crash, partition, revocation, and fault-injection
acceptance is completed against the closed Control Plane and greenfield
Cloud v2 deployment.
- Box provides and proves hard byte and inode quotas for Workspace, Skill,
root, home, and tmp storage.
- Plugin installation writable data receives an operator-owned hard total
disk quota.
- Plugin and future networked Box workloads have tenant-safe egress/SSRF
policy.
- Plugin Runtime adds jitter, global restart concurrency limiting, and a
Runtime-level circuit breaker, then passes systemic-failure injection.
- M0 rolls Core and Plugin Runtime together until authenticated Runtime takeover
or an owner lease/fencing protocol exists.
- Account registration moves New API provisioning to an in-process,
transactional outbox reconciler keyed by Space Account UUID.
- EPay and Stripe fulfillment bind provider identity, amount, currency,
channel/session, payment status, and expiry to the locked order transaction;
provider transaction IDs are unique and replay-safe.
- Payment-session creation and stale `processing` orders gain in-process
reconciliation.
- Cloud v2 subscription service periods are stored immutably and included in
recognized-revenue reporting.
- Production migration Job, backup/rollback, PostgreSQL credential/network
boundaries, horizontal-replica fault injection, Workspace release/export,
deletion, and restore semantics are completed.
The following product work is also deliberately deferred: WebUI configuration
for tenant-provided remote E2B, multi-replica lease storage and sharding rules,
artifact signing/garbage collection, custom roles, SSO, and SCIM.
These are explicit release gates, not hidden fallbacks. The current branches
are a verified isolation foundation on which the closed SaaS system can be
built; they are not a production Cloud activation.
These gates intentionally add no tenant-specific service. They are implemented
inside the existing Space, Core, Plugin Runtime, Box Runtime, and PostgreSQL
components to preserve the architecture goal: near-zero static cost for a new
Workspace.
docs/multi-tenant/workspace-multi-user-architecture.md
+12 -5
View File
@@ -148,7 +148,8 @@ M1 是 M0 的透明扩容，M2 是相同架构下的资源等级。外部 API
### 4.4 组件边界
- Core、Plugin Runtime 和 Box Runtime 可以由发布与故障域决定是否共用 Pod，但必须保持独立进程身份、容器和 security context。
- Core、Plugin Runtime 和 Box Runtime 必须保持独立进程身份、容器和 security context。M0 中 Core 与 Plugin Runtime
需要处于同一 rollout/restart unit；在实现受认证 takeover 或 owner lease/fencing 前，Core 不能单独重启后接管仍存活的 Runtime。
- Core 不能继承 nsjail、cgroup 或 mount namespace 所需的高权限。
- Plugin Runtime 与 Box Runtime 不合并为一个高权限进程。
- MVP 不新增 Runtime 专用数据库、Box 专用数据库、Kafka、Redis、租户级 scheduler 或 artifact service。
@@ -321,7 +322,9 @@ delegated issuers and keyset revision
### 7.2 DirectoryEvent 与目录新鲜度
Control Plane 通过 transactional outbox 发布 Account、Workspace 和 Membership 的版本化事件。
Core 使用 inbox 按 `event_id` 去重，以 aggregate revision 拒绝旧写，并追踪连续应用水位。
Core 使用 inbox 按 `event_id` 去重，以 aggregate revision 拒绝旧写，并追踪连续应用水位。启动时读取一个 PostgreSQL
`REPEATABLE READ` 事务内生成的签名全量 snapshot；运行时先消费携带当前 high-water 的签名事件页，再只请求该页涉及的 Workspace 签名增量。
增量响应不携带新的事件 cursor，因此即使其内容已包含并发提交的后续 revision，也不能跳过尚未消费的事件。
要求：
@@ -329,6 +332,10 @@ Core 使用 inbox 按 `event_id` 去重，以 aggregate revision 拒绝旧写，
- 重复、乱序、延迟、断流和全量 replay 都安全。
- 删除使用 tombstone。
- 新实例先导入带 high watermark 的 snapshot，再消费增量。
- 常态目录更新成本与本页发生变化的 Workspace 数量相关，不得为每个 `directory.changed` 重新读取和投影全部 Workspace。
- 每个 Core replica 独立保存进程内消费 cursor，以确保各自的 entitlement cache 都看到事件；共享 PostgreSQL 保存投影
high-water mark、全量 snapshot coverage 和 inbox。同一事件被多个 replica 消费时，第二个 replica 验证已有 receipt；
snapshot coverage 内缺少的 receipt 可以补写，coverage 之外缺失则失败关闭。只有本地 cursor 追平签名 high-water 后才续期 ready。
- projection 未就绪或落后于授权 lease 要求时，交互与自动化请求按策略失败关闭。
- SaaS pending Invitation、email 和 token hash 不进入 Core 投影。
@@ -630,8 +637,8 @@ installation 总磁盘配额需要可原子拒绝写入的 quota provider，不
- Runtime 重连执行实例范围 full reconciliation，清理 stale worker 并恢复 enabled installation。
- dependency preparation 失败记录在对应 installation，不启动半就绪 worker，也不阻塞其他 installation。
- desired semantics 要求 enabled installation 常驻，不做 idle eviction；是否按负载回收以后再决定。
- 当前 Supervisor 可在 Runtime 重连或 Core apply/reconcile 时恢复 desired state，但尚未为意外退出的 worker
实现带有界 backoff 的 completion callback；这项可靠性缺口在 Cloud 激活前必须补齐并验证。
- 当前 Supervisor 已在意外退出时通过 completion callback 和有界指数 backoff 恢复 enabled worker。
Cloud 激活前仍需加入 jitter、全局重启并发上限和 Runtime 级 circuit breaker，并验证系统性故障不会形成跨租户重启风暴。
真实 Linux/nsjail/cgroup 与受控 egress 的 Cloud 部署验证尚未完成，是生产激活门禁。
@@ -813,7 +820,7 @@ SaaS：
2. 普通业务写入贯穿 commit 的 generation-aware fence，以及与外部副作用同事务的 business outbox。
3. generation cutover 后稳定的 durable object identity 或原子对象引用迁移。
4. 所有 tenant-configurable outbound URL 的 SSRF 防护与 tenant-safe egress；Plugin Runtime 还需在真实 Linux/nsjail/cgroup v2 环境验证 namespace、资源限制和文件隔离。
5. Plugin Runtime 对意外退出的 enabled worker 实现 completion callback、有界 backoff 和自动恢复，并验证不会形成跨租户重启风暴。
5. Plugin Runtime 已实现意外退出 worker 的 completion callback、有界 backoff 和自动恢复；Cloud 激活前增加全局重启风暴抑制并完成故障注入验证。
6. Plugin installation data 的 production hard disk quota provider，能够在写入边界原子拒绝超额，不能以目录扫描代替。
7. Box Runtime 的 production hard quota provider，包括 Workspace、Skill、root/tmp/home 的 byte 与 inode quota；真实部署还必须在启动和重连时通过共享卷 marker challenge。
8. PostgreSQL runtime credential 的专用 endpoint 或 HBA/proxy 跨 database 隔离证明、生产 migration/rollback 流程，以及 legacy pgvector migration 失败后精确恢复 RLS/FORCE 并可安全重试的集成证据。
pyproject.toml
+1 -1
View File
@@ -70,7 +70,7 @@ dependencies = [
"chromadb>=1.0.0,<2.0.0",
"qdrant-client (>=1.15.1,<2.0.0)",
"pyseekdb==1.1.0.post3",
"langbot-plugin @ git+https://github.com/langbot-app/langbot-plugin-sdk.git@95a1805af2038c745de2c018a00db1305089a32e",
"langbot-plugin @ git+https://github.com/langbot-app/langbot-plugin-sdk.git@e7d946af4a6b1494fbe74627c1815ace19ac8991",
"asyncpg>=0.30.0",
"line-bot-sdk>=3.19.0",
"matrix-nio>=0.25.2",
src/langbot/pkg/api/http/service/user.py
+60 -2
View File
@@ -109,12 +109,15 @@ class UserService:
return await asyncio.to_thread(argon2.PasswordHasher().hash, password)
def _require_local_directory(self) -> None:
workspace_service = getattr(self.ap, 'workspace_service', None)
if workspace_service is not None and workspace_service.policy.multi_workspace_enabled:
if self._uses_control_plane_directory():
raise ControlPlaneDirectoryRequiredError(
'Cloud Accounts and directory changes are managed by the SaaS control plane'
)
def _uses_control_plane_directory(self) -> bool:
workspace_service = getattr(self.ap, 'workspace_service', None)
return bool(workspace_service is not None and workspace_service.policy.multi_workspace_enabled)
async def _verify_password(self, hashed_password: str, password: str) -> None:
async with self._password_hash_lock:
await asyncio.to_thread(argon2.PasswordHasher().verify, hashed_password, password)
@@ -427,6 +430,15 @@ class UserService:
expires_in: int = 0,
) -> user.User:
"""Create or update a Space user account (only if system not initialized or user exists)"""
if self._uses_control_plane_directory():
return await self._update_projected_space_user(
space_account_uuid=space_account_uuid,
email=email,
access_token=access_token,
refresh_token=refresh_token,
api_key=api_key,
expires_in=expires_in,
)
self._require_local_directory()
expires_at = datetime.datetime.now() + datetime.timedelta(seconds=expires_in) if expires_in > 0 else None
@@ -507,6 +519,52 @@ class UserService:
await self._update_space_provider_for_account(created_user, api_key)
return created_user
async def _update_projected_space_user(
self,
*,
space_account_uuid: str,
email: str,
access_token: str,
refresh_token: str,
api_key: str,
expires_in: int,
) -> user.User:
"""Attach OAuth credentials to an already projected Cloud Account."""
normalized_email = normalize_email(email)
expires_at = datetime.datetime.now() + datetime.timedelta(seconds=expires_in) if expires_in > 0 else None
async with self._create_user_lock:
projected = await self.get_user_by_space_account_uuid(space_account_uuid)
if (
projected is None
or projected.uuid != space_account_uuid
or projected.normalized_email != normalized_email
or projected.source != user.AccountSource.CLOUD_PROJECTION.value
or projected.account_type != 'space'
):
raise ControlPlaneDirectoryRequiredError('Space Account is not present in the verified Cloud directory')
self._require_active_account(projected)
await self._identity_execute(
sqlalchemy.update(user.User)
.where(
user.User.uuid == projected.uuid,
user.User.space_account_uuid == space_account_uuid,
user.User.source == user.AccountSource.CLOUD_PROJECTION.value,
)
.values(
space_access_token=access_token,
space_refresh_token=refresh_token,
space_api_key=api_key,
space_access_token_expires_at=expires_at,
),
f'space:{space_account_uuid}',
)
refreshed = await self.get_user_by_space_account_uuid(space_account_uuid)
if refreshed is None:
raise ControlPlaneDirectoryRequiredError('Space Account disappeared from the verified Cloud directory')
self._require_active_account(refreshed)
return refreshed
async def authenticate_space_user(
self, access_token: str, refresh_token: str, expires_in: int = 0
) -> typing.Tuple[str, user.User]:
src/langbot/pkg/box/connector.py
+1
View File
@@ -351,6 +351,7 @@ class BoxRuntimeConnector(ManagedRuntimeConnector):
make_connection_failed_callback=on_connect_failed,
additional_headers=self.get_control_headers(),
)
self._ctrl = ctrl
self._ctrl_task = asyncio.create_task(
ctrl.run(self._make_connection_callback(transport_name, connected, connect_error, self._generation))
)
src/langbot/pkg/box/service.py
+50 -3
View File
@@ -145,8 +145,38 @@ class BoxService:
self._available = False
self._connector_error = str(exc)
if self._cloud_managed:
await self._abort_failed_cloud_initialization()
raise
async def _abort_failed_cloud_initialization(self) -> None:
"""Close a connected Cloud transport before propagating readiness failure.
Connector initialization starts control and heartbeat tasks before Core
performs the stricter Cloud readiness challenge. If that challenge
fails, startup must remain fail-closed without leaving those tasks free
to schedule reconnect work while the application loop is unwinding.
"""
self._closing = True
self._available = False
reconnect_task = self._reconnect_task
self._reconnect_task = None
self._reconnecting = False
if reconnect_task is not None and reconnect_task is not asyncio.current_task():
reconnect_task.cancel()
await asyncio.gather(reconnect_task, return_exceptions=True)
connector = self._runtime_connector
if connector is None:
return
connector.runtime_disconnect_callback = None
try:
await connector.aclose()
except Exception:
# Cleanup failure must not replace the readiness error which caused
# Cloud startup to fail closed.
self.ap.logger.exception('Failed to close Box runtime after Cloud readiness validation failed')
async def _on_runtime_disconnect(self, connector: BoxRuntimeConnector) -> None:
"""Called by the connector when the Box runtime connection drops.
@@ -156,13 +186,28 @@ class BoxService:
"""
if not self._enabled or self._closing:
return
try:
loop = asyncio.get_running_loop()
except RuntimeError:
return
if loop.is_closed():
return
if self._reconnect_task is not None and not self._reconnect_task.done():
return # Another reconnect loop is already running
self._reconnecting = True
self._available = False
self._connector_error = 'Disconnected from Box runtime'
self.ap.logger.warning('Box runtime disconnected, sandbox features temporarily disabled.')
self._reconnect_task = asyncio.create_task(self._reconnect_loop(connector))
reconnect = self._reconnect_loop(connector)
try:
self._reconnect_task = loop.create_task(reconnect)
except RuntimeError:
# The loop may begin closing between get_running_loop() and task
# creation. Explicitly close the coroutine so shutdown emits no
# "coroutine was never awaited" warning.
reconnect.close()
self._reconnecting = False
self._reconnect_task = None
async def _reconnect_loop(self, connector: BoxRuntimeConnector) -> None:
"""Retry reconnection with exponential backoff (3s → 60s max)."""
@@ -173,9 +218,11 @@ class BoxService:
self.ap.logger.info(f'Attempting to reconnect to Box runtime in {delay}s...')
await asyncio.sleep(delay)
try:
connector.dispose()
await connector.initialize()
await connector.reconnect()
self._ensure_default_workspace()
await self._verify_cloud_runtime()
if not self._cloud_managed:
await self._purge_attachment_dirs()
self._available = True
self._connector_error = ''
skill_mgr = getattr(self.ap, 'skill_mgr', None)
src/langbot/pkg/cloud/__init__.py
+24
View File
@@ -2,10 +2,23 @@
from .bootstrap import (
CloudBootstrapError,
CloudManifestProvider,
CloudManifestRefreshService,
OpenSourceDeployment,
VerifiedCloudDeployment,
resolve_deployment,
)
from .directory import (
DirectoryDelta,
DirectoryEvent,
DirectoryEventBatch,
DirectoryMember,
DirectoryProjectionProvider,
DirectoryProjectionUnavailableError,
DirectorySnapshot,
DirectoryWorkspace,
)
from .directory_projection import DirectoryProjectionService
from .entitlements import (
EntitlementProvider,
EntitlementResolver,
@@ -16,6 +29,17 @@ from .entitlements import (
__all__ = [
'CloudBootstrapError',
'CloudManifestProvider',
'CloudManifestRefreshService',
'DirectoryDelta',
'DirectoryEvent',
'DirectoryEventBatch',
'DirectoryMember',
'DirectoryProjectionProvider',
'DirectoryProjectionService',
'DirectoryProjectionUnavailableError',
'DirectorySnapshot',
'DirectoryWorkspace',
'EntitlementProvider',
'EntitlementResolver',
'EntitlementSnapshot',
src/langbot/pkg/cloud/bootstrap.py
+83 -1
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
import asyncio
import dataclasses
import importlib.metadata
import inspect
@@ -7,9 +8,10 @@ import os
import threading
import time
from collections.abc import Awaitable, Callable
from typing import Any, Protocol
from typing import Any, Protocol, runtime_checkable
from ..workspace.policy import CloudWorkspacePolicy, SingleWorkspacePolicy
from .directory import DirectoryProjectionProvider
from .entitlements import EntitlementProvider, OpenSourceEntitlementProvider
@@ -26,6 +28,17 @@ class CloudRuntimeUnavailableError(CloudBootstrapError):
"""The verified Cloud receipt no longer admits runtime work."""
@runtime_checkable
class CloudManifestProvider(Protocol):
"""Closed adapter responsible for renewing the signed deployment receipt."""
async def refresh_manifest(self) -> VerifiedCloudDeployment:
"""Fetch, verify, and return the newest deployment receipt."""
async def aclose(self) -> None:
"""Release control-plane transport resources."""
@dataclasses.dataclass(frozen=True, slots=True)
class OpenSourceDeployment:
"""Default deployment selected when no closed bootstrap is installed."""
@@ -35,6 +48,8 @@ class OpenSourceDeployment:
entitlement_provider: OpenSourceEntitlementProvider = dataclasses.field(
default_factory=OpenSourceEntitlementProvider
)
directory_provider: None = None
manifest_provider: None = None
persistence_mode: str = 'oss_compat'
required_vector_backend: str | None = None
@@ -63,6 +78,8 @@ class VerifiedCloudDeployment:
capabilities: frozenset[str]
tenant_isolation_version: int
entitlement_provider: EntitlementProvider
directory_provider: DirectoryProjectionProvider
manifest_provider: CloudManifestProvider
verification_key_id: str
mode: str = dataclasses.field(default='cloud', init=False)
workspace_policy: CloudWorkspacePolicy = dataclasses.field(default_factory=CloudWorkspacePolicy, init=False)
@@ -89,6 +106,10 @@ class VerifiedCloudDeployment:
raise CloudBootstrapError('Verified Cloud Manifest does not grant multi_workspace_v2')
if not isinstance(self.entitlement_provider, EntitlementProvider):
raise CloudBootstrapError('Verified Cloud bootstrap did not provide an entitlement adapter')
if not isinstance(self.directory_provider, DirectoryProjectionProvider):
raise CloudBootstrapError('Verified Cloud bootstrap did not provide a directory adapter')
if not isinstance(self.manifest_provider, CloudManifestProvider):
raise CloudBootstrapError('Verified Cloud bootstrap did not provide a Manifest renewal adapter')
def validate_instance_config(self, config: dict[str, Any]) -> None:
if config.get('database', {}).get('use') != 'postgresql':
@@ -262,6 +283,67 @@ class DeploymentAdmissionGuard:
return deployment
class CloudManifestRefreshService:
"""Renew a short-lived verified Manifest before runtime admission expires."""
def __init__(
self,
admission: DeploymentAdmissionGuard,
provider: CloudManifestProvider,
logger: Any,
*,
wall_time: Callable[[], float] = time.time,
refresh_margin_seconds: int = 180,
maximum_sleep_seconds: int = 300,
) -> None:
if not isinstance(provider, CloudManifestProvider):
raise TypeError('Cloud Manifest refresh requires a CloudManifestProvider')
if refresh_margin_seconds < 120:
raise ValueError('Cloud Manifest refresh margin must be at least 120 seconds')
if maximum_sleep_seconds <= 0:
raise ValueError('Cloud Manifest refresh maximum sleep must be positive')
self.admission = admission
self.provider = provider
self.logger = logger
self._wall_time = wall_time
self.refresh_margin_seconds = refresh_margin_seconds
self.maximum_sleep_seconds = maximum_sleep_seconds
def next_refresh_delay(self) -> float:
deployment = self.admission.deployment
if not isinstance(deployment, VerifiedCloudDeployment):
return float(self.maximum_sleep_seconds)
remaining = deployment.expires_at - self._wall_time()
return max(
5.0,
min(
float(self.maximum_sleep_seconds),
remaining - self.refresh_margin_seconds,
),
)
async def refresh_once(self) -> VerifiedCloudDeployment:
candidate = await self.provider.refresh_manifest()
if not isinstance(candidate, VerifiedCloudDeployment):
raise CloudBootstrapError('Cloud Manifest provider returned an invalid deployment receipt')
self.admission.replace(candidate)
return candidate
async def run(self) -> None:
retry_delay = 5.0
while True:
try:
await asyncio.sleep(self.next_refresh_delay())
await self.refresh_once()
retry_delay = 5.0
except asyncio.CancelledError:
raise
except Exception:
self.logger.exception('Cloud Manifest refresh failed')
await asyncio.sleep(retry_delay)
retry_delay = min(retry_delay * 2, 30.0)
async def _invoke_provider(
loaded: Any,
*,
src/langbot/pkg/cloud/directory.py
+215
View File
@@ -0,0 +1,215 @@
from __future__ import annotations
import datetime
from collections.abc import Sequence
from typing import Any, Protocol, runtime_checkable
import pydantic
class DirectoryProjectionUnavailableError(RuntimeError):
"""Raised when the verified Cloud directory cannot safely admit work."""
class DirectoryMember(pydantic.BaseModel):
"""One account membership published by the SaaS control plane."""
model_config = pydantic.ConfigDict(frozen=True, extra='forbid')
membership_uuid: str = pydantic.Field(min_length=1, max_length=36)
account_uuid: str = pydantic.Field(min_length=1, max_length=36)
normalized_email: str = pydantic.Field(min_length=1, max_length=320)
display_name: str = pydantic.Field(min_length=1, max_length=255)
account_status: str = pydantic.Field(pattern=r'^(active|blocked|disabled|deleted)$')
role: str = pydantic.Field(pattern=r'^(owner|admin|member|developer|operator|viewer)$')
membership_status: str = pydantic.Field(pattern=r'^(active|invited|disabled|removed)$')
projection_revision: int = pydantic.Field(ge=1)
joined_at: datetime.datetime | None = None
@pydantic.field_validator('normalized_email')
@classmethod
def _normalize_email(cls, value: str) -> str:
normalized = value.strip().casefold()
if normalized != value:
raise ValueError('Directory email must already be normalized')
return normalized
class DirectoryWorkspace(pydantic.BaseModel):
"""One Workspace and its authoritative membership projection."""
model_config = pydantic.ConfigDict(frozen=True, extra='forbid')
uuid: str = pydantic.Field(min_length=1, max_length=36)
name: str = pydantic.Field(min_length=1, max_length=255)
slug: str = pydantic.Field(min_length=1, max_length=255)
type: str = pydantic.Field(pattern=r'^(personal|team)$')
status: str = pydantic.Field(pattern=r'^(provisioning|active|suspended|archived|deleted)$')
created_by_account_uuid: str = pydantic.Field(min_length=1, max_length=36)
projection_revision: int = pydantic.Field(ge=1)
execution_generation: int = pydantic.Field(ge=1)
members: tuple[DirectoryMember, ...] = ()
@pydantic.field_validator('members', mode='before')
@classmethod
def _copy_members(cls, value: Sequence[DirectoryMember] | None) -> tuple[DirectoryMember, ...]:
return tuple(value or ())
@pydantic.model_validator(mode='after')
def _validate_members(self) -> DirectoryWorkspace:
membership_uuids = [member.membership_uuid for member in self.members]
account_uuids = [member.account_uuid for member in self.members]
if len(membership_uuids) != len(set(membership_uuids)):
raise ValueError('Directory Workspace contains duplicate membership UUIDs')
if len(account_uuids) != len(set(account_uuids)):
raise ValueError('Directory Workspace contains duplicate account UUIDs')
if self.created_by_account_uuid not in set(account_uuids):
raise ValueError('Directory Workspace must include its creator')
return self
class DirectorySnapshot(pydantic.BaseModel):
"""Full signed directory state at one monotonic outbox cursor."""
model_config = pydantic.ConfigDict(frozen=True, extra='forbid')
instance_uuid: str = pydantic.Field(min_length=1, max_length=255)
cursor: int = pydantic.Field(ge=0)
generated_at: datetime.datetime
workspaces: tuple[DirectoryWorkspace, ...] = ()
@pydantic.field_validator('workspaces', mode='before')
@classmethod
def _copy_workspaces(cls, value: Sequence[DirectoryWorkspace] | None) -> tuple[DirectoryWorkspace, ...]:
return tuple(value or ())
@pydantic.model_validator(mode='after')
def _validate_workspaces(self) -> DirectorySnapshot:
workspace_uuids = [workspace.uuid for workspace in self.workspaces]
slugs = [workspace.slug for workspace in self.workspaces]
membership_uuids = [member.membership_uuid for workspace in self.workspaces for member in workspace.members]
if len(workspace_uuids) != len(set(workspace_uuids)):
raise ValueError('Directory snapshot contains duplicate Workspace UUIDs')
if len(slugs) != len(set(slugs)):
raise ValueError('Directory snapshot contains duplicate Workspace slugs')
if len(membership_uuids) != len(set(membership_uuids)):
raise ValueError('Directory snapshot contains duplicate membership UUIDs')
return self
class DirectoryDelta(pydantic.BaseModel):
"""Signed authoritative state for an explicitly requested Workspace set."""
model_config = pydantic.ConfigDict(frozen=True, extra='forbid')
instance_uuid: str = pydantic.Field(min_length=1, max_length=255)
requested_workspace_uuids: tuple[str, ...]
generated_at: datetime.datetime
workspaces: tuple[DirectoryWorkspace, ...] = ()
@pydantic.field_validator('requested_workspace_uuids', mode='before')
@classmethod
def _copy_requested_workspace_uuids(cls, value: Sequence[str]) -> tuple[str, ...]:
return tuple(value)
@pydantic.field_validator('workspaces', mode='before')
@classmethod
def _copy_workspaces(cls, value: Sequence[DirectoryWorkspace] | None) -> tuple[DirectoryWorkspace, ...]:
return tuple(value or ())
@pydantic.model_validator(mode='after')
def _validate_workspaces(self) -> DirectoryDelta:
requested = self.requested_workspace_uuids
if not requested or len(requested) > 100:
raise ValueError('Directory delta must request between 1 and 100 Workspaces')
if any(not workspace_uuid or len(workspace_uuid) > 36 for workspace_uuid in requested):
raise ValueError('Directory delta contains an invalid requested Workspace UUID')
if len(requested) != len(set(requested)):
raise ValueError('Directory delta contains duplicate requested Workspace UUIDs')
workspace_uuids = [workspace.uuid for workspace in self.workspaces]
slugs = [workspace.slug for workspace in self.workspaces]
membership_uuids = [member.membership_uuid for workspace in self.workspaces for member in workspace.members]
if len(workspace_uuids) != len(set(workspace_uuids)):
raise ValueError('Directory delta contains duplicate Workspace UUIDs')
if not set(workspace_uuids).issubset(set(requested)):
raise ValueError('Directory delta returned an unrequested Workspace')
if len(slugs) != len(set(slugs)):
raise ValueError('Directory delta contains duplicate Workspace slugs')
if len(membership_uuids) != len(set(membership_uuids)):
raise ValueError('Directory delta contains duplicate membership UUIDs')
return self
class DirectoryEvent(pydantic.BaseModel):
"""One signed control-plane outbox notification."""
model_config = pydantic.ConfigDict(frozen=True, extra='forbid')
cursor: int = pydantic.Field(ge=1)
uuid: str = pydantic.Field(min_length=1, max_length=36)
aggregate_uuid: str = pydantic.Field(min_length=1, max_length=36)
event_type: str = pydantic.Field(min_length=1, max_length=128)
revision: int = pydantic.Field(ge=1)
payload: dict[str, Any] = pydantic.Field(default_factory=dict)
created_at: datetime.datetime
class DirectoryEventBatch(pydantic.BaseModel):
"""Signed events returned after a caller-supplied directory cursor."""
model_config = pydantic.ConfigDict(frozen=True, extra='forbid')
instance_uuid: str = pydantic.Field(min_length=1, max_length=255)
after_cursor: int = pydantic.Field(ge=0)
cursor: int = pydantic.Field(ge=0)
high_water_cursor: int = pydantic.Field(ge=0)
events: tuple[DirectoryEvent, ...] = ()
@pydantic.field_validator('events', mode='before')
@classmethod
def _copy_events(cls, value: Sequence[DirectoryEvent] | None) -> tuple[DirectoryEvent, ...]:
return tuple(value or ())
@pydantic.model_validator(mode='after')
def _validate_events(self) -> DirectoryEventBatch:
if self.cursor < self.after_cursor:
raise ValueError('Directory event cursor rolled back')
if self.high_water_cursor < self.cursor:
raise ValueError('Directory event high-water mark rolled back')
event_cursors = [event.cursor for event in self.events]
event_uuids = [event.uuid for event in self.events]
if event_cursors != sorted(event_cursors) or len(event_cursors) != len(set(event_cursors)):
raise ValueError('Directory events must have strictly increasing cursors')
if len(event_uuids) != len(set(event_uuids)):
raise ValueError('Directory event batch contains duplicate UUIDs')
if any(cursor <= self.after_cursor or cursor > self.cursor for cursor in event_cursors):
raise ValueError('Directory event falls outside the requested cursor window')
if not self.events and (self.cursor != self.after_cursor or self.high_water_cursor != self.after_cursor):
raise ValueError('Empty Directory event batch cannot advance or trail the high-water mark')
if self.events and self.cursor != self.events[-1].cursor:
raise ValueError('Directory event batch cursor must equal its final event cursor')
return self
@runtime_checkable
class DirectoryProjectionProvider(Protocol):
"""Closed adapter that returns signature-verified control-plane data."""
async def fetch_snapshot(self, instance_uuid: str) -> DirectorySnapshot:
"""Fetch and verify an authoritative full snapshot."""
async def fetch_events(
self,
instance_uuid: str,
after_cursor: int,
limit: int,
) -> DirectoryEventBatch:
"""Fetch and verify directory events after one process-local cursor."""
async def fetch_workspaces(
self,
instance_uuid: str,
workspace_uuids: tuple[str, ...],
) -> DirectoryDelta:
"""Fetch and verify authoritative state for an explicit Workspace set."""
src/langbot/pkg/cloud/directory_projection.py
+881
View File
@@ -0,0 +1,881 @@
from __future__ import annotations
import asyncio
import datetime
import hashlib
import json
import time
from collections.abc import Callable, Iterable
from typing import TYPE_CHECKING, Any
import sqlalchemy
from sqlalchemy.dialects import postgresql, sqlite
from ..entity.persistence.cloud_directory import DirectoryProjectionInbox, DirectoryProjectionState
from ..entity.persistence.user import AccountSource, AccountStatus, User
from ..entity.persistence.workspace import (
MembershipRole,
MembershipStatus,
Workspace,
WorkspaceExecutionSource,
WorkspaceExecutionState,
WorkspaceExecutionStatus,
WorkspaceMembership,
WorkspaceSource,
WorkspaceStatus,
)
from .directory import (
DirectoryDelta,
DirectoryEvent,
DirectoryEventBatch,
DirectoryMember,
DirectoryProjectionProvider,
DirectoryProjectionUnavailableError,
DirectorySnapshot,
DirectoryWorkspace,
)
if TYPE_CHECKING:
from ..core.app import Application
_ROLE_MAP = {
'owner': MembershipRole.OWNER.value,
'admin': MembershipRole.ADMIN.value,
# Space deliberately exposes a smaller product role vocabulary. A regular
# SaaS member receives the Core developer role; operator/viewer can be
# introduced later without changing the signed directory contract.
'member': MembershipRole.DEVELOPER.value,
'developer': MembershipRole.DEVELOPER.value,
'operator': MembershipRole.OPERATOR.value,
'viewer': MembershipRole.VIEWER.value,
}
_ACCOUNT_STATUS_MAP = {
'active': AccountStatus.ACTIVE.value,
'blocked': AccountStatus.DISABLED.value,
'disabled': AccountStatus.DISABLED.value,
'deleted': AccountStatus.DELETED.value,
}
_MEMBERSHIP_STATUS_MAP = {
'active': MembershipStatus.ACTIVE.value,
'invited': MembershipStatus.DISABLED.value,
'disabled': MembershipStatus.DISABLED.value,
'removed': MembershipStatus.REMOVED.value,
}
_INCREMENTAL_PROJECTION_FINGERPRINT = hashlib.sha256(b'langbot-directory-incremental-v1').hexdigest()
class _DirectorySnapshotSuperseded(DirectoryProjectionUnavailableError):
"""A valid snapshot lost a race with a newer shared projection."""
class DirectoryProjectionService:
"""Project a verified SaaS directory into Core-owned tenant tables.
The closed adapter verifies transport signatures and returns immutable
models. Core owns database transactions, revision checks, execution fences,
and readiness. This keeps the ORM and PostgreSQL RLS boundary out of the
closed control-plane package.
"""
def __init__(
self,
ap: Application,
provider: DirectoryProjectionProvider,
instance_uuid: str,
*,
sync_interval_seconds: float = 5.0,
max_staleness_seconds: float = 60.0,
event_limit: int = 100,
monotonic_time: Callable[[], float] = time.monotonic,
) -> None:
if not isinstance(provider, DirectoryProjectionProvider):
raise TypeError('Cloud directory projection requires a DirectoryProjectionProvider')
if not instance_uuid.strip():
raise ValueError('Cloud directory projection requires an instance UUID')
if sync_interval_seconds <= 0:
raise ValueError('Directory sync interval must be positive')
if max_staleness_seconds <= sync_interval_seconds:
raise ValueError('Directory max staleness must exceed the sync interval')
if event_limit <= 0 or event_limit > 100:
raise ValueError('Directory event limit must be between 1 and 100')
self.ap = ap
self.provider = provider
self.instance_uuid = instance_uuid.strip()
self.sync_interval_seconds = float(sync_interval_seconds)
self.max_staleness_seconds = float(max_staleness_seconds)
self.event_limit = event_limit
self._monotonic_time = monotonic_time
self._last_success_monotonic: float | None = None
self._ready = False
# Every runtime replica must consume the event stream independently:
# entitlement snapshots live in the closed adapter's process memory.
# The database cursor remains the shared projection high-water mark,
# while this cursor tracks what this process has actually observed.
self._consumer_cursor: int | None = None
async def initialize(self) -> None:
"""Block Cloud startup until one full signed snapshot is committed."""
last_superseded: _DirectorySnapshotSuperseded | None = None
for _attempt in range(5):
snapshot = await self.provider.fetch_snapshot(self.instance_uuid)
try:
await self.apply_snapshot(snapshot)
except _DirectorySnapshotSuperseded as exc:
last_superseded = exc
continue
self._consumer_cursor = snapshot.cursor
return
raise DirectoryProjectionUnavailableError(
'Directory snapshot was repeatedly superseded by another runtime replica'
) from last_superseded
async def run(self) -> None:
"""Continuously refresh the directory and fail closed when it goes stale."""
delay = self.sync_interval_seconds
while True:
try:
await asyncio.sleep(delay)
await self.sync_once()
delay = self.sync_interval_seconds
except asyncio.CancelledError:
raise
except Exception:
self.ap.logger.exception('Cloud directory synchronization failed')
delay = min(max(delay * 2, self.sync_interval_seconds), self.max_staleness_seconds / 2)
async def sync_once(self) -> None:
cursor = self._consumer_cursor
if cursor is None:
await self.initialize()
return
batch = await self.provider.fetch_events(
self.instance_uuid,
cursor,
self.event_limit,
)
batch = DirectoryEventBatch.model_validate(batch.model_dump())
self._validate_batch(batch, expected_after_cursor=cursor)
if batch.events:
directory_revisions = self._directory_event_revisions(batch.events)
if directory_revisions:
requested_workspace_uuids = tuple(sorted(directory_revisions))
delta = await self.provider.fetch_workspaces(
self.instance_uuid,
requested_workspace_uuids,
)
await self.apply_delta(delta, batch)
else:
await self.apply_event_batch(batch)
self._consumer_cursor = batch.cursor
return
await self._touch_freshness(cursor)
def require_ready(self) -> None:
"""Fail synchronously at execution admission when projection is stale."""
last_success = self._last_success_monotonic
if not self._ready or last_success is None:
raise DirectoryProjectionUnavailableError('Cloud directory projection is not ready')
if self._monotonic_time() - last_success >= self.max_staleness_seconds:
raise DirectoryProjectionUnavailableError('Cloud directory projection is stale')
async def apply_snapshot(
self,
snapshot: DirectorySnapshot,
*,
events: Iterable[DirectoryEvent] = (),
) -> None:
"""Atomically apply one monotonic full snapshot and its event receipts."""
if not isinstance(snapshot, DirectorySnapshot):
raise DirectoryProjectionUnavailableError('Directory provider returned an invalid snapshot')
snapshot = DirectorySnapshot.model_validate(snapshot.model_dump())
if snapshot.instance_uuid != self.instance_uuid:
raise DirectoryProjectionUnavailableError('Directory snapshot targets another LangBot instance')
fingerprint = self._snapshot_fingerprint(snapshot)
now = self._utcnow()
lease_expires_at = now + datetime.timedelta(seconds=self.max_staleness_seconds)
event_list = tuple(DirectoryEvent.model_validate(event.model_dump()) for event in events)
directory_uow = getattr(self.ap.persistence_mgr, 'directory_projection_uow', None)
if not callable(directory_uow):
raise DirectoryProjectionUnavailableError('Directory projection persistence scope is unavailable')
async with directory_uow(self.instance_uuid) as uow:
session = uow.session
state_values = {
'instance_uuid': self.instance_uuid,
'cursor': snapshot.cursor,
'snapshot_coverage_cursor': snapshot.cursor,
'snapshot_fingerprint': fingerprint,
'last_applied_at': now,
'lease_expires_at': lease_expires_at,
}
dialect_name = self.ap.persistence_mgr.get_db_engine().dialect.name
if dialect_name == 'postgresql':
insert_state = postgresql.insert(DirectoryProjectionState)
elif dialect_name == 'sqlite':
insert_state = sqlite.insert(DirectoryProjectionState)
else: # pragma: no cover - Cloud supports PostgreSQL; tests use SQLite.
raise DirectoryProjectionUnavailableError('Directory projection database is unsupported')
await session.execute(
insert_state.values(**state_values).on_conflict_do_nothing(
index_elements=[DirectoryProjectionState.instance_uuid]
)
)
state = await session.scalar(
sqlalchemy.select(DirectoryProjectionState)
.where(DirectoryProjectionState.instance_uuid == self.instance_uuid)
.with_for_update()
)
if state is None: # pragma: no cover - insert/select are one transaction.
raise DirectoryProjectionUnavailableError('Directory projection state could not be locked')
if snapshot.cursor < state.cursor:
raise _DirectorySnapshotSuperseded('Directory snapshot cursor rolled back')
if snapshot.cursor == state.cursor and state.snapshot_fingerprint not in {
fingerprint,
_INCREMENTAL_PROJECTION_FINGERPRINT,
}:
raise DirectoryProjectionUnavailableError('Directory snapshot cursor has conflicting contents')
await self._record_events(session, event_list, now=now)
await self._apply_accounts(session, snapshot)
await self._apply_workspaces(session, snapshot)
await self._fence_absent_workspaces(session, snapshot)
state.cursor = snapshot.cursor
state.snapshot_coverage_cursor = snapshot.cursor
state.snapshot_fingerprint = fingerprint
state.last_applied_at = now
state.lease_expires_at = lease_expires_at
await self._mark_events_applied(session, event_list, now=now)
await session.flush()
self._record_success()
self._consumer_cursor = snapshot.cursor
async def apply_delta(self, delta: DirectoryDelta, batch: DirectoryEventBatch) -> None:
"""Apply only Workspaces named by directory events in one signed page."""
if not isinstance(delta, DirectoryDelta):
raise DirectoryProjectionUnavailableError('Directory provider returned an invalid delta')
if not isinstance(batch, DirectoryEventBatch):
raise DirectoryProjectionUnavailableError('Directory provider returned an invalid event batch')
delta = DirectoryDelta.model_validate(delta.model_dump())
batch = DirectoryEventBatch.model_validate(batch.model_dump())
self._validate_batch(batch, expected_after_cursor=batch.after_cursor)
if delta.instance_uuid != self.instance_uuid:
raise DirectoryProjectionUnavailableError('Directory delta targets another LangBot instance')
required_revisions = self._directory_event_revisions(batch.events)
requested = set(delta.requested_workspace_uuids)
if not required_revisions or requested != set(required_revisions):
raise DirectoryProjectionUnavailableError('Directory delta does not match its event batch')
returned = {workspace.uuid: workspace for workspace in delta.workspaces}
for workspace_uuid, workspace in returned.items():
if workspace.projection_revision < required_revisions[workspace_uuid]:
raise DirectoryProjectionUnavailableError(
'Directory Workspace delta is older than its signed event notification'
)
now = self._utcnow()
lease_expires_at = now + datetime.timedelta(seconds=self.max_staleness_seconds)
directory_uow = getattr(self.ap.persistence_mgr, 'directory_projection_uow', None)
if not callable(directory_uow):
raise DirectoryProjectionUnavailableError('Directory projection persistence scope is unavailable')
projection_caught_up = False
async with directory_uow(self.instance_uuid) as uow:
session = uow.session
state = await session.scalar(
sqlalchemy.select(DirectoryProjectionState)
.where(DirectoryProjectionState.instance_uuid == self.instance_uuid)
.with_for_update()
)
if state is None:
raise DirectoryProjectionUnavailableError('Directory projection state disappeared')
state_cursor = int(state.cursor)
if state_cursor < batch.after_cursor:
raise DirectoryProjectionUnavailableError('Directory projection state cursor rolled back')
# A different runtime replica may already have applied this page.
# In that case every receipt through the shared cursor must exist;
# this replica still fetched the delta and refreshed its own
# entitlement cache before advancing its process-local cursor.
await self._record_events(
session,
batch.events,
now=now,
allow_missing_through_cursor=int(state.snapshot_coverage_cursor),
reject_missing_through_cursor=state_cursor,
)
if state_cursor < batch.cursor:
projected_delta = DirectorySnapshot(
instance_uuid=self.instance_uuid,
cursor=batch.cursor,
generated_at=delta.generated_at,
workspaces=delta.workspaces,
)
await self._apply_accounts(session, projected_delta)
await self._apply_workspaces(session, projected_delta)
await self._fence_workspaces(
session,
{
workspace_uuid: required_revisions[workspace_uuid]
for workspace_uuid in requested - set(returned)
},
)
state.cursor = batch.cursor
# A per-Workspace delta cannot prove a full-directory
# fingerprint. Event receipts and entity revisions protect the
# incremental path; a later full snapshot replaces this marker.
state.snapshot_fingerprint = _INCREMENTAL_PROJECTION_FINGERPRINT
state.last_applied_at = now
state.lease_expires_at = lease_expires_at
await self._mark_events_applied(session, batch.events, now=now)
await session.flush()
projection_caught_up = batch.cursor == batch.high_water_cursor and int(state.cursor) == batch.cursor
if projection_caught_up:
self._record_success()
self._consumer_cursor = batch.cursor
async def apply_event_batch(self, batch: DirectoryEventBatch) -> None:
"""Advance non-directory events after the adapter refreshes local caches."""
if not isinstance(batch, DirectoryEventBatch):
raise DirectoryProjectionUnavailableError('Directory provider returned an invalid event batch')
batch = DirectoryEventBatch.model_validate(batch.model_dump())
self._validate_batch(batch, expected_after_cursor=batch.after_cursor)
if any(event.event_type == 'directory.changed' for event in batch.events):
raise DirectoryProjectionUnavailableError('Directory changes require an authoritative full snapshot')
now = self._utcnow()
lease_expires_at = now + datetime.timedelta(seconds=self.max_staleness_seconds)
directory_uow = getattr(self.ap.persistence_mgr, 'directory_projection_uow', None)
if not callable(directory_uow):
raise DirectoryProjectionUnavailableError('Directory projection persistence scope is unavailable')
projection_caught_up = False
async with directory_uow(self.instance_uuid) as uow:
state = await uow.session.scalar(
sqlalchemy.select(DirectoryProjectionState)
.where(DirectoryProjectionState.instance_uuid == self.instance_uuid)
.with_for_update()
)
if state is None:
raise DirectoryProjectionUnavailableError('Directory projection state disappeared')
if state.cursor < batch.after_cursor:
raise DirectoryProjectionUnavailableError('Directory projection state cursor rolled back')
await self._record_events(
uow.session,
batch.events,
now=now,
allow_missing_through_cursor=int(state.snapshot_coverage_cursor),
reject_missing_through_cursor=int(state.cursor),
)
state.cursor = max(int(state.cursor), batch.cursor)
state.last_applied_at = now
state.lease_expires_at = lease_expires_at
await self._mark_events_applied(uow.session, batch.events, now=now)
await uow.session.flush()
projection_caught_up = batch.cursor == batch.high_water_cursor and int(state.cursor) == batch.cursor
if projection_caught_up:
self._record_success()
self._consumer_cursor = batch.cursor
async def _touch_freshness(self, requested_cursor: int) -> None:
now = self._utcnow()
lease_expires_at = now + datetime.timedelta(seconds=self.max_staleness_seconds)
directory_uow = getattr(self.ap.persistence_mgr, 'directory_projection_uow', None)
if not callable(directory_uow):
raise DirectoryProjectionUnavailableError('Directory projection persistence scope is unavailable')
async with directory_uow(self.instance_uuid) as uow:
state = await uow.session.scalar(
sqlalchemy.select(DirectoryProjectionState)
.where(DirectoryProjectionState.instance_uuid == self.instance_uuid)
.with_for_update()
)
if state is None:
raise DirectoryProjectionUnavailableError('Directory projection state disappeared')
if state.cursor < requested_cursor:
raise DirectoryProjectionUnavailableError('Directory projection state cursor rolled back')
if state.cursor > requested_cursor:
raise DirectoryProjectionUnavailableError(
'This runtime replica has not consumed the shared directory high-water mark'
)
state.last_applied_at = now
state.lease_expires_at = lease_expires_at
await uow.session.flush()
self._record_success()
def _validate_batch(self, batch: DirectoryEventBatch, *, expected_after_cursor: int) -> None:
if batch.instance_uuid != self.instance_uuid:
raise DirectoryProjectionUnavailableError('Directory event batch targets another LangBot instance')
if batch.after_cursor != expected_after_cursor:
raise DirectoryProjectionUnavailableError('Directory event batch does not match the requested cursor')
supported_event_types = {'directory.changed', 'entitlement.changed'}
if any(event.event_type not in supported_event_types for event in batch.events):
raise DirectoryProjectionUnavailableError('Directory event batch contains an unsupported event type')
for event in batch.events:
if event.payload.get('workspace_uuid') != event.aggregate_uuid:
raise DirectoryProjectionUnavailableError('Directory event payload has a conflicting Workspace scope')
revision_key = 'directory_revision' if event.event_type == 'directory.changed' else 'entitlement_revision'
payload_revision = event.payload.get(revision_key)
if type(payload_revision) is not int or payload_revision != event.revision:
raise DirectoryProjectionUnavailableError('Directory event payload has a conflicting revision')
@staticmethod
def _directory_event_revisions(events: Iterable[DirectoryEvent]) -> dict[str, int]:
revisions: dict[str, int] = {}
for event in events:
if event.event_type == 'directory.changed':
revisions[event.aggregate_uuid] = max(revisions.get(event.aggregate_uuid, 0), event.revision)
return revisions
async def _record_events(
self,
session: Any,
events: tuple[DirectoryEvent, ...],
*,
now: datetime.datetime,
allow_missing_through_cursor: int = -1,
reject_missing_through_cursor: int | None = None,
) -> None:
for event in events:
fingerprint = self._fingerprint(event.model_dump(mode='json'))
existing = await session.scalar(
sqlalchemy.select(DirectoryProjectionInbox).where(
DirectoryProjectionInbox.instance_uuid == self.instance_uuid,
DirectoryProjectionInbox.event_uuid == event.uuid,
)
)
if existing is not None:
if existing.cursor != event.cursor or existing.fingerprint != fingerprint:
raise DirectoryProjectionUnavailableError('Directory event UUID has conflicting contents')
continue
if (
reject_missing_through_cursor is not None
and allow_missing_through_cursor < event.cursor <= reject_missing_through_cursor
):
raise DirectoryProjectionUnavailableError(
'Directory projection cursor advanced without a matching event receipt'
)
session.add(
DirectoryProjectionInbox(
instance_uuid=self.instance_uuid,
event_uuid=event.uuid,
cursor=event.cursor,
event_type=event.event_type,
revision=event.revision,
fingerprint=fingerprint,
received_at=now,
applied_at=None,
)
)
async def _mark_events_applied(
self,
session: Any,
events: Iterable[DirectoryEvent],
*,
now: datetime.datetime,
) -> None:
event_uuids = [event.uuid for event in events]
if not event_uuids:
return
inbox_rows = (
await session.scalars(
sqlalchemy.select(DirectoryProjectionInbox).where(
DirectoryProjectionInbox.instance_uuid == self.instance_uuid,
DirectoryProjectionInbox.event_uuid.in_(event_uuids),
)
)
).all()
if len(inbox_rows) != len(event_uuids):
raise DirectoryProjectionUnavailableError('Directory event receipt could not be persisted')
for row in inbox_rows:
row.applied_at = now
async def _apply_accounts(self, session: Any, snapshot: DirectorySnapshot) -> None:
selected: dict[str, DirectoryMember] = {}
emails: dict[str, str] = {}
for workspace in snapshot.workspaces:
for member in workspace.members:
email_owner = emails.setdefault(member.normalized_email, member.account_uuid)
if email_owner != member.account_uuid:
raise DirectoryProjectionUnavailableError(
'Directory snapshot maps one normalized email to multiple accounts'
)
previous = selected.get(member.account_uuid)
if previous is not None and self._account_projection(previous) != self._account_projection(member):
raise DirectoryProjectionUnavailableError('Directory snapshot has conflicting account projections')
if previous is None:
selected[member.account_uuid] = member
for account_uuid, member in selected.items():
account = await session.scalar(sqlalchemy.select(User).where(User.uuid == account_uuid))
email_account = await session.scalar(
sqlalchemy.select(User).where(User.normalized_email == member.normalized_email)
)
if email_account is not None and email_account.uuid != account_uuid:
raise DirectoryProjectionUnavailableError('Directory account email collides with another Core account')
if account is None:
account = User(
uuid=account_uuid,
user=member.display_name,
normalized_email=member.normalized_email,
password='',
status=_ACCOUNT_STATUS_MAP[member.account_status],
source=AccountSource.CLOUD_PROJECTION.value,
projection_revision=snapshot.cursor,
account_type='space',
space_account_uuid=account_uuid,
)
session.add(account)
continue
if account.source != AccountSource.CLOUD_PROJECTION.value:
raise DirectoryProjectionUnavailableError('Directory account UUID collides with a local Core account')
if account.projection_revision > snapshot.cursor:
raise DirectoryProjectionUnavailableError('Directory account revision rolled back')
projected_account = self._account_projection(member)
persisted_account = self._persisted_account_projection(account)
if account.projection_revision == snapshot.cursor and persisted_account != projected_account:
raise DirectoryProjectionUnavailableError('Directory account revision has conflicting contents')
if persisted_account == projected_account:
# A Workspace rename, role update, or another member's change
# must not revoke this Account's JWT. Account revisions advance
# only when the Account projection itself changes.
continue
account.user = member.display_name
account.normalized_email = member.normalized_email
account.status = _ACCOUNT_STATUS_MAP[member.account_status]
account.projection_revision = snapshot.cursor
account.account_type = 'space'
account.space_account_uuid = account_uuid
await session.flush()
async def _apply_workspaces(self, session: Any, snapshot: DirectorySnapshot) -> None:
for candidate in snapshot.workspaces:
workspace = await session.get(Workspace, candidate.uuid)
if workspace is None:
workspace = Workspace(
uuid=candidate.uuid,
instance_uuid=self.instance_uuid,
name=candidate.name,
slug=candidate.slug,
type=candidate.type,
status=candidate.status,
created_by_account_uuid=await self._projected_creator_uuid(session, candidate),
source=WorkspaceSource.CLOUD_PROJECTION.value,
projection_revision=candidate.projection_revision,
)
session.add(workspace)
await session.flush()
else:
self._validate_existing_workspace(workspace, candidate)
workspace.name = candidate.name
workspace.slug = candidate.slug
workspace.type = candidate.type
workspace.status = candidate.status
workspace.created_by_account_uuid = await self._projected_creator_uuid(session, candidate)
workspace.projection_revision = candidate.projection_revision
await self._apply_memberships(session, workspace, candidate)
await self._apply_execution_state(session, workspace, candidate)
async def _projected_creator_uuid(self, session: Any, candidate: DirectoryWorkspace) -> str | None:
creator = await session.scalar(sqlalchemy.select(User).where(User.uuid == candidate.created_by_account_uuid))
if creator is None:
if candidate.status == WorkspaceStatus.ACTIVE.value:
raise DirectoryProjectionUnavailableError('Active Directory Workspace creator is not projected')
return None
return candidate.created_by_account_uuid
def _validate_existing_workspace(self, workspace: Workspace, candidate: DirectoryWorkspace) -> None:
if workspace.instance_uuid != self.instance_uuid:
raise DirectoryProjectionUnavailableError('Directory Workspace belongs to another LangBot instance')
if workspace.source != WorkspaceSource.CLOUD_PROJECTION.value:
raise DirectoryProjectionUnavailableError('Directory Workspace UUID collides with a local Workspace')
if workspace.projection_revision > candidate.projection_revision:
raise DirectoryProjectionUnavailableError('Directory Workspace revision rolled back')
if workspace.projection_revision == candidate.projection_revision and self._workspace_projection(
workspace
) != self._candidate_workspace_projection(candidate):
raise DirectoryProjectionUnavailableError('Directory Workspace revision has conflicting contents')
async def _apply_memberships(
self,
session: Any,
workspace: Workspace,
candidate: DirectoryWorkspace,
) -> None:
existing = {
membership.account_uuid: membership
for membership in (
await session.scalars(
sqlalchemy.select(WorkspaceMembership).where(WorkspaceMembership.workspace_uuid == workspace.uuid)
)
).all()
}
included_accounts: set[str] = set()
for member in candidate.members:
included_accounts.add(member.account_uuid)
membership = existing.get(member.account_uuid)
joined_at = self._naive_utc(member.joined_at)
role = _ROLE_MAP[member.role]
status = _MEMBERSHIP_STATUS_MAP[member.membership_status]
if candidate.status != WorkspaceStatus.ACTIVE.value:
status = (
MembershipStatus.REMOVED.value
if candidate.status in {WorkspaceStatus.ARCHIVED.value, WorkspaceStatus.DELETED.value}
else MembershipStatus.DISABLED.value
)
if membership is None:
session.add(
WorkspaceMembership(
uuid=member.membership_uuid,
workspace_uuid=workspace.uuid,
account_uuid=member.account_uuid,
role=role,
status=status,
joined_at=joined_at,
projection_revision=member.projection_revision,
)
)
continue
if membership.uuid != member.membership_uuid:
raise DirectoryProjectionUnavailableError('Directory membership UUID changed for one account')
if membership.projection_revision > member.projection_revision:
raise DirectoryProjectionUnavailableError('Directory membership revision rolled back')
if membership.projection_revision == member.projection_revision and self._membership_projection(
membership
) != (role, status, self._datetime_fingerprint(joined_at)):
raise DirectoryProjectionUnavailableError('Directory membership revision has conflicting contents')
membership.role = role
membership.status = status
membership.joined_at = joined_at
membership.projection_revision = member.projection_revision
for account_uuid, membership in existing.items():
if account_uuid not in included_accounts:
membership.status = MembershipStatus.REMOVED.value
membership.projection_revision = max(
int(membership.projection_revision),
candidate.projection_revision,
)
await session.flush()
async def _apply_execution_state(
self,
session: Any,
workspace: Workspace,
candidate: DirectoryWorkspace,
) -> None:
active = candidate.status == WorkspaceStatus.ACTIVE.value
desired_state = WorkspaceExecutionStatus.ACTIVE.value if active else WorkspaceExecutionStatus.INACTIVE.value
execution = await session.get(WorkspaceExecutionState, workspace.uuid)
if execution is None:
session.add(
WorkspaceExecutionState(
workspace_uuid=workspace.uuid,
instance_uuid=self.instance_uuid,
active_generation=candidate.execution_generation,
state=desired_state,
write_fenced=not active,
source=WorkspaceExecutionSource.CLOUD.value,
desired_state_revision=candidate.projection_revision,
)
)
await session.flush()
return
if execution.instance_uuid != self.instance_uuid or execution.source != WorkspaceExecutionSource.CLOUD.value:
raise DirectoryProjectionUnavailableError('Directory execution state has an invalid owner')
if execution.active_generation > candidate.execution_generation:
raise DirectoryProjectionUnavailableError('Directory execution generation rolled back')
if execution.desired_state_revision > candidate.projection_revision:
raise DirectoryProjectionUnavailableError('Directory desired-state revision rolled back')
if execution.desired_state_revision == candidate.projection_revision and (
execution.active_generation != candidate.execution_generation
or execution.state != desired_state
or execution.write_fenced != (not active)
):
raise DirectoryProjectionUnavailableError(
'Directory execution state has conflicting contents at one revision'
)
execution.active_generation = candidate.execution_generation
execution.state = desired_state
execution.write_fenced = not active
execution.desired_state_revision = candidate.projection_revision
await session.flush()
async def _fence_absent_workspaces(self, session: Any, snapshot: DirectorySnapshot) -> None:
included = {workspace.uuid for workspace in snapshot.workspaces}
projected = (
await session.scalars(
sqlalchemy.select(Workspace).where(
Workspace.instance_uuid == self.instance_uuid,
Workspace.source == WorkspaceSource.CLOUD_PROJECTION.value,
)
)
).all()
for workspace in projected:
if workspace.uuid in included:
continue
workspace.status = WorkspaceStatus.ARCHIVED.value
await self._remove_workspace_memberships(session, workspace.uuid)
execution = await session.get(WorkspaceExecutionState, workspace.uuid)
if execution is not None:
execution.state = WorkspaceExecutionStatus.INACTIVE.value
execution.write_fenced = True
await session.flush()
async def _fence_workspaces(self, session: Any, workspace_revisions: dict[str, int]) -> None:
"""Fence requested Workspaces omitted from an authoritative delta."""
if not workspace_revisions:
return
projected = (
await session.scalars(
sqlalchemy.select(Workspace).where(
Workspace.instance_uuid == self.instance_uuid,
Workspace.source == WorkspaceSource.CLOUD_PROJECTION.value,
Workspace.uuid.in_(workspace_revisions),
)
)
).all()
for workspace in projected:
tombstone_revision = workspace_revisions[workspace.uuid]
if int(workspace.projection_revision) > tombstone_revision:
raise DirectoryProjectionUnavailableError('Directory Workspace tombstone revision rolled back')
memberships = (
await session.scalars(
sqlalchemy.select(WorkspaceMembership).where(WorkspaceMembership.workspace_uuid == workspace.uuid)
)
).all()
if any(int(membership.projection_revision) > tombstone_revision for membership in memberships):
raise DirectoryProjectionUnavailableError('Directory membership tombstone revision rolled back')
execution = await session.get(WorkspaceExecutionState, workspace.uuid)
if execution is not None and int(execution.desired_state_revision) > tombstone_revision:
raise DirectoryProjectionUnavailableError('Directory execution tombstone revision rolled back')
workspace.status = WorkspaceStatus.ARCHIVED.value
workspace.projection_revision = max(int(workspace.projection_revision), tombstone_revision)
await self._remove_workspace_memberships(
session,
workspace.uuid,
projection_revision=tombstone_revision,
memberships=memberships,
)
if execution is not None:
execution.state = WorkspaceExecutionStatus.INACTIVE.value
execution.write_fenced = True
execution.desired_state_revision = max(
int(execution.desired_state_revision),
tombstone_revision,
)
await session.flush()
async def _remove_workspace_memberships(
self,
session: Any,
workspace_uuid: str,
*,
projection_revision: int | None = None,
memberships: Iterable[WorkspaceMembership] | None = None,
) -> None:
if memberships is None:
memberships = (
await session.scalars(
sqlalchemy.select(WorkspaceMembership).where(WorkspaceMembership.workspace_uuid == workspace_uuid)
)
).all()
for membership in memberships:
membership.status = MembershipStatus.REMOVED.value
if projection_revision is not None:
membership.projection_revision = max(
int(membership.projection_revision),
projection_revision,
)
def _record_success(self) -> None:
self._last_success_monotonic = self._monotonic_time()
self._ready = True
@classmethod
def _snapshot_fingerprint(cls, snapshot: DirectorySnapshot) -> str:
workspaces = []
for workspace in sorted(snapshot.workspaces, key=lambda item: item.uuid):
data = workspace.model_dump(mode='json')
data['members'] = sorted(data['members'], key=lambda item: item['membership_uuid'])
workspaces.append(data)
return cls._fingerprint(
{
'instance_uuid': snapshot.instance_uuid,
'workspaces': workspaces,
}
)
@staticmethod
def _fingerprint(value: Any) -> str:
encoded = json.dumps(value, sort_keys=True, separators=(',', ':'), ensure_ascii=True).encode()
return hashlib.sha256(encoded).hexdigest()
@staticmethod
def _account_projection(member: DirectoryMember) -> tuple[str, str, str]:
return member.normalized_email, member.display_name, _ACCOUNT_STATUS_MAP[member.account_status]
@staticmethod
def _persisted_account_projection(account: User) -> tuple[str, str, str]:
return account.normalized_email, account.user, account.status
@staticmethod
def _workspace_projection(workspace: Workspace) -> tuple[Any, ...]:
return (
workspace.name,
workspace.slug,
workspace.type,
workspace.status,
workspace.created_by_account_uuid,
)
@staticmethod
def _candidate_workspace_projection(candidate: DirectoryWorkspace) -> tuple[Any, ...]:
return (
candidate.name,
candidate.slug,
candidate.type,
candidate.status,
candidate.created_by_account_uuid,
)
@classmethod
def _membership_projection(cls, membership: WorkspaceMembership) -> tuple[Any, ...]:
return (
membership.role,
membership.status,
cls._datetime_fingerprint(membership.joined_at),
)
@staticmethod
def _datetime_fingerprint(value: datetime.datetime | None) -> str | None:
if value is None:
return None
return DirectoryProjectionService._naive_utc(value).isoformat(timespec='microseconds')
@staticmethod
def _naive_utc(value: datetime.datetime | None) -> datetime.datetime | None:
if value is None:
return None
if value.tzinfo is None:
return value
return value.astimezone(datetime.UTC).replace(tzinfo=None)
@staticmethod
def _utcnow() -> datetime.datetime:
return datetime.datetime.now(datetime.UTC)
src/langbot/pkg/core/app.py
+22
View File
@@ -48,6 +48,7 @@ from ..skill import manager as skill_mgr
from ..workspace import service as workspace_service_module
from ..workspace import collaboration as workspace_collaboration_module
from ..cloud import bootstrap as cloud_bootstrap_module
from ..cloud import directory_projection as cloud_directory_projection_module
from ..cloud import entitlements as cloud_entitlements_module
from ..api.http.context import ExecutionContext, PrincipalContext, PrincipalType
@@ -132,8 +133,12 @@ class Application:
deployment_admission: cloud_bootstrap_module.DeploymentAdmissionGuard = None
manifest_refresh_service: cloud_bootstrap_module.CloudManifestRefreshService | None = None
entitlement_resolver: cloud_entitlements_module.EntitlementResolver | None = None
directory_projection_service: cloud_directory_projection_module.DirectoryProjectionService | None = None
vector_db_mgr: vectordb_mgr.VectorDBManager = None
http_ctrl: http_controller.HTTPController = None
@@ -189,6 +194,19 @@ class Application:
async def run(self):
try:
if self.directory_projection_service is not None:
self.task_mgr.create_task(
self.directory_projection_service.run(),
name='cloud-directory-projection',
scopes=[core_entities.LifecycleControlScope.APPLICATION],
)
if self.manifest_refresh_service is not None:
self.task_mgr.create_task(
self.manifest_refresh_service.run(),
name='cloud-manifest-refresh',
scopes=[core_entities.LifecycleControlScope.APPLICATION],
)
await self.plugin_connector.initialize_plugins()
# 后续可能会允许动态重启其他任务
@@ -374,6 +392,10 @@ class Application:
if self.plugin_connector is not None:
with contextlib.suppress(Exception):
await self.plugin_connector.aclose()
manifest_provider = getattr(self.deployment, 'manifest_provider', None)
if manifest_provider is not None:
with contextlib.suppress(Exception):
await manifest_provider.aclose()
if self.task_mgr is not None:
tasks = [wrapper.task for wrapper in self.task_mgr.tasks if not wrapper.task.done()]
src/langbot/pkg/core/stages/build_app.py
+19
View File
@@ -40,6 +40,7 @@ from ...survey import manager as survey_module
from ...workspace import service as workspace_service_module
from ...workspace import collaboration as workspace_collaboration_module
from ...cloud import bootstrap as cloud_bootstrap
from ...cloud.directory_projection import DirectoryProjectionService
from ...cloud.entitlements import EntitlementResolver
from ...api.http.context import ExecutionContext, PrincipalContext, PrincipalType
from ...api.http.authz import WorkspaceRequiredError
@@ -63,6 +64,15 @@ class BuildAppStage(stage.BootingStage):
constants.instance_id,
deployment,
)
ap.manifest_refresh_service = (
cloud_bootstrap.CloudManifestRefreshService(
ap.deployment_admission,
deployment.manifest_provider,
ap.logger,
)
if deployment.multi_workspace_enabled
else None
)
ap.entitlement_resolver = (
EntitlementResolver(
constants.instance_id,
@@ -137,6 +147,15 @@ class BuildAppStage(stage.BootingStage):
ap.persistence_mgr = persistence_mgr_inst
await persistence_mgr_inst.initialize()
if deployment.multi_workspace_enabled:
directory_projection_service = DirectoryProjectionService(
ap,
deployment.directory_provider,
constants.instance_id,
)
await directory_projection_service.initialize()
ap.directory_projection_service = directory_projection_service
workspace_policy = deployment.workspace_policy
workspace_service_inst = workspace_service_module.WorkspaceService(
ap,
src/langbot/pkg/entity/persistence/cloud_directory.py
+69
View File
@@ -0,0 +1,69 @@
from __future__ import annotations
import sqlalchemy
from .base import Base
class DirectoryProjectionState(Base):
"""Durable cursor and lease for one verified Cloud directory."""
__tablename__ = 'directory_projection_states'
instance_uuid = sqlalchemy.Column(sqlalchemy.String(255), primary_key=True)
cursor = sqlalchemy.Column(sqlalchemy.BigInteger, nullable=False, server_default='0')
snapshot_coverage_cursor = sqlalchemy.Column(sqlalchemy.BigInteger, nullable=False, server_default='0')
snapshot_fingerprint = sqlalchemy.Column(sqlalchemy.Text, nullable=False)
last_applied_at = sqlalchemy.Column(sqlalchemy.DateTime(timezone=True), nullable=False)
lease_expires_at = sqlalchemy.Column(sqlalchemy.DateTime(timezone=True), nullable=True)
__table_args__ = (
sqlalchemy.CheckConstraint('cursor >= 0', name='ck_directory_projection_state_cursor'),
sqlalchemy.CheckConstraint(
'snapshot_coverage_cursor >= 0 AND snapshot_coverage_cursor <= cursor',
name='ck_directory_projection_state_snapshot_coverage',
),
sqlalchemy.CheckConstraint(
'length(snapshot_fingerprint) = 64',
name='ck_directory_projection_state_fingerprint',
),
)
class DirectoryProjectionInbox(Base):
"""Idempotency ledger for signed control-plane directory events."""
__tablename__ = 'directory_projection_inbox'
instance_uuid = sqlalchemy.Column(sqlalchemy.String(255), primary_key=True)
event_uuid = sqlalchemy.Column(sqlalchemy.String(36), primary_key=True)
cursor = sqlalchemy.Column(sqlalchemy.BigInteger, nullable=False)
event_type = sqlalchemy.Column(sqlalchemy.String(128), nullable=False)
revision = sqlalchemy.Column(sqlalchemy.BigInteger, nullable=False)
fingerprint = sqlalchemy.Column(sqlalchemy.Text, nullable=False)
received_at = sqlalchemy.Column(
sqlalchemy.DateTime(timezone=True),
nullable=False,
server_default=sqlalchemy.func.now(),
)
applied_at = sqlalchemy.Column(sqlalchemy.DateTime(timezone=True), nullable=True)
__table_args__ = (
sqlalchemy.UniqueConstraint(
'instance_uuid',
'cursor',
name='uq_directory_projection_inbox_cursor',
),
sqlalchemy.Index(
'ix_directory_projection_inbox_pending',
'instance_uuid',
'applied_at',
'cursor',
),
sqlalchemy.CheckConstraint('cursor > 0', name='ck_directory_projection_inbox_cursor'),
sqlalchemy.CheckConstraint('revision > 0', name='ck_directory_projection_inbox_revision'),
sqlalchemy.CheckConstraint(
'length(fingerprint) = 64',
name='ck_directory_projection_inbox_fingerprint',
),
)
src/langbot/pkg/persistence/alembic/versions/0014_cloud_directory_projection.py
+268
View File
@@ -0,0 +1,268 @@
"""add the Cloud directory projection persistence boundary
Revision ID: 0014_cloud_directory
Revises: 0013_tenant_pgvector
Create Date: 2026-07-24
The open Core projector receives already-verified control-plane data and is the
only runtime path allowed to mutate projected Workspace directory rows. Its
transaction-local instance setting is intentionally distinct from both normal
Workspace scope and the read-only instance discovery scope.
"""
from __future__ import annotations
import sqlalchemy as sa
from alembic import op
revision = '0014_cloud_directory'
down_revision = '0013_tenant_pgvector'
branch_labels = None
depends_on = None
_STATE_TABLE = 'directory_projection_states'
_INBOX_TABLE = 'directory_projection_inbox'
_DIRECTORY_POLICY_NAME = 'langbot_directory_projection'
_TENANT_POLICY_NAME = 'langbot_workspace_isolation'
_LOCAL_WRITE_POLICY_NAME = 'langbot_workspace_local_directory_write'
_DIRECTORY_SETTING = 'langbot.directory_instance_uuid'
_TENANT_SETTING = 'langbot.workspace_uuid'
_PROJECTED_TENANT_TABLES = (
'workspaces',
'workspace_memberships',
'workspace_execution_states',
)
def _setting(name: str) -> str:
return f"NULLIF(current_setting('{name}', true), '')"
def _quote(conn: sa.Connection, identifier: str) -> str:
return conn.dialect.identifier_preparer.quote(identifier)
def _create_tables(conn: sa.Connection) -> None:
existing_tables = set(sa.inspect(conn).get_table_names())
if _STATE_TABLE not in existing_tables:
op.create_table(
_STATE_TABLE,
sa.Column('instance_uuid', sa.String(255), nullable=False),
sa.Column('cursor', sa.BigInteger(), server_default='0', nullable=False),
sa.Column('snapshot_coverage_cursor', sa.BigInteger(), server_default='0', nullable=False),
sa.Column('snapshot_fingerprint', sa.Text(), nullable=False),
sa.Column('last_applied_at', sa.DateTime(timezone=True), nullable=False),
sa.Column('lease_expires_at', sa.DateTime(timezone=True), nullable=True),
sa.CheckConstraint(
'cursor >= 0',
name='ck_directory_projection_state_cursor',
),
sa.CheckConstraint(
'snapshot_coverage_cursor >= 0 AND snapshot_coverage_cursor <= cursor',
name='ck_directory_projection_state_snapshot_coverage',
),
sa.CheckConstraint(
'length(snapshot_fingerprint) = 64',
name='ck_directory_projection_state_fingerprint',
),
sa.PrimaryKeyConstraint('instance_uuid'),
)
if _INBOX_TABLE not in existing_tables:
op.create_table(
_INBOX_TABLE,
sa.Column('instance_uuid', sa.String(255), nullable=False),
sa.Column('event_uuid', sa.String(36), nullable=False),
sa.Column('cursor', sa.BigInteger(), nullable=False),
sa.Column('event_type', sa.String(128), nullable=False),
sa.Column('revision', sa.BigInteger(), nullable=False),
sa.Column('fingerprint', sa.Text(), nullable=False),
sa.Column(
'received_at',
sa.DateTime(timezone=True),
server_default=sa.func.now(),
nullable=False,
),
sa.Column('applied_at', sa.DateTime(timezone=True), nullable=True),
sa.CheckConstraint(
'cursor > 0',
name='ck_directory_projection_inbox_cursor',
),
sa.CheckConstraint(
'revision > 0',
name='ck_directory_projection_inbox_revision',
),
sa.CheckConstraint(
'length(fingerprint) = 64',
name='ck_directory_projection_inbox_fingerprint',
),
sa.PrimaryKeyConstraint('instance_uuid', 'event_uuid'),
sa.UniqueConstraint(
'instance_uuid',
'cursor',
name='uq_directory_projection_inbox_cursor',
),
)
op.create_index(
'ix_directory_projection_inbox_pending',
_INBOX_TABLE,
['instance_uuid', 'applied_at', 'cursor'],
unique=False,
)
def _drop_policy(conn: sa.Connection, table_name: str, policy_name: str) -> None:
table = _quote(conn, table_name)
policy = _quote(conn, policy_name)
op.execute(sa.text(f'DROP POLICY IF EXISTS {policy} ON {table}'))
def _create_policy(
conn: sa.Connection,
table_name: str,
policy_name: str,
expression: str,
*,
command: str = 'ALL',
) -> None:
table = _quote(conn, table_name)
policy = _quote(conn, policy_name)
_drop_policy(conn, table_name, policy_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'))
if command == 'SELECT':
sql = f'CREATE POLICY {policy} ON {table} AS PERMISSIVE FOR SELECT TO PUBLIC USING ({expression})'
elif command == 'ALL':
sql = (
f'CREATE POLICY {policy} ON {table} AS PERMISSIVE FOR ALL TO PUBLIC '
f'USING ({expression}) WITH CHECK ({expression})'
)
else: # pragma: no cover - migration-local invariant.
raise AssertionError(f'Unsupported RLS policy command: {command}')
op.execute(sa.text(sql))
def _install_postgres_policies(conn: sa.Connection) -> None:
existing_tables = set(sa.inspect(conn).get_table_names())
required_tables = set(_PROJECTED_TENANT_TABLES) | {_STATE_TABLE, _INBOX_TABLE}
missing_tables = required_tables - existing_tables
if missing_tables:
raise RuntimeError(
f'Cannot enable Cloud directory projection RLS before all required tables exist: {sorted(missing_tables)!r}'
)
directory_setting = _setting(_DIRECTORY_SETTING)
tenant_setting = _setting(_TENANT_SETTING)
directory_expressions = {
'workspaces': (f"instance_uuid::text = {directory_setting} AND source = 'cloud_projection'"),
'workspace_memberships': (
'EXISTS ('
'SELECT 1 FROM workspaces AS directory_workspace '
'WHERE directory_workspace.uuid = workspace_memberships.workspace_uuid '
f'AND directory_workspace.instance_uuid::text = {directory_setting} '
"AND directory_workspace.source = 'cloud_projection'"
')'
),
'workspace_execution_states': (
f"instance_uuid::text = {directory_setting} AND source = 'cloud' AND EXISTS ("
'SELECT 1 FROM workspaces AS directory_workspace '
'WHERE directory_workspace.uuid = workspace_execution_states.workspace_uuid '
f'AND directory_workspace.instance_uuid::text = {directory_setting} '
"AND directory_workspace.source = 'cloud_projection'"
')'
),
_STATE_TABLE: f'instance_uuid::text = {directory_setting}',
_INBOX_TABLE: f'instance_uuid::text = {directory_setting}',
}
tenant_expressions = {
'workspaces': f'uuid::text = {tenant_setting}',
'workspace_memberships': f'workspace_uuid::text = {tenant_setting}',
'workspace_execution_states': f'workspace_uuid::text = {tenant_setting}',
}
local_write_expressions = {
'workspaces': f"uuid::text = {tenant_setting} AND source = 'local'",
'workspace_memberships': (
f'workspace_uuid::text = {tenant_setting} AND EXISTS ('
'SELECT 1 FROM workspaces AS local_workspace '
'WHERE local_workspace.uuid = workspace_memberships.workspace_uuid '
"AND local_workspace.source = 'local'"
')'
),
'workspace_execution_states': (
f'workspace_uuid::text = {tenant_setting} AND EXISTS ('
'SELECT 1 FROM workspaces AS local_workspace '
'WHERE local_workspace.uuid = workspace_execution_states.workspace_uuid '
"AND local_workspace.source = 'local'"
')'
),
}
for table_name in _PROJECTED_TENANT_TABLES:
_create_policy(
conn,
table_name,
_TENANT_POLICY_NAME,
tenant_expressions[table_name],
command='SELECT',
)
_create_policy(
conn,
table_name,
_LOCAL_WRITE_POLICY_NAME,
local_write_expressions[table_name],
)
_create_policy(
conn,
table_name,
_DIRECTORY_POLICY_NAME,
directory_expressions[table_name],
)
for table_name in (_STATE_TABLE, _INBOX_TABLE):
_create_policy(
conn,
table_name,
_DIRECTORY_POLICY_NAME,
directory_expressions[table_name],
)
def upgrade() -> None:
conn = op.get_bind()
_create_tables(conn)
if conn.dialect.name == 'postgresql':
_install_postgres_policies(conn)
def downgrade() -> None:
conn = op.get_bind()
existing_tables = set(sa.inspect(conn).get_table_names())
if conn.dialect.name == 'postgresql':
tenant_setting = _setting(_TENANT_SETTING)
tenant_columns = {
'workspaces': 'uuid',
'workspace_memberships': 'workspace_uuid',
'workspace_execution_states': 'workspace_uuid',
}
for table_name in _PROJECTED_TENANT_TABLES:
if table_name not in existing_tables:
continue
_drop_policy(conn, table_name, _DIRECTORY_POLICY_NAME)
_drop_policy(conn, table_name, _LOCAL_WRITE_POLICY_NAME)
_create_policy(
conn,
table_name,
_TENANT_POLICY_NAME,
f'{tenant_columns[table_name]}::text = {tenant_setting}',
)
for table_name in (_STATE_TABLE, _INBOX_TABLE):
if table_name not in existing_tables:
continue
_drop_policy(conn, table_name, _DIRECTORY_POLICY_NAME)
table = _quote(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'))
if _INBOX_TABLE in existing_tables:
op.drop_table(_INBOX_TABLE)
if _STATE_TABLE in existing_tables:
op.drop_table(_STATE_TABLE)
src/langbot/pkg/persistence/mgr.py
+90 -4
View File
@@ -27,10 +27,15 @@ from .tenant_uow import (
ACCOUNT_DISCOVERY_POLICY_NAME,
INSTANCE_DISCOVERY_POLICY_NAME,
INVITATION_DISCOVERY_POLICY_NAME,
LOCAL_DIRECTORY_WRITE_POLICY_NAME,
TENANT_POLICY_NAME,
TENANT_SETTING,
TENANT_TABLE_COLUMNS,
CrossScopeTransactionError,
DIRECTORY_INSTANCE_SETTING,
DIRECTORY_PROJECTED_TENANT_TABLES,
DIRECTORY_PROJECTION_POLICY_NAME,
DIRECTORY_PROJECTION_TABLE_COLUMNS,
PersistenceScope,
PersistenceScopeBoundary,
PersistenceScopeKind,
@@ -74,6 +79,8 @@ _ALEMBIC_TENANT_TABLES = {
'monitoring_embedding_calls',
'monitoring_feedback',
'langbot_vectors',
'directory_projection_states',
'directory_projection_inbox',
}
_PRE_WORKSPACE_ALEMBIC_REVISIONS = {
@@ -1345,6 +1352,7 @@ class PersistenceManager:
engine = self.get_db_engine()
if engine.dialect.name != 'postgresql':
raise RuntimeError('PostgreSQL tenant schema validation requires PostgreSQL')
rls_table_names = tuple(sorted(set(TENANT_TABLE_COLUMNS) | set(DIRECTORY_PROJECTION_TABLE_COLUMNS)))
table_query = sqlalchemy.text(
"""
@@ -1408,7 +1416,7 @@ class PersistenceManager:
await conn.execute(
table_query,
{
'table_names': tuple(TENANT_TABLE_COLUMNS),
'table_names': rls_table_names,
},
)
)
@@ -1419,7 +1427,7 @@ class PersistenceManager:
(
await conn.execute(
policy_query,
{'table_names': tuple(TENANT_TABLE_COLUMNS)},
{'table_names': rls_table_names},
)
)
.mappings()
@@ -1427,7 +1435,7 @@ class PersistenceManager:
)
by_table = {row['table_name']: row for row in rows}
missing_tables = set(TENANT_TABLE_COLUMNS) - set(by_table)
missing_tables = set(rls_table_names) - set(by_table)
if missing_tables:
raise RuntimeError(f'PostgreSQL tenant tables are missing: {sorted(missing_tables)!r}')
@@ -1472,7 +1480,7 @@ class PersistenceManager:
@staticmethod
def _expected_postgres_tenant_policies() -> dict[str, dict[str, dict[str, str | None]]]:
"""Return the exact PostgreSQL 16 policy expressions emitted by 0011."""
"""Return the exact PostgreSQL 16 policy expressions emitted by 0011/0014."""
def setting(name: str) -> str:
return f"NULLIF(current_setting('{name}'::text, true), ''::text)"
@@ -1488,6 +1496,39 @@ class PersistenceManager:
}
}
local_workspace_expression = (
f"(((uuid)::text = {setting(TENANT_SETTING)}) AND ((source)::text = 'local'::text))"
)
local_membership_expression = (
f'(((workspace_uuid)::text = {setting(TENANT_SETTING)}) AND (EXISTS ( SELECT 1\n'
' FROM workspaces local_workspace\n'
' WHERE (((local_workspace.uuid)::text = (workspace_memberships.workspace_uuid)::text) '
"AND ((local_workspace.source)::text = 'local'::text)))))"
)
local_execution_expression = (
f'(((workspace_uuid)::text = {setting(TENANT_SETTING)}) AND (EXISTS ( SELECT 1\n'
' FROM workspaces local_workspace\n'
' WHERE (((local_workspace.uuid)::text = (workspace_execution_states.workspace_uuid)::text) '
"AND ((local_workspace.source)::text = 'local'::text)))))"
)
for table_name, local_write_expression in {
'workspaces': local_workspace_expression,
'workspace_memberships': local_membership_expression,
'workspace_execution_states': local_execution_expression,
}.items():
tenant_column = TENANT_TABLE_COLUMNS[table_name]
tenant_expression = f'(({tenant_column})::text = {setting(TENANT_SETTING)})'
policies[table_name][TENANT_POLICY_NAME] = {
'command': 'r',
'using_expression': tenant_expression,
'check_expression': None,
}
policies[table_name][LOCAL_DIRECTORY_WRITE_POLICY_NAME] = {
'command': '*',
'using_expression': local_write_expression,
'check_expression': local_write_expression,
}
policies['workspace_memberships'][ACCOUNT_DISCOVERY_POLICY_NAME] = {
'command': 'r',
'using_expression': (
@@ -1517,6 +1558,48 @@ class PersistenceManager:
),
'check_expression': None,
}
directory_setting = setting(DIRECTORY_INSTANCE_SETTING)
workspace_expression = (
f"(((instance_uuid)::text = {directory_setting}) AND ((source)::text = 'cloud_projection'::text))"
)
membership_expression = (
'(EXISTS ( SELECT 1\n'
' FROM workspaces directory_workspace\n'
' WHERE (((directory_workspace.uuid)::text = (workspace_memberships.workspace_uuid)::text) '
f'AND ((directory_workspace.instance_uuid)::text = {directory_setting}) '
"AND ((directory_workspace.source)::text = 'cloud_projection'::text))))"
)
execution_expression = (
f'(((instance_uuid)::text = {directory_setting}) '
"AND ((source)::text = 'cloud'::text) "
'AND (EXISTS ( SELECT 1\n'
' FROM workspaces directory_workspace\n'
' WHERE (((directory_workspace.uuid)::text = (workspace_execution_states.workspace_uuid)::text) '
f'AND ((directory_workspace.instance_uuid)::text = {directory_setting}) '
"AND ((directory_workspace.source)::text = 'cloud_projection'::text)))))"
)
directory_tenant_expressions = {
'workspaces': workspace_expression,
'workspace_memberships': membership_expression,
'workspace_execution_states': execution_expression,
}
for table_name in DIRECTORY_PROJECTED_TENANT_TABLES:
expression = directory_tenant_expressions[table_name]
policies[table_name][DIRECTORY_PROJECTION_POLICY_NAME] = {
'command': '*',
'using_expression': expression,
'check_expression': expression,
}
for table_name, instance_column in DIRECTORY_PROJECTION_TABLE_COLUMNS.items():
expression = f'(({instance_column})::text = {directory_setting})'
policies[table_name] = {
DIRECTORY_PROJECTION_POLICY_NAME: {
'command': '*',
'using_expression': expression,
'check_expression': expression,
}
}
return policies
async def write_space_model_providers(self):
@@ -1729,6 +1812,9 @@ class PersistenceManager:
def instance_discovery_uow(self, instance_uuid: str) -> TenantUnitOfWork:
return self._scoped_uow(PersistenceScope.instance(instance_uuid))
def directory_projection_uow(self, instance_uuid: str) -> TenantUnitOfWork:
return self._scoped_uow(PersistenceScope.directory(instance_uuid))
def identity_discovery_uow(self, identity_digest: str) -> TenantUnitOfWork:
return self._scoped_uow(PersistenceScope.identity(identity_digest))
src/langbot/pkg/persistence/tenant_uow.py
+38 -1
View File
@@ -14,7 +14,9 @@ import sqlalchemy
import sqlalchemy.ext.asyncio as sqlalchemy_asyncio
import sqlalchemy.orm as sqlalchemy_orm
from pgvector.sqlalchemy import Vector
from sqlalchemy.dialects.postgresql.dml import OnConflictDoNothing as PostgreSQLOnConflictDoNothing
from sqlalchemy.dialects.postgresql.dml import OnConflictDoUpdate as PostgreSQLOnConflictDoUpdate
from sqlalchemy.dialects.sqlite.dml import OnConflictDoNothing as SQLiteOnConflictDoNothing
from sqlalchemy.dialects.sqlite.dml import OnConflictDoUpdate as SQLiteOnConflictDoUpdate
@@ -24,12 +26,15 @@ 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'
DIRECTORY_INSTANCE_SETTING = 'langbot.directory_instance_uuid'
TENANT_POLICY_NAME = 'langbot_workspace_isolation'
LOCAL_DIRECTORY_WRITE_POLICY_NAME = 'langbot_workspace_local_directory_write'
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'
DIRECTORY_PROJECTION_POLICY_NAME = 'langbot_directory_projection'
# Keep this contract explicit. A new tenant-owned table must be added to both
# this runtime list and the corresponding Alembic migration before release.
@@ -67,6 +72,19 @@ TENANT_TABLE_COLUMNS: dict[str, str] = {
'langbot_vectors': 'workspace_uuid',
}
DIRECTORY_PROJECTION_TABLE_COLUMNS: dict[str, str] = {
'directory_projection_states': 'instance_uuid',
'directory_projection_inbox': 'instance_uuid',
}
DIRECTORY_PROJECTED_TENANT_TABLES = frozenset(
{
'workspaces',
'workspace_memberships',
'workspace_execution_states',
}
)
class PersistenceScopeKind(enum.StrEnum):
WORKSPACE = 'workspace'
@@ -75,6 +93,7 @@ class PersistenceScopeKind(enum.StrEnum):
INVITATION_DISCOVERY = 'invitation_discovery'
INSTANCE_DISCOVERY = 'instance_discovery'
IDENTITY_DISCOVERY = 'identity_discovery'
DIRECTORY_PROJECTION = 'directory_projection'
@dataclasses.dataclass(frozen=True, slots=True)
@@ -108,6 +127,14 @@ class PersistenceScope:
def identity(cls, identity_digest: str) -> PersistenceScope:
return cls._one(PersistenceScopeKind.IDENTITY_DISCOVERY, IDENTITY_DIGEST_SETTING, identity_digest)
@classmethod
def directory(cls, instance_uuid: str) -> PersistenceScope:
return cls._one(
PersistenceScopeKind.DIRECTORY_PROJECTION,
DIRECTORY_INSTANCE_SETTING,
instance_uuid,
)
@classmethod
def _one(cls, kind: PersistenceScopeKind, setting: str, value: str) -> PersistenceScope:
return cls(kind, ((setting, cls._normalize(value, kind.value)),))
@@ -188,7 +215,9 @@ _ALLOWED_SCOPED_STATEMENT_TYPES = (
sqlalchemy.sql.selectable.SelectBase,
)
_ALLOWED_SCOPED_POST_VALUES_TYPES = (
PostgreSQLOnConflictDoNothing,
PostgreSQLOnConflictDoUpdate,
SQLiteOnConflictDoNothing,
SQLiteOnConflictDoUpdate,
)
@@ -853,6 +882,7 @@ class TenantScopedAsyncSession(sqlalchemy_asyncio.AsyncSession):
INVITATION_HASH_SETTING,
INSTANCE_SETTING,
IDENTITY_DIGEST_SETTING,
DIRECTORY_INSTANCE_SETTING,
}:
self._reject_transaction_escape('unknown tenant scope configuration')
self._require_owner_task()
@@ -1230,7 +1260,14 @@ class TenantUnitOfWork:
# It must still never switch directly to another Workspace.
workspace_to_discovery = (
active_scope.scope.kind == PersistenceScopeKind.WORKSPACE
and self.scope.kind != PersistenceScopeKind.WORKSPACE
and self.scope.kind
in {
PersistenceScopeKind.ACCOUNT_DISCOVERY,
PersistenceScopeKind.API_KEY_DISCOVERY,
PersistenceScopeKind.INVITATION_DISCOVERY,
PersistenceScopeKind.INSTANCE_DISCOVERY,
PersistenceScopeKind.IDENTITY_DISCOVERY,
}
)
if not workspace_to_discovery:
raise CrossScopeTransactionError(
src/langbot/pkg/plugin/connector.py
+287 -177
View File
@@ -2,6 +2,7 @@
from __future__ import annotations
import asyncio
import contextlib
import contextvars
import hashlib
import io
@@ -73,6 +74,10 @@ _GITHUB_ASSET_HOSTS = frozenset(
}
)
_HTTP_REDIRECT_STATUSES = frozenset({301, 302, 303, 307, 308})
_CONNECT_TIMEOUT_SEC = 30.0
_HEARTBEAT_INTERVAL_SEC = 20.0
_HEARTBEAT_FAILURE_THRESHOLD = 3
_RECONNECT_MAX_DELAY_SEC = 60.0
class PluginRuntimeNotConnectedError(RuntimeError):
@@ -139,6 +144,24 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
self._installation_failures: dict[str, dict[str, str]] = {}
self._state_lock = asyncio.Lock()
self._control_token = str(os.environ.get(PLUGIN_RUNTIME_CONTROL_TOKEN_ENV) or '').strip()
self._transport_task: asyncio.Task | None = None
self._reconnect_task: asyncio.Task | None = None
self._generation = 0
self._connected = asyncio.Event()
def _runtime_handler(self) -> handler.RuntimeConnectionHandler:
runtime_handler = getattr(self, 'handler', None)
if runtime_handler is None:
raise PluginRuntimeNotConnectedError('Plugin runtime is not connected')
return runtime_handler
def _runtime_available(self) -> bool:
runtime_handler = getattr(self, 'handler', None)
if runtime_handler is None:
return False
# Unit-level and explicitly injected handlers do not own a transport.
# A managed transport must also have completed its handshake.
return self._transport_task is None or self._connected.is_set()
def _load_worker_policy(self) -> PluginWorkerPolicy:
"""Validate the instance policy without consulting plugin manifests."""
@@ -351,8 +374,9 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
desired_states: list[PluginInstallationDesiredState] = []
for setting in settings:
binding = self._binding_from_setting(execution_context, setting)
if hasattr(self, 'handler'):
self.handler.register_installation_binding(
runtime_handler = getattr(self, 'handler', None)
if runtime_handler is not None:
runtime_handler.register_installation_binding(
binding,
plugin_author=setting.plugin_author,
plugin_name=setting.plugin_name,
@@ -371,7 +395,8 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
*,
artifact_package: bytes | None = None,
) -> dict[str, Any]:
result = await self.handler.apply_plugin_installation(
runtime_handler = self._runtime_handler()
result = await runtime_handler.apply_plugin_installation(
desired.binding,
artifact_package=artifact_package,
enabled=desired.enabled,
@@ -396,7 +421,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
f'Durable plugin artifact {desired.binding.artifact_digest} is missing for '
f'installation {desired.binding.installation_uuid}'
)
repaired = await self.handler.apply_plugin_installation(
repaired = await runtime_handler.apply_plugin_installation(
desired.binding,
artifact_package=persisted_package,
enabled=desired.enabled,
@@ -559,6 +584,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
async def _prepare_connected_runtime(self) -> None:
"""Handshake follow-up: pin OSS compatibility, then replay authority."""
runtime_handler = self._runtime_handler()
workspace_service = getattr(self.ap, 'workspace_service', None)
if workspace_service is None:
raise RuntimeError('Plugin Runtime requires the Workspace projection service')
@@ -580,15 +606,15 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
# One fully-bound action releases the SDK's deliberately retained
# pre-v4 data/plugins/debug compatibility path. Shared mode never
# creates this bridge.
with self.handler.installation_scope(bridge):
await self.handler.list_plugins()
with runtime_handler.installation_scope(bridge):
await runtime_handler.list_plugins()
desired_states = await self._load_workspace_desired_states(execution_context)
self._workspace_installations[execution_context.workspace_uuid] = {
state.binding.installation_uuid for state in desired_states
}
self._known_desired_states.update({state.binding.installation_uuid: state for state in desired_states})
result = await self.handler.reconcile_plugin_installations(tuple(self._known_desired_states.values()))
result = await runtime_handler.reconcile_plugin_installations(tuple(self._known_desired_states.values()))
await self._repair_reconcile_missing_artifacts(self._known_desired_states, result)
self._record_reconcile_failures(self._known_desired_states, result)
@@ -604,8 +630,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
Runtime reconcile.
"""
if not hasattr(self, 'handler'):
raise PluginRuntimeNotConnectedError('Plugin runtime is not connected')
runtime_handler = self._runtime_handler()
async with self._state_lock:
all_states: dict[str, PluginInstallationDesiredState] = {}
workspace_installations: dict[str, set[str]] = {}
@@ -619,12 +644,12 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
if state.binding.installation_uuid in all_states:
raise ValueError('Duplicate plugin installation UUID across projected Workspaces')
all_states[state.binding.installation_uuid] = state
result = await self.handler.reconcile_plugin_installations(tuple(all_states.values()))
result = await runtime_handler.reconcile_plugin_installations(tuple(all_states.values()))
await self._repair_reconcile_missing_artifacts(all_states, result)
self._record_reconcile_failures(all_states, result)
for installation_uuid, previous in tuple(self._known_desired_states.items()):
if installation_uuid not in all_states:
self.handler.unregister_installation_binding(previous.binding)
runtime_handler.unregister_installation_binding(previous.binding)
self._known_desired_states = all_states
self._workspace_installations = workspace_installations
return result
@@ -652,6 +677,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
async def _synchronize_workspace(self, execution_context: ExecutionContext) -> None:
if not self.is_enable_plugin or not hasattr(self, 'handler'):
return
runtime_handler = self._runtime_handler()
desired_states = await self._load_workspace_desired_states(execution_context)
desired_by_uuid = {state.binding.installation_uuid: state for state in desired_states}
async with self._state_lock:
@@ -659,8 +685,8 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
for installation_uuid in previous_ids - set(desired_by_uuid):
previous = self._known_desired_states.get(installation_uuid)
if previous is not None:
await self.handler.remove_plugin_installation(previous.binding)
self.handler.unregister_installation_binding(previous.binding)
await runtime_handler.remove_plugin_installation(previous.binding)
runtime_handler.unregister_installation_binding(previous.binding)
self._known_desired_states.pop(installation_uuid, None)
self._installation_failures.pop(installation_uuid, None)
@@ -755,48 +781,152 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
await notify_disconnect()
return False
self.handler = handler.RuntimeConnectionHandler(
connection,
disconnect_callback,
self.ap,
)
self.handler_task = asyncio.create_task(self.handler.run())
_ = await self.handler.ping()
# Push the configured marketplace (Space) URL to the runtime so it
# downloads plugins from the same Space LangBot is bound to, rather
# than relying on the runtime's own env/default.
space_url = self.ap.instance_config.data.get('space', {}).get('url', '').rstrip('/')
try:
if self.runtime_identity is None or self.worker_policy is None: # pragma: no cover
raise RuntimeError('Plugin Runtime identity or worker policy was not loaded')
await self.handler.set_runtime_config(
runtime_identity=self.runtime_identity,
worker_policy=self.worker_policy,
runtime_profile=self.runtime_profile,
cloud_service_url=space_url or None,
runtime_handler = handler.RuntimeConnectionHandler(
connection,
disconnect_callback,
self.ap,
)
if space_url:
self.ap.logger.info(f'Pushed marketplace URL to plugin runtime: {space_url}')
except Exception as e:
self.ap.logger.warning(f'Failed to bind plugin runtime config: {e}')
raise
await self._prepare_connected_runtime()
self.ap.logger.info('Connected to instance-scoped plugin runtime.')
await self.handler_task
self.handler = runtime_handler
self.handler_task = asyncio.create_task(runtime_handler.run())
try:
await runtime_handler.ping()
if self.runtime_identity is None or self.worker_policy is None: # pragma: no cover
raise RuntimeError('Plugin Runtime identity or worker policy was not loaded')
space_url = self.ap.instance_config.data.get('space', {}).get('url', '').rstrip('/')
await runtime_handler.set_runtime_config(
runtime_identity=self.runtime_identity,
worker_policy=self.worker_policy,
runtime_profile=self.runtime_profile,
cloud_service_url=space_url or None,
)
if space_url:
self.ap.logger.info(f'Pushed marketplace URL to plugin runtime: {space_url}')
await self._prepare_connected_runtime()
if generation == self._generation and not self._closing:
connection_ready = True
self._connected.set()
self.ap.logger.info('Connected to instance-scoped plugin runtime.')
await self.handler_task
except asyncio.CancelledError:
raise
except Exception as exc:
if not self._connected.is_set():
connect_errors.append(exc)
self._connected.set()
finally:
if generation == self._generation and not self._closing:
self._connected.clear()
if getattr(self, 'handler', None) is runtime_handler:
del self.handler
await notify_disconnect()
task_coro: typing.Coroutine
task_coro: typing.Coroutine[Any, Any, Any]
if platform.get_platform() == 'docker' or platform.use_websocket_to_connect_plugin_runtime():
self.ap.logger.info('use websocket to connect to plugin runtime')
control_headers = self._control_headers(allow_generate=False)
ws_url = self.ap.instance_config.data.get('plugin', {}).get(
'runtime_ws_url',
'ws://langbot_plugin_runtime:5400/control/ws',
)
if platform.get_platform() == 'docker' or platform.use_websocket_to_connect_plugin_runtime(): # use websocket
self.ap.logger.info('use websocket to connect to plugin runtime')
control_headers = self._control_headers(allow_generate=False)
ws_url = self.ap.instance_config.data.get('plugin', {}).get(
'runtime_ws_url', 'ws://langbot_plugin_runtime:5400/control/ws'
async def connection_failed(
ctrl: ws_client_controller.WebSocketClientController,
exc: Exception | None = None,
) -> None:
del ctrl
connect_errors.append(exc or RuntimeError('WebSocket connection failed'))
self._connected.set()
self.ctrl = ws_client_controller.WebSocketClientController(
ws_url=ws_url,
make_connection_failed_callback=connection_failed,
additional_headers=control_headers,
)
task_coro = self.ctrl.run(new_connection_callback)
elif platform.get_platform() == 'win32':
# Windows cannot use the stdio subprocess transport, so launch
# a managed runtime and authenticate its WebSocket controller.
self.ap.logger.info('(windows) use cmd to launch plugin runtime and communicate via ws')
control_headers = self._control_headers(allow_generate=True)
await self._start_runtime_subprocess(
'-m',
'langbot_plugin.cli.__init__',
'rt',
env_overrides={PLUGIN_RUNTIME_CONTROL_TOKEN_ENV: self._control_token},
)
ws_url = 'ws://localhost:5400/control/ws'
async def connection_failed(
ctrl: ws_client_controller.WebSocketClientController,
exc: Exception | None = None,
) -> None:
del ctrl
connect_errors.append(exc or RuntimeError('WebSocket connection failed'))
self._connected.set()
self.ctrl = ws_client_controller.WebSocketClientController(
ws_url=ws_url,
make_connection_failed_callback=connection_failed,
additional_headers=control_headers,
)
task_coro = self.ctrl.run(new_connection_callback)
else:
self.ap.logger.info('use stdio to connect to plugin runtime')
self.ctrl = stdio_client_controller.StdioClientController(
command=sys.executable,
args=['-m', 'langbot_plugin.cli.__init__', 'rt', '-s'],
env=os.environ.copy(),
capture_stderr=False,
)
task_coro = self.ctrl.run(new_connection_callback)
self._transport_task = asyncio.create_task(task_coro)
try:
await asyncio.wait_for(self._connected.wait(), timeout=_CONNECT_TIMEOUT_SEC)
except asyncio.TimeoutError as exc:
await self._stop_transport()
raise PluginRuntimeNotConnectedError('Plugin runtime did not become ready within 30 seconds') from exc
if connect_errors:
await self._stop_transport()
raise PluginRuntimeNotConnectedError(f'Plugin runtime connection failed: {connect_errors[-1]}')
if self.heartbeat_task is None or self.heartbeat_task.done():
self.heartbeat_task = asyncio.create_task(self.heartbeat_loop())
def schedule_reconnect(self) -> None:
if self._closing or not self.is_enable_plugin:
return
if self._reconnect_task is not None and not self._reconnect_task.done():
return
self._reconnect_task = asyncio.create_task(self._reconnect_loop())
async def _reconnect_loop(self) -> None:
delay = 1.0
try:
while not self._closing:
try:
await self.initialize()
return
except Exception as exc:
self.ap.logger.warning(f'Plugin runtime reconnection failed: {exc}; retrying in {delay:.0f}s')
await asyncio.sleep(delay)
delay = min(delay * 2, _RECONNECT_MAX_DELAY_SEC)
finally:
self._reconnect_task = None
async def _stop_transport(self) -> None:
self._connected.clear()
runtime_handler = getattr(self, 'handler', None)
if runtime_handler is not None:
with contextlib.suppress(Exception):
await runtime_handler.close()
if getattr(self, 'handler', None) is runtime_handler:
del self.handler
tasks = [
task
for task in (
getattr(self, 'handler_task', None),
self._transport_task,
)
if task is not None and task is not asyncio.current_task()
]
@@ -812,73 +942,20 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
with contextlib.suppress(Exception):
await close_ctrl()
async def make_connection_failed_callback(
ctrl: ws_client_controller.WebSocketClientController,
exc: Exception = None,
) -> None:
if exc is not None:
self.ap.logger.error(f'Failed to connect to plugin runtime({ws_url}): {exc}')
else:
self.ap.logger.error(f'Failed to connect to plugin runtime({ws_url}), trying to reconnect...')
await self.runtime_disconnect_callback(self)
self.ctrl = ws_client_controller.WebSocketClientController(
ws_url=ws_url,
make_connection_failed_callback=make_connection_failed_callback,
additional_headers=control_headers,
)
task = self.ctrl.run(new_connection_callback)
elif platform.get_platform() == 'win32':
# Due to Windows's lack of supports for both stdio and subprocess:
# See also: https://docs.python.org/zh-cn/3.13/library/asyncio-platforms.html
# We have to launch runtime via cmd but communicate via ws.
self.ap.logger.info('(windows) use cmd to launch plugin runtime and communicate via ws')
control_headers = self._control_headers(allow_generate=True)
await self._start_runtime_subprocess(
'-m',
'langbot_plugin.cli.__init__',
'rt',
env_overrides={PLUGIN_RUNTIME_CONTROL_TOKEN_ENV: self._control_token},
)
ws_url = 'ws://localhost:5400/control/ws'
async def make_connection_failed_callback(
ctrl: ws_client_controller.WebSocketClientController,
exc: Exception = None,
) -> None:
if exc is not None:
self.ap.logger.error(f'(windows) Failed to connect to plugin runtime({ws_url}): {exc}')
else:
self.ap.logger.error(
f'(windows) Failed to connect to plugin runtime({ws_url}), trying to reconnect...'
)
await self.runtime_disconnect_callback(self)
self.ctrl = ws_client_controller.WebSocketClientController(
ws_url=ws_url,
make_connection_failed_callback=make_connection_failed_callback,
additional_headers=control_headers,
)
task = self.ctrl.run(new_connection_callback)
else: # stdio
self.ap.logger.info('use stdio to connect to plugin runtime')
# cmd: lbp rt -s
python_path = sys.executable
env = os.environ.copy()
self.ctrl = stdio_client_controller.StdioClientController(
command=python_path,
args=['-m', 'langbot_plugin.cli.__init__', 'rt', '-s'],
env=env,
)
task = self.ctrl.run(new_connection_callback)
if self.heartbeat_task is None:
self.heartbeat_task = asyncio.create_task(self.heartbeat_loop())
asyncio.create_task(task)
async def aclose(self) -> None:
self._closing = True
self._generation += 1
reconnect_task = self._reconnect_task
self._reconnect_task = None
if reconnect_task is not None and reconnect_task is not asyncio.current_task():
reconnect_task.cancel()
await asyncio.gather(reconnect_task, return_exceptions=True)
if self.heartbeat_task is not None:
self.heartbeat_task.cancel()
await asyncio.gather(self.heartbeat_task, return_exceptions=True)
self.heartbeat_task = None
await self._stop_transport()
await self._close_managed_subprocess()
async def initialize_plugins(self):
pass
@@ -1001,12 +1078,13 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
component_kind: typing.Literal['tool', 'command'],
include_plugins: list[str] | None,
) -> InstallationBinding:
runtime_handler = self._runtime_handler()
for binding in await self._operation_bindings(include_plugins=include_plugins):
with self.handler.installation_scope(binding):
with runtime_handler.installation_scope(binding):
components = (
await self.handler.list_tools(include_plugins=include_plugins)
await runtime_handler.list_tools(include_plugins=include_plugins)
if component_kind == 'tool'
else await self.handler.list_commands(include_plugins=include_plugins)
else await runtime_handler.list_commands(include_plugins=include_plugins)
)
for component in components:
manifest = ComponentManifest.model_validate(component)
@@ -1467,6 +1545,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
install_info: dict[str, Any],
task_context: taskmgr.TaskContext | None = None,
) -> None:
runtime_handler = self._runtime_handler()
execution_context = await self._current_execution_context()
plugin_author = str(install_info.get('plugin_author') or '')
plugin_name = str(install_info.get('plugin_name') or '')
@@ -1521,7 +1600,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
except Exception:
await self._delete_artifact_if_unreferenced(execution_context, artifact_digest)
raise
self.handler.register_installation_binding(
runtime_handler.register_installation_binding(
binding,
plugin_author=plugin_author,
plugin_name=plugin_name,
@@ -1538,8 +1617,8 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
if previous_digest is not None and not previous_was_durable and self.runtime_profile == 'oss_dev':
bridge = self._legacy_oss_bridge_binding(execution_context)
try:
with self.handler.installation_scope(bridge):
async for _ in self.handler.delete_plugin(plugin_author, plugin_name):
with runtime_handler.installation_scope(bridge):
async for _ in runtime_handler.delete_plugin(plugin_author, plugin_name):
pass
except Exception as exc:
self.ap.logger.debug(f'Legacy OSS plugin cleanup skipped: {exc}')
@@ -1570,6 +1649,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
delete_data: bool = False,
task_context: taskmgr.TaskContext | None = None,
) -> dict[str, Any]:
runtime_handler = self._runtime_handler()
execution_context, setting = await self._setting_for_plugin(plugin_author, plugin_name)
binding = self._binding_from_setting(execution_context, setting)
is_legacy_oss = self.runtime_profile == 'oss_dev' and (
@@ -1578,11 +1658,11 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
)
if is_legacy_oss:
bridge = self._legacy_oss_bridge_binding(execution_context)
with self.handler.installation_scope(bridge):
async for _ in self.handler.delete_plugin(plugin_author, plugin_name):
with runtime_handler.installation_scope(bridge):
async for _ in runtime_handler.delete_plugin(plugin_author, plugin_name):
pass
await self.handler.remove_plugin_installation(binding)
self.handler.unregister_installation_binding(binding)
await runtime_handler.remove_plugin_installation(binding)
runtime_handler.unregister_installation_binding(binding)
async def delete(execute):
await execute(
@@ -1629,11 +1709,12 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
if not self.is_enable_plugin or not self._runtime_available():
return []
runtime_handler = self._runtime_handler()
plugins: list[dict[str, Any]] = []
seen_plugin_ids: set[str] = set()
for binding in await self._operation_bindings():
with self.handler.installation_scope(binding):
scoped_plugins = await self.handler.list_plugins()
with runtime_handler.installation_scope(binding):
scoped_plugins = await runtime_handler.list_plugins()
for plugin in scoped_plugins:
metadata = plugin.get('manifest', {}).get('manifest', {}).get('metadata', {})
plugin_id = f'{metadata.get("author", "")}/{metadata.get("name", "")}'
@@ -1685,11 +1766,13 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
return plugins
async def get_plugin_info(self, author: str, plugin_name: str) -> dict[str, Any]:
runtime_handler = self._runtime_handler()
binding = await self._target_binding(author, plugin_name)
with self.handler.installation_scope(binding):
return await self.handler.get_plugin_info(author, plugin_name)
with runtime_handler.installation_scope(binding):
return await runtime_handler.get_plugin_info(author, plugin_name)
async def set_plugin_config(self, plugin_author: str, plugin_name: str, config: dict[str, Any]) -> dict[str, Any]:
runtime_handler = self._runtime_handler()
execution_context, setting = await self._setting_for_plugin(plugin_author, plugin_name)
next_revision = setting.runtime_revision + 1
statement = (
@@ -1716,7 +1799,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
runtime_revision=next_revision,
artifact_digest=setting.artifact_digest,
)
self.handler.register_installation_binding(
runtime_handler.register_installation_binding(
binding,
plugin_author=plugin_author,
plugin_name=plugin_name,
@@ -1731,22 +1814,24 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
)
if is_legacy_oss:
bridge = self._legacy_oss_bridge_binding(execution_context)
with self.handler.installation_scope(bridge):
await self.handler.set_plugin_config(plugin_author, plugin_name, config)
with runtime_handler.installation_scope(bridge):
await runtime_handler.set_plugin_config(plugin_author, plugin_name, config)
else:
await self._apply_desired_state(desired)
self._known_desired_states[binding.installation_uuid] = desired
return {}
async def get_plugin_icon(self, plugin_author: str, plugin_name: str) -> dict[str, Any]:
runtime_handler = self._runtime_handler()
binding = await self._target_binding(plugin_author, plugin_name)
with self.handler.installation_scope(binding):
return await self.handler.get_plugin_icon(plugin_author, plugin_name)
with runtime_handler.installation_scope(binding):
return await runtime_handler.get_plugin_icon(plugin_author, plugin_name)
async def get_plugin_readme(self, plugin_author: str, plugin_name: str, language: str = 'en') -> str:
runtime_handler = self._runtime_handler()
binding = await self._target_binding(plugin_author, plugin_name)
with self.handler.installation_scope(binding):
return await self.handler.get_plugin_readme(plugin_author, plugin_name, language)
with runtime_handler.installation_scope(binding):
return await runtime_handler.get_plugin_readme(plugin_author, plugin_name, language)
async def get_plugin_logs(
self,
@@ -1755,14 +1840,16 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
limit: int = 200,
level: str | None = None,
) -> list[dict[str, Any]]:
runtime_handler = self._runtime_handler()
binding = await self._target_binding(plugin_author, plugin_name)
with self.handler.installation_scope(binding):
return await self.handler.get_plugin_logs(plugin_author, plugin_name, limit, level)
with runtime_handler.installation_scope(binding):
return await runtime_handler.get_plugin_logs(plugin_author, plugin_name, limit, level)
async def get_plugin_assets(self, plugin_author: str, plugin_name: str, filepath: str) -> dict[str, Any]:
runtime_handler = self._runtime_handler()
binding = await self._target_binding(plugin_author, plugin_name)
with self.handler.installation_scope(binding):
return await self.handler.get_plugin_assets(plugin_author, plugin_name, filepath)
with runtime_handler.installation_scope(binding):
return await runtime_handler.get_plugin_assets(plugin_author, plugin_name, filepath)
async def handle_page_api(
self,
@@ -1773,9 +1860,10 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
method: str,
body: Any = None,
) -> dict[str, Any]:
runtime_handler = self._runtime_handler()
binding = await self._target_binding(plugin_author, plugin_name)
with self.handler.installation_scope(binding):
return await self.handler.handle_page_api(plugin_author, plugin_name, page_id, endpoint, method, body)
with runtime_handler.installation_scope(binding):
return await runtime_handler.handle_page_api(plugin_author, plugin_name, page_id, endpoint, method, body)
async def get_debug_info(self) -> dict[str, Any]:
"""Get debug information including debug key and WS URL"""
@@ -1800,11 +1888,12 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
event_ctx._response_sources = []
return event_ctx
runtime_handler = self._runtime_handler()
emitted_plugins: list[Any] = []
response_sources: list[dict[str, Any]] = []
for binding in await self._operation_bindings(include_plugins=bound_plugins):
with self.handler.installation_scope(binding):
result = await self.handler.emit_event(
with runtime_handler.installation_scope(binding):
result = await runtime_handler.emit_event(
event_ctx.model_dump(serialize_as_any=False),
include_plugins=bound_plugins,
)
@@ -1821,18 +1910,19 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
if not self.is_enable_plugin or not self._runtime_available():
return
try:
runtime_handler = self._runtime_handler()
plugin_ref = diagnostic.get('plugin') if isinstance(diagnostic, dict) else None
if isinstance(plugin_ref, dict):
author = plugin_ref.get('author') or plugin_ref.get('plugin_author')
name = plugin_ref.get('name') or plugin_ref.get('plugin_name')
if author and name:
binding = await self._target_binding(str(author), str(name))
with self.handler.installation_scope(binding):
await self.handler.notify_plugin_diagnostic(diagnostic)
with runtime_handler.installation_scope(binding):
await runtime_handler.notify_plugin_diagnostic(diagnostic)
return
for binding in await self._operation_bindings():
with self.handler.installation_scope(binding):
await self.handler.notify_plugin_diagnostic(diagnostic)
with runtime_handler.installation_scope(binding):
await runtime_handler.notify_plugin_diagnostic(diagnostic)
except Exception as e:
self.ap.logger.debug(f'Plugin diagnostic forwarding skipped: {e}')
@@ -1840,11 +1930,12 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
if not self.is_enable_plugin or not self._runtime_available():
return []
runtime_handler = self._runtime_handler()
tools: list[ComponentManifest] = []
seen: set[tuple[str, str]] = set()
for binding in await self._operation_bindings(include_plugins=bound_plugins):
with self.handler.installation_scope(binding):
scoped = await self.handler.list_tools(include_plugins=bound_plugins)
with runtime_handler.installation_scope(binding):
scoped = await runtime_handler.list_tools(include_plugins=bound_plugins)
for raw_tool in scoped:
tool = ComponentManifest.model_validate(raw_tool)
key = (str(tool.owner), tool.metadata.name)
@@ -1878,8 +1969,9 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
component_kind='tool',
include_plugins=bound_plugins,
)
with self.handler.installation_scope(binding):
return await self.handler.call_tool(
runtime_handler = self._runtime_handler()
with runtime_handler.installation_scope(binding):
return await runtime_handler.call_tool(
tool_name,
parameters,
session.model_dump(serialize_as_any=True),
@@ -1892,11 +1984,12 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
if not self.is_enable_plugin or not self._runtime_available():
return []
runtime_handler = self._runtime_handler()
commands: list[ComponentManifest] = []
seen: set[tuple[str, str]] = set()
for binding in await self._operation_bindings(include_plugins=bound_plugins):
with self.handler.installation_scope(binding):
scoped = await self.handler.list_commands(include_plugins=bound_plugins)
with runtime_handler.installation_scope(binding):
scoped = await runtime_handler.list_commands(include_plugins=bound_plugins)
for raw_command in scoped:
command = ComponentManifest.model_validate(raw_command)
key = (str(command.owner), command.metadata.name)
@@ -1925,8 +2018,9 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
component_kind='command',
include_plugins=bound_plugins,
)
with self.handler.installation_scope(binding):
gen = self.handler.execute_command(
runtime_handler = self._runtime_handler()
with runtime_handler.installation_scope(binding):
gen = runtime_handler.execute_command(
command_ctx.model_dump(serialize_as_any=True),
include_plugins=bound_plugins,
)
@@ -1944,9 +2038,15 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
if not self.is_enable_plugin or not self._runtime_available():
return {'results': []}
runtime_handler = self._runtime_handler()
binding = await self._target_binding(plugin_author, plugin_name)
with self.handler.installation_scope(binding):
return await self.handler.retrieve_knowledge(plugin_author, plugin_name, retriever_name, retrieval_context)
with runtime_handler.installation_scope(binding):
return await runtime_handler.retrieve_knowledge(
plugin_author,
plugin_name,
retriever_name,
retrieval_context,
)
def dispose(self):
"""Best-effort synchronous compatibility wrapper; prefer ``aclose``."""
@@ -1997,41 +2097,47 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
context_data: IngestionContext data.
"""
plugin_author, plugin_name = self._parse_plugin_id(plugin_id)
runtime_handler = self._runtime_handler()
binding = await self._target_binding(plugin_author, plugin_name)
with self.handler.installation_scope(binding):
return await self.handler.rag_ingest_document(plugin_author, plugin_name, context_data)
with runtime_handler.installation_scope(binding):
return await runtime_handler.rag_ingest_document(plugin_author, plugin_name, context_data)
async def call_rag_delete_document(self, plugin_id: str, document_id: str, kb_id: str) -> bool:
plugin_author, plugin_name = self._parse_plugin_id(plugin_id)
runtime_handler = self._runtime_handler()
binding = await self._target_binding(plugin_author, plugin_name)
with self.handler.installation_scope(binding):
return await self.handler.rag_delete_document(plugin_author, plugin_name, document_id, kb_id)
with runtime_handler.installation_scope(binding):
return await runtime_handler.rag_delete_document(plugin_author, plugin_name, document_id, kb_id)
async def get_rag_creation_schema(self, plugin_id: str) -> dict[str, Any]:
plugin_author, plugin_name = self._parse_plugin_id(plugin_id)
runtime_handler = self._runtime_handler()
binding = await self._target_binding(plugin_author, plugin_name)
with self.handler.installation_scope(binding):
return await self.handler.get_rag_creation_schema(plugin_author, plugin_name)
with runtime_handler.installation_scope(binding):
return await runtime_handler.get_rag_creation_schema(plugin_author, plugin_name)
async def get_rag_retrieval_schema(self, plugin_id: str) -> dict[str, Any]:
plugin_author, plugin_name = self._parse_plugin_id(plugin_id)
runtime_handler = self._runtime_handler()
binding = await self._target_binding(plugin_author, plugin_name)
with self.handler.installation_scope(binding):
return await self.handler.get_rag_retrieval_schema(plugin_author, plugin_name)
with runtime_handler.installation_scope(binding):
return await runtime_handler.get_rag_retrieval_schema(plugin_author, plugin_name)
async def rag_on_kb_create(self, plugin_id: str, kb_id: str, config: dict[str, Any]) -> dict[str, Any]:
"""Notify plugin about KB creation."""
plugin_author, plugin_name = self._parse_plugin_id(plugin_id)
runtime_handler = self._runtime_handler()
binding = await self._target_binding(plugin_author, plugin_name)
with self.handler.installation_scope(binding):
return await self.handler.rag_on_kb_create(plugin_author, plugin_name, kb_id, config)
with runtime_handler.installation_scope(binding):
return await runtime_handler.rag_on_kb_create(plugin_author, plugin_name, kb_id, config)
async def rag_on_kb_delete(self, plugin_id: str, kb_id: str) -> dict[str, Any]:
"""Notify plugin about KB deletion."""
plugin_author, plugin_name = self._parse_plugin_id(plugin_id)
runtime_handler = self._runtime_handler()
binding = await self._target_binding(plugin_author, plugin_name)
with self.handler.installation_scope(binding):
return await self.handler.rag_on_kb_delete(plugin_author, plugin_name, kb_id)
with runtime_handler.installation_scope(binding):
return await runtime_handler.rag_on_kb_delete(plugin_author, plugin_name, kb_id)
async def call_rag_retrieve(self, plugin_id: str, retrieval_context: dict[str, Any]) -> dict[str, Any]:
"""Call plugin to retrieve knowledge.
@@ -2041,9 +2147,10 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
retrieval_context: RetrievalContext data.
"""
plugin_author, plugin_name = self._parse_plugin_id(plugin_id)
runtime_handler = self._runtime_handler()
binding = await self._target_binding(plugin_author, plugin_name)
with self.handler.installation_scope(binding):
return await self.handler.retrieve_knowledge(plugin_author, plugin_name, '', retrieval_context)
with runtime_handler.installation_scope(binding):
return await runtime_handler.retrieve_knowledge(plugin_author, plugin_name, '', retrieval_context)
async def list_knowledge_engines(self) -> list[dict[str, Any]]:
"""List all available Knowledge Engines from plugins.
@@ -2053,11 +2160,12 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
if not self.is_enable_plugin or not self._runtime_available():
return []
runtime_handler = self._runtime_handler()
engines: list[dict[str, Any]] = []
seen: set[tuple[str, str]] = set()
for binding in await self._operation_bindings():
with self.handler.installation_scope(binding):
scoped = await self.handler.list_knowledge_engines()
with runtime_handler.installation_scope(binding):
scoped = await runtime_handler.list_knowledge_engines()
for engine in scoped:
key = (str(engine.get('plugin_id', '')), str(engine.get('name', '')))
if key not in seen:
@@ -2069,11 +2177,12 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
"""List all available parsers from plugins."""
if not self.is_enable_plugin or not self._runtime_available():
return []
runtime_handler = self._runtime_handler()
parsers: list[dict[str, Any]] = []
seen: set[tuple[str, str]] = set()
for binding in await self._operation_bindings():
with self.handler.installation_scope(binding):
scoped = await self.handler.list_parsers()
with runtime_handler.installation_scope(binding):
scoped = await runtime_handler.list_parsers()
for parser in scoped:
key = (str(parser.get('plugin_id', '')), str(parser.get('name', '')))
if key not in seen:
@@ -2084,6 +2193,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
async def call_parser(self, plugin_id: str, context_data: dict[str, Any], file_bytes: bytes) -> dict[str, Any]:
"""Call plugin to parse a document."""
plugin_author, plugin_name = self._parse_plugin_id(plugin_id)
runtime_handler = self._runtime_handler()
binding = await self._target_binding(plugin_author, plugin_name)
with self.handler.installation_scope(binding):
return await self.handler.parse_document(plugin_author, plugin_name, context_data, file_bytes)
with runtime_handler.installation_scope(binding):
return await runtime_handler.parse_document(plugin_author, plugin_name, context_data, file_bytes)
src/langbot/pkg/provider/runners/localagent.py
+3 -2
View File
@@ -224,12 +224,13 @@ class LocalAgentRunner(runner.RequestRunner):
) -> list[modelmgr_requester.RuntimeLLMModel]:
"""Build ordered list of models to try: primary model + fallback models."""
candidates = []
execution_context = get_query_execution_context(query)
# Primary model
if query.use_llm_model_uuid:
try:
primary = await self.ap.model_mgr.get_model_by_uuid(
get_query_execution_context(query),
execution_context,
query.use_llm_model_uuid,
)
candidates.append(primary)
@@ -241,7 +242,7 @@ class LocalAgentRunner(runner.RequestRunner):
for fb_uuid in fallback_uuids:
try:
fb_model = await self.ap.model_mgr.get_model_by_uuid(
get_query_execution_context(query),
execution_context,
fb_uuid,
)
candidates.append(fb_model)
src/langbot/pkg/provider/tools/loaders/mcp.py
+49 -4
View File
@@ -33,7 +33,13 @@ from ....workspace.errors import WorkspaceError, WorkspaceInvariantError
import langbot_plugin.api.entities.builtin.resource.tool as resource_tool
import langbot_plugin.api.entities.builtin.provider.message as provider_message
from ....entity.persistence import mcp as persistence_mcp
from .mcp_stdio import BoxStdioSessionRuntime, MCPServerBoxConfig, MCPSessionErrorPhase, _ColdStartRetry # noqa: F401
from .mcp_stdio import (
BoxStdioSessionRuntime,
MCPServerBoxConfig as MCPServerBoxConfig, # noqa: F401 - public re-export
MCPSessionErrorPhase,
_ColdStartRetry,
_get_default_memory_mb,
)
from .mcp_policy import require_stdio_mcp_enabled, stdio_mcp_enabled
# Synthesized LLM tools for MCP resources (not from server tools/list).
@@ -321,6 +327,26 @@ class RuntimeMCPSession:
self._box_stdio_runtime = BoxStdioSessionRuntime(self)
self.box_config = self._box_stdio_runtime.config
def _parse_tool_call_timeout(self, value: typing.Any) -> float:
"""Return a safe tool-call timeout; zero explicitly disables it."""
try:
timeout = -1 if isinstance(value, bool) else float(value)
if timeout > 0:
# Validate the exact conversion used for each call here, so a
# finite-but-enormous manual config cannot fail at invocation.
timedelta(seconds=timeout)
except (TypeError, ValueError, OverflowError):
timeout = -1
if not math.isfinite(timeout) or timeout < 0:
self.ap.logger.warning(
f'Invalid MCP tool call timeout {value!r} for {self.server_name}; '
f'using {MCP_TOOL_CALL_TIMEOUT_DEFAULT_SECONDS:g} seconds'
)
return MCP_TOOL_CALL_TIMEOUT_DEFAULT_SECONDS
return timeout
async def _assert_execution_active(self) -> None:
"""Fail closed when this long-lived session belongs to a stale placement."""
@@ -723,7 +749,11 @@ class RuntimeMCPSession:
self.status = MCPSessionStatus.CONNECTING
self.error_message = None
self.error_phase = None
await asyncio.sleep(1)
try:
await self._sleep_with_execution_fence(1)
except WorkspaceError as fence_error:
self._stop_for_stale_execution(fence_error)
return
continue
# Explicitly disabled Box is a deliberate refusal, not a
# transient failure. Surface it immediately without log
@@ -1082,7 +1112,12 @@ class RuntimeMCPSession:
try:
await self._assert_execution_active()
result = await self.session.call_tool(tool_name, arguments)
read_timeout = timedelta(seconds=self.tool_call_timeout_sec) if self.tool_call_timeout_sec > 0 else None
result = await self.session.call_tool(
tool_name,
arguments,
read_timeout_seconds=read_timeout,
)
await self._assert_execution_active()
except Exception as e:
if self._is_tool_call_timeout(e):
@@ -2144,7 +2179,15 @@ class MCPLoader(loader.ToolLoader):
async def shutdown(self):
"""关闭所有工具"""
self.ap.logger.info('Shutting down all MCP sessions...')
for key, session in list(self.sessions.items()):
hosted_tasks = [task for task in self._hosted_mcp_tasks if not task.done()]
for task in hosted_tasks:
task.cancel()
if hosted_tasks:
await asyncio.gather(*hosted_tasks, return_exceptions=True)
self._hosted_mcp_tasks.clear()
async def shutdown_session(session: RuntimeMCPSession) -> None:
try:
await session.shutdown()
self.ap.logger.debug(f'Shutdown MCP session: {session.server_name}')
@@ -2152,5 +2195,7 @@ class MCPLoader(loader.ToolLoader):
self.ap.logger.error(
f'Error shutting down MCP session {session.server_name}: {e}\n{traceback.format_exc()}'
)
await asyncio.gather(*(shutdown_session(session) for session in list(self.sessions.values())))
self.sessions.clear()
self.ap.logger.info('All MCP sessions shutdown complete')
src/langbot/pkg/workspace/collaboration.py
+16 -7
View File
@@ -25,7 +25,7 @@ from ..entity.persistence.workspace import (
WorkspaceStatus,
)
from .entities import WorkspaceExecutionBinding
from .errors import WorkspaceNotFoundError
from .errors import WorkspaceExecutionUnavailableError, WorkspaceInvariantError, WorkspaceNotFoundError
from .policy import CloudWorkspacePolicy, SingleWorkspacePolicy
from .service import WorkspaceService
@@ -231,13 +231,22 @@ class WorkspaceCollaborationService:
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,
try:
accesses.append(
await self.resolve_account_workspace(
account_uuid,
workspace_uuid,
session=workspace_uow.session,
)
)
except (
WorkspaceExecutionUnavailableError,
WorkspaceInvariantError,
WorkspaceNotFoundError,
) as exc:
self.ap.logger.warning(
f'Skipping inactive Workspace discovery projection {workspace_uuid!r}: {exc}'
)
)
accesses.sort(key=lambda access: (access.workspace.created_at, access.workspace.uuid))
return accesses
src/langbot/pkg/workspace/service.py
+14
View File
@@ -177,6 +177,7 @@ class WorkspaceService:
"""
self._require_deployment_admission()
self._require_directory_projection()
tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
if session is None and workspace_uuid is not None and callable(tenant_uow):
@@ -247,6 +248,7 @@ class WorkspaceService:
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_directory_projection()
self._require_deployment_admission()
return binding
@@ -500,3 +502,15 @@ class WorkspaceService:
guard = getattr(self.ap, 'deployment_admission', None)
if guard is not None:
guard.require_active()
def _require_directory_projection(self) -> None:
deployment = getattr(self.ap, 'deployment', None)
if deployment is None or not getattr(deployment, 'multi_workspace_enabled', False):
return
projection = getattr(self.ap, 'directory_projection_service', None)
if projection is None:
raise WorkspaceExecutionUnavailableError('Cloud directory projection is unavailable')
try:
projection.require_ready()
except RuntimeError as exc:
raise WorkspaceExecutionUnavailableError(str(exc)) from exc
tests/integration/persistence/test_migrations.py
+19
View File
@@ -15,6 +15,7 @@ from sqlalchemy.ext.asyncio import create_async_engine
from langbot.pkg.entity.persistence.base import Base
from langbot.pkg.persistence.alembic_runner import (
run_alembic_downgrade,
run_alembic_upgrade,
run_alembic_stamp,
get_alembic_current,
@@ -171,6 +172,24 @@ class TestSQLiteMigrationUpgrade:
)
assert 'embedding_dimension' in columns
@pytest.mark.asyncio
async def test_directory_projection_upgrade_downgrade_round_trip(self, sqlite_engine):
await run_alembic_stamp(sqlite_engine, '0013_tenant_pgvector')
await run_alembic_upgrade(sqlite_engine, 'head')
async with sqlite_engine.connect() as conn:
tables = await conn.run_sync(lambda sync_conn: set(sqlalchemy.inspect(sync_conn).get_table_names()))
assert {'directory_projection_states', 'directory_projection_inbox'} <= tables
await run_alembic_downgrade(sqlite_engine, '0013_tenant_pgvector')
async with sqlite_engine.connect() as conn:
tables = await conn.run_sync(lambda sync_conn: set(sqlalchemy.inspect(sync_conn).get_table_names()))
assert 'directory_projection_states' not in tables
assert 'directory_projection_inbox' not in tables
await run_alembic_upgrade(sqlite_engine, 'head')
assert await get_alembic_current(sqlite_engine) == _get_script_head()
class TestSQLiteMigrationFreshDatabase:
"""Tests for fresh database workflow."""
tests/integration/persistence/test_migrations_postgres.py
+208 -3
View File
@@ -31,7 +31,13 @@ from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine, create_async_engine
from sqlalchemy import text
from langbot.pkg.cloud.directory import DirectoryEvent, DirectoryEventBatch
from langbot.pkg.cloud.directory_projection import DirectoryProjectionService
from langbot.pkg.entity.persistence.base import Base
from langbot.pkg.entity.persistence.cloud_directory import (
DirectoryProjectionInbox,
DirectoryProjectionState,
)
from langbot.pkg.entity.persistence.user import User
from langbot.pkg.persistence.mgr import PersistenceManager, PersistenceMode
from langbot.pkg.persistence.tenant_uow import (
@@ -82,6 +88,17 @@ from langbot.pkg.rag.knowledge.kbmgr import RAGManager
from .resource_migration_support import TENANT_TABLES, create_legacy_resource_schema
class _NoopDirectoryProjectionProvider:
async def fetch_snapshot(self, instance_uuid: str):
raise AssertionError(f'Unexpected snapshot fetch for {instance_uuid}')
async def fetch_events(self, instance_uuid: str, after_cursor: int, limit: int):
raise AssertionError(f'Unexpected event fetch for {instance_uuid} after {after_cursor} with limit {limit}')
async def fetch_workspaces(self, instance_uuid: str, workspace_uuids: tuple[str, ...]):
raise AssertionError(f'Unexpected delta fetch for {instance_uuid}: {workspace_uuids!r}')
def _get_script_head() -> str:
"""Resolve the current Alembic head revision from the script directory.
@@ -648,6 +665,10 @@ class TestPostgreSQLTenantRuntime:
instance_uuid = 'cloud-runtime-persistence-test'
workspace_a = '10000000-0000-0000-0000-000000000001'
workspace_b = '20000000-0000-0000-0000-000000000002'
workspace_other = '30000000-0000-0000-0000-000000000003'
workspace_local = '40000000-0000-0000-0000-000000000004'
directory_account = '50000000-0000-0000-0000-000000000005'
workspace_projected = '60000000-0000-0000-0000-000000000006'
role_suffix = uuid.uuid4().hex[:12]
runtime_role = f'lb_runtime_{role_suffix}'
bypass_role = f'lb_bypass_{role_suffix}'
@@ -730,17 +751,22 @@ class TestPostgreSQLTenantRuntime:
assert {row['relname'] for row in rls_rows} == set(TENANT_TABLE_COLUMNS)
assert all(row['relrowsecurity'] and row['relforcerowsecurity'] and row['has_policy'] for row in rls_rows)
for workspace_uuid, slug in ((workspace_a, 'workspace-a'), (workspace_b, 'workspace-b')):
for workspace_uuid, slug, target_instance, source in (
(workspace_a, 'workspace-a', instance_uuid, 'cloud_projection'),
(workspace_b, 'workspace-b', instance_uuid, 'cloud_projection'),
(workspace_other, 'workspace-other', 'other-instance', 'cloud_projection'),
(workspace_local, 'workspace-local', instance_uuid, 'local'),
):
async with TenantUnitOfWork(release_engine, workspace_uuid) as uow:
await uow.execute(
sa.insert(Workspace).values(
uuid=workspace_uuid,
instance_uuid=instance_uuid,
instance_uuid=target_instance,
name=slug,
slug=slug,
type='team',
status='active',
source='cloud_projection',
source=source,
projection_revision=0,
)
)
@@ -751,6 +777,30 @@ class TestPostgreSQLTenantRuntime:
value=slug,
)
)
async with release_engine.begin() as conn:
await conn.execute(
sa.insert(User).values(
uuid=directory_account,
user='directory@example.com',
normalized_email='directory@example.com',
password='closed-directory',
account_type='space',
status='active',
source='cloud_projection',
projection_revision=1,
)
)
async with TenantUnitOfWork(release_engine, workspace_local) as uow:
uow.session.add(
WorkspaceMembership(
uuid=str(uuid.uuid4()),
workspace_uuid=workspace_local,
account_uuid=directory_account,
role='owner',
status='active',
projection_revision=1,
)
)
await create_role(runtime_role)
runtime_engine = create_async_engine(role_url(runtime_role), pool_size=1, max_overflow=0)
@@ -844,6 +894,161 @@ class TestPostgreSQLTenantRuntime:
cloud_manager.create_tables.assert_not_awaited()
cloud_manager._run_alembic_migrations.assert_not_awaited()
directory_fingerprint = hashlib.sha256(b'directory-snapshot').hexdigest()
async with cloud_manager.directory_projection_uow(instance_uuid) as directory:
assert set((await directory.session.scalars(sa.select(Workspace.uuid))).all()) == {
workspace_a,
workspace_b,
}
directory.session.add(
Workspace(
uuid=workspace_projected,
instance_uuid=instance_uuid,
name='workspace-projected',
slug='workspace-projected',
type='team',
status='active',
source='cloud_projection',
projection_revision=1,
)
)
await directory.session.flush()
directory.session.add(
WorkspaceMembership(
uuid=str(uuid.uuid4()),
workspace_uuid=workspace_projected,
account_uuid=directory_account,
role='owner',
status='active',
projection_revision=1,
)
)
directory.session.add(
WorkspaceExecutionState(
workspace_uuid=workspace_projected,
instance_uuid=instance_uuid,
active_generation=1,
state='active',
write_fenced=False,
source='cloud',
desired_state_revision=1,
)
)
directory.session.add(
DirectoryProjectionState(
instance_uuid=instance_uuid,
cursor=1,
snapshot_fingerprint=directory_fingerprint,
last_applied_at=datetime.datetime.now(datetime.UTC),
)
)
directory.session.add(
DirectoryProjectionInbox(
instance_uuid=instance_uuid,
event_uuid=str(uuid.uuid4()),
cursor=1,
event_type='workspace.changed',
revision=1,
fingerprint=directory_fingerprint,
)
)
await directory.session.flush()
assert await directory.session.scalar(sa.select(sa.func.count()).select_from(WorkspaceMembership)) == 1
assert (await directory.session.execute(sa.select(WorkspaceMetadata))).all() == []
async with cloud_manager.tenant_uow(workspace_projected) as tenant:
workspace_update = await tenant.session.execute(
sa.update(Workspace)
.where(Workspace.uuid == workspace_projected)
.values(name='tenant-must-not-update-cloud')
)
membership_update = await tenant.session.execute(
sa.update(WorkspaceMembership)
.where(WorkspaceMembership.workspace_uuid == workspace_projected)
.values(role='admin')
)
execution_update = await tenant.session.execute(
sa.update(WorkspaceExecutionState)
.where(WorkspaceExecutionState.workspace_uuid == workspace_projected)
.values(write_fenced=True)
)
assert workspace_update.rowcount == 0
assert membership_update.rowcount == 0
assert execution_update.rowcount == 0
async with cloud_manager.tenant_uow(workspace_local) as tenant:
local_update = await tenant.session.execute(
sa.update(Workspace).where(Workspace.uuid == workspace_local).values(name='tenant-can-update-local')
)
assert local_update.rowcount == 1
async with cloud_manager.directory_projection_uow(instance_uuid) as directory:
local_update = await directory.session.execute(
sa.update(Workspace)
.where(Workspace.uuid == workspace_local)
.values(name='directory-must-not-update-local')
)
assert local_update.rowcount == 0
projection_application = SimpleNamespace(
persistence_mgr=cloud_manager,
logger=logging.getLogger('postgres-directory-projection-concurrency'),
)
first_projection = DirectoryProjectionService(
projection_application,
_NoopDirectoryProjectionProvider(),
instance_uuid,
)
second_projection = DirectoryProjectionService(
projection_application,
_NoopDirectoryProjectionProvider(),
instance_uuid,
)
concurrent_event = DirectoryEvent(
cursor=2,
uuid='70000000-0000-0000-0000-000000000007',
aggregate_uuid=workspace_projected,
event_type='entitlement.changed',
revision=2,
payload={
'workspace_uuid': workspace_projected,
'entitlement_revision': 2,
},
created_at=datetime.datetime.now(datetime.UTC),
)
concurrent_batch = DirectoryEventBatch(
instance_uuid=instance_uuid,
after_cursor=1,
cursor=2,
high_water_cursor=2,
events=[concurrent_event],
)
await asyncio.gather(
first_projection.apply_event_batch(concurrent_batch),
second_projection.apply_event_batch(concurrent_batch),
)
async with cloud_manager.directory_projection_uow(instance_uuid) as directory:
projection_state = await directory.session.get(DirectoryProjectionState, instance_uuid)
concurrent_receipts = (
await directory.session.scalars(
sa.select(DirectoryProjectionInbox).where(
DirectoryProjectionInbox.event_uuid == concurrent_event.uuid
)
)
).all()
assert projection_state.cursor == 2
assert len(concurrent_receipts) == 1
assert concurrent_receipts[0].applied_at is not None
async with cloud_manager.tenant_uow(workspace_a) as tenant:
assert (await tenant.session.execute(sa.select(DirectoryProjectionState))).all() == []
assert (await tenant.session.execute(sa.select(DirectoryProjectionInbox))).all() == []
async with cloud_manager.instance_discovery_uow(instance_uuid) as discovery:
assert (await discovery.session.execute(sa.select(DirectoryProjectionState))).all() == []
update_result = await discovery.session.execute(sa.update(DirectoryProjectionState).values(cursor=2))
assert update_result.rowcount == 0
with pytest.raises(TransactionRollbackOnlyError, match='after-commit work was cancelled'):
async with cloud_manager.tenant_uow(workspace_a):
duplicate_statement = sa.insert(WorkspaceMetadata).values(
tests/unit_tests/api/service/test_user_service.py
+61 -2
View File
@@ -19,8 +19,11 @@ import datetime
from unittest.mock import AsyncMock, Mock
from types import SimpleNamespace
from langbot.pkg.api.http.service.user import UserService
from langbot.pkg.entity.persistence.user import AccountStatus, User
from langbot.pkg.api.http.service.user import (
ControlPlaneDirectoryRequiredError,
UserService,
)
from langbot.pkg.entity.persistence.user import AccountSource, AccountStatus, User
from langbot.pkg.entity.errors.account import AccountEmailMismatchError
@@ -563,6 +566,62 @@ class TestUserServiceCreateOrUpdateSpaceUser:
ap.persistence_mgr.execute_async.assert_called()
assert updated_user.space_account_uuid == 'existing-space-uuid'
async def test_cloud_login_updates_only_the_projected_space_account(self):
projected = SimpleNamespace(
uuid='projected-space-uuid',
user='Cloud Owner',
normalized_email='owner@example.com',
password='',
account_type='space',
status=AccountStatus.ACTIVE.value,
source=AccountSource.CLOUD_PROJECTION.value,
projection_revision=7,
space_account_uuid='projected-space-uuid',
)
persistence = SimpleNamespace(execute_async=AsyncMock())
ap = SimpleNamespace(
persistence_mgr=persistence,
workspace_service=SimpleNamespace(policy=SimpleNamespace(multi_workspace_enabled=True)),
)
service = UserService(ap)
service.get_user_by_space_account_uuid = AsyncMock(side_effect=[projected, projected])
result = await service.create_or_update_space_user(
space_account_uuid='projected-space-uuid',
email='OWNER@example.com',
access_token='access-token',
refresh_token='refresh-token',
api_key='api-key',
expires_in=3600,
)
assert result is projected
persistence.execute_async.assert_awaited_once()
async def test_cloud_login_never_creates_an_unprojected_account(self):
persistence = SimpleNamespace(execute_async=AsyncMock())
ap = SimpleNamespace(
persistence_mgr=persistence,
workspace_service=SimpleNamespace(policy=SimpleNamespace(multi_workspace_enabled=True)),
)
service = UserService(ap)
service.get_user_by_space_account_uuid = AsyncMock(return_value=None)
with pytest.raises(
ControlPlaneDirectoryRequiredError,
match='verified Cloud directory',
):
await service.create_or_update_space_user(
space_account_uuid='unknown-space-uuid',
email='unknown@example.com',
access_token='access-token',
refresh_token='refresh-token',
api_key='api-key',
expires_in=3600,
)
persistence.execute_async.assert_not_awaited()
async def test_create_or_update_new_space_user_first_init(self):
"""Creates new Space user on first initialization."""
# Setup
tests/unit_tests/box/test_box_connector.py
+136 -2
View File
@@ -6,6 +6,7 @@ from unittest.mock import AsyncMock, Mock
import pytest
from langbot.pkg.box import connector as connector_module
from langbot_plugin.box.client import ActionRPCBoxClient
from langbot_plugin.box.errors import BoxRuntimeUnavailableError
from langbot_plugin.box.security import (
@@ -121,6 +122,139 @@ def test_box_runtime_connector_dispose_terminates_subprocess(monkeypatch: pytest
assert connector._ctrl_task is None
@pytest.mark.asyncio
async def test_box_runtime_connector_cleans_partial_transport_on_connect_failure(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setattr('langbot.pkg.utils.platform.get_platform', lambda: 'linux')
monkeypatch.setattr('langbot.pkg.utils.platform.standalone_box', False)
connector = BoxRuntimeConnector(make_app(Mock()))
connector._start_local_stdio = AsyncMock(side_effect=RuntimeError('bind failed'))
connector._stop_transport = AsyncMock()
connector._close_managed_subprocess = AsyncMock()
with pytest.raises(RuntimeError, match='bind failed'):
await connector.initialize()
assert connector._stop_transport.await_count == 2
connector._close_managed_subprocess.assert_awaited_once()
@pytest.mark.asyncio
async def test_box_runtime_connector_starts_heartbeat_after_reconnect(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setattr('langbot.pkg.utils.platform.get_platform', lambda: 'linux')
monkeypatch.setattr('langbot.pkg.utils.platform.standalone_box', False)
connector = BoxRuntimeConnector(make_app(Mock()))
connector._start_local_stdio = AsyncMock(side_effect=[RuntimeError('bind failed'), None])
connector._stop_transport = AsyncMock()
connector._close_managed_subprocess = AsyncMock()
with pytest.raises(RuntimeError, match='bind failed'):
await connector.initialize()
assert connector._heartbeat_task is None
await connector.reconnect()
assert connector._heartbeat_task is not None
assert not connector._heartbeat_task.done()
await connector.aclose()
@pytest.mark.asyncio
async def test_box_stdio_connection_does_not_capture_unconsumed_stderr(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setattr('langbot.pkg.utils.platform.get_platform', lambda: 'linux')
monkeypatch.setattr('langbot.pkg.utils.platform.standalone_box', False)
created = {}
class FakeHandler:
def __init__(self, connection):
self.release = asyncio.Event()
async def call_action(self, action, data):
return None
async def run(self):
await self.release.wait()
async def close(self):
self.release.set()
class FakeController:
def __init__(self, **kwargs):
created.update(kwargs)
self.process = SimpleNamespace(returncode=0)
async def run(self, callback):
await callback(object())
async def close(self):
return None
monkeypatch.setattr(connector_module, 'Handler', FakeHandler)
monkeypatch.setattr(
'langbot_plugin.runtime.io.controllers.stdio.client.StdioClientController',
FakeController,
)
connector = BoxRuntimeConnector(make_app(Mock()))
await connector.initialize()
assert created['capture_stderr'] is False
assert connector._handler is not None
await connector.aclose()
@pytest.mark.asyncio
async def test_box_disconnect_notifies_once_and_clears_handler(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setattr('langbot.pkg.utils.platform.get_platform', lambda: 'linux')
monkeypatch.setattr('langbot.pkg.utils.platform.standalone_box', False)
disconnect = AsyncMock()
class FakeHandler:
def __init__(self, connection):
pass
async def call_action(self, action, data):
return None
async def run(self):
return None
async def close(self):
return None
class FakeController:
def __init__(self, **kwargs):
self.process = SimpleNamespace(returncode=0)
async def run(self, callback):
await callback(object())
async def close(self):
return None
monkeypatch.setattr(connector_module, 'Handler', FakeHandler)
monkeypatch.setattr(
'langbot_plugin.runtime.io.controllers.stdio.client.StdioClientController',
FakeController,
)
connector = BoxRuntimeConnector(make_app(Mock()), runtime_disconnect_callback=disconnect)
await connector.initialize()
await asyncio.sleep(0)
disconnect.assert_awaited_once_with(connector)
assert connector._handler is None
await connector.aclose()
def test_box_runtime_connector_builds_host_control_headers(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv(BOX_CONTROL_TOKEN_ENV, _CONTROL_TOKEN)
connector = BoxRuntimeConnector(make_app(Mock(), runtime_endpoint='http://box-runtime:5410'))
@@ -200,7 +334,7 @@ async def test_local_stdio_injects_generated_token_and_trusted_instance(
)
connector = BoxRuntimeConnector(make_app(Mock()))
def fake_callback(_transport_name, connected, _connect_error):
def fake_callback(_transport_name, connected, _connect_error, _generation):
async def callback(_connection):
connected.set()
@@ -231,7 +365,7 @@ async def test_websocket_controller_receives_control_headers(monkeypatch: pytest
)
connector = BoxRuntimeConnector(make_app(Mock(), runtime_endpoint='http://box-runtime:5410'))
def fake_callback(_transport_name, connected, _connect_error):
def fake_callback(_transport_name, connected, _connect_error, _generation):
async def callback(_connection):
connected.set()
tests/unit_tests/box/test_box_service.py
+91
View File
@@ -253,6 +253,46 @@ async def test_box_service_without_explicit_client_initializes_internal_connecto
connector.initialize.assert_awaited_once()
@pytest.mark.asyncio
async def test_cloud_initialize_validation_failure_closes_connector_and_cancels_reconnect(
monkeypatch: pytest.MonkeyPatch,
):
logger = Mock()
app = make_app(logger)
app.deployment = SimpleNamespace(multi_workspace_enabled=True)
app.instance_config.data['box'].update(
{
'backend': 'nsjail',
'admission': {'required': True, 'workspace_quota_mb': 32},
}
)
connector = Mock()
connector.client = Mock(spec=BoxRuntimeClient)
connector.initialize = AsyncMock()
connector.aclose = AsyncMock()
connector.runtime_disconnect_callback = Mock()
monkeypatch.setattr('langbot.pkg.box.service.BoxRuntimeConnector', Mock(return_value=connector))
service = BoxService(app)
service._ensure_default_workspace = Mock()
readiness_error = BoxValidationError('Cloud Box nsjail isolation readiness failed')
service._verify_cloud_runtime = AsyncMock(side_effect=readiness_error)
reconnect_task = asyncio.create_task(asyncio.Event().wait())
service._reconnect_task = reconnect_task
service._reconnecting = True
with pytest.raises(BoxValidationError) as exc_info:
await service.initialize()
assert exc_info.value is readiness_error
assert service.available is False
assert service._closing is True
assert service._reconnecting is False
assert service._reconnect_task is None
assert reconnect_task.cancelled()
assert connector.runtime_disconnect_callback is None
connector.aclose.assert_awaited_once()
class TestSharesFilesystemWithBox:
"""``shares_filesystem_with_box`` must reflect the real LangBot<->Box
filesystem topology, which is derived from the connector transport:
@@ -1788,6 +1828,57 @@ class TestBoxDisabledByConfig:
assert service._reconnecting is False
@pytest.mark.asyncio
async def test_disconnect_callback_does_not_schedule_on_closing_event_loop(
monkeypatch: pytest.MonkeyPatch,
):
service = BoxService(make_app(Mock()), client=Mock(spec=BoxRuntimeClient))
closed_loop = Mock()
closed_loop.is_closed.return_value = True
monkeypatch.setattr('langbot.pkg.box.service.asyncio.get_running_loop', Mock(return_value=closed_loop))
await service._on_runtime_disconnect(connector=Mock())
closed_loop.create_task.assert_not_called()
assert service._reconnect_task is None
assert service._reconnecting is False
@pytest.mark.asyncio
async def test_disconnect_callback_closes_reconnect_coroutine_when_task_creation_races_with_loop_close(
monkeypatch: pytest.MonkeyPatch,
):
service = BoxService(make_app(Mock()), client=Mock(spec=BoxRuntimeClient))
loop = Mock()
loop.is_closed.return_value = False
loop.create_task.side_effect = RuntimeError('event loop is closed')
monkeypatch.setattr('langbot.pkg.box.service.asyncio.get_running_loop', Mock(return_value=loop))
async def reconnect():
await asyncio.Event().wait()
reconnect_coroutine = reconnect()
service._reconnect_loop = Mock(return_value=reconnect_coroutine)
await service._on_runtime_disconnect(connector=Mock())
loop.create_task.assert_called_once_with(reconnect_coroutine)
assert reconnect_coroutine.cr_frame is None
assert service._reconnect_task is None
assert service._reconnecting is False
def test_disconnect_callback_does_not_schedule_without_running_event_loop():
service = BoxService(make_app(Mock()), client=Mock(spec=BoxRuntimeClient))
callback = service._on_runtime_disconnect(connector=Mock())
with pytest.raises(StopIteration):
callback.send(None)
assert service._reconnect_task is None
assert service._reconnecting is False
class TestBuildSkillExtraMounts:
"""Robustness of skill mount construction against a stale skill cache.
tests/unit_tests/cloud/test_bootstrap.py
+61
View File
@@ -7,6 +7,7 @@ import pytest
from langbot.pkg.cloud.bootstrap import (
CloudBootstrapError,
CloudManifestRefreshService,
CloudRuntimeUnavailableError,
DeploymentAdmissionGuard,
OpenSourceDeployment,
@@ -33,7 +34,38 @@ class _Entitlements:
)
class _Directory:
async def fetch_snapshot(self, instance_uuid: str):
del instance_uuid
raise AssertionError('not used by bootstrap contract tests')
async def fetch_events(self, instance_uuid: str, after_cursor: int, limit: int):
del instance_uuid, after_cursor, limit
raise AssertionError('not used by bootstrap contract tests')
async def fetch_workspaces(self, instance_uuid: str, workspace_uuids: tuple[str, ...]):
del instance_uuid, workspace_uuids
raise AssertionError('not used by bootstrap contract tests')
class _Manifest:
def __init__(self):
self.candidate = None
self.closed = False
async def refresh_manifest(self):
if self.candidate is None:
raise AssertionError('no refreshed Manifest was configured')
return self.candidate
async def aclose(self) -> None:
self.closed = True
class _Provider:
def __init__(self):
self.manifest_provider = _Manifest()
def bootstrap(self, *, instance_uuid: str, instance_config: dict):
del instance_config
return VerifiedCloudDeployment(
@@ -45,6 +77,8 @@ class _Provider:
capabilities=frozenset({'multi_workspace_v2'}),
tenant_isolation_version=2,
entitlement_provider=_Entitlements(),
directory_provider=_Directory(),
manifest_provider=self.manifest_provider,
verification_key_id='root-2026',
)
@@ -289,3 +323,30 @@ async def test_deployment_admission_accepts_only_monotonic_non_conflicting_renew
conflicting = dataclasses.replace(renewed, manifest_jti='different')
with pytest.raises(CloudRuntimeUnavailableError, match='conflicting'):
guard.replace(conflicting)
async def test_manifest_refresh_replaces_receipt_before_short_ttl_expires():
wall = [1_000.0]
provider = _Provider()
current = dataclasses.replace(
provider.bootstrap(instance_uuid='instance-a', instance_config={}),
expires_at=1_300,
)
guard = DeploymentAdmissionGuard('instance-a', current, wall_time=lambda: wall[0])
renewed = dataclasses.replace(
current,
manifest_jti='manifest-renewed',
manifest_generation=current.manifest_generation + 1,
expires_at=2_000,
)
provider.manifest_provider.candidate = renewed
service = CloudManifestRefreshService(
guard,
provider.manifest_provider,
SimpleNamespace(exception=lambda *_: None),
wall_time=lambda: wall[0],
)
assert service.next_refresh_delay() == 120
assert await service.refresh_once() is renewed
assert guard.deployment is renewed
tests/unit_tests/cloud/test_directory_projection.py
+817
View File
@@ -0,0 +1,817 @@
from __future__ import annotations
import datetime
import logging
from types import SimpleNamespace
import pytest
import sqlalchemy
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
from langbot.pkg.cloud.directory import (
DirectoryDelta,
DirectoryEvent,
DirectoryEventBatch,
DirectoryMember,
DirectoryProjectionUnavailableError,
DirectorySnapshot,
DirectoryWorkspace,
)
from langbot.pkg.cloud.directory_projection import DirectoryProjectionService
from langbot.pkg.entity.persistence.base import Base
from langbot.pkg.entity.persistence.cloud_directory import DirectoryProjectionInbox, DirectoryProjectionState
from langbot.pkg.entity.persistence.user import User
from langbot.pkg.entity.persistence.workspace import (
Workspace,
WorkspaceExecutionState,
WorkspaceMembership,
)
from langbot.pkg.persistence.tenant_uow import PersistenceScope, TenantUnitOfWork
pytestmark = pytest.mark.asyncio
INSTANCE_UUID = 'instance-directory-test'
WORKSPACE_UUID = '10000000-0000-0000-0000-000000000001'
SECOND_WORKSPACE_UUID = '10000000-0000-0000-0000-000000000002'
ACCOUNT_UUID = '20000000-0000-0000-0000-000000000001'
MEMBERSHIP_UUID = '30000000-0000-0000-0000-000000000001'
SECOND_MEMBERSHIP_UUID = '30000000-0000-0000-0000-000000000002'
class _Persistence:
def __init__(self, engine):
self.engine = engine
def directory_projection_uow(self, instance_uuid: str) -> TenantUnitOfWork:
return TenantUnitOfWork(
self.engine,
scope=PersistenceScope.directory(instance_uuid),
)
def get_db_engine(self):
return self.engine
class _Provider:
def __init__(
self,
snapshots: list[DirectorySnapshot],
batches: list[DirectoryEventBatch] | None = None,
deltas: list[DirectoryDelta] | None = None,
) -> None:
self.snapshots = snapshots
self.batches = list(batches or [])
self.deltas = list(deltas or [])
self.snapshot_calls = 0
self.delta_calls = 0
self.after_cursors: list[int] = []
async def fetch_snapshot(self, instance_uuid: str) -> DirectorySnapshot:
assert instance_uuid == INSTANCE_UUID
snapshot = self.snapshots[min(self.snapshot_calls, len(self.snapshots) - 1)]
self.snapshot_calls += 1
return snapshot
async def fetch_events(
self,
instance_uuid: str,
after_cursor: int,
limit: int,
) -> DirectoryEventBatch:
assert instance_uuid == INSTANCE_UUID
assert limit == 100
self.after_cursors.append(after_cursor)
if self.batches:
return self.batches.pop(0)
return DirectoryEventBatch(
instance_uuid=instance_uuid,
after_cursor=after_cursor,
cursor=after_cursor,
high_water_cursor=after_cursor,
events=[],
)
async def fetch_workspaces(
self,
instance_uuid: str,
workspace_uuids: tuple[str, ...],
) -> DirectoryDelta:
assert instance_uuid == INSTANCE_UUID
assert workspace_uuids == (WORKSPACE_UUID,)
delta = self.deltas[min(self.delta_calls, len(self.deltas) - 1)]
self.delta_calls += 1
return delta
@pytest.fixture
async def projection_context(tmp_path):
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "directory-projection.db"}')
async with engine.begin() as connection:
await connection.run_sync(Base.metadata.create_all)
application = SimpleNamespace(
persistence_mgr=_Persistence(engine),
logger=logging.getLogger('test-directory-projection'),
)
session_factory = async_sessionmaker(engine, expire_on_commit=False)
yield application, session_factory
await engine.dispose()
def _member(*, revision: int = 1, role: str = 'member') -> DirectoryMember:
return DirectoryMember(
membership_uuid=MEMBERSHIP_UUID,
account_uuid=ACCOUNT_UUID,
normalized_email='owner@example.com',
display_name='Workspace Owner',
account_status='active',
role=role,
membership_status='active',
projection_revision=revision,
joined_at=datetime.datetime(2026, 7, 24, tzinfo=datetime.UTC),
)
def _workspace(
*,
revision: int = 1,
status: str = 'active',
name: str = 'Personal Workspace',
role: str = 'member',
execution_generation: int = 1,
) -> DirectoryWorkspace:
return DirectoryWorkspace(
uuid=WORKSPACE_UUID,
name=name,
slug='personal-workspace',
type='personal',
status=status,
created_by_account_uuid=ACCOUNT_UUID,
projection_revision=revision,
execution_generation=execution_generation,
members=[_member(revision=revision, role=role)],
)
def _snapshot(
cursor: int,
*,
workspaces: list[DirectoryWorkspace] | None = None,
) -> DirectorySnapshot:
return DirectorySnapshot(
instance_uuid=INSTANCE_UUID,
cursor=cursor,
generated_at=datetime.datetime(2026, 7, 24, 12, cursor % 60, tzinfo=datetime.UTC),
workspaces=[_workspace(revision=max(cursor, 1))] if workspaces is None else workspaces,
)
def _delta(
*,
workspaces: list[DirectoryWorkspace] | None = None,
requested_workspace_uuids: tuple[str, ...] = (WORKSPACE_UUID,),
) -> DirectoryDelta:
return DirectoryDelta(
instance_uuid=INSTANCE_UUID,
requested_workspace_uuids=requested_workspace_uuids,
generated_at=datetime.datetime(2026, 7, 24, 12, 30, tzinfo=datetime.UTC),
workspaces=[_workspace(revision=2)] if workspaces is None else workspaces,
)
async def test_initial_snapshot_projects_core_owned_rows(projection_context):
application, session_factory = projection_context
service = DirectoryProjectionService(
application,
_Provider([_snapshot(1)]),
INSTANCE_UUID,
)
await service.initialize()
service.require_ready()
async with session_factory() as session:
account = await session.scalar(sqlalchemy.select(User))
workspace = await session.get(Workspace, WORKSPACE_UUID)
membership = await session.scalar(sqlalchemy.select(WorkspaceMembership))
execution = await session.get(WorkspaceExecutionState, WORKSPACE_UUID)
state = await session.get(DirectoryProjectionState, INSTANCE_UUID)
assert account is not None
assert account.uuid == ACCOUNT_UUID
assert account.source == 'cloud_projection'
assert account.account_type == 'space'
assert account.password == ''
assert workspace is not None
assert workspace.source == 'cloud_projection'
assert membership is not None
assert membership.role == 'developer'
assert execution is not None
assert execution.source == 'cloud'
assert execution.write_fenced is False
assert state is not None
assert state.cursor == 1
assert state.snapshot_coverage_cursor == 1
async def test_same_cursor_equivocation_and_rollback_fail_closed(projection_context):
application, _session_factory = projection_context
service = DirectoryProjectionService(
application,
_Provider([_snapshot(2)]),
INSTANCE_UUID,
)
await service.initialize()
conflicting = _snapshot(
2,
workspaces=[_workspace(revision=2, name='Conflicting Name')],
)
with pytest.raises(DirectoryProjectionUnavailableError, match='conflicting contents'):
await service.apply_snapshot(conflicting)
with pytest.raises(DirectoryProjectionUnavailableError, match='rolled back'):
await service.apply_snapshot(_snapshot(1))
async def test_execution_generation_cannot_change_without_directory_revision(projection_context):
application, _session_factory = projection_context
service = DirectoryProjectionService(
application,
_Provider([_snapshot(1)]),
INSTANCE_UUID,
)
await service.initialize()
conflicting = _snapshot(
2,
workspaces=[_workspace(revision=1, execution_generation=2)],
)
with pytest.raises(DirectoryProjectionUnavailableError, match='execution state has conflicting contents'):
await service.apply_snapshot(conflicting)
async def test_event_limit_matches_single_delta_request_limit(projection_context):
application, _session_factory = projection_context
with pytest.raises(ValueError, match='between 1 and 100'):
DirectoryProjectionService(
application,
_Provider([_snapshot(1)]),
INSTANCE_UUID,
event_limit=101,
)
async def test_archived_and_absent_workspaces_are_execution_fenced(projection_context):
application, session_factory = projection_context
service = DirectoryProjectionService(
application,
_Provider([_snapshot(1)]),
INSTANCE_UUID,
)
await service.initialize()
await service.apply_snapshot(
_snapshot(
2,
workspaces=[_workspace(revision=2, status='archived')],
)
)
async with session_factory() as session:
workspace = await session.get(Workspace, WORKSPACE_UUID)
execution = await session.get(WorkspaceExecutionState, WORKSPACE_UUID)
membership = await session.scalar(sqlalchemy.select(WorkspaceMembership))
assert workspace.status == 'archived'
assert execution.state == 'inactive'
assert execution.write_fenced is True
assert membership.status == 'removed'
await service.apply_snapshot(_snapshot(3, workspaces=[]))
async with session_factory() as session:
workspace = await session.get(Workspace, WORKSPACE_UUID)
execution = await session.get(WorkspaceExecutionState, WORKSPACE_UUID)
assert workspace.status == 'archived'
assert execution.state == 'inactive'
assert execution.write_fenced is True
async def test_event_poll_fetches_workspace_delta_and_records_receipt(projection_context):
application, session_factory = projection_context
event = DirectoryEvent(
cursor=2,
uuid='40000000-0000-0000-0000-000000000001',
aggregate_uuid=WORKSPACE_UUID,
event_type='directory.changed',
revision=2,
payload={'workspace_uuid': WORKSPACE_UUID, 'directory_revision': 2},
created_at=datetime.datetime(2026, 7, 24, 12, 2, tzinfo=datetime.UTC),
)
batch = DirectoryEventBatch(
instance_uuid=INSTANCE_UUID,
after_cursor=1,
cursor=2,
high_water_cursor=2,
events=[event],
)
provider = _Provider(
[_snapshot(1)],
[batch],
[_delta(workspaces=[_workspace(revision=2, name='Updated Workspace')])],
)
service = DirectoryProjectionService(application, provider, INSTANCE_UUID)
await service.initialize()
await service.sync_once()
async with session_factory() as session:
account = await session.scalar(sqlalchemy.select(User).where(User.uuid == ACCOUNT_UUID))
workspace = await session.get(Workspace, WORKSPACE_UUID)
state = await session.get(DirectoryProjectionState, INSTANCE_UUID)
inbox = await session.scalar(sqlalchemy.select(DirectoryProjectionInbox))
assert account.projection_revision == 1
assert workspace.name == 'Updated Workspace'
assert state.cursor == 2
assert inbox.event_uuid == event.uuid
assert inbox.applied_at is not None
assert provider.snapshot_calls == 1
assert provider.delta_calls == 1
async def test_directory_delta_does_not_skip_unfetched_event_cursors(projection_context):
application, session_factory = projection_context
event = DirectoryEvent(
cursor=2,
uuid='40000000-0000-0000-0000-000000000003',
aggregate_uuid=WORKSPACE_UUID,
event_type='directory.changed',
revision=2,
payload={'workspace_uuid': WORKSPACE_UUID, 'directory_revision': 2},
created_at=datetime.datetime(2026, 7, 24, 12, 2, tzinfo=datetime.UTC),
)
batch = DirectoryEventBatch(
instance_uuid=INSTANCE_UUID,
after_cursor=1,
cursor=2,
high_water_cursor=3,
events=[event],
)
provider = _Provider(
[_snapshot(1)],
[batch],
[_delta(workspaces=[_workspace(revision=3, name='Latest Workspace')])],
)
service = DirectoryProjectionService(application, provider, INSTANCE_UUID)
await service.initialize()
await service.sync_once()
async with session_factory() as session:
workspace = await session.get(Workspace, WORKSPACE_UUID)
state = await session.get(DirectoryProjectionState, INSTANCE_UUID)
assert workspace.name == 'Latest Workspace'
assert state.cursor == 2
async def test_entitlement_event_advances_cursor_without_refetching_directory(projection_context):
application, session_factory = projection_context
event = DirectoryEvent(
cursor=2,
uuid='40000000-0000-0000-0000-000000000002',
aggregate_uuid=WORKSPACE_UUID,
event_type='entitlement.changed',
revision=2,
payload={'workspace_uuid': WORKSPACE_UUID, 'entitlement_revision': 2},
created_at=datetime.datetime(2026, 7, 24, 12, 2, tzinfo=datetime.UTC),
)
batch = DirectoryEventBatch(
instance_uuid=INSTANCE_UUID,
after_cursor=1,
cursor=2,
high_water_cursor=2,
events=[event],
)
provider = _Provider([_snapshot(1)], [batch])
service = DirectoryProjectionService(application, provider, INSTANCE_UUID)
await service.initialize()
await service.sync_once()
await service.apply_snapshot(_snapshot(2, workspaces=[_workspace(revision=1)]))
async with session_factory() as session:
state = await session.get(DirectoryProjectionState, INSTANCE_UUID)
inbox = await session.scalar(sqlalchemy.select(DirectoryProjectionInbox))
assert provider.snapshot_calls == 1
assert state.cursor == 2
assert inbox.event_uuid == event.uuid
assert inbox.applied_at is not None
async def test_missing_workspace_in_delta_fences_only_requested_workspace(projection_context):
application, session_factory = projection_context
event = DirectoryEvent(
cursor=2,
uuid='40000000-0000-0000-0000-000000000004',
aggregate_uuid=WORKSPACE_UUID,
event_type='directory.changed',
revision=2,
payload={'workspace_uuid': WORKSPACE_UUID, 'directory_revision': 2},
created_at=datetime.datetime(2026, 7, 24, 12, 2, tzinfo=datetime.UTC),
)
batch = DirectoryEventBatch(
instance_uuid=INSTANCE_UUID,
after_cursor=1,
cursor=2,
high_water_cursor=2,
events=[event],
)
provider = _Provider([_snapshot(1)], [batch], [_delta(workspaces=[])])
service = DirectoryProjectionService(application, provider, INSTANCE_UUID)
await service.initialize()
await service.sync_once()
async with session_factory() as session:
workspace = await session.get(Workspace, WORKSPACE_UUID)
execution = await session.get(WorkspaceExecutionState, WORKSPACE_UUID)
membership = await session.scalar(sqlalchemy.select(WorkspaceMembership))
assert workspace.status == 'archived'
assert workspace.projection_revision == 2
assert execution.state == 'inactive'
assert execution.write_fenced is True
assert execution.desired_state_revision == 2
assert membership.status == 'removed'
assert membership.projection_revision == 2
stale_same_cursor_snapshot = _snapshot(
2,
workspaces=[_workspace(revision=2, status='active')],
)
with pytest.raises(DirectoryProjectionUnavailableError, match='conflicting contents'):
await service.apply_snapshot(stale_same_cursor_snapshot)
async with session_factory() as session:
workspace = await session.get(Workspace, WORKSPACE_UUID)
execution = await session.get(WorkspaceExecutionState, WORKSPACE_UUID)
assert workspace.status == 'archived'
assert execution.write_fenced is True
async def test_each_replica_consumes_events_with_its_own_cursor(projection_context):
application, session_factory = projection_context
event = DirectoryEvent(
cursor=2,
uuid='40000000-0000-0000-0000-000000000005',
aggregate_uuid=WORKSPACE_UUID,
event_type='directory.changed',
revision=2,
payload={'workspace_uuid': WORKSPACE_UUID, 'directory_revision': 2},
created_at=datetime.datetime(2026, 7, 24, 12, 2, tzinfo=datetime.UTC),
)
batch = DirectoryEventBatch(
instance_uuid=INSTANCE_UUID,
after_cursor=1,
cursor=2,
high_water_cursor=2,
events=[event],
)
updated = _delta(workspaces=[_workspace(revision=2, name='Replica-safe Workspace')])
first_provider = _Provider([_snapshot(1)], [batch], [updated])
second_provider = _Provider([_snapshot(1)], [batch], [updated])
first = DirectoryProjectionService(application, first_provider, INSTANCE_UUID)
second = DirectoryProjectionService(application, second_provider, INSTANCE_UUID)
await first.initialize()
await second.initialize()
await first.sync_once()
await second.sync_once()
await second.sync_once()
async with session_factory() as session:
workspace = await session.get(Workspace, WORKSPACE_UUID)
state = await session.get(DirectoryProjectionState, INSTANCE_UUID)
inbox_rows = (await session.scalars(sqlalchemy.select(DirectoryProjectionInbox))).all()
assert workspace.name == 'Replica-safe Workspace'
assert state.cursor == 2
assert len(inbox_rows) == 1
assert first_provider.after_cursors == [1]
assert second_provider.after_cursors == [1, 2]
async def test_snapshot_coverage_allows_lagging_replica_to_replay_receipts(projection_context):
application, session_factory = projection_context
event_two = DirectoryEvent(
cursor=2,
uuid='40000000-0000-0000-0000-000000000008',
aggregate_uuid=WORKSPACE_UUID,
event_type='directory.changed',
revision=2,
payload={'workspace_uuid': WORKSPACE_UUID, 'directory_revision': 2},
created_at=datetime.datetime(2026, 7, 24, 12, 2, tzinfo=datetime.UTC),
)
event_three = DirectoryEvent(
cursor=3,
uuid='40000000-0000-0000-0000-000000000009',
aggregate_uuid=WORKSPACE_UUID,
event_type='directory.changed',
revision=3,
payload={'workspace_uuid': WORKSPACE_UUID, 'directory_revision': 3},
created_at=datetime.datetime(2026, 7, 24, 12, 3, tzinfo=datetime.UTC),
)
batch_two = DirectoryEventBatch(
instance_uuid=INSTANCE_UUID,
after_cursor=1,
cursor=2,
high_water_cursor=3,
events=[event_two],
)
batch_three = DirectoryEventBatch(
instance_uuid=INSTANCE_UUID,
after_cursor=2,
cursor=3,
high_water_cursor=3,
events=[event_three],
)
lagging_provider = _Provider(
[_snapshot(1)],
[batch_two, batch_three],
[
_delta(workspaces=[_workspace(revision=2, name='Intermediate Workspace')]),
_delta(workspaces=[_workspace(revision=3, name='Current Workspace')]),
],
)
lagging = DirectoryProjectionService(application, lagging_provider, INSTANCE_UUID)
leading = DirectoryProjectionService(
application,
_Provider([_snapshot(3, workspaces=[_workspace(revision=3, name='Current Workspace')])]),
INSTANCE_UUID,
)
await lagging.initialize()
await leading.initialize()
await lagging.sync_once()
await lagging.sync_once()
lagging.require_ready()
async with session_factory() as session:
workspace = await session.get(Workspace, WORKSPACE_UUID)
state = await session.get(DirectoryProjectionState, INSTANCE_UUID)
inbox_rows = (
await session.scalars(sqlalchemy.select(DirectoryProjectionInbox).order_by(DirectoryProjectionInbox.cursor))
).all()
assert workspace.name == 'Current Workspace'
assert state.cursor == 3
assert state.snapshot_coverage_cursor == 3
assert [row.cursor for row in inbox_rows] == [2, 3]
assert lagging_provider.after_cursors == [1, 2]
async def test_initialize_retries_snapshot_superseded_by_another_replica(projection_context):
application, _session_factory = projection_context
leading = DirectoryProjectionService(
application,
_Provider([_snapshot(2)]),
INSTANCE_UUID,
)
await leading.initialize()
racing_provider = _Provider([_snapshot(1), _snapshot(2)])
racing = DirectoryProjectionService(application, racing_provider, INSTANCE_UUID)
await racing.initialize()
await racing.sync_once()
racing.require_ready()
assert racing_provider.snapshot_calls == 2
assert racing_provider.after_cursors == [2]
async def test_page_below_signed_high_water_does_not_renew_readiness(projection_context):
application, _session_factory = projection_context
monotonic = [10.0]
event = DirectoryEvent(
cursor=2,
uuid='40000000-0000-0000-0000-000000000010',
aggregate_uuid=WORKSPACE_UUID,
event_type='directory.changed',
revision=2,
payload={'workspace_uuid': WORKSPACE_UUID, 'directory_revision': 2},
created_at=datetime.datetime(2026, 7, 24, 12, 2, tzinfo=datetime.UTC),
)
batch = DirectoryEventBatch(
instance_uuid=INSTANCE_UUID,
after_cursor=1,
cursor=2,
high_water_cursor=3,
events=[event],
)
service = DirectoryProjectionService(
application,
_Provider([_snapshot(1)], [batch], [_delta()]),
INSTANCE_UUID,
sync_interval_seconds=5,
max_staleness_seconds=60,
monotonic_time=lambda: monotonic[0],
)
await service.initialize()
await service.sync_once()
monotonic[0] = 70.0
with pytest.raises(DirectoryProjectionUnavailableError, match='stale'):
service.require_ready()
async def test_empty_page_cannot_hide_shared_projection_progress(projection_context):
application, _session_factory = projection_context
lagging = DirectoryProjectionService(
application,
_Provider([_snapshot(1)]),
INSTANCE_UUID,
)
leading = DirectoryProjectionService(
application,
_Provider([_snapshot(2)]),
INSTANCE_UUID,
)
await lagging.initialize()
await leading.initialize()
with pytest.raises(DirectoryProjectionUnavailableError, match='has not consumed'):
await lagging.sync_once()
async def test_event_payload_revision_mismatch_fails_closed(projection_context):
application, _session_factory = projection_context
event = DirectoryEvent(
cursor=2,
uuid='40000000-0000-0000-0000-000000000006',
aggregate_uuid=WORKSPACE_UUID,
event_type='directory.changed',
revision=2,
payload={'workspace_uuid': WORKSPACE_UUID, 'directory_revision': 3},
created_at=datetime.datetime(2026, 7, 24, 12, 2, tzinfo=datetime.UTC),
)
batch = DirectoryEventBatch(
instance_uuid=INSTANCE_UUID,
after_cursor=1,
cursor=2,
high_water_cursor=2,
events=[event],
)
service = DirectoryProjectionService(
application,
_Provider([_snapshot(1)], [batch], [_delta()]),
INSTANCE_UUID,
)
await service.initialize()
with pytest.raises(DirectoryProjectionUnavailableError, match='conflicting revision'):
await service.sync_once()
async def test_workspace_delta_older_than_event_fails_closed(projection_context):
application, _session_factory = projection_context
event = DirectoryEvent(
cursor=2,
uuid='40000000-0000-0000-0000-000000000007',
aggregate_uuid=WORKSPACE_UUID,
event_type='directory.changed',
revision=3,
payload={'workspace_uuid': WORKSPACE_UUID, 'directory_revision': 3},
created_at=datetime.datetime(2026, 7, 24, 12, 2, tzinfo=datetime.UTC),
)
batch = DirectoryEventBatch(
instance_uuid=INSTANCE_UUID,
after_cursor=1,
cursor=2,
high_water_cursor=2,
events=[event],
)
service = DirectoryProjectionService(
application,
_Provider([_snapshot(1)], [batch], [_delta(workspaces=[_workspace(revision=2)])]),
INSTANCE_UUID,
)
await service.initialize()
with pytest.raises(DirectoryProjectionUnavailableError, match='older'):
await service.sync_once()
async def test_stale_workspace_tombstone_revision_fails_closed(projection_context):
application, session_factory = projection_context
event = DirectoryEvent(
cursor=2,
uuid='40000000-0000-0000-0000-000000000011',
aggregate_uuid=WORKSPACE_UUID,
event_type='directory.changed',
revision=2,
payload={'workspace_uuid': WORKSPACE_UUID, 'directory_revision': 2},
created_at=datetime.datetime(2026, 7, 24, 12, 2, tzinfo=datetime.UTC),
)
batch = DirectoryEventBatch(
instance_uuid=INSTANCE_UUID,
after_cursor=1,
cursor=2,
high_water_cursor=2,
events=[event],
)
service = DirectoryProjectionService(
application,
_Provider(
[_snapshot(1, workspaces=[_workspace(revision=3)])],
[batch],
[_delta(workspaces=[])],
),
INSTANCE_UUID,
)
await service.initialize()
with pytest.raises(DirectoryProjectionUnavailableError, match='tombstone revision rolled back'):
await service.sync_once()
async with session_factory() as session:
workspace = await session.get(Workspace, WORKSPACE_UUID)
membership = await session.scalar(sqlalchemy.select(WorkspaceMembership))
execution = await session.get(WorkspaceExecutionState, WORKSPACE_UUID)
assert workspace.status == 'active'
assert membership.status == 'active'
assert execution.write_fenced is False
async def test_conflicting_account_projection_across_workspaces_fails_closed(projection_context):
application, _session_factory = projection_context
conflicting_member = DirectoryMember(
membership_uuid=SECOND_MEMBERSHIP_UUID,
account_uuid=ACCOUNT_UUID,
normalized_email='owner@example.com',
display_name='Workspace Owner',
account_status='blocked',
role='member',
membership_status='active',
projection_revision=99,
joined_at=datetime.datetime(2026, 7, 24, tzinfo=datetime.UTC),
)
second_workspace = DirectoryWorkspace(
uuid=SECOND_WORKSPACE_UUID,
name='Second Workspace',
slug='second-workspace',
type='team',
status='active',
created_by_account_uuid=ACCOUNT_UUID,
projection_revision=99,
execution_generation=1,
members=[conflicting_member],
)
snapshot = _snapshot(
1,
workspaces=[_workspace(revision=1), second_workspace],
)
service = DirectoryProjectionService(application, _Provider([snapshot]), INSTANCE_UUID)
with pytest.raises(DirectoryProjectionUnavailableError, match='conflicting account projections'):
await service.initialize()
async def test_account_projection_revision_advances_when_account_fields_change(projection_context):
application, session_factory = projection_context
service = DirectoryProjectionService(
application,
_Provider([_snapshot(1)]),
INSTANCE_UUID,
)
await service.initialize()
blocked_member = _member(revision=2).model_copy(update={'account_status': 'blocked'})
blocked_workspace = _workspace(revision=2).model_copy(update={'members': (blocked_member,)})
await service.apply_snapshot(_snapshot(2, workspaces=[blocked_workspace]))
async with session_factory() as session:
account = await session.scalar(sqlalchemy.select(User).where(User.uuid == ACCOUNT_UUID))
assert account.status == 'disabled'
assert account.projection_revision == 2
async def test_readiness_expires_on_monotonic_staleness(projection_context):
application, _session_factory = projection_context
monotonic = [10.0]
service = DirectoryProjectionService(
application,
_Provider([_snapshot(1)]),
INSTANCE_UUID,
sync_interval_seconds=5,
max_staleness_seconds=60,
monotonic_time=lambda: monotonic[0],
)
await service.initialize()
service.require_ready()
monotonic[0] = 70.0
with pytest.raises(DirectoryProjectionUnavailableError, match='stale'):
service.require_ready()
async def test_snapshot_for_another_instance_is_rejected(projection_context):
application, _session_factory = projection_context
other = _snapshot(1).model_copy(update={'instance_uuid': 'other-instance'})
service = DirectoryProjectionService(application, _Provider([other]), INSTANCE_UUID)
with pytest.raises(DirectoryProjectionUnavailableError, match='another LangBot instance'):
await service.initialize()
tests/unit_tests/persistence/test_tenant_uow.py
+53
View File
@@ -2,6 +2,7 @@ from __future__ import annotations
import asyncio
import contextvars
import datetime
from types import SimpleNamespace
import pytest
@@ -15,6 +16,10 @@ from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, registry, rel
from sqlalchemy.sql import quoted_name
from langbot.pkg.core.task_boundary import create_detached_task
from langbot.pkg.entity.persistence.cloud_directory import (
DirectoryProjectionInbox,
DirectoryProjectionState,
)
from langbot.pkg.entity.persistence.user import User
from langbot.pkg.persistence.mgr import PersistenceManager, PersistenceMode
from langbot.pkg.persistence.tenant_uow import (
@@ -159,6 +164,54 @@ async def test_manager_reuses_one_session_and_rejects_cross_workspace(tmp_path)
await engine.dispose()
async def test_directory_projection_uow_is_instance_scoped_and_not_workspace_nestable(tmp_path) -> None:
engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "directory-projection-uow.db"}')
manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME)
manager.db = SimpleNamespace(get_engine=lambda: engine)
fingerprint = 'a' * 64
try:
async with engine.begin() as conn:
await conn.run_sync(DirectoryProjectionState.__table__.create)
await conn.run_sync(DirectoryProjectionInbox.__table__.create)
async with manager.directory_projection_uow('instance-a') as uow:
assert manager.current_scope() is not None
assert manager.current_scope().kind is PersistenceScopeKind.DIRECTORY_PROJECTION
assert manager.current_scope().settings == (('langbot.directory_instance_uuid', 'instance-a'),)
manager.require_current_session(PersistenceScopeKind.DIRECTORY_PROJECTION)
uow.session.add(
DirectoryProjectionState(
instance_uuid='instance-a',
cursor=1,
snapshot_fingerprint=fingerprint,
last_applied_at=datetime.datetime.now(datetime.UTC),
)
)
uow.session.add(
DirectoryProjectionInbox(
instance_uuid='instance-a',
event_uuid='00000000-0000-0000-0000-000000000001',
cursor=1,
event_type='workspace.changed',
revision=1,
fingerprint=fingerprint,
)
)
with pytest.raises(CrossScopeTransactionError, match='while directory_projection scope is active'):
async with manager.tenant_uow('workspace-a'):
pass
async with manager.tenant_scope('workspace-a'):
with pytest.raises(CrossScopeTransactionError, match='while workspace scope is active'):
async with manager.directory_projection_uow('instance-a'):
pass
with pytest.raises(ValueError, match='must not be empty'):
manager.directory_projection_uow(' ')
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)
tests/unit_tests/plugin/test_connector_ping.py
+152 -1
View File
@@ -7,6 +7,7 @@ from unittest.mock import AsyncMock, Mock
import pytest
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.plugin import connector as connector_module
from langbot.pkg.plugin.connector import PluginRuntimeConnector, PluginRuntimeNotConnectedError
from langbot_plugin.runtime.security import (
PLUGIN_RUNTIME_CONTROL_TOKEN_ENV,
@@ -17,7 +18,22 @@ from langbot_plugin.runtime.security import (
def make_connector() -> PluginRuntimeConnector:
app = SimpleNamespace(
logger=Mock(),
instance_config=SimpleNamespace(data={'plugin': {'enable': True}, 'space': {'url': ''}}),
instance_config=SimpleNamespace(
data={
'plugin': {
'enable': True,
'worker': {
'max_cpus': 1.0,
'max_memory_mb': 512,
'max_pids': 128,
'max_open_files': 256,
'max_file_size_mb': 512,
'require_hard_limits': False,
},
},
'space': {'url': ''},
}
),
)
return PluginRuntimeConnector(app, AsyncMock())
@@ -41,6 +57,141 @@ async def test_ping_plugin_runtime_delegates_to_connected_handler():
connector.handler.ping.assert_awaited_once()
@pytest.mark.asyncio
async def test_stop_transport_tolerates_handler_callback_removing_attribute():
connector = make_connector()
class Handler:
async def close(self):
del connector.handler
connector.handler = Handler()
await connector._stop_transport()
assert not hasattr(connector, 'handler')
@pytest.mark.asyncio
async def test_stdio_runtime_connection_does_not_capture_unconsumed_stderr(
monkeypatch: pytest.MonkeyPatch,
):
connector = make_connector()
connector._prepare_connected_runtime = AsyncMock()
created = {}
monkeypatch.setattr(connector_module.constants, 'instance_id', 'instance-a')
class FakeRuntimeHandler:
def __init__(self, connection, disconnect_callback, ap):
self.release = asyncio.Event()
async def ping(self):
return None
async def set_runtime_config(self, **kwargs):
return None
async def run(self):
await self.release.wait()
async def close(self):
self.release.set()
class FakeController:
def __init__(self, **kwargs):
created.update(kwargs)
async def run(self, callback):
await callback(object())
async def close(self):
return None
monkeypatch.setattr(connector_module.platform, 'get_platform', lambda: 'linux')
monkeypatch.setattr(
connector_module.platform,
'use_websocket_to_connect_plugin_runtime',
lambda: False,
)
monkeypatch.setattr(
connector_module.handler,
'RuntimeConnectionHandler',
FakeRuntimeHandler,
)
monkeypatch.setattr(
connector_module.stdio_client_controller,
'StdioClientController',
FakeController,
)
await connector.initialize()
assert created['capture_stderr'] is False
assert connector._connected.is_set()
connector._prepare_connected_runtime.assert_awaited_once()
await connector.aclose()
@pytest.mark.asyncio
async def test_runtime_disconnect_notifies_once_and_clears_handler(
monkeypatch: pytest.MonkeyPatch,
):
disconnect = AsyncMock()
connector = PluginRuntimeConnector(make_connector().ap, disconnect)
connector._prepare_connected_runtime = AsyncMock()
monkeypatch.setattr(connector_module.constants, 'instance_id', 'instance-a')
class FakeRuntimeHandler:
def __init__(self, connection, disconnect_callback, ap):
self.disconnect_callback = disconnect_callback
async def ping(self):
return None
async def set_runtime_config(self, **kwargs):
return None
async def run(self):
await self.disconnect_callback(self)
async def close(self):
return None
class FakeController:
def __init__(self, **kwargs):
pass
async def run(self, callback):
await callback(object())
async def close(self):
return None
monkeypatch.setattr(connector_module.platform, 'get_platform', lambda: 'linux')
monkeypatch.setattr(
connector_module.platform,
'use_websocket_to_connect_plugin_runtime',
lambda: False,
)
monkeypatch.setattr(
connector_module.handler,
'RuntimeConnectionHandler',
FakeRuntimeHandler,
)
monkeypatch.setattr(
connector_module.stdio_client_controller,
'StdioClientController',
FakeController,
)
await connector.initialize()
await asyncio.sleep(0)
disconnect.assert_awaited_once_with(connector)
assert not hasattr(connector, 'handler')
await connector.aclose()
@pytest.mark.asyncio
async def test_disabled_connector_validates_workspace_without_runtime_handler():
app = SimpleNamespace(
tests/unit_tests/provider/test_localagent_no_duplicate.py
+13 -1
View File
@@ -10,6 +10,7 @@ import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
import langbot_plugin.api.entities.builtin.provider.message as provider_message
import langbot_plugin.api.entities.builtin.provider.session as provider_session
from langbot.pkg.api.http.context import ExecutionContext, PrincipalContext, PrincipalType
from langbot.pkg.provider.runners.localagent import LocalAgentRunner
@@ -97,7 +98,7 @@ def make_query() -> pipeline_query.Query:
adapter = AsyncMock()
adapter.is_stream_output_supported = AsyncMock(return_value=False)
return pipeline_query.Query.model_construct(
query = pipeline_query.Query.model_construct(
query_id='no-dup-query',
launcher_type=provider_session.LauncherTypes.PERSON,
launcher_id=12345,
@@ -124,6 +125,17 @@ def make_query() -> pipeline_query.Query:
use_llm_model_uuid='test-model-uuid',
variables={},
)
object.__setattr__(
query,
'_execution_context',
ExecutionContext(
instance_uuid='instance-test',
workspace_uuid='workspace-test',
placement_generation=1,
trigger_principal=PrincipalContext(PrincipalType.SYSTEM),
),
)
return query
def _make_app(provider) -> SimpleNamespace:
tests/unit_tests/provider/test_mcp_box_integration.py
+3 -4
View File
@@ -829,7 +829,7 @@ class TestGetRuntimeInfoDict:
assert ap.box_service.available is True
@pytest.mark.asyncio
async def test_enabled_box_timeout_does_not_exhaust_mcp_retry_budget(self, mcp_module, monkeypatch):
async def test_enabled_box_timeout_does_not_exhaust_mcp_retry_budget(self, mcp_module):
ap = _make_ap()
ap.box_service.available = False
ap.box_service.enabled = True
@@ -854,14 +854,13 @@ class TestGetRuntimeInfoDict:
raise RuntimeError('Box runtime is not available after 1 seconds')
session._lifecycle_loop = lifecycle
sleep = AsyncMock()
monkeypatch.setattr(mcp_module.asyncio, 'sleep', sleep)
session._sleep_with_execution_fence = AsyncMock()
await session._lifecycle_loop_with_retry()
assert attempts == 2
assert session.retry_count == 0
sleep.assert_awaited_once_with(1)
session._sleep_with_execution_fence.assert_awaited_once_with(1)
@pytest.mark.asyncio
async def test_disabled_box_still_stops_mcp_retry_loop(self, mcp_module):
tests/unit_tests/provider/test_mcp_remote_transport.py
+16 -2
View File
@@ -13,7 +13,7 @@ from aiohttp import web
from mcp import types as mcp_types
from langbot.pkg.api.http.context import ExecutionContext
from langbot.pkg.provider.tools.loaders.mcp import RuntimeMCPSession
from langbot.pkg.provider.tools.loaders.mcp import MCPToolCallTimeoutError, RuntimeMCPSession
TEST_EXECUTION_CONTEXT = ExecutionContext(
@@ -159,7 +159,21 @@ def _session(
timeout: float = 2,
tool_call_timeout_sec: float = 300,
) -> RuntimeMCPSession:
app = cast(Any, SimpleNamespace(logger=Mock()))
app = cast(
Any,
SimpleNamespace(
logger=Mock(),
workspace_service=SimpleNamespace(
get_execution_binding=AsyncMock(
return_value=SimpleNamespace(
instance_uuid=TEST_EXECUTION_CONTEXT.instance_uuid,
workspace_uuid=TEST_EXECUTION_CONTEXT.workspace_uuid,
placement_generation=TEST_EXECUTION_CONTEXT.placement_generation,
)
)
),
),
)
return RuntimeMCPSession(
'remote-transport-test',
{
tests/unit_tests/provider/test_mcp_resources.py
+49 -1
View File
@@ -124,6 +124,7 @@ async def test_invoke_mcp_tool_uses_configurable_request_timeout():
},
True,
_app(),
TEST_EXECUTION_CONTEXT,
)
session.session = SimpleNamespace(call_tool=AsyncMock(return_value=_tool_result()))
@@ -148,6 +149,7 @@ async def test_invoke_mcp_tool_zero_timeout_disables_request_deadline():
},
True,
_app(),
TEST_EXECUTION_CONTEXT,
)
session.session = SimpleNamespace(call_tool=AsyncMock(return_value=_tool_result()))
@@ -171,6 +173,7 @@ async def test_invoke_mcp_tool_timeout_is_not_retried_and_session_remains_usable
},
True,
_app(),
TEST_EXECUTION_CONTEXT,
)
timeout = McpError(
mcp_types.ErrorData(
@@ -205,6 +208,7 @@ def test_invalid_tool_call_timeout_falls_back_to_default(invalid_timeout):
},
True,
ap,
TEST_EXECUTION_CONTEXT,
)
assert session.tool_call_timeout_sec == MCP_TOOL_CALL_TIMEOUT_DEFAULT_SECONDS
@@ -525,7 +529,11 @@ async def test_mcp_tool_result_is_discarded_when_generation_changes_during_call(
with pytest.raises(WorkspaceGenerationMismatchError):
await session.invoke_mcp_tool('side_effecting_tool', {})
session.session.call_tool.assert_awaited_once_with('side_effecting_tool', {})
session.session.call_tool.assert_awaited_once_with(
'side_effecting_tool',
{},
read_timeout_seconds=timedelta(seconds=MCP_TOOL_CALL_TIMEOUT_DEFAULT_SECONDS),
)
@pytest.mark.asyncio
@@ -567,3 +575,43 @@ async def test_mcp_idle_lifecycle_stops_without_retry_after_generation_bump():
assert session.status == MCPSessionStatus.ERROR
assert session.error_message == 'Workspace execution binding is stale'
assert session._shutdown_event.is_set()
@pytest.mark.asyncio
async def test_mcp_loader_shutdown_cancels_startup_tasks_and_closes_sessions_concurrently():
loader = MCPLoader(_app())
hosted_cancelled = asyncio.Event()
async def pending_host():
try:
await asyncio.Event().wait()
finally:
hosted_cancelled.set()
hosted_task = asyncio.create_task(pending_host())
await asyncio.sleep(0)
loader._hosted_mcp_tasks = [hosted_task]
started: set[str] = set()
all_started = asyncio.Event()
class Session:
def __init__(self, name: str):
self.name = name
self.server_name = name
async def shutdown(self):
started.add(self.name)
if len(started) == 2:
all_started.set()
await all_started.wait()
loader.sessions = {'one': Session('one'), 'two': Session('two')}
await asyncio.wait_for(loader.shutdown(), timeout=1)
assert hosted_cancelled.is_set()
assert hosted_task.cancelled()
assert started == {'one', 'two'}
assert loader._hosted_mcp_tasks == []
assert loader.sessions == {}
tests/unit_tests/provider/test_skill_tools.py
+1 -1
View File
@@ -508,7 +508,7 @@ class TestSkillToolLoader:
ap.skill_mgr = SimpleNamespace(skills={'demo': _make_skill_data(name='demo')})
ap.box_service = SimpleNamespace(
available=False,
get_status=AsyncMock(return_value={'backend': {'available': True}}),
get_backend_status=AsyncMock(return_value={'backend': {'available': True}}),
)
loader = SkillToolLoader(ap)
tests/unit_tests/provider/test_tool_manager_native.py
+15
View File
@@ -206,6 +206,21 @@ async def test_native_tool_loader_exposes_all_tools_when_box_available():
assert await loader.has_tool(tool_name) is True
@pytest.mark.asyncio
async def test_native_tool_loader_refreshes_after_box_recovers():
box_service = SimpleNamespace(
available=False,
get_backend_status=AsyncMock(return_value={'backend': {'available': True}}),
)
loader = NativeToolLoader(SimpleNamespace(box_service=box_service, logger=Mock()))
await loader.initialize()
assert await loader.get_tools() == []
box_service.available = True
assert [tool.name for tool in await loader.get_tools()] == ['exec', 'read', 'write', 'edit', 'glob', 'grep']
@pytest.mark.asyncio
async def test_native_tool_loader_rechecks_admission_at_the_final_invoke_boundary():
box_service = SimpleNamespace(
uv.lock
Generated
+3 -3
View File
@@ -2115,7 +2115,7 @@ requires-dist = [
{ name = "ebooklib", specifier = ">=0.18" },
{ name = "gewechat-client", specifier = ">=0.1.5" },
{ name = "html2text", specifier = ">=2024.2.26" },
{ name = "langbot-plugin", git = "https://github.com/langbot-app/langbot-plugin-sdk.git?rev=95a1805af2038c745de2c018a00db1305089a32e" },
{ name = "langbot-plugin", git = "https://github.com/langbot-app/langbot-plugin-sdk.git?rev=e7d946af4a6b1494fbe74627c1815ace19ac8991" },
{ name = "langchain", specifier = ">=1.3.9" },
{ name = "langchain-core", specifier = ">=1.3.3" },
{ name = "langchain-text-splitters", specifier = ">=1.1.2" },
@@ -2180,8 +2180,8 @@ dev = [
[[package]]
name = "langbot-plugin"
version = "0.4.15"
source = { git = "https://github.com/langbot-app/langbot-plugin-sdk.git?rev=95a1805af2038c745de2c018a00db1305089a32e#95a1805af2038c745de2c018a00db1305089a32e" }
version = "0.4.18"
source = { git = "https://github.com/langbot-app/langbot-plugin-sdk.git?rev=e7d946af4a6b1494fbe74627c1815ace19ac8991#e7d946af4a6b1494fbe74627c1815ace19ac8991" }
dependencies = [
{ name = "aiofiles" },
{ name = "aiohttp" },