diff --git a/docs/multi-tenant/cloud-runtime-soak-gate.md b/docs/multi-tenant/cloud-runtime-soak-gate.md index 7118ffa07..35151d1bb 100644 --- a/docs/multi-tenant/cloud-runtime-soak-gate.md +++ b/docs/multi-tenant/cloud-runtime-soak-gate.md @@ -81,6 +81,9 @@ uv run python scripts/cloud_runtime_soak.py \ - 健康接口缺少 event-loop monitor、monitor 未持续运行，或其 sample counter 回退。 - blocking executor rejection counter 增长。 - Plugin Runtime restart circuit 的累计打开次数增长。 +- Core 目录 active Workspace、最近 snapshot/delta Workspace 或 membership 基数 + 超过各自配置上限，或 PostgreSQL `checked_out` 超过配置 pool 容量；相关 current/max + 指标只出现一半或 max 非法也失败。 负载结束后的冷却尾段还必须满足： diff --git a/docs/multi-tenant/cloud-v2-pending-verification.md b/docs/multi-tenant/cloud-v2-pending-verification.md index df768a301..ec5a25302 100644 --- a/docs/multi-tenant/cloud-v2-pending-verification.md +++ b/docs/multi-tenant/cloud-v2-pending-verification.md @@ -16,8 +16,9 @@ ## 1. 当前已形成的交付基线 -- LangBot Core 全量 `2811 passed, 33 skipped`，Plugin SDK 全量 - `1317 passed`；两仓 Ruff、`git diff --check` 已通过。 +- LangBot Core 全量 `2833 passed, 33 skipped`，Plugin SDK 全量 + `1325 passed`，闭源适配器 `40 passed`，Space Go 全量测试通过；三仓格式、 + 静态检查和 `git diff --check` 已通过。 - Plugin Runtime 和 Box Runtime 的公开健康接口、event-loop lag 与有界 blocking executor 指标已经过真实进程短时验证。 - 仓库 Dockerfile 构建的 Linux/cgroup v2 短时探针已证明 CPU、memory、 @@ -25,12 +26,29 @@ - PostgreSQL 16 + RLS 的 1,000 Workspace 真实启动测试，以及 5,000 Workspace 三代替换合成探针已通过。 - Core 已精确钉住 Plugin SDK 提交 - `c7893f8f94e83cb09ebcc98e25490fa699b0fb69`。最终验证必须使用包含该提交的 + `7c0b9827ed8597a1c84151b83fcf6307934fd944`。最终验证必须使用包含该提交的 Core、Plugin Runtime 和 Box Runtime 镜像，不能混用旧 SDK。 - 独立资源复核已经移除 Cloud MCP 每会话 5 秒查询执行绑定的轮询，改由签名目录 投影提交后向一个合并回收任务发布代次变化；工具与资源调用前后仍使用数据库 execution fence。Plugin restart 冷却等待者、MCP 投影回收、消息聚合 buffer/scope 均已纳入健康快照和 soak 归零门禁。 +- 单实例目录现在有一致的操作容量契约：Space 在注册事务内通过 PostgreSQL + advisory lock 串行执行 active Workspace check-and-create；Space 全量快照只返回 + active Workspace，并在查询阶段限制 Workspace/membership 数量；闭源适配器限制 + 解压后的 HTTP 响应字节和签名目录基数；Core 在持有目录投影行锁的事务内再次 + COUNT active Workspace，超限时整批回滚且不推进 cursor。任何一层都不截断权威数据。 +- Core Cloud PostgreSQL pool 的 `pool_size + max_overflow` 有绝对上限 100， + Cloud runtime 连接默认强制 60 秒 statement/idle-transaction timeout 和 5 秒 + lock timeout；pool 使用量、超时累计数与目录 active/max 基数进入 `/healthz`。 + Box Runtime 的 session、process、admission record、RPC 文件和 completed retention + 配置也有不可被实例配置放大的绝对上限。 +- Space 的 concurrent registration 容量准入已在一次性 PostgreSQL 16 上真实执行： + 两个 Account 同时争用最后一个 Workspace 槽位时，精确一个事务成功、一个事务 + 得到 capacity error，最终 active Workspace 数为 1。active-only snapshot 与 + archived delta tombstone 的同一真实 PostgreSQL 集成流程也通过。 +- Core Cloud manager 已连接一次性 PostgreSQL 16，并从 `pg_settings` 读回 + `statement_timeout=60000ms`、`lock_timeout=5000ms` 和 + `idle_in_transaction_session_timeout=60000ms`；测试结束后引擎已显式 dispose。 以上结果是进入生产候选验证的前提，不是 SaaS 上线批准。 @@ -107,6 +125,12 @@ - Core、Plugin Runtime、Box Runtime 的不可变镜像 digest； - LangBot 与 SDK commit； - `data/config.yaml` 的非敏感摘要和所有环境变量覆写； +- Space 的 `CLOUD_V2_MAX_DIRECTORY_WORKSPACES` 必须与 Core + `cloud.directory.max_active_workspaces` 和 + `cloud.directory.max_snapshot_workspaces` 一致；Space 的 membership 上限必须与 + Core `cloud.directory.max_snapshot_memberships` 一致； +- Core 的 PostgreSQL pool、statement/lock/idle-transaction timeout，以及 Plugin + Runtime/Box Runtime 的全部实例级资源上限； - PostgreSQL migration revision； - Cloud Adapter、Space control plane 和 workload 的版本。 @@ -176,6 +200,10 @@ PostgreSQL/pgvector 和代表性 Workspace 配置分布，测量： - 单实例可批准的 Workspace、活跃 Bot、plugin worker 和 sandbox 上限。 容量上限必须写入生产配置与告警，不能只保留在测试报告中。 +代码中的默认值和绝对上限只是失控配置的最后防线，不等于生产容量结论。V-08 必须 +根据最终镜像的真实曲线把 Space 与 Core 的匹配上限调到已验证容量以内；如果最终 +批准值高于当前默认 1,000 active Workspace，必须重新执行目录启动、故障恢复和 +24 小时门禁。 ### V-09：24 小时资源 soak @@ -193,8 +221,12 @@ PostgreSQL/pgvector 和代表性 Workspace 配置分布，测量： 最后至少保留 30 分钟无测试流量冷却。任一健康失败、OOM/memory pressure、 PID limit、blocking executor rejection、超阈值 CPU throttling/event-loop lag、 +目录 `active_workspaces > max_active_workspaces`、数据库 pool 使用量超过配置容量、 冷却尾段内存持续增长，或 Plugin restart `gate_waiters`、MCP 投影回收、 消息聚合 buffer/scope 等临时 gauge 不回落都判为失败。 +标准 soak 工具已自动比较目录 active/最近批次与各自配置上限，并比较 PostgreSQL +`checked_out` 与配置 pool 容量；名为 `core` 的标准 endpoint 缺少任一容量指标、 +current/max 只出现一半、数值非法或任一样本越界都会直接失败。 必须归档： diff --git a/docs/multi-tenant/implementation-decisions.md b/docs/multi-tenant/implementation-decisions.md index d0007142d..a74c13f32 100644 --- a/docs/multi-tenant/implementation-decisions.md +++ b/docs/multi-tenant/implementation-decisions.md @@ -270,3 +270,36 @@ This log records implementation choices made while delivering the Workspace arch - 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. + +## 2026-07-29 + +### Directory capacity is one instance-level admission contract + +- Decision: Space and Core share an explicitly configured operational ceiling for active + Workspace and snapshot membership cardinality. Space serializes new personal-Workspace + creation across replicas with a PostgreSQL transaction advisory lock and rejects the + registration transaction before the limit is crossed. Core independently counts the + projected active Workspace set while holding the per-instance directory projection row + lock. Exceeding any limit rolls back the complete projection and does not advance its + cursor; neither side truncates authoritative data. +- Snapshot boundary: A full snapshot is current desired state and contains active + Workspaces only. Archived Workspace revisions are retained as bounded, targeted deltas + so Core can validate and apply monotonic tombstones without every bootstrap carrying + unbounded history. +- Memory boundary: Space bounds database result cardinality before signing. The closed + adapter bounds decompressed response bytes before JSON/JWS parsing and validates + Workspace/membership cardinality before entitlement-cache fan-out. Core schema models + have absolute list ceilings, validate duplicates in one pass, and bulk-read Accounts in + bounded chunks rather than issuing two serial queries per Account. Manifest and + entitlement responses have smaller endpoint-specific byte ceilings; entitlement refresh + validates and releases batches of at most 16 raw responses instead of retaining the + entire directory fan-out. +- Operations boundary: Core health exposes aggregate active/max directory cardinality and + PostgreSQL pool occupancy/timeouts. The default active Workspace ceiling is 1,000, with + a hard code ceiling of 5,000; this is a safety stop, not a production capacity claim. + Space and Core configuration must match, and the approved production value comes from + the real V-08 capacity curve plus the V-09 24-hour soak. +- Reason: One LangBot instance should admit new tenants at the cheapest data-only boundary, + while still preventing legitimate registration growth or a malformed control-plane + response from causing an unbounded startup allocation, database connection storm, or + CPU spike. diff --git a/docs/multi-tenant/runtime-resource-audit-2026-07-28.md b/docs/multi-tenant/runtime-resource-audit-2026-07-28.md index ba25e00ef..236731fe3 100644 --- a/docs/multi-tenant/runtime-resource-audit-2026-07-28.md +++ b/docs/multi-tenant/runtime-resource-audit-2026-07-28.md @@ -17,7 +17,7 @@ 1. 在最终 Cloud 部署权限和 cgroup 拓扑下重复 nsjail、namespace 和 delegated cgroup v2 的 CPU、内存、swap、PID、文件句柄验证。本轮一次性 Linux 容器已经证明代码路径可工作，但普通容器和仅 `--privileged` 的 private cgroup namespace 都不满足条件。 2. 为 Cloud Box 提供并验证硬文件系统 quota provider。普通 nsjail bind mount 不能证明总字节数和 inode 硬配额，当前严格 readiness 按设计会失败关闭。 3. 使用最终生产配置分布继续做容量测试，并据此确定单实例 Workspace placement 上限。本轮真实 PostgreSQL 16 + RLS 启动测试已经覆盖 1,000 个各带 Provider、三类 Model、Bot、Pipeline、KnowledgeBase、MCP 和 Plugin setting 的 Workspace，启动加载耗时和 SQL 次数保持线性；5,000 Workspace 的合成三代替换探针也证明旧运行时会释放。仓库已新增可同时采集 Core/Plugin/Box HTTP、进程树和 cgroup v2 的 24 小时门禁工具，并在受 CPU、memory、swap、PID 硬限制的 Linux 容器中完成短时自检；但最终生产候选拓扑的 24 小时运行仍未执行。测试中的 fake adapter/requester/Plugin handler 仍不能替代真实平台 SDK、外部连接池和插件进程的容量数据；合法活跃租户本身仍会线性占用内存。 -SDK 已先行发布到分支提交 `c7893f8f94e83cb09ebcc98e25490fa699b0fb69`，本提交集中的 LangBot +SDK 已先行发布到分支提交 `7c0b9827ed8597a1c84151b83fcf6307934fd944`，本提交集中的 LangBot `pyproject.toml` 和 `uv.lock` 已精确钉住该提交。最终镜像仍需按待验证清单记录并核对实际安装版本。 ## 覆盖范围 @@ -99,6 +99,17 @@ SDK 已先行发布到分支提交 `c7893f8f94e83cb09ebcc98e25490fa699b0fb69`， - Dify 待提交表单、用户 Space OAuth state 和 Cloud launch JTI replay cache 改为带 revision 校验的最小过期堆；Space credits 使用按时间有序的 LRU/TTL 队列。原实现会在每次攻击者可触发的请求上扫描整个历史缓存并在满容量时再次线性寻找最旧项；现在过期回收为摊销 `O(log N)` 或仅消费已过期前缀，旧堆项会忽略并按活跃状态的有界倍数压缩。 - Cloud launch JTI cache 达到 4,096 个仍有效 token 时失败关闭，不再为了接纳新 token 淘汰仍有效的 replay 记录；否则攻击者可以在容量满后重放被提前遗忘的合法签名。 - Entitlement resolver 现在跟随 Cloud directory 的权威 Workspace 活跃集合。全量目录投影会丢弃已 fenced/removed Workspace 的历史 entitlement snapshot，delta 批量更新不会对每个变化重复扫描；provider 请求进行中发生目录撤销时，返回前的第二次 active fence 会阻止旧结果重新写回缓存。 +- Cloud directory 的签名响应、Workspace、membership 和实例 active Workspace + 均新增可配置操作上限与绝对上限。闭源适配器按流读取响应并在 JSON/JWS 解析前 + 拒绝超过 32 MiB 默认值的解压后正文；Manifest/entitlement/event endpoint 再 + 分别限制为 256 KiB、256 KiB 和 2 MiB。entitlement 刷新最多并发并驻留 16 个 + 原始响应，逐批验证成小型 snapshot 后释放；Core 在目录投影行锁保护的事务内检查最终 + active 数量，超限时回滚 Workspace、Account、membership、inbox 和 cursor，不会 + 截断权威数据或让并发副本各自越过最后一个容量槽。 +- Space 全量目录只投影 active Workspace，历史 archived Workspace 只在最多 100 个 + 目标的签名 delta 中作为 tombstone 返回。注册创建新个人 Workspace 前通过 + PostgreSQL transaction advisory lock 串行执行全局 active 数量准入；达到上限 + 返回 503，避免多个 Space 副本同时观察到最后一个空位。 ### CPU 和事件循环保护 @@ -141,12 +152,23 @@ SDK 已先行发布到分支提交 `c7893f8f94e83cb09ebcc98e25490fa699b0fb69`， - 仓库 Docker Compose/Kubernetes 示例显式下发 Core、Plugin Runtime 和 Box Runtime 的 blocking executor 上限；Kubernetes Box readiness probe 从仅报告进程存活的 `/healthz` 改为 `/readyz`，使 backend 或 managed-mode 隔离检查失败时不会把 Pod 加入就绪流量。 - 相同 digest 的已验证代码和依赖环境可只读共享，每个 installation 的 home/tmp/data 和进程独立。 - Box session、managed process、completed process、admission record 和 RPC 文件均有实例级上限；Cloud entitlement 仍限制每个合资格 Workspace 一个 `global` session、零 managed process。 +- Box Runtime 对上述实例级配置再增加不可放大的硬上限：session 5,000、managed + process 1,024、completed process 10,000、admission record 250,000、RPC 单文件 + 100 MiB、completed retention 86,400 秒。初始化与远程 INIT 对错误类型、负数和 + 超上限均失败关闭，错误动态更新不会留下部分生效的 limit。 ### PostgreSQL - Cloud 强制 PostgreSQL 业务库、共享 pgvector 和允许的固定向量维度。 - pgvector Cloud 模式复用业务数据库的同一个 AsyncEngine，不创建第二个连接池。 - `database.postgresql` 新增并校验 `pool_size`、`max_overflow`、`pool_timeout_seconds`、`pool_recycle_seconds`；默认最大连接数为 `10 + 10`。 +- `pool_size + max_overflow` 的绝对上限为 100，timeout/recycle 也有绝对上限； + Cloud runtime 的 asyncpg 连接默认设置 60 秒 statement timeout、5 秒 lock + timeout 和 60 秒 idle-in-transaction timeout，并分别限制最大 300/60/300 秒。 + 一次性 release migration 不继承这些短 runtime timeout。 +- `/healthz` 输出 pool 配置容量、checked-in/out、overflow、pool admission timeout + 累计数和 SQL timeout 配置；目录同时输出 active/max 与最近批次 + Workspace/membership 数，供生产 soak 和告警核对。 - Application shutdown 显式 dispose 业务引擎；standalone pgvector 仅关闭自己拥有的引擎。 - PersistenceManager 提供统一异步 shutdown；Cloud 常驻进程的启动失败、正常停机和一次性 release migration 的成功/异常路径都会释放数据库引擎。真实 PG catalog 测试还覆盖了“入口已经关闭后测试再次复用 manager 会重开 pool”的第二生命周期，严格资源告警模式下无 asyncpg socket/transport 遗留。 @@ -155,6 +177,10 @@ SDK 已先行发布到分支提交 `c7893f8f94e83cb09ebcc98e25490fa699b0fb69`， - 优先 fail closed 或淘汰最老的 idle cache，不允许攻击者控制的历史 key 无限驻留。 - 插件依赖准备选择实例级串行化，以稳定 CPU/磁盘峰值；代价是批量安装耗时增加。 - PostgreSQL 使用一个显式有界共享连接池；未拆分 pgvector pool。 +- 单实例目录的默认 active/full-snapshot Workspace 上限均为 1,000，membership + 上限为 20,000，签名响应上限为 32 MiB；绝对上限分别为 5,000、100,000 和 + 64 MiB。Space 和 Core 必须配置为同一操作上限，生产值只能根据 V-08 容量曲线 + 向下调整或在重跑全部门禁后提高。 - 第三方 runner 采用 1 MiB 单结果、16 MiB 单流总量和 100,000 个同步/异步事件的统一实例级安全上限；超限请求失败关闭。 - S3 默认允许 16 个并发阻塞调用，最大配置值 128；在没有独立 worker service 的前提下限制线程池排队和上游连接压力。 - MCP 生命周期默认并发 16、最大 128；该限制统一约束实例启动时的 session host 峰值和 shutdown 批次，不允许租户配置单独放大。 @@ -171,19 +197,22 @@ SDK 已先行发布到分支提交 `c7893f8f94e83cb09ebcc98e25490fa699b0fb69`， | --- | --- | | LangBot Ruff + `git diff --check` | 通过 | | Plugin SDK Ruff + `git diff --check` | 通过 | -| LangBot 全量测试（使用远端精确钉住的新 SDK，含 unit/integration/Box/E2E） | `2811 passed, 33 skipped` | -| Plugin SDK 全量测试 | `1317 passed` | +| LangBot 全量测试（使用远端精确钉住的新 SDK，含 unit/integration/Box/E2E） | `2833 passed, 33 skipped` | +| Plugin SDK 全量测试 | `1325 passed` | +| Space Go 全量测试与闭源 Cloud Adapter 测试 | Go `go test ./...` 通过；Adapter `40 passed` | +| Space PostgreSQL 16 Cloud v2 目录与并发容量准入 | 通过；两个注册并发争用最后一个槽位时 `1 success / 1 capacity rejection / 1 active Workspace` | +| Core PostgreSQL 16 Cloud runtime server timeout | 真实连接从 `pg_settings` 读回 `60000ms / 5000ms / 60000ms` 的 statement/lock/idle-transaction timeout，并显式 dispose | | 真实 PostgreSQL 16 + pgvector 迁移/RLS/发布测试（严格资源告警） | `22 passed` | | 真实 PostgreSQL 16 + RLS populated Cloud 启动容量 | 500 Workspace `6.178s / CPU 3.026s`；当前 1,000 Workspace 复跑 `12.109s / CPU 5.967s` | | 较早 Core Dockerfile Linux 镜像构建与 `regex` 导入 | 通过，image SHA `8893a14053df`；该镜像使用旧 SDK pin，已失效，最终候选必须重建 | | `ResourceWarning` + `PytestUnraisableExceptionWarning` 全量门禁 | Core 与 SDK 均通过，并已固化到 pytest 配置 | | Plugin SDK Box 专项测试（含全局扫描回归保护） | `669 passed` | | Docker Compose 渲染、Compose/Kubernetes YAML 解析与 diff 检查 | 通过 | -| Cloud soak 门禁解析/采样/判定单元测试 | `25 passed` | +| Cloud soak 门禁解析/采样/判定单元测试 | `27 passed` | | Core/Plugin SDK event-loop monitor 专项测试 | 两仓各 `7 passed`，包含真实 50 ms scheduler stall | | Cloud soak Linux 硬限制短时自检 | 通过；CPU `0.5`、memory+swap `256 MiB`、PID `128` 均从 cgroup v2 读回，冷却尾段 verdict `pass` | | Core 双阶段历史 churn 资源探针（使用精确钉住的新 SDK） | audit 通过，`11.895s` | -| Core 5,000 个 populated Workspace 三代容量探针（使用精确钉住的新 SDK） | audit 通过，最大替换耗时比 `1.435` | +| Core 5,000 个 populated Workspace 三代容量探针（使用精确钉住的新 SDK） | 当前复跑通过，最大替换耗时比 `1.382` | | Plugin SDK 双阶段资源探针 | audit 通过，`9.407s` | 两个仓库新增了可重复执行的历史 churn 探针，Core 另有 populated Workspace 三代替换探针： @@ -220,9 +249,9 @@ Populated Workspace audit 为 5,000 个 Workspace 各加载一个 Provider、LLM - 三个阶段的活跃 provider/model、pipeline、bot、knowledge 和 MCP registry 均精确维持 `5,000`，不存在按历史 generation 增长。 - 到第三阶段，前两代的 requester、Bot adapter 和 MCP session 各 `10,000` 个全部收到确定性关闭；weak reference 断言旧代对象可被回收。 -- event-loop task、线程和文件描述符保持 `1 / 1 / 6`；使用远端精确钉住 SDK 的当前复跑中，第三阶段相对第二阶段 RSS 增长 `1,261,568 bytes`，tracemalloc current 仅增长 `510 bytes`。 -- 初始/第一次替换/第二次替换分别耗时 `1.804s / 2.420s / 2.588s`，最大替换耗时比为 `1.435`，未随历史代次出现 CPU 退化。 -- macOS RSS sample 从初始的 `154,484,736` 增至第一阶段 `368,050,176`、第二阶段 `388,939,776` 和第三阶段 `390,201,344 bytes`；第二次替换只比第一次替换增加约 1.20 MiB，但“合法活跃租户资源的线性容量”仍必须作为 placement 容量输入。这里使用轻量 fake adapter/requester，不应把第一阶段约 204 MiB 增量外推为生产每租户成本。 +- event-loop task、线程和文件描述符保持 `1 / 1 / 6`；使用远端精确钉住 SDK 的当前复跑中，第三阶段相对第二阶段 RSS 增长 `1,245,184 bytes`，tracemalloc current 仅增长 `2,061 bytes`。 +- 初始/第一次替换/第二次替换分别耗时 `1.660s / 2.151s / 2.296s`，最大替换耗时比为 `1.382`，未随历史代次出现 CPU 退化。 +- macOS RSS sample 从初始的 `154,648,576` 增至第一阶段 `368,181,248`、第二阶段 `389,087,232` 和第三阶段 `390,332,416 bytes`；第二次替换只比第一次替换增加约 1.19 MiB，但“合法活跃租户资源的线性容量”仍必须作为 placement 容量输入。这里使用轻量 fake adapter/requester，不应把第一阶段约 204 MiB 增量外推为生产每租户成本。 Plugin SDK audit 每个阶段执行 25,000 次 loopback RPC、5,000 次安装 binding 激活/撤销、10,000 个 Workspace generation 更新和 2,500 次带 Workspace 上下文的 Box session 创建/删除。第一、第二阶段的保留状态完全一致： diff --git a/pyproject.toml b/pyproject.toml index 06015bd44..2864c76c5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -71,7 +71,7 @@ dependencies = [ "chromadb>=1.0.0,<2.0.0", "qdrant-client (>=1.15.1,<2.0.0)", "pyseekdb==1.1.0.post3", - "langbot-plugin @ git+https://github.com/langbot-app/langbot-plugin-sdk.git@c7893f8f94e83cb09ebcc98e25490fa699b0fb69", + "langbot-plugin @ git+https://github.com/langbot-app/langbot-plugin-sdk.git@7c0b9827ed8597a1c84151b83fcf6307934fd944", "asyncpg>=0.30.0", "line-bot-sdk>=3.19.0", "matrix-nio>=0.25.2", diff --git a/scripts/cloud_runtime_soak.py b/scripts/cloud_runtime_soak.py index f0eafdd89..527108d29 100644 --- a/scripts/cloud_runtime_soak.py +++ b/scripts/cloud_runtime_soak.py @@ -93,6 +93,25 @@ TRANSIENT_GAUGE_SUFFIXES = ( '.resources.closing_session_tasks', '.resources.background_tasks', ) +CAPACITY_GAUGE_PAIRS = ( + ( + '.resources.directory.active_workspaces', + '.resources.directory.max_active_workspaces', + ), + ( + '.resources.directory.last_batch_workspaces', + '.resources.directory.max_snapshot_workspaces', + ), + ( + '.resources.directory.last_batch_memberships', + '.resources.directory.max_snapshot_memberships', + ), + ( + '.resources.database_pool.checked_out', + '.resources.database_pool.configured_capacity', + ), +) +REQUIRED_CORE_CAPACITY_GAUGES = frozenset(suffix for pair in CAPACITY_GAUGE_PAIRS for suffix in pair) @dataclasses.dataclass(frozen=True, slots=True) @@ -712,6 +731,17 @@ def evaluate_gate( | set(state.baseline_metrics or ()) | set(state.last_metrics or ()) ) + if state.target.name == 'core': + for suffix in sorted(REQUIRED_CORE_CAPACITY_GAUGES): + matching_metrics = sorted(key for key in metric_keys if key.endswith(suffix)) + if not matching_metrics: + failures.append(f'{target_id} did not expose required capacity gauge {suffix}') + continue + for metric in matching_metrics: + if any(metric not in sample.metrics for sample in tail_samples): + failures.append( + f'{target_id} did not expose required capacity gauge {metric} throughout the tail' + ) for suffix in REJECTION_SUFFIXES: if thresholds.allow_rejections: break @@ -746,6 +776,24 @@ def evaluate_gate( growth = _counter_delta(first, last, metric) if growth is not None and growth > thresholds.max_transient_gauge_growth: failures.append(f'{target_id} transient gauge {metric} grew by {growth:g} during the idle tail') + for current_suffix, capacity_suffix in CAPACITY_GAUGE_PAIRS: + for current_metric in sorted(key for key in metric_keys if key.endswith(current_suffix)): + capacity_metric = current_metric[: -len(current_suffix)] + capacity_suffix + observed = [ + (sample.metrics[current_metric], sample.metrics.get(capacity_metric)) + for sample in tail_samples + if current_metric in sample.metrics + ] + if any(capacity is None for _current, capacity in observed): + failures.append(f'{target_id} exposed {current_metric} without matching {capacity_metric}') + continue + if any( + current < 0 or capacity is None or capacity <= 0 or current > capacity + for current, capacity in observed + ): + failures.append( + f'{target_id} exceeded or invalidated capacity pair {current_metric} <= {capacity_metric}' + ) loop_running_metrics = sorted(key for key in metric_keys if key.endswith('.event_loop.running')) if not loop_running_metrics: diff --git a/src/langbot/pkg/cloud/bootstrap.py b/src/langbot/pkg/cloud/bootstrap.py index 01c2382b7..8e9ee1886 100644 --- a/src/langbot/pkg/cloud/bootstrap.py +++ b/src/langbot/pkg/cloud/bootstrap.py @@ -11,7 +11,7 @@ from collections.abc import Awaitable, Callable from typing import Any, Protocol, runtime_checkable from ..workspace.policy import CloudWorkspacePolicy, SingleWorkspacePolicy -from .directory import DirectoryProjectionProvider +from .directory import DirectoryProjectionProvider, directory_projection_limits_from_config from .entitlements import EntitlementProvider, OpenSourceEntitlementProvider @@ -112,6 +112,10 @@ class VerifiedCloudDeployment: raise CloudBootstrapError('Verified Cloud bootstrap did not provide a Manifest renewal adapter') def validate_instance_config(self, config: dict[str, Any]) -> None: + try: + directory_projection_limits_from_config(config) + except (TypeError, ValueError) as exc: + raise CloudBootstrapError(f'Cloud directory limits are invalid: {exc}') from exc if config.get('database', {}).get('use') != 'postgresql': raise CloudBootstrapError('Cloud runtime requires database.use=postgresql') if config.get('vdb', {}).get('use') != self.required_vector_backend: diff --git a/src/langbot/pkg/cloud/directory.py b/src/langbot/pkg/cloud/directory.py index ae9f03ee1..3232e6bed 100644 --- a/src/langbot/pkg/cloud/directory.py +++ b/src/langbot/pkg/cloud/directory.py @@ -7,10 +7,86 @@ from typing import Any, Protocol, runtime_checkable import pydantic +DEFAULT_MAX_ACTIVE_WORKSPACES = 1_000 +HARD_MAX_ACTIVE_WORKSPACES = 5_000 +DEFAULT_MAX_SNAPSHOT_WORKSPACES = 1_000 +HARD_MAX_SNAPSHOT_WORKSPACES = 5_000 +DEFAULT_MAX_SNAPSHOT_MEMBERSHIPS = 20_000 +HARD_MAX_SNAPSHOT_MEMBERSHIPS = 100_000 +DEFAULT_MAX_CONTROL_PLANE_RESPONSE_BYTES = 32 * 1024 * 1024 +HARD_MAX_CONTROL_PLANE_RESPONSE_BYTES = 64 * 1024 * 1024 + + class DirectoryProjectionUnavailableError(RuntimeError): """Raised when the verified Cloud directory cannot safely admit work.""" +class DirectoryProjectionLimits(pydantic.BaseModel): + """Instance-owned cardinality limits for verified Cloud directory data. + + These are operational safety limits, not subscription entitlements. Core + fails the complete projection transaction when a limit is exceeded instead + of truncating an authoritative directory and accidentally hiding tenants. + The closed adapter consumes the same limits before priming entitlement + caches, and additionally bounds the HTTP response buffered for signature + verification. + """ + + model_config = pydantic.ConfigDict(frozen=True, extra='forbid') + + max_active_workspaces: int = pydantic.Field( + default=DEFAULT_MAX_ACTIVE_WORKSPACES, + ge=1, + le=HARD_MAX_ACTIVE_WORKSPACES, + ) + max_snapshot_workspaces: int = pydantic.Field( + default=DEFAULT_MAX_SNAPSHOT_WORKSPACES, + ge=1, + le=HARD_MAX_SNAPSHOT_WORKSPACES, + ) + max_snapshot_memberships: int = pydantic.Field( + default=DEFAULT_MAX_SNAPSHOT_MEMBERSHIPS, + ge=1, + le=HARD_MAX_SNAPSHOT_MEMBERSHIPS, + ) + max_response_bytes: int = pydantic.Field( + default=DEFAULT_MAX_CONTROL_PLANE_RESPONSE_BYTES, + ge=1024 * 1024, + le=HARD_MAX_CONTROL_PLANE_RESPONSE_BYTES, + ) + + @pydantic.field_validator( + 'max_active_workspaces', + 'max_snapshot_workspaces', + 'max_snapshot_memberships', + 'max_response_bytes', + mode='before', + ) + @classmethod + def _reject_boolean_limits(cls, value: object) -> object: + if isinstance(value, bool): + raise ValueError('must be an integer') + return value + + @pydantic.model_validator(mode='after') + def _validate_workspace_limits(self) -> DirectoryProjectionLimits: + if self.max_snapshot_workspaces < self.max_active_workspaces: + raise ValueError('max_snapshot_workspaces must be greater than or equal to max_active_workspaces') + return self + + +def directory_projection_limits_from_config(config: dict[str, Any]) -> DirectoryProjectionLimits: + """Parse typed Cloud directory limits from the instance configuration.""" + + cloud_config = config.get('cloud', {}) + if not isinstance(cloud_config, dict): + raise ValueError('cloud must be a mapping') + directory_config = cloud_config.get('directory', {}) + if not isinstance(directory_config, dict): + raise ValueError('cloud.directory must be a mapping') + return DirectoryProjectionLimits.model_validate(directory_config) + + class DirectoryMember(pydantic.BaseModel): """One account membership published by the SaaS control plane.""" @@ -48,7 +124,10 @@ class DirectoryWorkspace(pydantic.BaseModel): 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, ...] = () + members: tuple[DirectoryMember, ...] = pydantic.Field( + default=(), + max_length=HARD_MAX_SNAPSHOT_MEMBERSHIPS, + ) @pydantic.field_validator('members', mode='before') @classmethod @@ -57,13 +136,16 @@ class DirectoryWorkspace(pydantic.BaseModel): @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): + membership_uuids: set[str] = set() + account_uuids: set[str] = set() + for member in self.members: + if member.membership_uuid in membership_uuids: + raise ValueError('Directory Workspace contains duplicate membership UUIDs') + if member.account_uuid in account_uuids: + raise ValueError('Directory Workspace contains duplicate account UUIDs') + membership_uuids.add(member.membership_uuid) + account_uuids.add(member.account_uuid) + if self.created_by_account_uuid not in account_uuids: raise ValueError('Directory Workspace must include its creator') return self @@ -76,7 +158,10 @@ class DirectorySnapshot(pydantic.BaseModel): instance_uuid: str = pydantic.Field(min_length=1, max_length=255) cursor: int = pydantic.Field(ge=0) generated_at: datetime.datetime - workspaces: tuple[DirectoryWorkspace, ...] = () + workspaces: tuple[DirectoryWorkspace, ...] = pydantic.Field( + default=(), + max_length=HARD_MAX_SNAPSHOT_WORKSPACES, + ) @pydantic.field_validator('workspaces', mode='before') @classmethod @@ -85,15 +170,20 @@ class DirectorySnapshot(pydantic.BaseModel): @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') + workspace_uuids: set[str] = set() + slugs: set[str] = set() + membership_uuids: set[str] = set() + for workspace in self.workspaces: + if workspace.uuid in workspace_uuids: + raise ValueError('Directory snapshot contains duplicate Workspace UUIDs') + if workspace.slug in slugs: + raise ValueError('Directory snapshot contains duplicate Workspace slugs') + workspace_uuids.add(workspace.uuid) + slugs.add(workspace.slug) + for member in workspace.members: + if member.membership_uuid in membership_uuids: + raise ValueError('Directory snapshot contains duplicate membership UUIDs') + membership_uuids.add(member.membership_uuid) return self @@ -127,17 +217,23 @@ class DirectoryDelta(pydantic.BaseModel): 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') + requested_set = set(requested) + workspace_uuids: set[str] = set() + slugs: set[str] = set() + membership_uuids: set[str] = set() + for workspace in self.workspaces: + if workspace.uuid in workspace_uuids: + raise ValueError('Directory delta contains duplicate Workspace UUIDs') + if workspace.uuid not in requested_set: + raise ValueError('Directory delta returned an unrequested Workspace') + if workspace.slug in slugs: + raise ValueError('Directory delta contains duplicate Workspace slugs') + workspace_uuids.add(workspace.uuid) + slugs.add(workspace.slug) + for member in workspace.members: + if member.membership_uuid in membership_uuids: + raise ValueError('Directory delta contains duplicate membership UUIDs') + membership_uuids.add(member.membership_uuid) return self diff --git a/src/langbot/pkg/cloud/directory_projection.py b/src/langbot/pkg/cloud/directory_projection.py index ded071ca6..3d0f82e44 100644 --- a/src/langbot/pkg/cloud/directory_projection.py +++ b/src/langbot/pkg/cloud/directory_projection.py @@ -29,6 +29,7 @@ from .directory import ( DirectoryEvent, DirectoryEventBatch, DirectoryMember, + DirectoryProjectionLimits, DirectoryProjectionProvider, DirectoryProjectionUnavailableError, DirectorySnapshot, @@ -65,6 +66,7 @@ _MEMBERSHIP_STATUS_MAP = { 'removed': MembershipStatus.REMOVED.value, } _INCREMENTAL_PROJECTION_FINGERPRINT = hashlib.sha256(b'langbot-directory-incremental-v1').hexdigest() +_ACCOUNT_QUERY_CHUNK_SIZE = 500 class _DirectorySnapshotSuperseded(DirectoryProjectionUnavailableError): @@ -89,6 +91,7 @@ class DirectoryProjectionService: sync_interval_seconds: float = 5.0, max_staleness_seconds: float = 60.0, event_limit: int = 100, + limits: DirectoryProjectionLimits | None = None, monotonic_time: Callable[[], float] = time.monotonic, ) -> None: if not isinstance(provider, DirectoryProjectionProvider): @@ -101,15 +104,21 @@ class DirectoryProjectionService: 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') + if limits is not None and not isinstance(limits, DirectoryProjectionLimits): + raise TypeError('Directory projection limits must be a DirectoryProjectionLimits value') 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.limits = limits or DirectoryProjectionLimits() self._monotonic_time = monotonic_time self._last_success_monotonic: float | None = None self._ready = False + self._active_workspace_count = 0 + self._last_batch_workspace_count = 0 + self._last_batch_membership_count = 0 # 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, @@ -184,6 +193,90 @@ class DirectoryProjectionService: if self._monotonic_time() - last_success >= self.max_staleness_seconds: raise DirectoryProjectionUnavailableError('Cloud directory projection is stale') + def resource_snapshot(self) -> dict[str, int]: + """Return aggregate, tenant-free cardinality gauges for health checks.""" + + return { + 'active_workspaces': self._active_workspace_count, + 'max_active_workspaces': self.limits.max_active_workspaces, + 'last_batch_workspaces': self._last_batch_workspace_count, + 'last_batch_memberships': self._last_batch_membership_count, + 'max_snapshot_workspaces': self.limits.max_snapshot_workspaces, + 'max_snapshot_memberships': self.limits.max_snapshot_memberships, + } + + def _validate_batch_capacity( + self, + workspaces: tuple[DirectoryWorkspace, ...], + *, + full_snapshot: bool, + ) -> tuple[int, int]: + workspace_count = len(workspaces) + if full_snapshot and workspace_count > self.limits.max_snapshot_workspaces: + raise DirectoryProjectionUnavailableError( + 'Directory snapshot Workspace capacity exceeded ' + f'({workspace_count} > {self.limits.max_snapshot_workspaces})' + ) + + active_count = 0 + membership_count = 0 + for workspace in workspaces: + if workspace.status == WorkspaceStatus.ACTIVE.value: + active_count += 1 + membership_count += len(workspace.members) + if membership_count > self.limits.max_snapshot_memberships: + raise DirectoryProjectionUnavailableError( + 'Directory membership capacity exceeded ' + f'({membership_count} > {self.limits.max_snapshot_memberships})' + ) + if active_count > self.limits.max_active_workspaces: + raise DirectoryProjectionUnavailableError( + f'Directory active Workspace capacity exceeded ({active_count} > {self.limits.max_active_workspaces})' + ) + return workspace_count, membership_count + + async def _enforce_active_workspace_capacity(self, session: Any) -> int: + """Count the committed candidate state while holding the projection lock. + + Full snapshots can validate their own active count before doing any + database work. Incremental deltas cannot know the instance total, so + every projection path also checks the database after applying fences. + The caller holds the per-instance DirectoryProjectionState row lock; + concurrent replicas therefore cannot race two individually-admitted + deltas above the instance ceiling. + """ + + active_count = int( + ( + await session.scalar( + sqlalchemy.select(sqlalchemy.func.count()) + .select_from(Workspace) + .where( + Workspace.instance_uuid == self.instance_uuid, + Workspace.source == WorkspaceSource.CLOUD_PROJECTION.value, + Workspace.status == WorkspaceStatus.ACTIVE.value, + ) + ) + ) + or 0 + ) + if active_count > self.limits.max_active_workspaces: + raise DirectoryProjectionUnavailableError( + f'Projected active Workspace capacity exceeded ({active_count} > {self.limits.max_active_workspaces})' + ) + return active_count + + def _record_batch_cardinality( + self, + *, + active_workspaces: int, + workspaces: int, + memberships: int, + ) -> None: + self._active_workspace_count = active_workspaces + self._last_batch_workspace_count = workspaces + self._last_batch_membership_count = memberships + async def apply_snapshot( self, snapshot: DirectorySnapshot, @@ -194,6 +287,10 @@ class DirectoryProjectionService: if not isinstance(snapshot, DirectorySnapshot): raise DirectoryProjectionUnavailableError('Directory provider returned an invalid snapshot') + workspace_count, membership_count = self._validate_batch_capacity( + snapshot.workspaces, + full_snapshot=True, + ) snapshot = DirectorySnapshot.model_validate(snapshot.model_dump()) if snapshot.instance_uuid != self.instance_uuid: raise DirectoryProjectionUnavailableError('Directory snapshot targets another LangBot instance') @@ -244,9 +341,10 @@ class DirectoryProjectionService: 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) + accounts_by_uuid = await self._apply_accounts(session, snapshot) + await self._apply_workspaces(session, snapshot, accounts_by_uuid=accounts_by_uuid) await self._fence_absent_workspaces(session, snapshot) + active_workspace_count = await self._enforce_active_workspace_capacity(session) state.cursor = snapshot.cursor state.snapshot_coverage_cursor = snapshot.cursor @@ -260,6 +358,11 @@ class DirectoryProjectionService: await self._reconcile_entitlement_snapshot_set(snapshot) self._publish_runtime_execution_projection(snapshot.workspaces) + self._record_batch_cardinality( + active_workspaces=active_workspace_count, + workspaces=workspace_count, + memberships=membership_count, + ) self._record_success() self._consumer_cursor = snapshot.cursor @@ -270,6 +373,10 @@ class DirectoryProjectionService: raise DirectoryProjectionUnavailableError('Directory provider returned an invalid delta') if not isinstance(batch, DirectoryEventBatch): raise DirectoryProjectionUnavailableError('Directory provider returned an invalid event batch') + workspace_count, membership_count = self._validate_batch_capacity( + delta.workspaces, + full_snapshot=False, + ) delta = DirectoryDelta.model_validate(delta.model_dump()) batch = DirectoryEventBatch.model_validate(batch.model_dump()) self._validate_batch(batch, expected_after_cursor=batch.after_cursor) @@ -325,8 +432,12 @@ class DirectoryProjectionService: generated_at=delta.generated_at, workspaces=delta.workspaces, ) - await self._apply_accounts(session, projected_delta) - await self._apply_workspaces(session, projected_delta) + accounts_by_uuid = await self._apply_accounts(session, projected_delta) + await self._apply_workspaces( + session, + projected_delta, + accounts_by_uuid=accounts_by_uuid, + ) await self._fence_workspaces( session, { @@ -340,6 +451,7 @@ class DirectoryProjectionService: # incremental path; a later full snapshot replaces this marker. state.snapshot_fingerprint = _INCREMENTAL_PROJECTION_FINGERPRINT + active_workspace_count = await self._enforce_active_workspace_capacity(session) state.last_applied_at = now state.lease_expires_at = lease_expires_at await self._mark_events_applied(session, batch.events, now=now) @@ -354,6 +466,11 @@ class DirectoryProjectionService: returned.values(), affected_workspace_uuids=requested, ) + self._record_batch_cardinality( + active_workspaces=active_workspace_count, + workspaces=workspace_count, + memberships=membership_count, + ) if projection_caught_up: self._record_success() self._consumer_cursor = batch.cursor @@ -580,7 +697,7 @@ class DirectoryProjectionService: for row in inbox_rows: row.applied_at = now - async def _apply_accounts(self, session: Any, snapshot: DirectorySnapshot) -> None: + async def _apply_accounts(self, session: Any, snapshot: DirectorySnapshot) -> dict[str, User]: selected: dict[str, DirectoryMember] = {} emails: dict[str, str] = {} for workspace in snapshot.workspaces: @@ -596,11 +713,35 @@ class DirectoryProjectionService: if previous is None: selected[member.account_uuid] = member + # Fetch existing UUID and email owners in bounded batches. The previous + # two SELECTs per unique account made a large but valid directory + # snapshot produce tens of thousands of serial round trips during + # startup. The configured membership ceiling bounds the materialized + # maps, while batching stays below PostgreSQL parameter limits. + accounts_by_uuid: dict[str, User] = {} + accounts_by_email: dict[str, User] = {} + selected_items = list(selected.items()) + for start in range(0, len(selected_items), _ACCOUNT_QUERY_CHUNK_SIZE): + chunk = selected_items[start : start + _ACCOUNT_QUERY_CHUNK_SIZE] + account_uuids = [account_uuid for account_uuid, _member in chunk] + normalized_emails = [member.normalized_email for _account_uuid, member in chunk] + rows = ( + await session.scalars( + sqlalchemy.select(User).where( + sqlalchemy.or_( + User.uuid.in_(account_uuids), + User.normalized_email.in_(normalized_emails), + ) + ) + ) + ).all() + for account in rows: + accounts_by_uuid[account.uuid] = account + accounts_by_email[account.normalized_email] = account + 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) - ) + account = accounts_by_uuid.get(account_uuid) + email_account = accounts_by_email.get(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: @@ -616,6 +757,8 @@ class DirectoryProjectionService: space_account_uuid=account_uuid, ) session.add(account) + accounts_by_uuid[account_uuid] = account + accounts_by_email[member.normalized_email] = account continue if account.source != AccountSource.CLOUD_PROJECTION.value: raise DirectoryProjectionUnavailableError('Directory account UUID collides with a local Core account') @@ -637,8 +780,15 @@ class DirectoryProjectionService: account.account_type = 'space' account.space_account_uuid = account_uuid await session.flush() + return accounts_by_uuid - async def _apply_workspaces(self, session: Any, snapshot: DirectorySnapshot) -> None: + async def _apply_workspaces( + self, + session: Any, + snapshot: DirectorySnapshot, + *, + accounts_by_uuid: dict[str, User], + ) -> None: for candidate in snapshot.workspaces: workspace = await session.get(Workspace, candidate.uuid) if workspace is None: @@ -649,7 +799,7 @@ class DirectoryProjectionService: slug=candidate.slug, type=candidate.type, status=candidate.status, - created_by_account_uuid=await self._projected_creator_uuid(session, candidate), + created_by_account_uuid=self._projected_creator_uuid(candidate, accounts_by_uuid), source=WorkspaceSource.CLOUD_PROJECTION.value, projection_revision=candidate.projection_revision, ) @@ -661,14 +811,18 @@ class DirectoryProjectionService: 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.created_by_account_uuid = self._projected_creator_uuid(candidate, accounts_by_uuid) 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)) + @staticmethod + def _projected_creator_uuid( + candidate: DirectoryWorkspace, + accounts_by_uuid: dict[str, User], + ) -> str | None: + creator = accounts_by_uuid.get(candidate.created_by_account_uuid) if creator is None: if candidate.status == WorkspaceStatus.ACTIVE.value: raise DirectoryProjectionUnavailableError('Active Directory Workspace creator is not projected') @@ -801,9 +955,24 @@ class DirectoryProjectionService: included = {workspace.uuid for workspace in snapshot.workspaces} projected = ( await session.scalars( - sqlalchemy.select(Workspace).where( + sqlalchemy.select(Workspace) + .outerjoin( + WorkspaceExecutionState, + WorkspaceExecutionState.workspace_uuid == Workspace.uuid, + ) + .where( Workspace.instance_uuid == self.instance_uuid, Workspace.source == WorkspaceSource.CLOUD_PROJECTION.value, + sqlalchemy.or_( + Workspace.status.not_in( + ( + WorkspaceStatus.ARCHIVED.value, + WorkspaceStatus.DELETED.value, + ) + ), + WorkspaceExecutionState.state == WorkspaceExecutionStatus.ACTIVE.value, + WorkspaceExecutionState.write_fenced == sqlalchemy.false(), + ), ) ) ).all() diff --git a/src/langbot/pkg/core/app.py b/src/langbot/pkg/core/app.py index de49de979..3913857fc 100644 --- a/src/langbot/pkg/core/app.py +++ b/src/langbot/pkg/core/app.py @@ -268,11 +268,23 @@ class Application: } ) + directory_stats = {} + directory_snapshot = getattr(self.directory_projection_service, 'resource_snapshot', None) + if callable(directory_snapshot): + directory_stats = directory_snapshot() + + database_stats = {} + database_snapshot = getattr(self.persistence_mgr, 'get_resource_stats', None) + if callable(database_snapshot): + database_stats = database_snapshot() + return { 'asyncio_tasks': asyncio_tasks, 'event_loop': self.event_loop_monitor.snapshot(), 'blocking_executor': (self.blocking_executor.snapshot() if self.blocking_executor is not None else {}), 'application_tasks': task_stats, + 'database_pool': database_stats, + 'directory': directory_stats, 'query_pool': query_pool_stats, 'models': model_stats, 'runtimes': runtime_stats, diff --git a/src/langbot/pkg/core/stages/build_app.py b/src/langbot/pkg/core/stages/build_app.py index 16dbd9462..813a6c4e6 100644 --- a/src/langbot/pkg/core/stages/build_app.py +++ b/src/langbot/pkg/core/stages/build_app.py @@ -42,6 +42,7 @@ from ...workspace import collaboration as workspace_collaboration_module from ...workspace import invitation_delivery as invitation_delivery_module from ...cloud import bootstrap as cloud_bootstrap from ...cloud import launch as cloud_launch_module +from ...cloud.directory import directory_projection_limits_from_config from ...cloud.directory_projection import DirectoryProjectionService from ...cloud.entitlements import EntitlementResolver from ...api.http.context import ExecutionContext, PrincipalContext, PrincipalType @@ -154,6 +155,7 @@ class BuildAppStage(stage.BootingStage): ap, deployment.directory_provider, constants.instance_id, + limits=directory_projection_limits_from_config(ap.instance_config.data), ) await directory_projection_service.initialize() ap.directory_projection_service = directory_projection_service diff --git a/src/langbot/pkg/core/stages/load_config.py b/src/langbot/pkg/core/stages/load_config.py index 37f0bc566..e68e92497 100644 --- a/src/langbot/pkg/core/stages/load_config.py +++ b/src/langbot/pkg/core/stages/load_config.py @@ -14,6 +14,25 @@ from ..bootutils import config _RUNTIME_POLICY_DEFAULTS = { + 'cloud': { + 'directory': { + 'max_active_workspaces': 1000, + 'max_snapshot_workspaces': 1000, + 'max_snapshot_memberships': 20000, + 'max_response_bytes': 33554432, + } + }, + 'database': { + 'postgresql': { + 'pool_size': 10, + 'max_overflow': 10, + 'pool_timeout_seconds': 30, + 'pool_recycle_seconds': 1800, + 'statement_timeout_ms': 60000, + 'lock_timeout_ms': 5000, + 'idle_in_transaction_session_timeout_ms': 60000, + } + }, 'system': { 'blocking_executor': { 'max_workers': bounded_executor.DEFAULT_MAX_WORKERS, @@ -49,7 +68,7 @@ def _complete_runtime_policy_defaults(cfg: dict) -> dict: The historic config loader intentionally does not deep-complete the whole template. These fields are different: their native env overrides must retain boolean/numeric types on upgraded instances, so their defaults must - exist before ``PLUGIN__...`` and ``MCP__...`` are parsed. + exist before ``CLOUD__...``, ``PLUGIN__...`` and ``MCP__...`` are parsed. """ def merge(target: dict, defaults: dict, path: tuple[str, ...] = ()) -> None: diff --git a/src/langbot/pkg/persistence/database.py b/src/langbot/pkg/persistence/database.py index eb361edcb..2eee09a12 100644 --- a/src/langbot/pkg/persistence/database.py +++ b/src/langbot/pkg/persistence/database.py @@ -39,6 +39,7 @@ class BaseDatabaseManager(abc.ABC): ) -> None: self.ap = ap self.url_override = url_override + self.persistence_mode: str | None = None @abc.abstractmethod async def initialize(self) -> None: diff --git a/src/langbot/pkg/persistence/databases/postgresql.py b/src/langbot/pkg/persistence/databases/postgresql.py index 3e99e2129..75e340f7f 100644 --- a/src/langbot/pkg/persistence/databases/postgresql.py +++ b/src/langbot/pkg/persistence/databases/postgresql.py @@ -7,6 +7,14 @@ from .. import database from ..postgresql_url import normalize_asyncpg_url +MAX_POOL_CONNECTIONS = 100 +MAX_POOL_TIMEOUT_SECONDS = 300 +MAX_POOL_RECYCLE_SECONDS = 86_400 +MAX_STATEMENT_TIMEOUT_MS = 300_000 +MAX_LOCK_TIMEOUT_MS = 60_000 +MAX_IDLE_TRANSACTION_TIMEOUT_MS = 300_000 + + @database.manager_class('postgresql') class PostgreSQLDatabaseManager(database.BaseDatabaseManager): """PostgreSQL database manager""" @@ -18,14 +26,16 @@ class PostgreSQLDatabaseManager(database.BaseDatabaseManager): default: int, *, minimum: int, + maximum: int, ) -> int: value = config.get(name, default) - if isinstance(value, bool) or not isinstance(value, int) or value < minimum: + if isinstance(value, bool) or not isinstance(value, int) or not minimum <= value <= maximum: comparator = 'non-negative' if minimum == 0 else 'positive' - raise ValueError(f'database.postgresql.{name} must be a {comparator} integer') + raise ValueError(f'database.postgresql.{name} must be a {comparator} integer no greater than {maximum}') return value async def initialize(self) -> None: + self._pool_timeouts_total = 0 postgresql_config = self.ap.instance_config.data.get('database', {}).get('postgresql', {}) if not isinstance(postgresql_config, dict): raise ValueError('database.postgresql must be an object') @@ -53,31 +63,105 @@ class PostgreSQLDatabaseManager(database.BaseDatabaseManager): port=postgresql_config.get('port', 5432), database=postgresql_config.get('database', 'postgres'), ) + self.pool_size = self._pool_integer( + postgresql_config, + 'pool_size', + 10, + minimum=1, + maximum=MAX_POOL_CONNECTIONS, + ) + self.max_overflow = self._pool_integer( + postgresql_config, + 'max_overflow', + 10, + minimum=0, + maximum=MAX_POOL_CONNECTIONS, + ) + if self.pool_size + self.max_overflow > MAX_POOL_CONNECTIONS: + raise ValueError(f'database.postgresql pool_size + max_overflow must not exceed {MAX_POOL_CONNECTIONS}') + self.pool_timeout_seconds = self._pool_integer( + postgresql_config, + 'pool_timeout_seconds', + 30, + minimum=1, + maximum=MAX_POOL_TIMEOUT_SECONDS, + ) + self.pool_recycle_seconds = self._pool_integer( + postgresql_config, + 'pool_recycle_seconds', + 1800, + minimum=1, + maximum=MAX_POOL_RECYCLE_SECONDS, + ) + connect_args = {} + self.statement_timeout_ms = 0 + self.lock_timeout_ms = 0 + self.idle_transaction_timeout_ms = 0 + if self.persistence_mode == 'cloud_runtime': + self.statement_timeout_ms = self._pool_integer( + postgresql_config, + 'statement_timeout_ms', + 60_000, + minimum=1, + maximum=MAX_STATEMENT_TIMEOUT_MS, + ) + self.lock_timeout_ms = self._pool_integer( + postgresql_config, + 'lock_timeout_ms', + 5_000, + minimum=1, + maximum=MAX_LOCK_TIMEOUT_MS, + ) + self.idle_transaction_timeout_ms = self._pool_integer( + postgresql_config, + 'idle_in_transaction_session_timeout_ms', + 60_000, + minimum=1, + maximum=MAX_IDLE_TRANSACTION_TIMEOUT_MS, + ) + connect_args = { + 'server_settings': { + 'statement_timeout': str(self.statement_timeout_ms), + 'lock_timeout': str(self.lock_timeout_ms), + 'idle_in_transaction_session_timeout': str(self.idle_transaction_timeout_ms), + } + } self.engine = sqlalchemy_asyncio.create_async_engine( engine_url, - pool_size=self._pool_integer( - postgresql_config, - 'pool_size', - 10, - minimum=1, - ), - max_overflow=self._pool_integer( - postgresql_config, - 'max_overflow', - 10, - minimum=0, - ), - pool_timeout=self._pool_integer( - postgresql_config, - 'pool_timeout_seconds', - 30, - minimum=1, - ), - pool_recycle=self._pool_integer( - postgresql_config, - 'pool_recycle_seconds', - 1800, - minimum=1, - ), + pool_size=self.pool_size, + max_overflow=self.max_overflow, + pool_timeout=self.pool_timeout_seconds, + pool_recycle=self.pool_recycle_seconds, pool_pre_ping=True, + **({'connect_args': connect_args} if connect_args else {}), ) + + def resource_stats(self) -> dict[str, int]: + """Return aggregate pool gauges without exposing connection details.""" + + pool = self.engine.pool + + def read(name: str) -> int: + method = getattr(pool, name, None) + if not callable(method): + return 0 + try: + return int(method()) + except Exception: + return 0 + + return { + 'configured_size': self.pool_size, + 'configured_max_overflow': self.max_overflow, + 'configured_capacity': self.pool_size + self.max_overflow, + 'statement_timeout_ms': self.statement_timeout_ms, + 'lock_timeout_ms': self.lock_timeout_ms, + 'idle_in_transaction_session_timeout_ms': self.idle_transaction_timeout_ms, + 'timeouts_total': self._pool_timeouts_total, + 'checked_in': read('checkedin'), + 'checked_out': read('checkedout'), + 'overflow': max(read('overflow'), 0), + } + + def record_pool_timeout(self) -> None: + self._pool_timeouts_total += 1 diff --git a/src/langbot/pkg/persistence/mgr.py b/src/langbot/pkg/persistence/mgr.py index 749a3d04e..9565dfdd0 100644 --- a/src/langbot/pkg/persistence/mgr.py +++ b/src/langbot/pkg/persistence/mgr.py @@ -159,6 +159,7 @@ class PersistenceManager: for manager in database.preregistered_managers: if manager.name == database_type: self.db = manager(self.ap, url_override=self._database_url_override) + self.db.persistence_mode = self.mode.value await self.db.initialize() selected_manager = self.db break @@ -195,6 +196,17 @@ class PersistenceManager: if engine is not None: await engine.dispose() + def get_resource_stats(self) -> dict[str, int]: + """Return database-manager-owned aggregate resource gauges.""" + + resource_stats = getattr(getattr(self, 'db', None), 'resource_stats', None) + if not callable(resource_stats): + return {} + try: + return resource_stats() + except Exception: + return {} + @contextlib.asynccontextmanager async def _release_migration_lock(self) -> typing.AsyncIterator[None]: """Serialize the complete PostgreSQL migration and validation window.""" @@ -1865,11 +1877,13 @@ class PersistenceManager: return active.session def _scoped_uow(self, scope: PersistenceScope) -> TenantUnitOfWork: + on_pool_timeout = getattr(getattr(self, 'db', None), 'record_pool_timeout', None) return TenantUnitOfWork( self.get_db_engine(), scope=scope, active_transaction=self._active_transaction, active_scope=self._active_scope, + on_pool_timeout=(on_pool_timeout if callable(on_pool_timeout) else None), ) def _get_active_transaction(self) -> ActiveScopedTransaction | None: diff --git a/src/langbot/pkg/persistence/tenant_uow.py b/src/langbot/pkg/persistence/tenant_uow.py index 9f56f2aff..27db2d64a 100644 --- a/src/langbot/pkg/persistence/tenant_uow.py +++ b/src/langbot/pkg/persistence/tenant_uow.py @@ -1220,6 +1220,7 @@ class TenantUnitOfWork: scope: PersistenceScope | None = None, active_transaction: ActiveTransactionVar | None = None, active_scope: ActivePersistenceScopeVar | None = None, + on_pool_timeout: typing.Callable[[], None] | None = None, ) -> None: if (workspace_uuid is None) == (scope is None): raise ValueError('TenantUnitOfWork requires exactly one Workspace or persistence scope') @@ -1229,6 +1230,7 @@ class TenantUnitOfWork: self.workspace_uuid = self.scope.settings[0][1] if self.scope.kind == PersistenceScopeKind.WORKSPACE else None self._active_transaction = active_transaction self._active_scope = active_scope + self._on_pool_timeout = on_pool_timeout self._session: sqlalchemy_asyncio.AsyncSession | None = None self._transaction: sqlalchemy_asyncio.AsyncSessionTransaction | None = None self._active_state: ActiveScopedTransaction | None = None @@ -1322,7 +1324,9 @@ class TenantUnitOfWork: self._context_token = self._active_transaction.set(state) self._database_operation_token = _DATABASE_OPERATION_TRANSACTION.set(state) self._owns_transaction = True - except BaseException: + except BaseException as exc: + if isinstance(exc, sqlalchemy.exc.TimeoutError) and self._on_pool_timeout is not None: + self._on_pool_timeout() if self._database_operation_token is not None: _DATABASE_OPERATION_TRANSACTION.reset(self._database_operation_token) self._database_operation_token = None diff --git a/src/langbot/templates/config.yaml b/src/langbot/templates/config.yaml index 84356b63c..857e33cfe 100644 --- a/src/langbot/templates/config.yaml +++ b/src/langbot/templates/config.yaml @@ -50,6 +50,22 @@ concurrency: # Hard admission limits for queued + running pipeline queries. pending_queries: 1000 pending_queries_per_workspace: 100 +cloud: + # Operational safety ceilings for the one logical Cloud instance. These + # are not subscription entitlements. An authoritative directory update + # that would exceed them is rejected atomically rather than truncated. + directory: + # Tune downward from the measured production capacity curve. Core has + # an absolute safety ceiling of 5,000 active Workspaces. + max_active_workspaces: 1000 + # Full snapshots contain current Workspaces only. Archived tombstones + # are delivered through bounded per-Workspace deltas. + max_snapshot_workspaces: 1000 + # Aggregate memberships accepted in one signed snapshot or delta. + max_snapshot_memberships: 20000 + # Signed control-plane envelope buffered by the closed adapter before + # JSON/JWS verification (32 MiB; absolute maximum 64 MiB). + max_response_bytes: 33554432 proxy: http: '' https: '' @@ -135,6 +151,11 @@ database: max_overflow: 10 pool_timeout_seconds: 30 pool_recycle_seconds: 1800 + # Applied only to Cloud runtime connections. The one-shot release + # migration uses its operator connection without these short limits. + statement_timeout_ms: 60000 + lock_timeout_ms: 5000 + idle_in_transaction_session_timeout_ms: 60000 cloud_migration: # `langbot migrate --cloud` reads an operator-only PostgreSQL DSN from # this environment variable. The operator role must differ from the diff --git a/tests/unit_tests/cloud/test_bootstrap.py b/tests/unit_tests/cloud/test_bootstrap.py index e28869b14..77ccefa5d 100644 --- a/tests/unit_tests/cloud/test_bootstrap.py +++ b/tests/unit_tests/cloud/test_bootstrap.py @@ -178,6 +178,34 @@ async def test_cloud_runtime_config_is_fail_closed(field, value, message): ) +@pytest.mark.parametrize( + ('directory_config', 'message'), + [ + ({'max_active_workspaces': 0}, 'greater than or equal to 1'), + ({'max_active_workspaces': True}, 'must be an integer'), + ( + { + 'max_active_workspaces': 10, + 'max_snapshot_workspaces': 9, + }, + 'max_snapshot_workspaces', + ), + ({'max_response_bytes': 64 * 1024 * 1024 + 1}, 'less than or equal to'), + ], +) +async def test_cloud_directory_capacity_contract_is_fail_closed(directory_config, message): + config = _cloud_config() + config['cloud'] = {'directory': directory_config} + + with pytest.raises(CloudBootstrapError, match=message): + await resolve_deployment( + instance_uuid='instance-a', + instance_config=config, + entry_points=lambda: _EntryPoints([_EntryPoint(_Provider())]), + now=1_000, + ) + + @pytest.mark.parametrize( ('pgvector_config', 'message'), [ diff --git a/tests/unit_tests/cloud/test_directory_projection.py b/tests/unit_tests/cloud/test_directory_projection.py index f19e9fe0f..ec11b9f74 100644 --- a/tests/unit_tests/cloud/test_directory_projection.py +++ b/tests/unit_tests/cloud/test_directory_projection.py @@ -14,6 +14,7 @@ from langbot.pkg.cloud.directory import ( DirectoryEvent, DirectoryEventBatch, DirectoryMember, + DirectoryProjectionLimits, DirectoryProjectionUnavailableError, DirectorySnapshot, DirectoryWorkspace, @@ -274,6 +275,187 @@ async def test_event_limit_matches_single_delta_request_limit(projection_context ) +async def test_snapshot_capacity_rejects_atomically_before_projection(projection_context): + application, session_factory = projection_context + service = DirectoryProjectionService( + application, + _Provider([_snapshot(1)]), + INSTANCE_UUID, + limits=DirectoryProjectionLimits( + max_active_workspaces=1, + max_snapshot_workspaces=1, + max_snapshot_memberships=1, + ), + ) + second_workspace = DirectoryWorkspace( + uuid=SECOND_WORKSPACE_UUID, + name='Second Workspace', + slug='second-workspace', + type='personal', + status='active', + created_by_account_uuid='20000000-0000-0000-0000-000000000002', + projection_revision=1, + execution_generation=1, + members=[ + DirectoryMember( + membership_uuid=SECOND_MEMBERSHIP_UUID, + account_uuid='20000000-0000-0000-0000-000000000002', + normalized_email='second@example.com', + display_name='Second Owner', + account_status='active', + role='owner', + membership_status='active', + projection_revision=1, + ) + ], + ) + + with pytest.raises(DirectoryProjectionUnavailableError, match='Workspace capacity exceeded'): + await service.apply_snapshot( + _snapshot( + 1, + workspaces=[ + _workspace(), + second_workspace, + ], + ) + ) + + async with session_factory() as session: + assert await session.scalar(sqlalchemy.select(sqlalchemy.func.count()).select_from(Workspace)) == 0 + assert await session.get(DirectoryProjectionState, INSTANCE_UUID) is None + + +async def test_incremental_capacity_rolls_back_without_advancing_cursor(projection_context): + application, session_factory = projection_context + service = DirectoryProjectionService( + application, + _Provider([_snapshot(1)]), + INSTANCE_UUID, + limits=DirectoryProjectionLimits( + max_active_workspaces=1, + max_snapshot_workspaces=1, + max_snapshot_memberships=2, + ), + ) + await service.initialize() + + second_account_uuid = '20000000-0000-0000-0000-000000000002' + second_workspace = DirectoryWorkspace( + uuid=SECOND_WORKSPACE_UUID, + name='Second Workspace', + slug='second-workspace', + type='personal', + status='active', + created_by_account_uuid=second_account_uuid, + projection_revision=2, + execution_generation=1, + members=[ + DirectoryMember( + membership_uuid=SECOND_MEMBERSHIP_UUID, + account_uuid=second_account_uuid, + normalized_email='second@example.com', + display_name='Second Owner', + account_status='active', + role='owner', + membership_status='active', + projection_revision=2, + ) + ], + ) + event = DirectoryEvent( + cursor=2, + uuid='40000000-0000-0000-0000-000000000002', + aggregate_uuid=SECOND_WORKSPACE_UUID, + event_type='directory.changed', + revision=2, + payload={ + 'workspace_uuid': SECOND_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], + ) + delta = DirectoryDelta( + instance_uuid=INSTANCE_UUID, + requested_workspace_uuids=[SECOND_WORKSPACE_UUID], + generated_at=datetime.datetime(2026, 7, 24, 12, 2, tzinfo=datetime.UTC), + workspaces=[second_workspace], + ) + + with pytest.raises(DirectoryProjectionUnavailableError, match='Projected active Workspace capacity exceeded'): + await service.apply_delta(delta, batch) + + async with session_factory() as session: + assert await session.get(Workspace, SECOND_WORKSPACE_UUID) is None + state = await session.get(DirectoryProjectionState, INSTANCE_UUID) + assert state is not None + assert state.cursor == 1 + assert service.resource_snapshot()['active_workspaces'] == 1 + + +async def test_account_projection_reads_large_directory_in_bounded_batches(projection_context): + application, session_factory = projection_context + service = DirectoryProjectionService( + application, + _Provider([_snapshot(1)]), + INSTANCE_UUID, + ) + workspaces = [] + for number in range(501): + account_uuid = f'20000000-0000-0000-0000-{number:012d}' + workspaces.append( + DirectoryWorkspace( + uuid=f'10000000-0000-0000-0000-{number:012d}', + name=f'Workspace {number}', + slug=f'workspace-{number}', + type='personal', + status='active', + created_by_account_uuid=account_uuid, + projection_revision=1, + execution_generation=1, + members=[ + DirectoryMember( + membership_uuid=f'30000000-0000-0000-0000-{number:012d}', + account_uuid=account_uuid, + normalized_email=f'owner-{number}@example.com', + display_name=f'Owner {number}', + account_status='active', + role='owner', + membership_status='active', + projection_revision=1, + ) + ], + ) + ) + snapshot = _snapshot(1, workspaces=workspaces) + user_selects = 0 + + def count_user_selects(_connection, _cursor, statement, _parameters, _context, _executemany): + nonlocal user_selects + if statement.lstrip().upper().startswith('SELECT') and 'users' in statement: + user_selects += 1 + + engine = application.persistence_mgr.engine + sqlalchemy.event.listen(engine.sync_engine, 'before_cursor_execute', count_user_selects) + try: + async with application.persistence_mgr.directory_projection_uow(INSTANCE_UUID) as uow: + projected = await service._apply_accounts(uow.session, snapshot) + finally: + sqlalchemy.event.remove(engine.sync_engine, 'before_cursor_execute', count_user_selects) + + assert len(projected) == 501 + assert user_selects == 2 + async with session_factory() as session: + assert await session.scalar(sqlalchemy.select(sqlalchemy.func.count()).select_from(User)) == 501 + + async def test_archived_and_absent_workspaces_are_execution_fenced(projection_context): application, session_factory = projection_context service = DirectoryProjectionService( diff --git a/tests/unit_tests/core/test_app_shutdown.py b/tests/unit_tests/core/test_app_shutdown.py index 7f17ab825..e284f5000 100644 --- a/tests/unit_tests/core/test_app_shutdown.py +++ b/tests/unit_tests/core/test_app_shutdown.py @@ -88,6 +88,18 @@ async def test_runtime_resource_stats_are_aggregate_and_constant_time() -> None: app.pipeline_mgr = SimpleNamespace(_pipelines_by_key={}) app.rag_mgr = SimpleNamespace(knowledge_bases={}) app.plugin_connector = SimpleNamespace(_known_desired_states={'installation': object()}) + app.persistence_mgr = SimpleNamespace( + get_resource_stats=lambda: { + 'configured_capacity': 20, + 'checked_out': 3, + } + ) + app.directory_projection_service = SimpleNamespace( + resource_snapshot=lambda: { + 'active_workspaces': 10, + 'max_active_workspaces': 1000, + } + ) app.tool_mgr = SimpleNamespace( mcp_tool_loader=SimpleNamespace( _sessions={}, @@ -113,6 +125,14 @@ async def test_runtime_resource_stats_are_aggregate_and_constant_time() -> None: 'total': 5, 'completed': 2, } + assert stats['database_pool'] == { + 'configured_capacity': 20, + 'checked_out': 3, + } + assert stats['directory'] == { + 'active_workspaces': 10, + 'max_active_workspaces': 1000, + } assert stats['query_pool'] == { 'queued': 1, 'cached': 0, diff --git a/tests/unit_tests/core/test_load_config.py b/tests/unit_tests/core/test_load_config.py index f2f93ba9a..b95bc25ba 100644 --- a/tests/unit_tests/core/test_load_config.py +++ b/tests/unit_tests/core/test_load_config.py @@ -64,6 +64,20 @@ class TestApplyEnvOverridesToConfig: assert result['concurrency']['pipeline'] == 10 assert isinstance(result['concurrency']['pipeline'], int) + def test_cloud_directory_limit_override_keeps_integer_type_on_upgraded_config(self): + load_config = get_load_config_module() + cfg = load_config._complete_runtime_policy_defaults({}) + + with patch.dict( + os.environ, + {'CLOUD__DIRECTORY__MAX_ACTIVE_WORKSPACES': '250'}, + clear=True, + ): + result = load_config._apply_env_overrides_to_config(cfg) + + assert result['cloud']['directory']['max_active_workspaces'] == 250 + assert isinstance(result['cloud']['directory']['max_active_workspaces'], int) + def test_override_int_value_invalid_conversion(self): """Test that invalid int conversion keeps string value.""" load_config = get_load_config_module() diff --git a/tests/unit_tests/persistence/test_postgresql_database.py b/tests/unit_tests/persistence/test_postgresql_database.py index 63a3b751b..ad903ff95 100644 --- a/tests/unit_tests/persistence/test_postgresql_database.py +++ b/tests/unit_tests/persistence/test_postgresql_database.py @@ -121,6 +121,43 @@ async def test_postgresql_manager_applies_explicit_bounded_pool_options(monkeypa } +@pytest.mark.asyncio +async def test_cloud_postgresql_manager_applies_bounded_server_timeouts(monkeypatch) -> None: + captured_options = None + + def create_engine(_url, **options): + nonlocal captured_options + captured_options = options + return object() + + monkeypatch.setattr(postgresql.sqlalchemy_asyncio, 'create_async_engine', create_engine) + ap = SimpleNamespace( + instance_config=SimpleNamespace( + data={ + 'database': { + 'postgresql': { + 'statement_timeout_ms': 45_000, + 'lock_timeout_ms': 4_000, + 'idle_in_transaction_session_timeout_ms': 55_000, + } + } + } + ) + ) + + manager = postgresql.PostgreSQLDatabaseManager(ap) + manager.persistence_mode = 'cloud_runtime' + await manager.initialize() + + assert captured_options['connect_args'] == { + 'server_settings': { + 'statement_timeout': '45000', + 'lock_timeout': '4000', + 'idle_in_transaction_session_timeout': '55000', + } + } + + @pytest.mark.asyncio @pytest.mark.parametrize( ('name', 'value'), @@ -128,8 +165,12 @@ async def test_postgresql_manager_applies_explicit_bounded_pool_options(monkeypa ('pool_size', 0), ('pool_size', True), ('max_overflow', -1), + ('pool_size', 101), + ('max_overflow', 101), ('pool_timeout_seconds', 0), + ('pool_timeout_seconds', 301), ('pool_recycle_seconds', '1800'), + ('pool_recycle_seconds', 86401), ], ) async def test_postgresql_manager_rejects_invalid_pool_options(name, value) -> None: @@ -139,6 +180,45 @@ async def test_postgresql_manager_rejects_invalid_pool_options(name, value) -> N await postgresql.PostgreSQLDatabaseManager(ap).initialize() +@pytest.mark.asyncio +async def test_postgresql_manager_rejects_combined_pool_capacity_above_hard_ceiling() -> None: + ap = SimpleNamespace( + instance_config=SimpleNamespace( + data={ + 'database': { + 'postgresql': { + 'pool_size': 60, + 'max_overflow': 41, + } + } + } + ) + ) + + with pytest.raises(ValueError, match=r'pool_size \+ max_overflow'): + await postgresql.PostgreSQLDatabaseManager(ap).initialize() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ('name', 'value'), + [ + ('statement_timeout_ms', 0), + ('statement_timeout_ms', 300_001), + ('lock_timeout_ms', 60_001), + ('idle_in_transaction_session_timeout_ms', True), + ('idle_in_transaction_session_timeout_ms', 300_001), + ], +) +async def test_cloud_postgresql_manager_rejects_unsafe_server_timeouts(name, value) -> None: + ap = SimpleNamespace(instance_config=SimpleNamespace(data={'database': {'postgresql': {name: value}}})) + + with pytest.raises(ValueError, match=rf'database\.postgresql\.{name}'): + manager = postgresql.PostgreSQLDatabaseManager(ap) + manager.persistence_mode = 'cloud_runtime' + await manager.initialize() + + @pytest.mark.asyncio async def test_postgresql_manager_rejects_non_postgresql_url_without_echoing_secret() -> None: ap = SimpleNamespace( diff --git a/tests/unit_tests/persistence/test_tenant_uow.py b/tests/unit_tests/persistence/test_tenant_uow.py index 75634b968..7be62bf2e 100644 --- a/tests/unit_tests/persistence/test_tenant_uow.py +++ b/tests/unit_tests/persistence/test_tenant_uow.py @@ -26,6 +26,7 @@ from langbot.pkg.persistence.tenant_uow import ( CrossScopeTransactionError, PersistenceScopeKind, ScopedSessionTransactionError, + TenantScopedAsyncSession, TenantScopedSyncSession, TenantScopeRequiredError, TenantUnitOfWork, @@ -37,6 +38,33 @@ from langbot.pkg.persistence.tenant_uow import ( pytestmark = pytest.mark.asyncio +async def test_tenant_uow_reports_pool_timeout_during_transaction_admission(monkeypatch) -> None: + engine = create_async_engine('sqlite+aiosqlite:///:memory:') + pool_timeouts = 0 + + def record_pool_timeout() -> None: + nonlocal pool_timeouts + pool_timeouts += 1 + + async def fail_transaction_start(self, capability): + del self, capability + raise sa.exc.TimeoutError('pool exhausted') + + monkeypatch.setattr(TenantScopedAsyncSession, '_start_owned_transaction', fail_transaction_start) + try: + with pytest.raises(sa.exc.TimeoutError, match='pool exhausted'): + async with TenantUnitOfWork( + engine, + '10000000-0000-0000-0000-000000000001', + on_pool_timeout=record_pool_timeout, + ): + pass + finally: + await engine.dispose() + + assert pool_timeouts == 1 + + def _on_conflict_statement(*, update_value, update_key='value', index_element=None): table = sa.table('conflict_rows', sa.column('id'), sa.column('value')) if index_element is None: diff --git a/tests/unit_tests/utils/test_cloud_runtime_soak.py b/tests/unit_tests/utils/test_cloud_runtime_soak.py index a3c5cbd51..c10b14314 100644 --- a/tests/unit_tests/utils/test_cloud_runtime_soak.py +++ b/tests/unit_tests/utils/test_cloud_runtime_soak.py @@ -461,6 +461,58 @@ def test_evaluate_gate_detects_stuck_mcp_projection_cleanup() -> None: assert any('mcp_projection_reconcile_active above zero' in failure for failure in result.failures) +def test_evaluate_gate_detects_directory_and_database_capacity_violation() -> None: + prefix = 'body.resources' + state = _state( + 'endpoint', + [ + _sample( + 0, + **{ + f'{prefix}.directory.active_workspaces': 1000, + f'{prefix}.directory.max_active_workspaces': 1000, + f'{prefix}.database_pool.checked_out': 20, + f'{prefix}.database_pool.configured_capacity': 20, + }, + ), + _sample( + 60, + **{ + f'{prefix}.directory.active_workspaces': 1001, + f'{prefix}.directory.max_active_workspaces': 1000, + f'{prefix}.database_pool.checked_out': 21, + f'{prefix}.database_pool.configured_capacity': 20, + }, + ), + ], + ) + + result = soak.evaluate_gate( + [state], + analysis_start_seconds=0, + thresholds=_thresholds(), + ) + + assert any('directory.active_workspaces' in failure for failure in result.failures) + assert any('database_pool.checked_out' in failure for failure in result.failures) + + +def test_evaluate_gate_requires_capacity_gauges_from_named_core_endpoint() -> None: + state = _state( + 'endpoint', + [_sample(0, **{'http.ok': 1}), _sample(60, **{'http.ok': 1})], + ) + state.target = soak.Target(name='core', kind='endpoint', location='/core') + + result = soak.evaluate_gate( + [state], + analysis_start_seconds=0, + thresholds=_thresholds(require_event_loop_metrics=False), + ) + + assert any('required capacity gauge' in failure for failure in result.failures) + + def test_evaluate_gate_detects_event_loop_stall_and_sustained_lag() -> None: prefix = 'body.resources.event_loop' samples = [ diff --git a/uv.lock b/uv.lock index 074887671..c233e1d3c 100644 --- a/uv.lock +++ b/uv.lock @@ -2116,7 +2116,7 @@ requires-dist = [ { name = "ebooklib", specifier = ">=0.18" }, { name = "gewechat-client", specifier = ">=0.1.5" }, { name = "html2text", specifier = ">=2024.2.26" }, - { name = "langbot-plugin", git = "https://github.com/langbot-app/langbot-plugin-sdk.git?rev=c7893f8f94e83cb09ebcc98e25490fa699b0fb69" }, + { name = "langbot-plugin", git = "https://github.com/langbot-app/langbot-plugin-sdk.git?rev=7c0b9827ed8597a1c84151b83fcf6307934fd944" }, { name = "langchain", specifier = ">=1.3.9" }, { name = "langchain-core", specifier = ">=1.3.3" }, { name = "langchain-text-splitters", specifier = ">=1.1.2" }, @@ -2183,7 +2183,7 @@ dev = [ [[package]] name = "langbot-plugin" version = "0.4.18" -source = { git = "https://github.com/langbot-app/langbot-plugin-sdk.git?rev=c7893f8f94e83cb09ebcc98e25490fa699b0fb69#c7893f8f94e83cb09ebcc98e25490fa699b0fb69" } +source = { git = "https://github.com/langbot-app/langbot-plugin-sdk.git?rev=7c0b9827ed8597a1c84151b83fcf6307934fd944#7c0b9827ed8597a1c84151b83fcf6307934fd944" } dependencies = [ { name = "aiofiles" }, { name = "aiohttp" },