From 98f45aa88efb307315053987a1fa67520c2ff30d Mon Sep 17 00:00:00 2001 From: Junyan Qin Date: Fri, 24 Jul 2026 19:11:33 +0800 Subject: [PATCH] feat(tenancy): connect cloud workspace control plane --- docs/multi-tenant/implementation-checklist.md | 18 +- docs/multi-tenant/implementation-decisions.md | 38 +- .../pending-architecture-decisions.md | 25 +- docs/multi-tenant/verification-report.md | 403 ++++---- .../workspace-multi-user-architecture.md | 17 +- pyproject.toml | 2 +- src/langbot/pkg/api/http/service/user.py | 62 +- src/langbot/pkg/box/connector.py | 1 + src/langbot/pkg/box/service.py | 53 +- src/langbot/pkg/cloud/__init__.py | 24 + src/langbot/pkg/cloud/bootstrap.py | 84 +- src/langbot/pkg/cloud/directory.py | 215 +++++ src/langbot/pkg/cloud/directory_projection.py | 881 ++++++++++++++++++ src/langbot/pkg/core/app.py | 22 + src/langbot/pkg/core/stages/build_app.py | 19 + .../pkg/entity/persistence/cloud_directory.py | 69 ++ .../0014_cloud_directory_projection.py | 268 ++++++ src/langbot/pkg/persistence/mgr.py | 94 +- src/langbot/pkg/persistence/tenant_uow.py | 39 +- src/langbot/pkg/plugin/connector.py | 464 +++++---- .../pkg/provider/runners/localagent.py | 5 +- src/langbot/pkg/provider/tools/loaders/mcp.py | 53 +- src/langbot/pkg/workspace/collaboration.py | 23 +- src/langbot/pkg/workspace/service.py | 14 + .../persistence/test_migrations.py | 19 + .../persistence/test_migrations_postgres.py | 211 ++++- .../api/service/test_user_service.py | 63 +- tests/unit_tests/box/test_box_connector.py | 138 ++- tests/unit_tests/box/test_box_service.py | 91 ++ tests/unit_tests/cloud/test_bootstrap.py | 61 ++ .../cloud/test_directory_projection.py | 817 ++++++++++++++++ .../unit_tests/persistence/test_tenant_uow.py | 53 ++ .../unit_tests/plugin/test_connector_ping.py | 153 ++- .../provider/test_localagent_no_duplicate.py | 14 +- .../provider/test_mcp_box_integration.py | 7 +- .../provider/test_mcp_remote_transport.py | 18 +- .../unit_tests/provider/test_mcp_resources.py | 50 +- tests/unit_tests/provider/test_skill_tools.py | 2 +- .../provider/test_tool_manager_native.py | 15 + uv.lock | 6 +- 40 files changed, 4159 insertions(+), 452 deletions(-) create mode 100644 src/langbot/pkg/cloud/directory.py create mode 100644 src/langbot/pkg/cloud/directory_projection.py create mode 100644 src/langbot/pkg/entity/persistence/cloud_directory.py create mode 100644 src/langbot/pkg/persistence/alembic/versions/0014_cloud_directory_projection.py create mode 100644 tests/unit_tests/cloud/test_directory_projection.py diff --git a/docs/multi-tenant/implementation-checklist.md b/docs/multi-tenant/implementation-checklist.md index 6f86ed91b..c74c3dd57 100644 --- a/docs/multi-tenant/implementation-checklist.md +++ b/docs/multi-tenant/implementation-checklist.md @@ -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. diff --git a/docs/multi-tenant/implementation-decisions.md b/docs/multi-tenant/implementation-decisions.md index ff98eac12..3a1a1a3d5 100644 --- a/docs/multi-tenant/implementation-decisions.md +++ b/docs/multi-tenant/implementation-decisions.md @@ -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. diff --git a/docs/multi-tenant/pending-architecture-decisions.md b/docs/multi-tenant/pending-architecture-decisions.md index b47b46e9e..cf4c895f4 100644 --- a/docs/multi-tenant/pending-architecture-decisions.md +++ b/docs/multi-tenant/pending-architecture-decisions.md @@ -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. 本轮明确不做的事情 diff --git a/docs/multi-tenant/verification-report.md b/docs/multi-tenant/verification-report.md index b22974010..02f1b76d3 100644 --- a/docs/multi-tenant/verification-report.md +++ b/docs/multi-tenant/verification-report.md @@ -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. diff --git a/docs/multi-tenant/workspace-multi-user-architecture.md b/docs/multi-tenant/workspace-multi-user-architecture.md index ff50def2e..b01388139 100644 --- a/docs/multi-tenant/workspace-multi-user-architecture.md +++ b/docs/multi-tenant/workspace-multi-user-architecture.md @@ -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 并可安全重试的集成证据。 diff --git a/pyproject.toml b/pyproject.toml index 30ee73b2f..8bf3f4569 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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", diff --git a/src/langbot/pkg/api/http/service/user.py b/src/langbot/pkg/api/http/service/user.py index 516dd60eb..e5f1d5d1f 100644 --- a/src/langbot/pkg/api/http/service/user.py +++ b/src/langbot/pkg/api/http/service/user.py @@ -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]: diff --git a/src/langbot/pkg/box/connector.py b/src/langbot/pkg/box/connector.py index fa329ce03..2ef990d0c 100644 --- a/src/langbot/pkg/box/connector.py +++ b/src/langbot/pkg/box/connector.py @@ -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)) ) diff --git a/src/langbot/pkg/box/service.py b/src/langbot/pkg/box/service.py index 92afeb610..6d803b4d5 100644 --- a/src/langbot/pkg/box/service.py +++ b/src/langbot/pkg/box/service.py @@ -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) diff --git a/src/langbot/pkg/cloud/__init__.py b/src/langbot/pkg/cloud/__init__.py index 41e4e4877..ada0c48cf 100644 --- a/src/langbot/pkg/cloud/__init__.py +++ b/src/langbot/pkg/cloud/__init__.py @@ -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', diff --git a/src/langbot/pkg/cloud/bootstrap.py b/src/langbot/pkg/cloud/bootstrap.py index 9655bf680..01c2382b7 100644 --- a/src/langbot/pkg/cloud/bootstrap.py +++ b/src/langbot/pkg/cloud/bootstrap.py @@ -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, *, diff --git a/src/langbot/pkg/cloud/directory.py b/src/langbot/pkg/cloud/directory.py new file mode 100644 index 000000000..ae9f03ee1 --- /dev/null +++ b/src/langbot/pkg/cloud/directory.py @@ -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.""" diff --git a/src/langbot/pkg/cloud/directory_projection.py b/src/langbot/pkg/cloud/directory_projection.py new file mode 100644 index 000000000..ca4388afd --- /dev/null +++ b/src/langbot/pkg/cloud/directory_projection.py @@ -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) diff --git a/src/langbot/pkg/core/app.py b/src/langbot/pkg/core/app.py index bd5a3098c..8cc97f537 100644 --- a/src/langbot/pkg/core/app.py +++ b/src/langbot/pkg/core/app.py @@ -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()] diff --git a/src/langbot/pkg/core/stages/build_app.py b/src/langbot/pkg/core/stages/build_app.py index 76a073ed6..43881bb83 100644 --- a/src/langbot/pkg/core/stages/build_app.py +++ b/src/langbot/pkg/core/stages/build_app.py @@ -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, diff --git a/src/langbot/pkg/entity/persistence/cloud_directory.py b/src/langbot/pkg/entity/persistence/cloud_directory.py new file mode 100644 index 000000000..11e41cdfb --- /dev/null +++ b/src/langbot/pkg/entity/persistence/cloud_directory.py @@ -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', + ), + ) diff --git a/src/langbot/pkg/persistence/alembic/versions/0014_cloud_directory_projection.py b/src/langbot/pkg/persistence/alembic/versions/0014_cloud_directory_projection.py new file mode 100644 index 000000000..86b95bb49 --- /dev/null +++ b/src/langbot/pkg/persistence/alembic/versions/0014_cloud_directory_projection.py @@ -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) diff --git a/src/langbot/pkg/persistence/mgr.py b/src/langbot/pkg/persistence/mgr.py index e075e9cc8..3c11935de 100644 --- a/src/langbot/pkg/persistence/mgr.py +++ b/src/langbot/pkg/persistence/mgr.py @@ -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)) diff --git a/src/langbot/pkg/persistence/tenant_uow.py b/src/langbot/pkg/persistence/tenant_uow.py index 428186d8b..9f56f2aff 100644 --- a/src/langbot/pkg/persistence/tenant_uow.py +++ b/src/langbot/pkg/persistence/tenant_uow.py @@ -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( diff --git a/src/langbot/pkg/plugin/connector.py b/src/langbot/pkg/plugin/connector.py index 129d181e6..8b65d20c6 100644 --- a/src/langbot/pkg/plugin/connector.py +++ b/src/langbot/pkg/plugin/connector.py @@ -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) diff --git a/src/langbot/pkg/provider/runners/localagent.py b/src/langbot/pkg/provider/runners/localagent.py index 4fa90a3d4..dc8f82dcb 100644 --- a/src/langbot/pkg/provider/runners/localagent.py +++ b/src/langbot/pkg/provider/runners/localagent.py @@ -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) diff --git a/src/langbot/pkg/provider/tools/loaders/mcp.py b/src/langbot/pkg/provider/tools/loaders/mcp.py index 581bf0d6b..69dafb4b0 100644 --- a/src/langbot/pkg/provider/tools/loaders/mcp.py +++ b/src/langbot/pkg/provider/tools/loaders/mcp.py @@ -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') diff --git a/src/langbot/pkg/workspace/collaboration.py b/src/langbot/pkg/workspace/collaboration.py index de63a769f..3742ca4d6 100644 --- a/src/langbot/pkg/workspace/collaboration.py +++ b/src/langbot/pkg/workspace/collaboration.py @@ -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 diff --git a/src/langbot/pkg/workspace/service.py b/src/langbot/pkg/workspace/service.py index 549f86e83..963d26f30 100644 --- a/src/langbot/pkg/workspace/service.py +++ b/src/langbot/pkg/workspace/service.py @@ -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 diff --git a/tests/integration/persistence/test_migrations.py b/tests/integration/persistence/test_migrations.py index 62813471f..c1a455fb1 100644 --- a/tests/integration/persistence/test_migrations.py +++ b/tests/integration/persistence/test_migrations.py @@ -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.""" diff --git a/tests/integration/persistence/test_migrations_postgres.py b/tests/integration/persistence/test_migrations_postgres.py index c45086ace..87caee109 100644 --- a/tests/integration/persistence/test_migrations_postgres.py +++ b/tests/integration/persistence/test_migrations_postgres.py @@ -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( diff --git a/tests/unit_tests/api/service/test_user_service.py b/tests/unit_tests/api/service/test_user_service.py index 8a580a41d..eb1516520 100644 --- a/tests/unit_tests/api/service/test_user_service.py +++ b/tests/unit_tests/api/service/test_user_service.py @@ -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 diff --git a/tests/unit_tests/box/test_box_connector.py b/tests/unit_tests/box/test_box_connector.py index 9fc5452ff..367b2e362 100644 --- a/tests/unit_tests/box/test_box_connector.py +++ b/tests/unit_tests/box/test_box_connector.py @@ -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() diff --git a/tests/unit_tests/box/test_box_service.py b/tests/unit_tests/box/test_box_service.py index 66334c662..e299e3aa9 100644 --- a/tests/unit_tests/box/test_box_service.py +++ b/tests/unit_tests/box/test_box_service.py @@ -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. diff --git a/tests/unit_tests/cloud/test_bootstrap.py b/tests/unit_tests/cloud/test_bootstrap.py index ec2caeeb3..e28869b14 100644 --- a/tests/unit_tests/cloud/test_bootstrap.py +++ b/tests/unit_tests/cloud/test_bootstrap.py @@ -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 diff --git a/tests/unit_tests/cloud/test_directory_projection.py b/tests/unit_tests/cloud/test_directory_projection.py new file mode 100644 index 000000000..0e24490e1 --- /dev/null +++ b/tests/unit_tests/cloud/test_directory_projection.py @@ -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() diff --git a/tests/unit_tests/persistence/test_tenant_uow.py b/tests/unit_tests/persistence/test_tenant_uow.py index a4dd8c76a..75634b968 100644 --- a/tests/unit_tests/persistence/test_tenant_uow.py +++ b/tests/unit_tests/persistence/test_tenant_uow.py @@ -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) diff --git a/tests/unit_tests/plugin/test_connector_ping.py b/tests/unit_tests/plugin/test_connector_ping.py index e14d57ba9..cec36884d 100644 --- a/tests/unit_tests/plugin/test_connector_ping.py +++ b/tests/unit_tests/plugin/test_connector_ping.py @@ -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( diff --git a/tests/unit_tests/provider/test_localagent_no_duplicate.py b/tests/unit_tests/provider/test_localagent_no_duplicate.py index f047dfcbc..7e75b9f42 100644 --- a/tests/unit_tests/provider/test_localagent_no_duplicate.py +++ b/tests/unit_tests/provider/test_localagent_no_duplicate.py @@ -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: diff --git a/tests/unit_tests/provider/test_mcp_box_integration.py b/tests/unit_tests/provider/test_mcp_box_integration.py index a56b52bad..e8f5cd7fb 100644 --- a/tests/unit_tests/provider/test_mcp_box_integration.py +++ b/tests/unit_tests/provider/test_mcp_box_integration.py @@ -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): diff --git a/tests/unit_tests/provider/test_mcp_remote_transport.py b/tests/unit_tests/provider/test_mcp_remote_transport.py index 7dc9f9046..7abe8afea 100644 --- a/tests/unit_tests/provider/test_mcp_remote_transport.py +++ b/tests/unit_tests/provider/test_mcp_remote_transport.py @@ -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', { diff --git a/tests/unit_tests/provider/test_mcp_resources.py b/tests/unit_tests/provider/test_mcp_resources.py index e1b28afe0..afaabd7d6 100644 --- a/tests/unit_tests/provider/test_mcp_resources.py +++ b/tests/unit_tests/provider/test_mcp_resources.py @@ -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 == {} diff --git a/tests/unit_tests/provider/test_skill_tools.py b/tests/unit_tests/provider/test_skill_tools.py index 4711a4c24..e96156496 100644 --- a/tests/unit_tests/provider/test_skill_tools.py +++ b/tests/unit_tests/provider/test_skill_tools.py @@ -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) diff --git a/tests/unit_tests/provider/test_tool_manager_native.py b/tests/unit_tests/provider/test_tool_manager_native.py index 83c24cdf4..8ef217a1e 100644 --- a/tests/unit_tests/provider/test_tool_manager_native.py +++ b/tests/unit_tests/provider/test_tool_manager_native.py @@ -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( diff --git a/uv.lock b/uv.lock index f98708414..4f855da0f 100644 --- a/uv.lock +++ b/uv.lock @@ -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" },