mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-08-01 09:06:08 +00:00
feat(cloud): harden multi-tenant runtime resources
This commit is contained in:
Junyan Qin
@@ -17,6 +17,10 @@ services:
|
||||
# Shared with the langbot service and sent only as a WebSocket handshake
|
||||
# header. Generate with: openssl rand -hex 32
|
||||
- LANGBOT_PLUGIN_RUNTIME_CONTROL_TOKEN=${LANGBOT_PLUGIN_RUNTIME_CONTROL_TOKEN:-}
|
||||
# Process-wide admission for every asyncio.to_thread() call.
|
||||
- LANGBOT_BLOCKING_EXECUTOR_MAX_WORKERS=${LANGBOT_BLOCKING_EXECUTOR_MAX_WORKERS:-8}
|
||||
- LANGBOT_BLOCKING_EXECUTOR_MAX_PENDING=${LANGBOT_BLOCKING_EXECUTOR_MAX_PENDING:-128}
|
||||
- LANGBOT_BLOCKING_EXECUTOR_MAX_INFLIGHT_PER_SCOPE=${LANGBOT_BLOCKING_EXECUTOR_MAX_INFLIGHT_PER_SCOPE:-4}
|
||||
command: ["uv", "run", "--no-sync", "-m", "langbot_plugin.cli.__init__", "rt"]
|
||||
networks:
|
||||
- langbot_network
|
||||
@@ -49,10 +53,13 @@ services:
|
||||
# An empty value is accepted by Compose so Box can remain optional, but
|
||||
# the Box runtime itself fails closed when the profile is started.
|
||||
- LANGBOT_BOX_CONTROL_TOKEN=${LANGBOT_BOX_CONTROL_TOKEN:-}
|
||||
# Box has its own process-wide blocking-work budget.
|
||||
- LANGBOT_BLOCKING_EXECUTOR_MAX_WORKERS=${LANGBOT_BLOCKING_EXECUTOR_MAX_WORKERS:-8}
|
||||
- LANGBOT_BLOCKING_EXECUTOR_MAX_PENDING=${LANGBOT_BLOCKING_EXECUTOR_MAX_PENDING:-128}
|
||||
- LANGBOT_BLOCKING_EXECUTOR_MAX_INFLIGHT_PER_SCOPE=${LANGBOT_BLOCKING_EXECUTOR_MAX_INFLIGHT_PER_SCOPE:-4}
|
||||
# The Box runtime does NOT read box.local.* from config.yaml or env; it
|
||||
# receives its functional configuration from LangBot via the INIT RPC
|
||||
# action. LANGBOT_BOX_CONTROL_TOKEN is the intentional security-only
|
||||
# exception; do not add BOX__* here because those would be ignored.
|
||||
# action. Do not add BOX__* here because those would be ignored.
|
||||
# Launched through the same CLI entry point as the plugin runtime
|
||||
# (`langbot_plugin.cli.__init__ <subcommand>`). WebSocket is the default
|
||||
# control transport — mirrors `rt`, which also runs with no flag. Pass
|
||||
@@ -76,6 +83,11 @@ services:
|
||||
# Must match the value supplied to langbot_box. The token is sent only
|
||||
# in WebSocket handshake headers, never in URLs or action payloads.
|
||||
- LANGBOT_BOX_CONTROL_TOKEN=${LANGBOT_BOX_CONTROL_TOKEN:-}
|
||||
# Core process-wide blocking-work admission. These are native config
|
||||
# overrides and are persisted with the effective data/config.yaml.
|
||||
- SYSTEM__BLOCKING_EXECUTOR__MAX_WORKERS=${LANGBOT_BLOCKING_EXECUTOR_MAX_WORKERS:-8}
|
||||
- SYSTEM__BLOCKING_EXECUTOR__MAX_PENDING=${LANGBOT_BLOCKING_EXECUTOR_MAX_PENDING:-128}
|
||||
- SYSTEM__BLOCKING_EXECUTOR__MAX_INFLIGHT_PER_SCOPE=${LANGBOT_BLOCKING_EXECUTOR_MAX_INFLIGHT_PER_SCOPE:-4}
|
||||
# Unified env-override convention: SECTION__SUBSECTION__KEY overrides the
|
||||
# matching config.yaml field (see LoadConfigStage). These map onto
|
||||
# box.* and are forwarded to the Box runtime via INIT RPC.
|
||||
|
||||
+56
-4
@@ -91,6 +91,12 @@ metadata:
|
||||
data:
|
||||
TZ: "Asia/Shanghai"
|
||||
PLUGIN__RUNTIME_WS_URL: "ws://langbot-plugin-runtime:5400/control/ws"
|
||||
SYSTEM__BLOCKING_EXECUTOR__MAX_WORKERS: "8"
|
||||
SYSTEM__BLOCKING_EXECUTOR__MAX_PENDING: "128"
|
||||
SYSTEM__BLOCKING_EXECUTOR__MAX_INFLIGHT_PER_SCOPE: "4"
|
||||
LANGBOT_BLOCKING_EXECUTOR_MAX_WORKERS: "8"
|
||||
LANGBOT_BLOCKING_EXECUTOR_MAX_PENDING: "128"
|
||||
LANGBOT_BLOCKING_EXECUTOR_MAX_INFLIGHT_PER_SCOPE: "4"
|
||||
# Box sandbox runtime endpoint. LangBot connects to the Box runtime over
|
||||
# WebSocket. The hostname MUST match the langbot-box Service name. Note the
|
||||
# in-container default ("langbot_box") uses an underscore, which is an
|
||||
@@ -136,6 +142,21 @@ spec:
|
||||
secretKeyRef:
|
||||
name: langbot-plugin-runtime-control
|
||||
key: token
|
||||
- name: LANGBOT_BLOCKING_EXECUTOR_MAX_WORKERS
|
||||
valueFrom:
|
||||
configMapKeyRef:
|
||||
name: langbot-config
|
||||
key: LANGBOT_BLOCKING_EXECUTOR_MAX_WORKERS
|
||||
- name: LANGBOT_BLOCKING_EXECUTOR_MAX_PENDING
|
||||
valueFrom:
|
||||
configMapKeyRef:
|
||||
name: langbot-config
|
||||
key: LANGBOT_BLOCKING_EXECUTOR_MAX_PENDING
|
||||
- name: LANGBOT_BLOCKING_EXECUTOR_MAX_INFLIGHT_PER_SCOPE
|
||||
valueFrom:
|
||||
configMapKeyRef:
|
||||
name: langbot-config
|
||||
key: LANGBOT_BLOCKING_EXECUTOR_MAX_INFLIGHT_PER_SCOPE
|
||||
volumeMounts:
|
||||
- name: plugin-data
|
||||
mountPath: /app/data/plugins
|
||||
@@ -262,9 +283,23 @@ spec:
|
||||
secretKeyRef:
|
||||
name: langbot-box-control
|
||||
key: token
|
||||
- name: LANGBOT_BLOCKING_EXECUTOR_MAX_WORKERS
|
||||
valueFrom:
|
||||
configMapKeyRef:
|
||||
name: langbot-config
|
||||
key: LANGBOT_BLOCKING_EXECUTOR_MAX_WORKERS
|
||||
- name: LANGBOT_BLOCKING_EXECUTOR_MAX_PENDING
|
||||
valueFrom:
|
||||
configMapKeyRef:
|
||||
name: langbot-config
|
||||
key: LANGBOT_BLOCKING_EXECUTOR_MAX_PENDING
|
||||
- name: LANGBOT_BLOCKING_EXECUTOR_MAX_INFLIGHT_PER_SCOPE
|
||||
valueFrom:
|
||||
configMapKeyRef:
|
||||
name: langbot-config
|
||||
key: LANGBOT_BLOCKING_EXECUTOR_MAX_INFLIGHT_PER_SCOPE
|
||||
# The Box runtime does NOT read box.local.* / BOX__* from its own env;
|
||||
# it receives its functional configuration from LangBot via INIT.
|
||||
# LANGBOT_BOX_CONTROL_TOKEN is the security-only exception.
|
||||
volumeMounts:
|
||||
# Box workspace root — identical path on node, box, and sandbox
|
||||
# containers (see the IMPORTANT note above).
|
||||
@@ -290,7 +325,9 @@ spec:
|
||||
failureThreshold: 3
|
||||
readinessProbe:
|
||||
httpGet:
|
||||
path: /healthz
|
||||
# Unlike liveness, readiness validates the configured backend and
|
||||
# all strict managed-mode isolation guarantees.
|
||||
path: /readyz
|
||||
port: 5410
|
||||
initialDelaySeconds: 10
|
||||
periodSeconds: 5
|
||||
@@ -379,6 +416,21 @@ spec:
|
||||
secretKeyRef:
|
||||
name: langbot-plugin-runtime-control
|
||||
key: token
|
||||
- name: SYSTEM__BLOCKING_EXECUTOR__MAX_WORKERS
|
||||
valueFrom:
|
||||
configMapKeyRef:
|
||||
name: langbot-config
|
||||
key: SYSTEM__BLOCKING_EXECUTOR__MAX_WORKERS
|
||||
- name: SYSTEM__BLOCKING_EXECUTOR__MAX_PENDING
|
||||
valueFrom:
|
||||
configMapKeyRef:
|
||||
name: langbot-config
|
||||
key: SYSTEM__BLOCKING_EXECUTOR__MAX_PENDING
|
||||
- name: SYSTEM__BLOCKING_EXECUTOR__MAX_INFLIGHT_PER_SCOPE
|
||||
valueFrom:
|
||||
configMapKeyRef:
|
||||
name: langbot-config
|
||||
key: SYSTEM__BLOCKING_EXECUTOR__MAX_INFLIGHT_PER_SCOPE
|
||||
# Box (sandbox) runtime endpoint. Connects LangBot to the langbot-box
|
||||
# Service over WebSocket. Remove this (and the langbot-box Deployment)
|
||||
# and set BOX__ENABLED=false if you do not want the sandbox.
|
||||
@@ -426,7 +478,7 @@ spec:
|
||||
# Liveness probe to restart container if it becomes unresponsive
|
||||
livenessProbe:
|
||||
httpGet:
|
||||
path: /
|
||||
path: /healthz
|
||||
port: 5300
|
||||
initialDelaySeconds: 60
|
||||
periodSeconds: 10
|
||||
@@ -435,7 +487,7 @@ spec:
|
||||
# Readiness probe to know when container is ready to accept traffic
|
||||
readinessProbe:
|
||||
httpGet:
|
||||
path: /
|
||||
path: /healthz
|
||||
port: 5300
|
||||
initialDelaySeconds: 30
|
||||
periodSeconds: 5
|
||||
|
||||
@@ -0,0 +1,106 @@
|
||||
# LangBot Cloud 24 小时资源 Soak 门禁
|
||||
|
||||
`scripts/cloud_runtime_soak.py` 是生产候选拓扑的最终资源稳定性门禁。它不替代单元测试、历史 churn 探针或 nsjail 隔离测试；它把以下三类证据按同一时间轴采集并给出可机读的 pass/fail：
|
||||
|
||||
- Core、Plugin Runtime 和 Box Runtime 的 HTTP liveness/readiness。
|
||||
- 三个 Python 进程的 event-loop recent max/p95 调度延迟。
|
||||
- Linux `/proc` 进程树的 current RSS、累计 CPU、线程、文件描述符和子进程数。
|
||||
- cgroup v2 的 `memory.current/peak/events`、swap、CPU usage/throttling、PID current/events 和实际硬限制。
|
||||
|
||||
生产批准必须使用 cgroup 证据。`--pid` 只适合本地诊断，因为进程指标无法证明 OOM kill、PID limit 或 CPU throttling。
|
||||
|
||||
## 运行位置
|
||||
|
||||
建议把采集器放在独立的 node agent 或监控 sidecar 中，并只读挂载三个目标容器的 cgroup 路径。不要把采集器和样本文件放进被测容器自己的 cgroup/数据卷，否则采集器的 CPU、内存和 page cache 会污染目标数据。
|
||||
|
||||
Kubernetes/containerd 生成的 cgroup 路径不是稳定 API。每次生产候选部署都必须从实际 pod/container ID 解析，不能从 pod 名猜路径。传入的每个目录都必须至少可读：
|
||||
|
||||
- `memory.current`、`memory.events`
|
||||
- `cpu.stat`、`cpu.max`
|
||||
- `pids.current`、`pids.events`
|
||||
|
||||
最终门禁应加 `--require-hard-limits`。该选项要求每个目标 cgroup 都能观察到有限的 CPU quota、memory、swap 和 PID 上限；任一值为 `max` 都失败。
|
||||
|
||||
## 标准 24 小时命令
|
||||
|
||||
```bash
|
||||
uv run python scripts/cloud_runtime_soak.py \
|
||||
--duration 24h \
|
||||
--startup-grace 5m \
|
||||
--sample-interval 15s \
|
||||
--cooldown 30m \
|
||||
--analysis-window 30m \
|
||||
--http-timeout 5s \
|
||||
--max-memory-growth-mib 64 \
|
||||
--max-memory-slope-mib-per-hour 32 \
|
||||
--max-tail-cpu-cores 0.5 \
|
||||
--max-throttled-period-ratio 0.25 \
|
||||
--max-event-loop-lag-ms 1000 \
|
||||
--max-event-loop-p95-lag-ms 250 \
|
||||
--require-hard-limits \
|
||||
--endpoint core=http://langbot:5300/healthz \
|
||||
--endpoint plugin=http://langbot-plugin-runtime:5400/healthz \
|
||||
--endpoint box=http://langbot-box:5410/readyz \
|
||||
--cgroup core=/host-cgroup/CURRENT_CORE_CONTAINER \
|
||||
--cgroup plugin=/host-cgroup/CURRENT_PLUGIN_RUNTIME_CONTAINER \
|
||||
--cgroup box=/host-cgroup/CURRENT_BOX_RUNTIME_CONTAINER \
|
||||
--samples-file artifacts/cloud-soak-samples.jsonl \
|
||||
--report-file artifacts/cloud-soak-report.json \
|
||||
--workload uv run python tests/load/cloud_candidate_workload.py
|
||||
```
|
||||
|
||||
`--duration` 是包含启动观察、负载和冷却期的最大墙钟时间。工作负载必须在截止时间前退出并至少留出 30 分钟冷却；否则门禁会终止负载并失败。若负载由外部系统控制，可以省略 `--workload`，但必须保证最后 `--analysis-window` 完全无测试流量，该窗口才可解释为空闲尾段。
|
||||
|
||||
工作负载命令的 stdout/stderr 会转发到采集器 stderr，不会混入 stdout 的最终 JSON 报告。命令以独立 process group 启动；超时或中断时整组收到 TERM，10 秒后仍未退出则收到 KILL。
|
||||
|
||||
凭据只能通过 workload 进程环境或 secret mount 注入，不能放在命令参数中。最终报告只记录可执行文件名和参数个数，不保存参数正文；采集器也拒绝带 userinfo、query 或 fragment 的健康 URL。
|
||||
|
||||
## 必须覆盖的负载
|
||||
|
||||
同一候选版本至少要覆盖：
|
||||
|
||||
1. 大批 Workspace 注册、成员邀请、登录和 entitlement 刷新。
|
||||
2. Plugin installation reconcile、依赖准备、正常调用、进程崩溃与重启。
|
||||
3. Dashboard/Embed/平台 WebSocket 建连、突发消息和批量断连。
|
||||
4. Box session、文件同步、并发 exec、managed-process 输出和清理。
|
||||
5. PostgreSQL pool 接近容量、事务超时和恢复。
|
||||
6. Core、Plugin Runtime、Box 分别收到 SIGTERM 后的优雅重启。
|
||||
|
||||
工作负载不能把 API 过载拒绝当作成功吞掉。默认情况下，Core health 中 blocking executor 的 global/scope rejection counter 只要增长，门禁即失败；只有专门验证“过载会正确返回 429”的独立测试才可以使用 `--allow-rejections`，该次运行不能作为生产批准证据。
|
||||
|
||||
## 判定规则
|
||||
|
||||
整个有效观察期内出现以下任一情况即失败：
|
||||
|
||||
- 健康接口请求失败、非 2xx、Core `code != 0`，或 Box `ready=false`。
|
||||
- `memory.events.high/max/oom/oom_kill/oom_group_kill` 增长。
|
||||
- `pids.events.max` 增长。
|
||||
- cgroup 单调计数器回退，表示目标很可能发生了未记录的重启或 cgroup 替换。
|
||||
- CPU throttled-period ratio 超过配置阈值。
|
||||
- 任一健康采样窗口的 event-loop recent max 超过 1 秒，或冷却尾段 recent p95 超过 250 ms。
|
||||
- 健康接口缺少 event-loop monitor、monitor 未持续运行，或其 sample counter 回退。
|
||||
- blocking executor rejection counter 增长。
|
||||
|
||||
负载结束后的冷却尾段还必须满足：
|
||||
|
||||
- `memory.current`/RSS 的稳健首尾增长和线性斜率不能同时超过阈值。
|
||||
- 平均 CPU 核数不超过 `--max-tail-cpu-cores`。
|
||||
- event-loop recent p95 不超过 `--max-event-loop-p95-lag-ms`。
|
||||
- blocking executor `pending` 至少回到过零；不能整个尾段持续积压。
|
||||
- telemetry、QueryPool、MCP host/dispatch、Box creating/closing/background 等临时 gauge 不能继续增长。
|
||||
|
||||
内存判定要求“增长量”和“斜率”同时越界，避免几 MiB allocator/page-cache 噪声在短窗口被外推成很大的每小时斜率。最终报告仍保留实际增长与斜率，人工审查时不能只看 verdict。
|
||||
|
||||
## 产物与退出码
|
||||
|
||||
- `--samples-file`：逐样本 JSONL，写入后立即 flush，供时序图和故障定位。
|
||||
- `--report-file`：最终汇总、阈值、资源硬限制、OOM/PID/throttle delta、尾段斜率和 workload 状态。
|
||||
- stdout：与 report 文件相同的最终 JSON；workload 日志只写 stderr。
|
||||
|
||||
退出码：
|
||||
|
||||
- `0`：全部门禁通过。
|
||||
- `1`：采样完成但资源门禁失败。
|
||||
- `2`：CLI 参数或目标配置错误。
|
||||
|
||||
必须保存原始 JSONL、最终报告、三个镜像 digest、LangBot/SDK commit、生产配置摘要和工作负载版本。滚动更新、节点迁移或镜像变化后，旧报告不能继续作为新候选版本的批准证据。
|
||||
@@ -0,0 +1,212 @@
|
||||
# Cloud v2 仍待验证事项
|
||||
|
||||
状态：`NOT APPROVED FOR SAAS ACTIVATION`
|
||||
|
||||
更新日期：2026-07-29
|
||||
|
||||
本文是 Cloud v2 首期上线前的剩余验证清单。它只记录尚不能由当前代码审查、
|
||||
单元测试、集成测试、合成容量探针或短时 Linux 容器实验替代的证据。
|
||||
|
||||
相关文档：
|
||||
|
||||
- [多租户架构决策](./pending-architecture-decisions.md)
|
||||
- [实现决策记录](./implementation-decisions.md)
|
||||
- [Runtime 资源安全审查](./runtime-resource-audit-2026-07-28.md)
|
||||
- [24 小时资源 Soak 门禁](./cloud-runtime-soak-gate.md)
|
||||
|
||||
## 1. 当前已形成的交付基线
|
||||
|
||||
- LangBot Core 和 Plugin SDK 的全量测试、Ruff、`git diff --check` 已通过。
|
||||
- Plugin Runtime 和 Box Runtime 的公开健康接口、event-loop lag 与有界
|
||||
blocking executor 指标已经过真实进程短时验证。
|
||||
- 仓库 Dockerfile 构建的 Linux/cgroup v2 短时探针已证明 CPU、memory、
|
||||
swap 和 PID 限制代码路径可工作。
|
||||
- PostgreSQL 16 + RLS 的 1,000 Workspace 真实启动测试，以及 5,000
|
||||
Workspace 三代替换合成探针已通过。
|
||||
- Core 已精确钉住 Plugin SDK 提交
|
||||
`67be7c332ded4c1afa131b28232590cf9728f817`。最终验证必须使用包含该提交的
|
||||
Core、Plugin Runtime 和 Box Runtime 镜像，不能混用旧 SDK。
|
||||
|
||||
以上结果是进入生产候选验证的前提，不是 SaaS 上线批准。
|
||||
|
||||
## 2. 尚有实现前置条件的阻断项
|
||||
|
||||
以下项目不是“再跑一次测试”即可关闭。必须先完成实现，再执行对应验收。
|
||||
|
||||
| 编号 | 阻断项 | 完成实现后的最低验收证据 |
|
||||
| --- | --- | --- |
|
||||
| B-01 | Plugin Runtime 缺少跨 installation 的重启风暴抑制，包括 jitter、全局重启并发上限和 Runtime 级 circuit breaker | 同时使大量 worker 因系统性故障退出，证明重启速率和并发受限、控制面仍可用，且单个故障 installation 不拖垮其他租户 |
|
||||
| B-02 | Cloud 插件缺少生产 egress policy | 证明插件只能访问允许的公网目标，不能访问 Core/Box/数据库、其他内部服务、loopback、link-local 或云 metadata endpoint |
|
||||
| B-03 | Plugin installation 与 Box Workspace/Skill/root/tmp/home 缺少真实的 byte 和 inode 硬配额 provider | 在写入边界原子拒绝超额；并发写入、重启和配额耗尽后不能越界，也不能用目录扫描或事后清理冒充硬配额 |
|
||||
| B-04 | 普通业务写入尚未具备贯穿 commit 的 generation-aware fence、同事务 business outbox，以及 generation cutover 后稳定的 durable-object 引用 | 在旧 generation 与新 generation 并发、事务提交竞态和重复投递下，旧 owner 不产生业务写入或外部副作用，outbox 可幂等恢复 |
|
||||
|
||||
任一 B 类项目未关闭时，不得把 24 小时 soak 的通过结果解释为可以上线。
|
||||
|
||||
## 3. 最终部署环境验证
|
||||
|
||||
### V-01：Plugin Runtime 与 Box 的 Linux 隔离
|
||||
|
||||
必须在最终 Cloud Pod security context、容器 runtime 和 cgroup 拓扑中验证，
|
||||
不能使用开发机或权限不同的一次性容器替代。
|
||||
|
||||
验证内容：
|
||||
|
||||
1. nsjail 可以建立 mount、PID、IPC、UTS namespace 和 private `/proc`。
|
||||
2. delegated cgroup v2 对每个插件进程和 sandbox 强制 CPU、memory、
|
||||
`memory.swap.max=0` 和 PID 上限。
|
||||
3. open files、process 和单文件大小 rlimit 生效。
|
||||
4. 插件不能枚举、读取或 signal Runtime 及其他 installation 的进程，
|
||||
不能读取其他 installation 的 home/tmp/data，也不能修改共享只读
|
||||
artifact/environment。
|
||||
5. 超额只杀死或拒绝当前 installation/sandbox；Runtime、其他租户和健康接口
|
||||
继续工作。
|
||||
6. 进程退出、取消、超时、generation 切换和 Pod SIGTERM 后，cgroup、nsjail
|
||||
目录、子进程和文件描述符均被回收。
|
||||
7. 硬限制或 namespace 能力缺失时 readiness 失败关闭，不允许降级成普通进程。
|
||||
|
||||
通过证据必须包含实际容器安全配置、cgroup 文件值、探针原始输出和失败注入结果。
|
||||
|
||||
### V-02：Box 持久卷与硬存储配额
|
||||
|
||||
在 B-03 的 quota provider 实现后，必须验证：
|
||||
|
||||
1. Core 与 Box Runtime 通过随机 marker challenge 证明使用同一共享持久卷。
|
||||
2. Workspace、Skill store、ephemeral root/tmp/home 的 byte 和 inode quota
|
||||
都在写入点生效。
|
||||
3. 并发写、压缩包展开、文件同步、重启恢复和删除重建不能绕过配额。
|
||||
4. 配额耗尽只影响目标 Workspace，其他 Workspace 仍能执行。
|
||||
5. 任一硬存储能力缺失时 `/readyz` 返回非 2xx，Pod 不进入就绪流量。
|
||||
|
||||
### V-03：PostgreSQL 与 pgvector 生产边界
|
||||
|
||||
必须在最终 PostgreSQL endpoint、凭据和网络策略下验证：
|
||||
|
||||
1. migrator 与 runtime 使用不同 role；runtime role 无 superuser、
|
||||
`BYPASSRLS`、DDL、对象所有权、role membership 或额外 schema 权限。
|
||||
2. runtime credential 只能连接目标 business database。需要专用
|
||||
cluster/endpoint，或经过测试的 HBA/proxy policy；database 内 catalog
|
||||
audit 本身不能证明这一点。
|
||||
3. release migration Job 的 advisory lock、失败重试、回滚和精确 Alembic head
|
||||
校验有效；应用启动角色不能执行 migration 或其他 DDL。
|
||||
4. 若使用 PgBouncer，transaction pooling、异常回滚和连接复用不会残留
|
||||
tenant context。
|
||||
5. 故意遗漏应用层 Workspace filter 时，RLS 仍阻止跨租户读写。
|
||||
6. 两个 Workspace 使用相同 `vector_id`、猜测其他 Workspace ID、后台任务和
|
||||
连接复用时，pgvector CRUD 均不能越权。
|
||||
7. dimension mismatch、extension/schema/ACL drift 或 runtime audit 失败时，
|
||||
Core 启动失败且不回退到其他向量后端。
|
||||
|
||||
### V-04：最终镜像和配置一致性
|
||||
|
||||
生产候选验证前必须固定：
|
||||
|
||||
- Core、Plugin Runtime、Box Runtime 的不可变镜像 digest；
|
||||
- LangBot 与 SDK commit；
|
||||
- `data/config.yaml` 的非敏感摘要和所有环境变量覆写；
|
||||
- PostgreSQL migration revision；
|
||||
- Cloud Adapter、Space control plane 和 workload 的版本。
|
||||
|
||||
滚动更新、节点迁移、配置变化或任一镜像 digest 变化后，旧验证报告失效。
|
||||
|
||||
## 4. 多租户行为与故障注入
|
||||
|
||||
### V-05：目录、entitlement 与 generation
|
||||
|
||||
在真实 Space control plane、闭源 Cloud Adapter 和 Core 之间验证：
|
||||
|
||||
1. 注册自动创建个人 Workspace、owner membership、Free subscription、
|
||||
entitlement snapshot 和 outbox 事件，且重试不重复创建。
|
||||
2. 邀请、成员变更、套餐变更和 Workspace 撤销只影响目标 Workspace。
|
||||
3. 全量快照与增量事件覆盖乱序、重复、断点续传、缺页、签名错误、
|
||||
high-water gap、snapshot coverage 和消费者重启。
|
||||
4. Core、Plugin Runtime 或 Box 重启后，权威 desired state 可恢复；
|
||||
本地进程表和缓存不是唯一真相。
|
||||
5. generation/revision 切换期间，旧 callback、RPC、WebSocket、Box relay、
|
||||
plugin worker 和缓存写入全部失败关闭。
|
||||
6. 为未来多副本预留的 replica-local cursor 语义通过故障注入：
|
||||
一个副本追平不能使另一个副本跳过本地 cache 刷新。
|
||||
|
||||
### V-06：套餐、Box 与 stdio MCP
|
||||
|
||||
1. Free/非 Pro Workspace 不会自动获得 managed sandbox。
|
||||
2. 合资格 Workspace 最多只有一个持久 `global` sandbox，且新增 Workspace
|
||||
不创建专属 Runtime、Pod、PVC、database、schema、role 或连接池。
|
||||
3. Cloud 即使 `box.enabled=true`，也不能 create/update/test/start stdio MCP；
|
||||
旧记录和直接 API 调用同样失败关闭，且不创建 `mcp-shared` session。
|
||||
4. OSS 默认仍是单 Workspace、多用户，stdio MCP 保持兼容，多租户能力不会被
|
||||
未签名配置或普通环境变量开启。
|
||||
|
||||
### V-07：跨租户安全回归
|
||||
|
||||
至少使用两个恶意测试 Workspace 验证：
|
||||
|
||||
- 同 digest 插件只共享只读代码和依赖；进程、secret、日志及所有可写目录隔离。
|
||||
- 同 author/name/version 但 digest 不同的 artifact 不共享目录。
|
||||
- Plugin Host API、Box RPC、对象 key、WebSocket、RAG、storage、model/session
|
||||
cache 和平台回调不能接受调用方伪造的 Workspace scope。
|
||||
- 撤销 entitlement、删除 installation 或 generation 切换后，已有长连接和
|
||||
in-flight 请求不能继续访问旧权限。
|
||||
|
||||
## 5. 生产候选容量与 24 小时门禁
|
||||
|
||||
### V-08：真实容量曲线
|
||||
|
||||
现有 fake adapter/requester/Plugin handler 探针不能替代真实容量数据。必须使用
|
||||
计划上线的平台 SDK、外部 HTTP/WebSocket 连接池、真实插件进程、真实
|
||||
PostgreSQL/pgvector 和代表性 Workspace 配置分布，测量：
|
||||
|
||||
- 空 Workspace、活跃 Workspace、每个启用插件和每个 Pro sandbox 的边际
|
||||
RSS、线程、文件描述符、连接和 PostgreSQL pool 成本；
|
||||
- 启动、目录重放、批量 reconcile 和故障恢复的耗时与峰值；
|
||||
- 单实例可批准的 Workspace、活跃 Bot、plugin worker 和 sandbox 上限。
|
||||
|
||||
容量上限必须写入生产配置与告警，不能只保留在测试报告中。
|
||||
|
||||
### V-09：24 小时资源 soak
|
||||
|
||||
使用 [标准 24 小时命令](./cloud-runtime-soak-gate.md#标准-24-小时命令)，并强制
|
||||
`--require-hard-limits`。工作负载至少覆盖：
|
||||
|
||||
1. 注册、邀请、登录和 entitlement 刷新；
|
||||
2. plugin reconcile、依赖准备、调用、崩溃与重启；
|
||||
3. Dashboard/Embed/平台 WebSocket 建连、突发消息和断连；
|
||||
4. Box session、文件同步、并发 exec、输出与清理；
|
||||
5. PostgreSQL pool 接近容量、事务超时和恢复；
|
||||
6. Core、Plugin Runtime、Box 分别 SIGTERM 和恢复。
|
||||
|
||||
最后至少保留 30 分钟无测试流量冷却。任一健康失败、OOM/memory pressure、
|
||||
PID limit、blocking executor rejection、超阈值 CPU throttling/event-loop lag、
|
||||
冷却尾段内存持续增长或临时 gauge 不回落都判为失败。
|
||||
|
||||
必须归档：
|
||||
|
||||
- 原始 `cloud-soak-samples.jsonl`；
|
||||
- 最终 `cloud-soak-report.json`；
|
||||
- 三个镜像 digest、Core/SDK commit；
|
||||
- 生产配置摘要、数据库 migration revision 和 workload 版本；
|
||||
- 故障注入时间线及关联日志/trace。
|
||||
|
||||
## 6. 本轮不作为验收条件的后续事项
|
||||
|
||||
以下能力已明确暂缓，不能混入当前验证结果，也不能以“尚未验证”为理由临时发明方案：
|
||||
|
||||
- Workspace export、释放、delete、单 Workspace restore 和在线迁移；
|
||||
- Workspace 级 BYOK E2B WebUI 配置；
|
||||
- 多 Core/Plugin Runtime/Box replica 的 lease store 与调度实现；
|
||||
- PostgreSQL 多 shard、dedicated shard 和跨地域部署；
|
||||
- 多 CloudInstance、Cell Router 或 Workspace Placement。
|
||||
|
||||
这些事项需要后续单独决策。首期实现仍需保留稳定 UUID、generation fence、
|
||||
幂等事件和无副本地址泄漏的协议边界。
|
||||
|
||||
## 7. 关闭规则
|
||||
|
||||
每个 B/V 项只能通过以下方式关闭：
|
||||
|
||||
1. 记录被测 commit、镜像 digest、配置摘要和环境拓扑；
|
||||
2. 保存可复现命令、原始输出和失败注入证据；
|
||||
3. 由报告明确给出 pass/fail，不能只依赖日志中“看起来正常”；
|
||||
4. 任一生产候选输入变化后，重跑受影响的验证。
|
||||
|
||||
在 B-01 至 B-04 全部实现，且 V-01 至 V-09 均有当前生产候选版本的通过证据前，
|
||||
Cloud v2 状态保持 `NOT APPROVED FOR SAAS ACTIVATION`。
|
||||
@@ -0,0 +1,254 @@
|
||||
# LangBot Cloud Runtime 资源安全审查
|
||||
|
||||
日期：2026-07-28 至 2026-07-29
|
||||
|
||||
审查分支：
|
||||
|
||||
- LangBot：`feat/multi-tenants`，审查起点 `32abbb636f4455e965141d8d209b359dbfbb5aae`
|
||||
- Plugin SDK：`feat/multi-tenants`，审查起点 `0cddf3c2bea5939c67b71e488a719e9903c28d17`
|
||||
|
||||
## 结论
|
||||
|
||||
本轮已覆盖 LangBot Core、Plugin Runtime 和 Box Runtime 的主要常驻对象、后台任务、队列、网络客户端、进程生命周期及数据库连接池。本轮定位到的攻击者可控或历史累积状态均已补充容量、超时、淘汰或确定性清理边界；修正后的高基数探针没有观察到随历史请求继续增长的活跃缓存。该结论不等于证明任意生产负载下不存在资源问题，生产激活仍受下述 Linux 隔离、硬盘配额和 soak 门禁约束。
|
||||
|
||||
代码级审查、跨仓全量测试和仓库 Dockerfile 构建的 Linux/cgroup v2 探针已经通过，但当前状态仍不能作为 Cloud 生产激活批准。完整剩余清单见
|
||||
[Cloud v2 仍待验证事项](./cloud-v2-pending-verification.md)。其中与本轮资源审查直接相关、上线前还必须完成的项目包括：
|
||||
|
||||
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 已先行发布到分支提交 `67be7c332ded4c1afa131b28232590cf9728f817`，本提交集中的 LangBot
|
||||
`pyproject.toml` 和 `uv.lock` 已精确钉住该提交。最终镜像仍需按待验证清单记录并核对实际安装版本。
|
||||
|
||||
## 覆盖范围
|
||||
|
||||
### LangBot
|
||||
|
||||
- 启动、停机、全局任务管理和运行时配置。
|
||||
- PostgreSQL、Tenant UoW、RLS、迁移和共享 pgvector。
|
||||
- HTTP、MCP、WebSocket、上传下载、S3/本地存储、维护任务和异步用户任务。
|
||||
- QueryPool、Pipeline Controller、会话/对话、限流、第三方 Agent/LLM runner 和同步 SDK 桥接。
|
||||
- PlatformManager 以及 DingTalk、QQ、Lark、WeCom、WeComCS、WeChatPad、LINE、Kook、Satori、OpenClaw Weixin、Telegram、Discord、Matrix、HTTP/WebSocket 等适配器。
|
||||
- Plugin Runtime connector、插件包校验、Marketplace 下载、pip 安装输出和 desired-state reconcile。
|
||||
- Box connector、admission、session/process 生命周期和 RPC 文件。
|
||||
- RAG、向量后端、Skill、Storage、Telemetry 和日志缓存。
|
||||
|
||||
### Plugin SDK
|
||||
|
||||
- stdio/WebSocket transport、请求 waiter、action task 和文件传输。
|
||||
- Runtime control handler、Workspace/generation fence 和 EventContext。
|
||||
- 插件 artifact、Marketplace、pip 依赖安装、共享依赖环境、installation desired state、Supervisor 和 worker launcher。
|
||||
- nsjail 参数、cgroup/rlimit、进程注册 capability 和 Runtime shutdown。
|
||||
- Box admission、generation fence、session/process/reaper、RPC/relay WebSocket 和 nsjail backend。
|
||||
|
||||
## 主要修复
|
||||
|
||||
### 确定性生命周期
|
||||
|
||||
- 修复 `Application.shutdown()` 使用 `contextlib.suppress` 却未导入 `contextlib` 的问题。原行为会在实际资源关闭分支直接 `NameError`，阻断后续 Plugin、Box、HTTP 和数据库释放。
|
||||
- `make_app()` 在任一启动 stage 或初始化失败时会关闭尚未返回给 `main()` 的半构建 Application；Telemetry、Box、Tool、Platform、Vector 和 HTTP manager 在初始化前即挂到 Application，避免初始化中途失败后清理器无法发现已经创建的连接、会话、子进程或后台任务。
|
||||
- MCP streamable-HTTP session manager 现在随 Application shutdown 显式退出。
|
||||
- MCP loader 按 `(instance, workspace, generation)` 管理 host task 和 session；任务完成会从注册表移除，代次推进会取消旧 host task 并关闭旧 session，reload/shutdown 会先清空现有运行时，避免 completed task 和旧代连接永久驻留。
|
||||
- Platform bot reload/remove/shutdown 统一串行化，旧 bot、代理、adapter 任务和进程会先停止再从注册表移除。
|
||||
- Model provider requester 新增异步关闭契约；provider reload/remove、Workspace generation 替换、全量 reload 和 Application shutdown 都会确定性关闭旧 requester，允许第三方 requester 安全持有自己的 HTTP client 或连接池。
|
||||
- Plugin Runtime、Box Runtime、stdio transport、adapter 连接和共享 HTTP client 均补齐 close/cancel/await。
|
||||
- `Application.dispose()` 只允许一个可追踪 shutdown task；重复的信号、窗口关闭或调用方清理不会铺开多个并行停机流程。
|
||||
- Lark、微信、钉钉、企业微信和 QQ Official 的凭证交换后台任务统一进入 Application TaskManager，受全局/单 Workspace admission 约束并随应用停机取消；容量满时关闭尚未调度的 coroutine 并返回 429，不留下游离 task。
|
||||
- `TaskCapacityError` 已下沉到无 Application/controller 依赖的纯错误模块。原来的 HTTP 过载异常路径会在特定冷启动导入顺序下触发 TaskManager/controller 循环导入，把应返回的 429 变成框架 500。
|
||||
- S3 storage provider 在初始化失败或 Application shutdown 时关闭 botocore HTTP connection pool；Storage manager 在 provider 初始化前即挂到 Application，避免 bucket probe 失败后遗留 client。
|
||||
- 修复 Coze runner 每次请求创建 `aiohttp.ClientSession` 却不关闭的问题；现在 runner 的 `aclose()` 会确定性关闭底层 API client。
|
||||
- LINE SDK client 现在随 adapter 停止关闭；Plugin Runtime shutdown 回调只创建一个可追踪、可等待的后台任务，重复回调不会累积清理 task。
|
||||
- SDK `lbp publish` 现在用上下文管理器关闭插件包上传文件；原实现的成功、API 错误和 HTTP 错误返回路径都会遗留文件句柄。
|
||||
- Plugin worker、Box Docker/CLI backend、Box nsjail backend 和 nsjail 依赖安装在调用方取消时会 terminate/kill、读取管道并 `wait()` 回收子进程；Windows worker 的原生进程路径也进入相同的 `finally` 清理契约，避免取消安装、停机或超时后留下孤儿进程。
|
||||
- Box 服务入口现在从 Runtime initialize、aiohttp app/runner setup、端口 bind 到主循环共用一个外层清理边界；建站期间的非 `OSError` 也会关闭 Runtime 和 reaper。WebSocket 控制模式端口绑定失败会退出并交给编排器重启，不再在没有任何可用 RPC/health 端口时永久等待；stdio 模式仍允许仅 relay 绑定失败后继续控制通道。
|
||||
- Plugin artifact 解压、installation staging/activation/rollback/delete、共享依赖环境和 nsjail session 目录的关键文件系统变更均移出事件循环。已经开始的原子变更在取消时会先等待线程结束，再回滚临时目录、目标目录或旧 supervisor，不会让后台线程继续修改一个调用方已经认为清理完毕的路径。
|
||||
- Core 和 SDK 的阻塞 executor 新增独立的有界清理入口。普通工作在容量耗尽时仍快速拒绝；已经拥有资源的 close/unlink/rmtree/进程回收则等待有限 worker 槽位并保证原子操作结束后再传播取消，避免过载恰好导致清理任务被拒绝。SeekDB、Milvus、S3、LINE、WeChatPad、MCP staging 和 TBox 临时文件等关键路径已接入。
|
||||
|
||||
### 有界队列、缓存和历史状态
|
||||
|
||||
- QueryPool 同时限制全局与单 Workspace 的 queued/running query；调度后不再被过载淘汰；历史 scope counter 有上限。
|
||||
- Session、Conversation、WebSocket connection/proxy/message、rate-limit identity、task record/log、telemetry task、vector handle 和 adapter 私有队列均有容量或 LRU/TTL。
|
||||
- SessionManager 现在维护 Workspace 二级索引和带 revision 校验的最小过期堆。新会话只扫描目标 Workspace 的有界会话集，TTL 回收只消费已过期堆前缀，全局 idle 淘汰使用最小堆；高频命中产生的旧堆项按活跃会话的有界倍数压缩。原实现会在每个攻击者可制造的新 launcher id 上扫描并排序实例全部会话。
|
||||
- SDK 的 EventContext 和依赖准备锁使用 weak reference；generation、admission、installation、capability 和 completed-process 状态有上限。
|
||||
- 空 Workspace 不再预分配 Model generation scope、Plugin installation set 或 Box generation event；只有 Workspace 实际拥有对应运行时资源或等待任务时才创建这些对象。
|
||||
- Runtime RPC 文件同时限制单文件字节数和单连接未消费文件数量，连接关闭时清理连接拥有的临时文件。
|
||||
- Box Runtime 维护实例与 Workspace 到活跃 session 的二级索引；创建、删除、过期、撤销和 shutdown 共用同一清理路径，避免每个租户 RPC 都扫描实例中的全部 session。
|
||||
- Box Runtime 另外只索引“可过期 session”和“持有 managed process 的 session”。Cloud 的持久 `global` sandbox 不进入 TTL 索引，managed process 被禁用时进程索引为空；session 创建、周期 reaper、状态和 `/healthz` 因此不会随全部持久租户数线性扫描。测试把总 session 字典替换为禁止迭代的映射，仍能完成第二个持久 session 创建、reap、status 和 health。
|
||||
- Core MCP loader 同样维护 Workspace/generation 到 session、host task 的二级索引；请求、代次回收和动态配置不再扫描实例中的全部租户 MCP session，已完成 task 的 done callback 会同步移除所有索引。
|
||||
- Box admission 过期回收使用带 revision/generation 校验的最小堆，只访问已到期记录；重复续期产生的旧堆项会被忽略，堆大小超过活跃 grant 的有界倍数时压缩，不再在每次 RPC 上全表扫描所有租户 grant。
|
||||
- 旧 QQ message ID/object cache 和 stdio MCP Workspace copy lock 不再随历史请求无限增长。
|
||||
- LLM/Agent runner 的单次生成结果默认限制为 1 MiB，流式传输限制单事件 1 MiB、单请求累计 16 MiB，并限制最多 100,000 个流式事件，避免上游异常响应无限占用内存或 CPU。
|
||||
- Marketplace JSON 限制为 1 MiB、插件包限制为 64 MiB；pip stdout/stderr 各最多保留 1 MiB，超出部分继续 drain 但不驻留内存。
|
||||
- Plugin Runtime stdio/WebSocket 协议除 16 MiB 消息字节上限外，新增最多 4,096 个入站和出站碎片的对象数量上限；大消息的 UTF-8 编码、分片、拼接、JSON 编解码和 Pydantic 验证均在线程执行。WebSocket receive 异常使用稳定的 `ConnectionClosed` 类型，不再因库顶层未导出 `exceptions` 属性而在错误路径二次失败。
|
||||
- HTTPX response hook 和 aiohttp 有界读取统一把第三方响应限制为 10 MiB；JSON 解析及诊断文本转换在线程执行，错误正文最多保留 4 KiB，避免大 JSON 在共享事件循环集中解析。
|
||||
- 图片、data URI 和平台媒体默认限制为 10 MiB，Base64 在解码前先校验编码长度；Plugin binary storage 默认单值 10 MiB，并设置不可由错误配置绕过的 64 MiB 绝对上限。
|
||||
- Skill 文本单文件限制为 1 MiB，Plugin UI 文件限制为 4 MiB，host edit 文件限制为 1 MiB；Box Skill ZIP、插件 artifact 和 GitHub Skill archive 同时限制条目数、单文件、解压总量和压缩比。
|
||||
- SDK E2B 文件同步限制为最多 2,048 个目录项、1,024 个文件、单文件 10 MiB、总计 50 MiB；同步文件 IO 从事件循环移到线程。
|
||||
- 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 会阻止旧结果重新写回缓存。
|
||||
|
||||
### CPU 和事件循环保护
|
||||
|
||||
- 修复旧 QQ `repeat_seed('')` 空输入无限循环。
|
||||
- ZIP 校验/重打包、PIL、Base64、AES、JSON 解析、fsync probe、插件 artifact/依赖文件、Skill、S3、本地存储和维护目录扫描从事件循环移到线程。
|
||||
- 公开 Slack、QQ Official、HTTP Bot、公众号、WeCom/WeComCS 回调体显式限制为 1 MiB；JSON/XML 解码移出共享事件循环。QQ、DingTalk、Satori、WeCom AI 和 WeChatPad 网关帧同样设置 1 MiB 上限或在解码前拒绝超限消息；KOOK zlib 数据使用 10 MiB 解压后硬上限，阻断小压缩包制造的大内存解压。
|
||||
- Dashboard、Embed 和 Plugin Runtime 的协议 JSON 编解码在线程执行；Dashboard/Embed 在接收端提交 terminal error 后为发送端保留有界 drain 窗口并使用内部 sentinel 唤醒，不会因“任一方向结束即取消”在撤权错误帧发出前关闭连接。
|
||||
- 租户配置的敏感词、内容忽略和群响应正则统一使用声明为直接依赖的 `regex` 引擎：最多 64 个 pattern、单 pattern 1,024 字符、输入 1 MiB、单次总匹配 CPU 预算 50 ms，并在线程中执行。超时、非法正则和替换放大均失败关闭；灾难性 `(a+)+$` 回归在 1 ms 测试预算内被中断。
|
||||
- 原生 `read/write/edit/glob/grep` 文件工具移出事件循环并继承 Workspace 阻塞预算。目录列举、递归 walk、grep 文件/总字符、单行、pattern、结果和 regex CPU 均有硬上限；glob 只用固定大小最小堆保留最新 100 项，不再先把全部命中路径驻留内存。Box 内执行的 glob/grep 脚本同样限制命中集合、扫描量和正则时间。
|
||||
- Dify、DingTalk、QQ、WeCom 等客户端复用连接并在生命周期结束时关闭，响应体和下载有字节上限。
|
||||
- DashScope、TBox 等同步第三方 SDK 的调用和生成器迭代改为在线程执行；单个同步生成器最多消费 100,000 个事件。
|
||||
- Dashboard 和 Embed WebSocket 改为任一收发 task 结束即取消并等待另一方向，避免发送端退出后接收 task 永久阻塞；两方向 task 同时继承从认证结果或 RuntimeBot 得到的可信 Workspace 阻塞预算。
|
||||
- Plugin installation 生命周期全局串行化；不同租户的依赖 pip/nsjail 准备不会在安装高峰并发抢占 CPU。
|
||||
- S3 同步 SDK 使用线程执行，并通过实例级 semaphore 限制并发；默认 `storage.s3.max_concurrency=16`，可通过实例配置和环境变量覆写。
|
||||
- Box 子进程 stderr 以 64 KiB 块读取，日志最多每秒输出 4 个摘录并汇总抑制数量，避免无换行或刷屏输出制造无界缓冲与日志放大。
|
||||
- Plugin worker 日志单行最多保留 64 KiB；Box managed-process stdout relay 以固定 64 KiB 块读取，不再依赖换行符，避免超长无换行输出触发 `StreamReader` limit 或堵塞子进程。
|
||||
- Box generation fence 的代次更新改为只访问目标 Workspace 的 event 和 active-task 二级索引。原实现每次更新都会遍历全部 Workspace 的 fence/task 记录，10,000 个 Workspace 的第二阶段更新会退化为 O(N²) 并在 40 秒后仍未完成；修正后包括其他 SDK 高基数负载和本轮协议 offload 在内的当前完整双阶段探针耗时 `9.037s`。
|
||||
- Box session 枚举、旧 generation 回收和 admission 计数均通过 Workspace 索引执行；admission 过期回收通过最小堆执行，不再在每次 RPC 上产生 O(实例总 session/grant 数) 的扫描。
|
||||
- Model、Pipeline、RAG 和 Platform manager 均维护 Workspace 到运行时 key 的二级索引。Workspace generation 更新只清理目标 Workspace 的缓存和运行时，不再扫描实例内所有租户的 provider/model、pipeline、knowledge runtime 或 bot；回归测试使用禁止全局迭代的映射验证该边界。
|
||||
- Cloud heartbeat 直接读取已加载且有容量边界的 Pipeline、MCP、KnowledgeBase 和 Bot registry 计数，不再为每个活跃 Workspace 依次打开 Tenant UoW、执行四类 COUNT 查询；这消除了租户数增长后每日周期性形成的串行 SQL/CPU 尖峰。OSS 模式仍保留数据库统计语义。
|
||||
- Cloud 启动阶段先生成一份经过部署适配器和目录投影校验的 Workspace binding 快照，Model、Platform、Pipeline、RAG 和 Plugin 初始化共用该快照，初始化完成后立即释放；避免启动期间为每个 manager 重复执行整批租户发现和投影校验。
|
||||
- Platform、Pipeline 和 RAG 的资源加载在使用已验证启动快照时不再为每个 Bot/Pipeline/KnowledgeBase 重新查询同一个 execution binding；常规请求和动态更新路径仍保留数据库 generation fence。
|
||||
- MCP 初始 host 和 shutdown burst 由实例级 semaphore/批次限制；默认 `mcp.lifecycle_concurrency=16`，支持 `MCP__LIFECYCLE_CONCURRENCY` 覆写并硬性限制最大 128。初始加载不再先为每个 server 创建一个等待 semaphore 的 task，而是由一个可取消 dispatcher 每批最多物化 `lifecycle_concurrency` 个子 task；同时去掉了 ORM server/config 的双份临时列表，避免大量租户启动时集中占用 CPU、内存、socket 和文件句柄。
|
||||
- Core、Plugin Runtime、Box Runtime 和独立 Plugin worker 的默认 `asyncio.to_thread()` executor 统一改为硬有界线程池。默认最多同时运行 8 个阻塞调用、排队 128 个，达到容量后立即抛出 admission 错误，不再使用 Python 默认 `ThreadPoolExecutor` 的无界工作队列保留任意数量的请求对象、Future 和闭包。每个可信 Workspace 的 running + queued 默认再限制为 4，并强制配置值不超过 worker 数的一半，避免单个租户先提交一整批同步工作占满全部 worker/FIFO 队列。Core 使用 `system.blocking_executor.max_workers/max_pending/max_inflight_per_scope`，原生支持对应的 `SYSTEM__BLOCKING_EXECUTOR__*` 覆写；SDK 进程使用 `LANGBOT_BLOCKING_EXECUTOR_MAX_*`，并分别限制全局最大值为 64/4096。
|
||||
- Core、Plugin Runtime 和 Box Runtime 各自运行固定 1 秒间隔、仅保留最近 120 个样本的 event-loop lag monitor；健康快照输出 current/recent max/recent p95/进程期最大延迟和累计样本数。Plugin Runtime 的两个 WebSocket 端口现在都在免认证 `/healthz` 返回同一份无凭据、无租户标识的聚合 JSON，Box `/readyz` 同时附带资源快照。24 小时门禁默认拒绝缺失或停止的 monitor、超过 1 秒的 recent max、超过 250 ms 的尾段 recent p95 及 sample counter 回退。
|
||||
- Workspace 阻塞预算由服务端认证后的 `RequestContext`、公开 bot 的 RuntimeBot、公开对象 key 中经 binding fence 验证的 Workspace、Platform/TaskManager 的 ExecutionContext，以及 SDK 入站 ActionContext 建立，不接受调用方伪造的租户 header。公开 webhook、公开对象下载、Dashboard/Embed WebSocket、普通 HTTP handler、Platform adapter 和 detached tenant task 均已覆盖。容量拒绝在 Core HTTP 路径返回稳定的 429，health/debug counter 分开报告 global 与 scope rejection。
|
||||
- Argon2 密码 hash/verify 只允许一个实例级在途操作，额外并发立即返回容量错误而不是在 asyncio semaphore 中无限积累等待请求；该 CPU/内存密集工作同时使用独立的 `system:authentication` 阻塞作用域。Cloud 本身仍禁用本地密码登录。
|
||||
- WeCom 扩展 API 的无限客户端超时改为 120 秒；平台 webhook 的 AES、媒体 Base64 与同步 SDK 调用均移出共享事件循环。
|
||||
|
||||
### 插件和 Box 资源隔离
|
||||
|
||||
- Plugin worker 数量受 `max_workers`、`max_total_cpus / max_cpus` 和 `max_total_memory_mb / max_memory_mb` 的最小值约束。
|
||||
- Shared profile 强制 Linux 和 nsjail；Cloud 强制 `plugin.worker.require_hard_limits=true`；cgroup v2 delegation 不可用时拒绝启动。
|
||||
- 每个 worker 下发 CPU、memory、swap、PID cgroup 限制，以及 process、open-file、file-size rlimit；插件 manifest 不能提高限制。
|
||||
- Box nsjail 的 cgroup v2 路径现在同时设置 `memory.max` 和 `memory.swap.max=0`。修复前，48 MiB 沙盒可以把强制提交的 128 MiB 页面换出并正常退出，形成宿主 swap 抢占；修复后同一探针以 exit 137 被 cgroup 杀死。
|
||||
- 仓库 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。
|
||||
|
||||
### PostgreSQL
|
||||
|
||||
- Cloud 强制 PostgreSQL 业务库、共享 pgvector 和允许的固定向量维度。
|
||||
- pgvector Cloud 模式复用业务数据库的同一个 AsyncEngine，不创建第二个连接池。
|
||||
- `database.postgresql` 新增并校验 `pool_size`、`max_overflow`、`pool_timeout_seconds`、`pool_recycle_seconds`；默认最大连接数为 `10 + 10`。
|
||||
- Application shutdown 显式 dispose 业务引擎；standalone pgvector 仅关闭自己拥有的引擎。
|
||||
- PersistenceManager 提供统一异步 shutdown；Cloud 常驻进程的启动失败、正常停机和一次性 release migration 的成功/异常路径都会释放数据库引擎。真实 PG catalog 测试还覆盖了“入口已经关闭后测试再次复用 manager 会重开 pool”的第二生命周期，严格资源告警模式下无 asyncpg socket/transport 遗留。
|
||||
|
||||
## 本轮采用的默认决策
|
||||
|
||||
- 优先 fail closed 或淘汰最老的 idle cache，不允许攻击者控制的历史 key 无限驻留。
|
||||
- 插件依赖准备选择实例级串行化，以稳定 CPU/磁盘峰值；代价是批量安装耗时增加。
|
||||
- PostgreSQL 使用一个显式有界共享连接池；未拆分 pgvector pool。
|
||||
- 第三方 runner 采用 1 MiB 单结果、16 MiB 单流总量和 100,000 个同步/异步事件的统一实例级安全上限；超限请求失败关闭。
|
||||
- S3 默认允许 16 个并发阻塞调用，最大配置值 128；在没有独立 worker service 的前提下限制线程池排队和上游连接压力。
|
||||
- MCP 生命周期默认并发 16、最大 128；该限制统一约束实例启动时的 session host 峰值和 shutdown 批次，不允许租户配置单独放大。
|
||||
- Core 与 SDK 各进程的通用阻塞 executor 默认使用 8 个 worker、128 个 pending 槽位、每 Workspace 4 个在途槽位；它是实例/进程级共享背压，不由 Workspace 或插件 manifest 调高，单 Workspace 配置硬性不得超过 worker 的一半。生产值应按容器 CPU 和上游阻塞时延校准，不能把 pending 当吞吐配置无限放大。
|
||||
- 插件包下载上限 64 MiB，pip stdout/stderr 保留上限各 1 MiB；这不会限制安装进程实际输出，只限制父进程内存中的诊断副本。
|
||||
- 通用远程响应和媒体默认上限 10 MiB；错误诊断正文只保留 4 KiB。Plugin binary storage 默认 10 MiB、绝对上限 64 MiB；Skill 文本、Plugin UI 和 host edit 分别限制为 1 MiB、4 MiB 和 1 MiB。
|
||||
- Managed-process relay 保留 stdout 的原始换行，并按 64 KiB WebSocket frame 分块；不再承诺“一行对应一个 frame”。这是为无换行输出提供确定内存边界所需的协议收敛。
|
||||
- 本轮没有把 Pipeline、Model、KnowledgeBase 等合法租户资源改成 lazy runtime。该改动会改变启动和请求语义，留到 Workspace placement/释放机制一起设计。
|
||||
- 本轮没有为普通 nsjail 声称伪硬盘配额；严格 Cloud readiness 保持失败关闭。
|
||||
|
||||
## 验证结果
|
||||
|
||||
| 验证项 | 结果 |
|
||||
| --- | --- |
|
||||
| LangBot Ruff + `git diff --check` | 通过 |
|
||||
| Plugin SDK Ruff + `git diff --check` | 通过 |
|
||||
| LangBot 全量测试（使用本地新 SDK，含 unit/integration/Box/E2E） | `2803 passed, 33 skipped` |
|
||||
| Plugin SDK 全量测试 | `1301 passed` |
|
||||
| 真实 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 门禁解析/采样/判定单元测试 | `23 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.275s` |
|
||||
| Core 5,000 个 populated Workspace 三代容量探针（使用本地新 SDK） | audit 通过，最大替换耗时比 `1.378` |
|
||||
| Plugin SDK 双阶段资源探针 | audit 通过，`9.037s` |
|
||||
|
||||
两个仓库新增了可重复执行的历史 churn 探针，Core 另有 populated Workspace 三代替换探针：
|
||||
|
||||
```bash
|
||||
# LangBot Core
|
||||
PYTHONPATH=../langbot-plugin-sdk/src uv run python scripts/runtime_resource_probe.py --scale audit --json
|
||||
|
||||
# LangBot Core：5,000 个带代表性资源的 Workspace
|
||||
PYTHONPATH=../langbot-plugin-sdk/src uv run python scripts/workspace_runtime_capacity_probe.py --scale audit --json
|
||||
|
||||
# LangBot Core：真实 PostgreSQL 16 + RLS populated Workspace 启动
|
||||
TEST_POSTGRES_URL=postgresql+asyncpg://... \
|
||||
LANGBOT_PG_CAPACITY_WORKSPACES=1000 \
|
||||
uv run pytest \
|
||||
tests/integration/persistence/test_migrations_postgres.py::TestPostgreSQLTenantRuntime::test_populated_cloud_startup_is_linear_and_task_bounded \
|
||||
-q -W error::ResourceWarning --log-cli-level=INFO
|
||||
|
||||
# langbot-plugin-sdk
|
||||
uv run python scripts/runtime_resource_probe.py --scale audit --json
|
||||
```
|
||||
|
||||
Core audit 每个阶段执行 10,000 个空 Workspace 的真实 Model/Plugin manager 加载与 reconcile、25,000 次 Query、2,500 次 session churn、10,000 个限流身份、5,000 个 task 和 2,500 次 WebSocket churn。第一、第二阶段的保留状态完全一致：
|
||||
|
||||
- 20,000 个历史空 Workspace：Model scope/provider/LLM、Plugin Workspace set/installation 均为 `0`。
|
||||
- 50,000 个历史 Query：活跃 query cache `0`，历史 scope counter `100`。
|
||||
- 5,000 个会话身份：session cache `200`。
|
||||
- 20,000 个限流身份：rate-limit container `10,000`。
|
||||
- 10,000 个历史 task：task record `200`。
|
||||
- 5,000 次 WebSocket churn：conversation 与 stream index 均为 `200`。
|
||||
- event-loop task、线程和文件描述符保持 `1 / 1 / 6`；强制加载本地最新 SDK 的当前复跑中，第二阶段相对第一阶段 RSS 增长 `1,605,632 bytes`、tracemalloc current 增长 `344,735 bytes`，总耗时 `11.275s`。Session 淘汰改为 Workspace 索引和最小堆后，同一 audit 工作量相对此前 `16.150s` 明显下降。
|
||||
|
||||
Populated Workspace audit 为 5,000 个 Workspace 各加载一个 Provider、LLM、Embedding、Rerank、Pipeline、Bot、KnowledgeBase 和 MCP session，然后全部推进两个 generation：
|
||||
|
||||
- 三个阶段的活跃 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,277,952 bytes`，tracemalloc current 反而减少 `148 bytes`。
|
||||
- 初始/第一次替换/第二次替换分别耗时 `1.707s / 2.205s / 2.351s`，最大替换耗时比为 `1.378`，未随历史代次出现 CPU 退化。
|
||||
- macOS RSS sample 从初始的 `154,271,744` 增至第一阶段 `367,820,800`、第二阶段 `388,726,784` 和第三阶段 `390,004,736 bytes`；第二次替换只比第一次替换增加约 1.22 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 创建/删除。第一、第二阶段的保留状态完全一致：
|
||||
|
||||
- RPC waiter、stream queue、action task 和活跃 installation binding 均为 `0`。
|
||||
- installation watermark 为有界的 `5,000`；Workspace generation record 为有界的 `10,000`，没有等待者时 generation event 为 `0`。
|
||||
- generation active task/index、Box session、Box Workspace session index、creating/closing/background task 和 session lock 均为 `0`。
|
||||
- event-loop task 和文件描述符保持 `1 / 7`；当前复跑第二阶段相对第一阶段 RSS peak 增长 `2,654,208 bytes`、tracemalloc current 增长 `289,634 bytes`，总耗时 `9.037s`。耗时增加来自本轮把大协议消息的 JSON/Pydantic、UTF-8 编码、分片和拼接移入有界线程池；25,000 RPC/阶段的合成探针仍约为 5,500 RPC/s，结构状态和第二阶段 tracemalloc 增量保持平稳。
|
||||
|
||||
第二轮反向静态审查另外枚举了 Core 的 50 个显式 task 创建点和 204 个线程、阻塞调用及子进程调用点，以及 SDK 的 28 个显式 task 创建点和 62 个线程、阻塞调用及子进程调用点。显式 task 均具有持有者、完成回调或 `finally` 回收路径；所有生产入口在第一次 `asyncio.to_thread()` 前安装有界默认 executor。Core、Plugin Runtime 和 Box 的公开 `/healthz`（Box `/readyz` 亦同）会输出各自的 aggregate runtime/resource counter 和 event-loop lag，供 soak 对比活跃量、pending、累计 capacity rejection 与调度延迟；不输出 debug key、控制 token、租户或插件身份。Plugin Runtime 的授权 debug info 复用同一资源快照，避免公开/私有指标语义漂移。
|
||||
|
||||
真实 PostgreSQL populated 启动门禁会先通过 release migration 创建最新 schema，再用无 `BYPASSRLS` 的临时 Cloud Runtime 角色启动。每个 Workspace 都含九类代表性资源，测试会走实际的 instance discovery、tenant UoW、启动 binding 快照和 Model/Platform/Pipeline/RAG/MCP/Plugin 加载路径：
|
||||
|
||||
- 500 Workspace：启动加载 `6.178s`，进程 CPU `3.026s`。
|
||||
- 当前 1,000 Workspace 复跑：启动加载 `12.109s`，进程 CPU `5.967s`；相对此前 500 Workspace 的墙钟比为 `1.960`。
|
||||
- `model_providers`、`llm_models`、`embedding_models`、`rerank_models`、`bots`、`legacy_pipelines`、`knowledge_bases`、`mcp_servers`、`plugin_settings` 九张表的 SELECT 次数均精确等于 Workspace 数，没有重复的全租户发现或超线性资源扫描。
|
||||
- MCP host dispatcher、host task 和临时 Runtime 角色/asyncpg 连接在测试结束后均清空；严格 `ResourceWarning` 模式通过。
|
||||
|
||||
探针要求第二阶段的结构状态与第一阶段精确相等，并对第二阶段 RSS 与 tracemalloc 增长设置失败阈值。macOS 的 RSS 来源是 `getrusage` peak，因此这里验证的是峰值增量边界而非“当前 RSS 回落”；最终 Linux 24 小时 soak 仍需采集 current RSS/PSS 和 cgroup `memory.current`。
|
||||
|
||||
LangBot 全量测试的 33 个 skip 中，22 个是默认全量运行未提供 PostgreSQL/pgvector 而跳过的集成用例，10 个是未提供 Valkey，另 1 个是可选环境的 collection skip；真实 PostgreSQL 相关路径已由上表单独运行覆盖。Plugin SDK 的 22 个 warning 为现有 Pydantic v2 deprecation 与 aiohttp AppKey 建议；没有失败、未关闭资源或资源上限降级。Core 当前全量产生 192 个既有第三方/兼容性 warning；`ResourceWarning` 和 `PytestUnraisableExceptionWarning` 仍由 pytest 配置提升为错误，本轮没有此类泄漏告警。
|
||||
|
||||
Linux Runtime 探针使用上述镜像并只读挂载本地最新 SDK 源码：
|
||||
|
||||
- 普通容器：nsjail binary 可执行，但 namespace、mount、network 与 cgroup v2 检查均为 `false`，严格 readiness 按预期失败关闭。
|
||||
- `--privileged` + private cgroup namespace：namespace、mount、network 通过，但 cgroup v2 delegation 为 `false`，仍按预期不能进入 Cloud ready。
|
||||
- 一次性容器内建立可写 delegated cgroup 子树后：Plugin 与 Box cgroup 探针均为 `true`，nsjail namespace、mount、network 和 cgroup v2 均通过；硬文件系统与 inode quota 继续报告 `false`。
|
||||
- `cpus=0.1` 的 1.0 秒 process-CPU busy loop 实际耗时 `9.13s`；`memory_mb=48` 下逐页提交 128 MiB 以 exit `137` 终止；`pids_limit=8` 下批量 fork 返回 `EAGAIN`。这些结果验证了 CPU、memory+swap 和 PID 的实际内核执行路径。
|
||||
- 新增 `scripts/cloud_runtime_soak.py` 后，在同一 Linux 镜像的独立容器中设置 `--cpus 0.5 --memory 256m --memory-swap 256m --pids-limit 128`，工具从目标 cgroup 读回 quota `50000/100000 usec`、memory `268435456 bytes`、swap `0 bytes` 和 PID `128`。最终复跑中，32 MiB 子负载退出后的 4 秒冷却尾段 `memory.current` 稳健增长和斜率均为 `0`，平均 CPU `0.00132 cores`，OOM、memory pressure、PID max 和 throttle delta 均为 `0`，最终 verdict 为 `pass`。这只是采集器/判定器自检，不替代最终 24 小时生产候选运行。
|
||||
- 本地实际启动 Plugin Runtime 后，控制端口与 debug 端口的公开 `/healthz` 均返回相同聚合 JSON，event-loop monitor 为 running，且正文不含 debug key。采集器显式绕过进程级 HTTP proxy 后，对控制端口执行 6 秒短时 endpoint gate：无失败，观测到的 recent max/p95 均为 `2.233 ms`，verdict 为 `pass`。
|
||||
- 本地实际启动 Box Runtime（未创建 sandbox session）后，`/healthz` 与 `/readyz` 均返回 event-loop、blocking executor、session/process/task 聚合快照；monitor 为 running、样本持续增长，两个端点观测到的 recent max 均为 `2.265 ms`。SIGINT 后 aiohttp、Runtime、reaper 与 monitor 走统一清理路径并正常退出。
|
||||
|
||||
## 上线配置与监控门禁
|
||||
|
||||
最终 24 小时命令、运行位置、阈值语义、负载矩阵和产物要求见 [LangBot Cloud 24 小时资源 Soak 门禁](./cloud-runtime-soak-gate.md)。该工具默认把任一健康失败、OOM/memory pressure、PID limit、CPU throttling 超阈值、blocking executor rejection、冷却尾段内存持续增长或空闲 CPU 过高判为失败；生产运行必须使用 `--require-hard-limits`。
|
||||
|
||||
至少需要监控并告警：
|
||||
|
||||
- Core/Plugin Runtime/Box Runtime 的 RSS、CPU throttling、OOM、PID 数和 event-loop lag。
|
||||
- 各进程 blocking executor 的 running、pending、inflight、active scopes、`global_rejected_total` 和 `scope_rejected_total`；pending 持续不归零或 rejection 增长都应告警。
|
||||
- QueryPool、WebSocket、session、task、plugin worker、Box session 的当前量、容量拒绝和淘汰计数。
|
||||
- Plugin crash/restart 频率、dependency prepare 耗时和失败率。
|
||||
- PostgreSQL pool checked-out/overflow/wait timeout、事务耗时和连接错误。
|
||||
- 临时文件、artifact、dependency environment、Box Workspace volume 的字节数和 inode。
|
||||
|
||||
生产 soak 应覆盖租户突发登录、批量插件 reconcile、插件崩溃重启、WebSocket 断连、Box 并发执行、PG pool 饱和和应用 SIGTERM；持续运行至少 24 小时，并验证负载停止后 RSS、task、socket、文件和子进程数量回到稳定基线。
|
||||
+2
-1
@@ -39,6 +39,7 @@ dependencies = [
|
||||
"quart>=0.20.0",
|
||||
"quart-cors>=0.8.0",
|
||||
"requests>=2.33.0",
|
||||
"regex>=2026.1.15",
|
||||
"slack-sdk>=3.35.0",
|
||||
"alembic>=1.15.0",
|
||||
"sqlalchemy[asyncio]>=2.0.40",
|
||||
@@ -70,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@0cddf3c2bea5939c67b71e488a719e9903c28d17",
|
||||
"langbot-plugin @ git+https://github.com/langbot-app/langbot-plugin-sdk.git@67be7c332ded4c1afa131b28232590cf9728f817",
|
||||
"asyncpg>=0.30.0",
|
||||
"line-bot-sdk>=3.19.0",
|
||||
"matrix-nio>=0.25.2",
|
||||
|
||||
@@ -13,6 +13,12 @@ testpaths = tests
|
||||
# Asyncio configuration
|
||||
asyncio_mode = auto
|
||||
|
||||
# Resource leaks are often reported during object finalization and wrapped by
|
||||
# pytest. Keep both forms fatal so --disable-warnings cannot hide them.
|
||||
filterwarnings =
|
||||
error::ResourceWarning
|
||||
error::pytest.PytestUnraisableExceptionWarning
|
||||
|
||||
# Output options
|
||||
addopts =
|
||||
-v
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,466 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Exercise long-lived Core registries and verify that they reach a plateau.
|
||||
|
||||
This probe is intentionally separate from the default test suite because the
|
||||
audit profile creates tens of thousands of historical identities. It uses the
|
||||
real admission, eviction, and cleanup code while replacing external platform
|
||||
objects that are irrelevant to registry retention.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import gc
|
||||
import json
|
||||
import time
|
||||
import tracemalloc
|
||||
from dataclasses import asdict, dataclass
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import psutil
|
||||
|
||||
from langbot.pkg.api.http.context import ExecutionContext
|
||||
|
||||
# Import the Application graph before taskmgr. The production boot path has
|
||||
# this same ordering; importing taskmgr first exposes its historical cycle
|
||||
# through HTTP route annotations.
|
||||
from langbot.pkg.core import app as _core_app # noqa: F401
|
||||
from langbot.pkg.core.taskmgr import AsyncTaskManager
|
||||
from langbot.pkg.pipeline.pool import QueryPool
|
||||
from langbot.pkg.pipeline.ratelimit.algos.fixedwin import FixedWindowAlgo
|
||||
from langbot.pkg.plugin.connector import PluginRuntimeConnector
|
||||
from langbot.pkg.platform.sources.websocket_adapter import (
|
||||
WebSocketMessage,
|
||||
WebSocketSession,
|
||||
)
|
||||
from langbot.pkg.provider.modelmgr.modelmgr import ModelManager
|
||||
from langbot.pkg.provider.session.sessionmgr import SessionManager
|
||||
from langbot_plugin.api.entities.builtin.provider.session import LauncherTypes
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ProbeScale:
|
||||
query_churn_per_phase: int
|
||||
session_churn_per_phase: int
|
||||
rate_limit_churn_per_phase: int
|
||||
task_churn_per_phase: int
|
||||
websocket_churn_per_phase: int
|
||||
empty_workspace_churn_per_phase: int
|
||||
|
||||
|
||||
SCALES = {
|
||||
'quick': ProbeScale(
|
||||
query_churn_per_phase=2_500,
|
||||
session_churn_per_phase=500,
|
||||
rate_limit_churn_per_phase=10_000,
|
||||
task_churn_per_phase=1_000,
|
||||
websocket_churn_per_phase=500,
|
||||
empty_workspace_churn_per_phase=1_000,
|
||||
),
|
||||
'audit': ProbeScale(
|
||||
query_churn_per_phase=25_000,
|
||||
session_churn_per_phase=2_500,
|
||||
rate_limit_churn_per_phase=10_000,
|
||||
task_churn_per_phase=5_000,
|
||||
websocket_churn_per_phase=2_500,
|
||||
empty_workspace_churn_per_phase=10_000,
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
class _ProbeQuery:
|
||||
"""Small weak-referenceable stand-in for SDK Query construction."""
|
||||
|
||||
def __init__(self, **values):
|
||||
self.__dict__.update(values)
|
||||
|
||||
|
||||
class _EmptyResult:
|
||||
def all(self) -> list:
|
||||
return []
|
||||
|
||||
|
||||
class _EmptyPluginRuntimeHandler:
|
||||
async def reconcile_plugin_installations(self, _states: tuple) -> dict:
|
||||
return {
|
||||
'applied': [],
|
||||
'removed': [],
|
||||
'missing_artifacts': [],
|
||||
'failed_installations': [],
|
||||
}
|
||||
|
||||
def unregister_installation_binding(self, _binding) -> None:
|
||||
raise AssertionError('An empty Workspace exposed an installation binding')
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ProcessSample:
|
||||
rss_bytes: int
|
||||
traced_current_bytes: int
|
||||
traced_peak_bytes: int
|
||||
asyncio_tasks: int
|
||||
threads: int
|
||||
open_fds: int | None
|
||||
|
||||
|
||||
def _sample_process() -> ProcessSample:
|
||||
gc.collect()
|
||||
process = psutil.Process()
|
||||
try:
|
||||
open_fds = process.num_fds()
|
||||
except (AttributeError, psutil.Error):
|
||||
open_fds = None
|
||||
traced_current, traced_peak = tracemalloc.get_traced_memory()
|
||||
return ProcessSample(
|
||||
rss_bytes=process.memory_info().rss,
|
||||
traced_current_bytes=traced_current,
|
||||
traced_peak_bytes=traced_peak,
|
||||
asyncio_tasks=len(asyncio.all_tasks()),
|
||||
threads=process.num_threads(),
|
||||
open_fds=open_fds,
|
||||
)
|
||||
|
||||
|
||||
def _execution_context(index: int, *, query_uuid: str | None = None) -> ExecutionContext:
|
||||
return ExecutionContext(
|
||||
instance_uuid='runtime-resource-probe',
|
||||
workspace_uuid=f'workspace-{index}',
|
||||
placement_generation=1,
|
||||
bot_uuid='probe-bot',
|
||||
pipeline_uuid='probe-pipeline',
|
||||
query_uuid=query_uuid,
|
||||
)
|
||||
|
||||
|
||||
class CoreRuntimeProbe:
|
||||
"""Own the same manager instances across two equal churn phases."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.query_pool = QueryPool(max_queries=100, max_queries_per_workspace=1)
|
||||
app = SimpleNamespace(
|
||||
event_loop=asyncio.get_running_loop(),
|
||||
persistence_mgr=None,
|
||||
instance_config=SimpleNamespace(
|
||||
data={
|
||||
'concurrency': {'session': 1},
|
||||
'system': {
|
||||
'session_retention': {
|
||||
'idle_ttl_seconds': 86_400,
|
||||
'max_entries': 200,
|
||||
'max_entries_per_workspace': 200,
|
||||
'max_conversations_per_session': 20,
|
||||
'max_messages_per_conversation': 100,
|
||||
},
|
||||
'task_retention': {
|
||||
'completed_limit': 200,
|
||||
'max_log_chars': 4_096,
|
||||
'max_active_user_tasks': 256,
|
||||
'max_active_user_tasks_per_workspace': 8,
|
||||
},
|
||||
},
|
||||
}
|
||||
),
|
||||
)
|
||||
self.session_manager = SessionManager(app)
|
||||
self.task_manager = AsyncTaskManager(app)
|
||||
self.rate_limit = FixedWindowAlgo(SimpleNamespace())
|
||||
self.websocket_session = WebSocketSession(
|
||||
'resource-probe',
|
||||
max_conversations=200,
|
||||
max_messages=100,
|
||||
)
|
||||
logger = SimpleNamespace(
|
||||
debug=lambda *_args, **_kwargs: None,
|
||||
info=lambda *_args, **_kwargs: None,
|
||||
warning=lambda *_args, **_kwargs: None,
|
||||
error=lambda *_args, **_kwargs: None,
|
||||
)
|
||||
self.empty_model_queries = 0
|
||||
|
||||
async def execute_empty(_statement):
|
||||
self.empty_model_queries += 1
|
||||
return _EmptyResult()
|
||||
|
||||
model_app = SimpleNamespace(
|
||||
logger=logger,
|
||||
persistence_mgr=SimpleNamespace(execute_async=execute_empty),
|
||||
)
|
||||
self.empty_model_manager = ModelManager(model_app)
|
||||
|
||||
async def runtime_disconnect_callback(_connector) -> None:
|
||||
return None
|
||||
|
||||
plugin_app = SimpleNamespace(
|
||||
instance_config=SimpleNamespace(data={'plugin': {'enable': True}}),
|
||||
deployment=SimpleNamespace(mode='cloud'),
|
||||
logger=logger,
|
||||
)
|
||||
self.empty_plugin_connector = PluginRuntimeConnector(
|
||||
plugin_app,
|
||||
runtime_disconnect_callback,
|
||||
)
|
||||
self.empty_plugin_connector.handler = _EmptyPluginRuntimeHandler()
|
||||
|
||||
async def validate_context(context):
|
||||
return context
|
||||
|
||||
async def load_desired_states(_context):
|
||||
return []
|
||||
|
||||
self.empty_plugin_connector._validate_execution_context = validate_context
|
||||
self.empty_plugin_connector._load_workspace_desired_states = load_desired_states
|
||||
|
||||
async def initialize(self) -> None:
|
||||
await self.rate_limit.initialize()
|
||||
|
||||
async def run_phase(self, scale: ProbeScale, phase: int) -> None:
|
||||
offsets = {
|
||||
'query': (phase - 1) * scale.query_churn_per_phase,
|
||||
'session': (phase - 1) * scale.session_churn_per_phase,
|
||||
'rate': (phase - 1) * scale.rate_limit_churn_per_phase,
|
||||
'task': (phase - 1) * scale.task_churn_per_phase,
|
||||
'websocket': (phase - 1) * scale.websocket_churn_per_phase,
|
||||
'empty_workspace': ((phase - 1) * scale.empty_workspace_churn_per_phase),
|
||||
}
|
||||
await self._churn_queries(offsets['query'], scale.query_churn_per_phase)
|
||||
await self._churn_sessions(offsets['session'], scale.session_churn_per_phase)
|
||||
await self._churn_rate_limits(offsets['rate'], scale.rate_limit_churn_per_phase)
|
||||
await self._churn_tasks(offsets['task'], scale.task_churn_per_phase)
|
||||
self._churn_websocket_history(
|
||||
offsets['websocket'],
|
||||
scale.websocket_churn_per_phase,
|
||||
)
|
||||
await self._churn_empty_workspaces(
|
||||
offsets['empty_workspace'],
|
||||
scale.empty_workspace_churn_per_phase,
|
||||
)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
async def _churn_queries(self, start: int, count: int) -> None:
|
||||
def make_query(**values):
|
||||
return _ProbeQuery(**values)
|
||||
|
||||
with patch(
|
||||
'langbot.pkg.pipeline.pool.pipeline_query.Query',
|
||||
side_effect=make_query,
|
||||
):
|
||||
for index in range(start, start + count):
|
||||
context = _execution_context(index)
|
||||
query = await self.query_pool.add_query(
|
||||
bot_uuid='probe-bot',
|
||||
launcher_type=LauncherTypes.PERSON,
|
||||
launcher_id=f'launcher-{index}',
|
||||
sender_id=f'sender-{index}',
|
||||
message_event=SimpleNamespace(),
|
||||
message_chain=SimpleNamespace(),
|
||||
adapter=None,
|
||||
pipeline_uuid='probe-pipeline',
|
||||
execution_context=context,
|
||||
)
|
||||
removed = await self.query_pool.remove_query(query)
|
||||
if not removed:
|
||||
raise AssertionError('Query cleanup failed')
|
||||
|
||||
async def _churn_sessions(self, start: int, count: int) -> None:
|
||||
for index in range(start, start + count):
|
||||
workspace_index = index % 100
|
||||
context = _execution_context(
|
||||
workspace_index,
|
||||
query_uuid=f'session-query-{index}',
|
||||
)
|
||||
query = SimpleNamespace(
|
||||
launcher_type=LauncherTypes.PERSON,
|
||||
launcher_id=f'launcher-{index}',
|
||||
sender_id=f'sender-{index}',
|
||||
bot_uuid='probe-bot',
|
||||
pipeline_uuid='probe-pipeline',
|
||||
query_uuid=context.query_uuid,
|
||||
_execution_context=context,
|
||||
)
|
||||
await self.session_manager.get_session(query)
|
||||
|
||||
async def _churn_rate_limits(self, start: int, count: int) -> None:
|
||||
for index in range(start, start + count):
|
||||
context = _execution_context(
|
||||
index % 1_000,
|
||||
query_uuid=f'rate-query-{index}',
|
||||
)
|
||||
query = SimpleNamespace(
|
||||
bot_uuid='probe-bot',
|
||||
pipeline_uuid='probe-pipeline',
|
||||
_execution_context=context,
|
||||
pipeline_config={
|
||||
'safety': {
|
||||
'rate-limit': {
|
||||
'window-length': 60,
|
||||
'limitation': 100_000,
|
||||
'strategy': 'drop',
|
||||
}
|
||||
}
|
||||
},
|
||||
)
|
||||
admitted = await self.rate_limit.require_access(
|
||||
query,
|
||||
LauncherTypes.PERSON,
|
||||
f'rate-identity-{index}',
|
||||
)
|
||||
if not admitted:
|
||||
raise AssertionError('Rate-limit registry rejected bounded churn')
|
||||
|
||||
async def _churn_tasks(self, start: int, count: int) -> None:
|
||||
async def complete_immediately() -> None:
|
||||
return None
|
||||
|
||||
for batch_start in range(start, start + count, 256):
|
||||
batch_size = min(256, start + count - batch_start)
|
||||
wrappers = [
|
||||
self.task_manager.create_task(
|
||||
complete_immediately(),
|
||||
name=f'resource-probe-{batch_start + offset}',
|
||||
)
|
||||
for offset in range(batch_size)
|
||||
]
|
||||
await asyncio.gather(*(wrapper.task for wrapper in wrappers))
|
||||
await asyncio.sleep(0)
|
||||
|
||||
def _churn_websocket_history(self, start: int, count: int) -> None:
|
||||
for index in range(start, start + count):
|
||||
conversation_key = f'conversation-{index}'
|
||||
response_id = f'response-{index}'
|
||||
indexes = self.websocket_session.get_stream_message_indexes(conversation_key)
|
||||
indexes[response_id] = 0
|
||||
self.websocket_session.append_message(
|
||||
conversation_key,
|
||||
WebSocketMessage(
|
||||
id=self.websocket_session.next_message_id(conversation_key),
|
||||
role='assistant',
|
||||
content='probe',
|
||||
message_chain=[],
|
||||
timestamp='1970-01-01T00:00:00+00:00',
|
||||
is_final=True,
|
||||
),
|
||||
)
|
||||
|
||||
async def _churn_empty_workspaces(self, start: int, count: int) -> None:
|
||||
for index in range(start, start + count):
|
||||
await self.empty_model_manager._load_workspace_models(_execution_context(index))
|
||||
await self.empty_plugin_connector.reconcile_projected_workspaces(
|
||||
_execution_context(index) for index in range(start, start + count)
|
||||
)
|
||||
|
||||
def retained_state(self) -> dict[str, int]:
|
||||
return {
|
||||
'query_cached': len(self.query_pool.cached_queries),
|
||||
'query_queued': len(self.query_pool.queries),
|
||||
'query_active_workspaces': len(self.query_pool.active_query_count_by_workspace),
|
||||
'query_scope_counters': len(self.query_pool.query_count_by_scope),
|
||||
'sessions': len(self.session_manager.session_list),
|
||||
'session_index': len(self.session_manager._session_index),
|
||||
'rate_limit_containers': len(self.rate_limit.containers),
|
||||
'task_records': len(self.task_manager.tasks),
|
||||
'websocket_conversations': len(self.websocket_session.message_lists),
|
||||
'websocket_stream_indexes': len(self.websocket_session.stream_message_indexes),
|
||||
'empty_model_scopes': len(self.empty_model_manager._scope_generations),
|
||||
'empty_model_providers': len(self.empty_model_manager.provider_dict),
|
||||
'empty_model_llms': len(self.empty_model_manager.llm_model_dict),
|
||||
'empty_plugin_workspace_sets': len(self.empty_plugin_connector._workspace_installations),
|
||||
'empty_plugin_installations': len(self.empty_plugin_connector._known_desired_states),
|
||||
}
|
||||
|
||||
def assert_bounded(self) -> None:
|
||||
state = self.retained_state()
|
||||
expected_maximums = {
|
||||
'query_cached': 0,
|
||||
'query_queued': 0,
|
||||
'query_active_workspaces': 0,
|
||||
'query_scope_counters': 100,
|
||||
'sessions': 200,
|
||||
'session_index': 200,
|
||||
'rate_limit_containers': 10_000,
|
||||
'task_records': 200,
|
||||
'websocket_conversations': 200,
|
||||
'websocket_stream_indexes': 200,
|
||||
'empty_model_scopes': 0,
|
||||
'empty_model_providers': 0,
|
||||
'empty_model_llms': 0,
|
||||
'empty_plugin_workspace_sets': 0,
|
||||
'empty_plugin_installations': 0,
|
||||
}
|
||||
violations = {key: (state[key], maximum) for key, maximum in expected_maximums.items() if state[key] > maximum}
|
||||
if violations:
|
||||
raise AssertionError(f'Core retained-state limits failed: {violations}')
|
||||
|
||||
|
||||
async def _run(args: argparse.Namespace) -> dict:
|
||||
scale = SCALES[args.scale]
|
||||
tracemalloc.start()
|
||||
started_at = time.monotonic()
|
||||
probe = CoreRuntimeProbe()
|
||||
await probe.initialize()
|
||||
|
||||
baseline = _sample_process()
|
||||
await probe.run_phase(scale, 1)
|
||||
probe.assert_bounded()
|
||||
phase_one = _sample_process()
|
||||
state_one = probe.retained_state()
|
||||
|
||||
await probe.run_phase(scale, 2)
|
||||
probe.assert_bounded()
|
||||
phase_two = _sample_process()
|
||||
state_two = probe.retained_state()
|
||||
|
||||
if state_two != state_one:
|
||||
raise AssertionError(f'Core retained state did not plateau: phase_one={state_one}, phase_two={state_two}')
|
||||
traced_growth = phase_two.traced_current_bytes - phase_one.traced_current_bytes
|
||||
rss_growth = phase_two.rss_bytes - phase_one.rss_bytes
|
||||
max_traced_growth = int(args.max_traced_growth_mib * 1024 * 1024)
|
||||
max_rss_growth = int(args.max_rss_growth_mib * 1024 * 1024)
|
||||
if traced_growth > max_traced_growth:
|
||||
raise AssertionError(f'Second-phase traced memory grew by {traced_growth} bytes (limit {max_traced_growth})')
|
||||
if rss_growth > max_rss_growth:
|
||||
raise AssertionError(f'Second-phase RSS grew by {rss_growth} bytes (limit {max_rss_growth})')
|
||||
|
||||
return {
|
||||
'component': 'langbot-core',
|
||||
'scale': args.scale,
|
||||
'work_per_phase': asdict(scale),
|
||||
'elapsed_seconds': round(time.monotonic() - started_at, 3),
|
||||
'samples': {
|
||||
'baseline': asdict(baseline),
|
||||
'phase_one': asdict(phase_one),
|
||||
'phase_two': asdict(phase_two),
|
||||
},
|
||||
'second_phase_growth': {
|
||||
'rss_bytes': rss_growth,
|
||||
'traced_current_bytes': traced_growth,
|
||||
},
|
||||
'retained_state': {
|
||||
'phase_one': state_one,
|
||||
'phase_two': state_two,
|
||||
},
|
||||
'passed': True,
|
||||
}
|
||||
|
||||
|
||||
def _parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument('--scale', choices=tuple(SCALES), default='quick')
|
||||
parser.add_argument('--max-traced-growth-mib', type=float, default=8.0)
|
||||
parser.add_argument('--max-rss-growth-mib', type=float, default=64.0)
|
||||
parser.add_argument('--json', action='store_true', help='Print compact JSON')
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = _parse_args()
|
||||
result = asyncio.run(_run(args))
|
||||
if args.json:
|
||||
print(json.dumps(result, sort_keys=True))
|
||||
else:
|
||||
print(json.dumps(result, indent=2, sort_keys=True))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,561 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Measure populated Workspace runtime replacement cost and retention.
|
||||
|
||||
Unlike ``runtime_resource_probe.py``, which stresses historical request keys
|
||||
and empty tenants, this probe keeps one representative Provider, LLM,
|
||||
Embedding model, Rerank model, Pipeline, Bot, and Knowledge Base per Workspace.
|
||||
It then advances every Workspace to a new placement generation and verifies
|
||||
that old runtime objects are closed and collectible while active registry
|
||||
cardinality remains constant.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import gc
|
||||
import json
|
||||
import time
|
||||
import tracemalloc
|
||||
import weakref
|
||||
from dataclasses import asdict, dataclass
|
||||
from types import SimpleNamespace
|
||||
|
||||
import psutil
|
||||
|
||||
from langbot.pkg.api.http.context import ExecutionContext
|
||||
|
||||
# Match the production import order; importing a leaf manager first exposes a
|
||||
# historical annotation cycle that the application graph resolves.
|
||||
from langbot.pkg.core import app as _core_app # noqa: F401
|
||||
from langbot.pkg.entity.persistence import bot as persistence_bot
|
||||
from langbot.pkg.entity.persistence import model as persistence_model
|
||||
from langbot.pkg.entity.persistence import pipeline as persistence_pipeline
|
||||
from langbot.pkg.entity.persistence import rag as persistence_rag
|
||||
from langbot.pkg.pipeline.pipelinemgr import PipelineManager
|
||||
from langbot.pkg.platform.botmgr import PlatformManager
|
||||
from langbot.pkg.provider.modelmgr import requester
|
||||
from langbot.pkg.provider.modelmgr.modelmgr import ModelManager
|
||||
from langbot.pkg.provider.tools.loaders.mcp import MCPLoader
|
||||
from langbot.pkg.rag.knowledge.kbmgr import RAGManager
|
||||
from langbot.pkg.workspace.entities import WorkspaceExecutionBinding
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ProbeScale:
|
||||
workspaces: int
|
||||
|
||||
|
||||
SCALES = {
|
||||
'quick': ProbeScale(workspaces=250),
|
||||
'audit': ProbeScale(workspaces=5_000),
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ProcessSample:
|
||||
rss_bytes: int
|
||||
traced_current_bytes: int
|
||||
traced_peak_bytes: int
|
||||
asyncio_tasks: int
|
||||
threads: int
|
||||
open_fds: int | None
|
||||
|
||||
|
||||
class _ProbeLogger:
|
||||
def debug(self, *_args, **_kwargs) -> None:
|
||||
return None
|
||||
|
||||
def info(self, *_args, **_kwargs) -> None:
|
||||
return None
|
||||
|
||||
def warning(self, *_args, **_kwargs) -> None:
|
||||
return None
|
||||
|
||||
def error(self, *_args, **_kwargs) -> None:
|
||||
return None
|
||||
|
||||
|
||||
class _ProbeWorkspaceService:
|
||||
instance_uuid = 'runtime-capacity-probe'
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.generations: dict[str, int] = {}
|
||||
self.binding_lookups = 0
|
||||
|
||||
async def get_execution_binding(
|
||||
self,
|
||||
workspace_uuid: str,
|
||||
*,
|
||||
expected_generation: int | None = None,
|
||||
) -> WorkspaceExecutionBinding:
|
||||
self.binding_lookups += 1
|
||||
generation = self.generations[workspace_uuid]
|
||||
if expected_generation is not None and expected_generation != generation:
|
||||
raise AssertionError(f'stale probe generation {expected_generation} != {generation}')
|
||||
return WorkspaceExecutionBinding(
|
||||
instance_uuid=self.instance_uuid,
|
||||
workspace_uuid=workspace_uuid,
|
||||
placement_generation=generation,
|
||||
write_fenced=False,
|
||||
state='active',
|
||||
)
|
||||
|
||||
|
||||
class _ProbeRequester(requester.ProviderAPIRequester):
|
||||
name = 'capacity-probe'
|
||||
closed = 0
|
||||
|
||||
async def invoke_llm(
|
||||
self,
|
||||
query,
|
||||
model,
|
||||
messages,
|
||||
funcs=None,
|
||||
extra_args=None,
|
||||
remove_think=False,
|
||||
):
|
||||
return None
|
||||
|
||||
async def aclose(self) -> None:
|
||||
type(self).closed += 1
|
||||
|
||||
|
||||
class _ProbeAdapter:
|
||||
killed = 0
|
||||
|
||||
def __init__(self, _config, _logger) -> None:
|
||||
self.listeners = []
|
||||
|
||||
def register_listener(self, event_type, listener) -> None:
|
||||
self.listeners.append((event_type, listener))
|
||||
|
||||
async def kill(self) -> None:
|
||||
type(self).killed += 1
|
||||
|
||||
|
||||
class _ProbeMCPSession:
|
||||
closed = 0
|
||||
|
||||
def __init__(self, server_name: str) -> None:
|
||||
self.server_name = server_name
|
||||
|
||||
async def shutdown(self) -> None:
|
||||
type(self).closed += 1
|
||||
|
||||
|
||||
def _sample_process() -> ProcessSample:
|
||||
gc.collect()
|
||||
process = psutil.Process()
|
||||
try:
|
||||
open_fds = process.num_fds()
|
||||
except (AttributeError, psutil.Error):
|
||||
open_fds = None
|
||||
traced_current, traced_peak = tracemalloc.get_traced_memory()
|
||||
return ProcessSample(
|
||||
rss_bytes=process.memory_info().rss,
|
||||
traced_current_bytes=traced_current,
|
||||
traced_peak_bytes=traced_peak,
|
||||
asyncio_tasks=len(asyncio.all_tasks()),
|
||||
threads=process.num_threads(),
|
||||
open_fds=open_fds,
|
||||
)
|
||||
|
||||
|
||||
class PopulatedWorkspaceProbe:
|
||||
def __init__(self) -> None:
|
||||
_ProbeRequester.closed = 0
|
||||
_ProbeAdapter.killed = 0
|
||||
_ProbeMCPSession.closed = 0
|
||||
self.workspace_service = _ProbeWorkspaceService()
|
||||
self.logger = _ProbeLogger()
|
||||
self.app = SimpleNamespace(
|
||||
logger=self.logger,
|
||||
workspace_service=self.workspace_service,
|
||||
persistence_mgr=SimpleNamespace(
|
||||
mode=SimpleNamespace(value='cloud_runtime'),
|
||||
),
|
||||
pipeline_config_meta_trigger={'name': 'trigger', 'stages': []},
|
||||
pipeline_config_meta_safety={'name': 'safety', 'stages': []},
|
||||
pipeline_config_meta_ai={'name': 'ai', 'stages': []},
|
||||
pipeline_config_meta_output={'name': 'output', 'stages': []},
|
||||
task_mgr=SimpleNamespace(
|
||||
cancel_by_scope=lambda *_args, **_kwargs: None,
|
||||
cancel_task=lambda *_args, **_kwargs: None,
|
||||
),
|
||||
)
|
||||
self.model_manager = ModelManager(self.app)
|
||||
self.model_manager.requester_dict = {
|
||||
_ProbeRequester.name: _ProbeRequester,
|
||||
}
|
||||
self.pipeline_manager = PipelineManager(self.app)
|
||||
self.pipeline_manager.stage_dict = {}
|
||||
self.rag_manager = RAGManager(self.app)
|
||||
self.mcp_loader = MCPLoader(self.app)
|
||||
self.platform_manager = PlatformManager(self.app)
|
||||
self.platform_manager.adapter_dict = {
|
||||
'capacity-probe': _ProbeAdapter,
|
||||
}
|
||||
self.generation_refs: dict[
|
||||
int,
|
||||
list[weakref.ReferenceType],
|
||||
] = {}
|
||||
|
||||
def _context(
|
||||
self,
|
||||
workspace_uuid: str,
|
||||
generation: int,
|
||||
*,
|
||||
bot_uuid: str | None = None,
|
||||
pipeline_uuid: str | None = None,
|
||||
) -> ExecutionContext:
|
||||
return ExecutionContext(
|
||||
instance_uuid=self.workspace_service.instance_uuid,
|
||||
workspace_uuid=workspace_uuid,
|
||||
placement_generation=generation,
|
||||
bot_uuid=bot_uuid,
|
||||
pipeline_uuid=pipeline_uuid,
|
||||
)
|
||||
|
||||
async def load_generation(self, workspaces: int, generation: int) -> None:
|
||||
for index in range(workspaces):
|
||||
workspace_uuid = f'workspace-{index}'
|
||||
provider_uuid = f'provider-{index}'
|
||||
llm_uuid = f'llm-{index}'
|
||||
embedding_uuid = f'embedding-{index}'
|
||||
rerank_uuid = f'rerank-{index}'
|
||||
pipeline_uuid = f'pipeline-{index}'
|
||||
bot_uuid = f'bot-{index}'
|
||||
kb_uuid = f'knowledge-{index}'
|
||||
mcp_server_name = f'mcp-{index}'
|
||||
self.workspace_service.generations[workspace_uuid] = generation
|
||||
context = self._context(workspace_uuid, generation)
|
||||
|
||||
runtime_provider = await self.model_manager.load_provider(
|
||||
context,
|
||||
persistence_model.ModelProvider(
|
||||
uuid=provider_uuid,
|
||||
workspace_uuid=workspace_uuid,
|
||||
name='Capacity Provider',
|
||||
requester=_ProbeRequester.name,
|
||||
base_url='https://capacity.invalid',
|
||||
api_keys=['probe'],
|
||||
),
|
||||
)
|
||||
await self.model_manager.cache_provider(context, runtime_provider)
|
||||
|
||||
runtime_llm = await self.model_manager.load_llm_model_with_provider(
|
||||
context,
|
||||
persistence_model.LLMModel(
|
||||
uuid=llm_uuid,
|
||||
workspace_uuid=workspace_uuid,
|
||||
name='Capacity LLM',
|
||||
provider_uuid=provider_uuid,
|
||||
abilities=['func_call'],
|
||||
extra_args={'temperature': 0.1},
|
||||
),
|
||||
runtime_provider,
|
||||
)
|
||||
await self.model_manager.cache_llm_model(context, runtime_llm)
|
||||
runtime_embedding = await self.model_manager.load_embedding_model_with_provider(
|
||||
context,
|
||||
persistence_model.EmbeddingModel(
|
||||
uuid=embedding_uuid,
|
||||
workspace_uuid=workspace_uuid,
|
||||
name='Capacity Embedding',
|
||||
provider_uuid=provider_uuid,
|
||||
extra_args={'dimensions': 1_024},
|
||||
),
|
||||
runtime_provider,
|
||||
)
|
||||
await self.model_manager.cache_embedding_model(
|
||||
context,
|
||||
runtime_embedding,
|
||||
)
|
||||
runtime_rerank = await self.model_manager.load_rerank_model_with_provider(
|
||||
context,
|
||||
persistence_model.RerankModel(
|
||||
uuid=rerank_uuid,
|
||||
workspace_uuid=workspace_uuid,
|
||||
name='Capacity Rerank',
|
||||
provider_uuid=provider_uuid,
|
||||
extra_args={},
|
||||
),
|
||||
runtime_provider,
|
||||
)
|
||||
await self.model_manager.cache_rerank_model(
|
||||
context,
|
||||
runtime_rerank,
|
||||
)
|
||||
|
||||
pipeline_context = self._context(
|
||||
workspace_uuid,
|
||||
generation,
|
||||
pipeline_uuid=pipeline_uuid,
|
||||
)
|
||||
await self.pipeline_manager.load_pipeline(
|
||||
pipeline_context,
|
||||
persistence_pipeline.LegacyPipeline(
|
||||
uuid=pipeline_uuid,
|
||||
workspace_uuid=workspace_uuid,
|
||||
name='Capacity Pipeline',
|
||||
description='',
|
||||
for_version='probe',
|
||||
is_default=True,
|
||||
stages=[],
|
||||
config={},
|
||||
extensions_preferences={},
|
||||
),
|
||||
_binding_validated=True,
|
||||
)
|
||||
runtime_pipeline = self.pipeline_manager._pipelines_by_key[
|
||||
(
|
||||
self.workspace_service.instance_uuid,
|
||||
workspace_uuid,
|
||||
pipeline_uuid,
|
||||
)
|
||||
]
|
||||
|
||||
runtime_kb = await self.rag_manager.load_knowledge_base(
|
||||
context,
|
||||
persistence_rag.KnowledgeBase(
|
||||
uuid=kb_uuid,
|
||||
workspace_uuid=workspace_uuid,
|
||||
name='Capacity Knowledge',
|
||||
description='',
|
||||
knowledge_engine_plugin_id=None,
|
||||
collection_id=kb_uuid,
|
||||
creation_settings={},
|
||||
retrieval_settings={},
|
||||
),
|
||||
_binding_validated=True,
|
||||
)
|
||||
|
||||
await self.mcp_loader._assert_execution_active(context)
|
||||
runtime_mcp = _ProbeMCPSession(mcp_server_name)
|
||||
self.mcp_loader._register_session(
|
||||
context,
|
||||
mcp_server_name,
|
||||
runtime_mcp,
|
||||
)
|
||||
|
||||
bot_context = self._context(
|
||||
workspace_uuid,
|
||||
generation,
|
||||
bot_uuid=bot_uuid,
|
||||
)
|
||||
runtime_bot = await self.platform_manager.load_bot(
|
||||
bot_context,
|
||||
persistence_bot.Bot(
|
||||
uuid=bot_uuid,
|
||||
workspace_uuid=workspace_uuid,
|
||||
name='Capacity Bot',
|
||||
description='',
|
||||
adapter='capacity-probe',
|
||||
adapter_config={},
|
||||
enable=True,
|
||||
use_pipeline_uuid=pipeline_uuid,
|
||||
pipeline_routing_rules=[],
|
||||
),
|
||||
_binding_validated=True,
|
||||
)
|
||||
|
||||
self.generation_refs.setdefault(generation, []).extend(
|
||||
(
|
||||
weakref.ref(runtime_provider),
|
||||
weakref.ref(runtime_llm),
|
||||
weakref.ref(runtime_embedding),
|
||||
weakref.ref(runtime_rerank),
|
||||
weakref.ref(runtime_pipeline),
|
||||
weakref.ref(runtime_kb),
|
||||
weakref.ref(runtime_mcp),
|
||||
weakref.ref(runtime_bot),
|
||||
)
|
||||
)
|
||||
|
||||
await asyncio.sleep(0)
|
||||
|
||||
def retained_state(self) -> dict[str, int]:
|
||||
return {
|
||||
'model_providers': len(self.model_manager.provider_dict),
|
||||
'llm_models': len(self.model_manager.llm_model_dict),
|
||||
'embedding_models': len(self.model_manager.embedding_model_dict),
|
||||
'rerank_models': len(self.model_manager.rerank_model_dict),
|
||||
'model_scopes': len(self.model_manager._scope_generations),
|
||||
'pipelines': len(self.pipeline_manager._pipelines_by_key),
|
||||
'pipeline_scopes': len(self.pipeline_manager._scope_generations),
|
||||
'knowledge_bases': len(self.rag_manager.knowledge_bases),
|
||||
'knowledge_scopes': len(self.rag_manager._scope_generations),
|
||||
'mcp_sessions': len(self.mcp_loader.sessions),
|
||||
'mcp_scopes': len(self.mcp_loader._scope_generations),
|
||||
'bots': len(self.platform_manager._bots_by_key),
|
||||
'bot_scopes': len(self.platform_manager._scope_generations),
|
||||
'requesters_closed': _ProbeRequester.closed,
|
||||
'adapters_killed': _ProbeAdapter.killed,
|
||||
'mcp_sessions_closed': _ProbeMCPSession.closed,
|
||||
'binding_lookups': self.workspace_service.binding_lookups,
|
||||
}
|
||||
|
||||
def assert_generation_state(
|
||||
self,
|
||||
workspaces: int,
|
||||
generation: int,
|
||||
) -> None:
|
||||
state = self.retained_state()
|
||||
cardinality_keys = (
|
||||
'model_providers',
|
||||
'llm_models',
|
||||
'embedding_models',
|
||||
'rerank_models',
|
||||
'model_scopes',
|
||||
'pipelines',
|
||||
'pipeline_scopes',
|
||||
'knowledge_bases',
|
||||
'knowledge_scopes',
|
||||
'mcp_sessions',
|
||||
'mcp_scopes',
|
||||
'bots',
|
||||
'bot_scopes',
|
||||
)
|
||||
invalid = {key: value for key in cardinality_keys if (value := state[key]) != workspaces}
|
||||
if invalid:
|
||||
raise AssertionError(f'populated Workspace cardinality mismatch: {invalid}')
|
||||
expected_retired = (generation - 1) * workspaces
|
||||
if state['requesters_closed'] != expected_retired:
|
||||
raise AssertionError(f'retired requester count {state["requesters_closed"]} != {expected_retired}')
|
||||
if state['adapters_killed'] != expected_retired:
|
||||
raise AssertionError(f'retired adapter count {state["adapters_killed"]} != {expected_retired}')
|
||||
if state['mcp_sessions_closed'] != expected_retired:
|
||||
raise AssertionError(f'retired MCP session count {state["mcp_sessions_closed"]} != {expected_retired}')
|
||||
|
||||
def assert_generation_collected(self, generation: int) -> None:
|
||||
gc.collect()
|
||||
references = self.generation_refs.pop(generation)
|
||||
retained = sum(reference() is not None for reference in references)
|
||||
if retained:
|
||||
raise AssertionError(f'{retained} generation-{generation} runtime objects remain reachable')
|
||||
|
||||
|
||||
async def _run(args: argparse.Namespace) -> dict:
|
||||
scale = SCALES[args.scale]
|
||||
tracemalloc.start()
|
||||
probe = PopulatedWorkspaceProbe()
|
||||
baseline = _sample_process()
|
||||
|
||||
phase_one_started = time.monotonic()
|
||||
await probe.load_generation(scale.workspaces, 1)
|
||||
phase_one_seconds = time.monotonic() - phase_one_started
|
||||
probe.assert_generation_state(scale.workspaces, 1)
|
||||
phase_one = _sample_process()
|
||||
phase_one_state = probe.retained_state()
|
||||
|
||||
phase_two_started = time.monotonic()
|
||||
await probe.load_generation(scale.workspaces, 2)
|
||||
phase_two_seconds = time.monotonic() - phase_two_started
|
||||
probe.assert_generation_state(scale.workspaces, 2)
|
||||
probe.assert_generation_collected(1)
|
||||
phase_two = _sample_process()
|
||||
phase_two_state = probe.retained_state()
|
||||
|
||||
phase_three_started = time.monotonic()
|
||||
await probe.load_generation(scale.workspaces, 3)
|
||||
phase_three_seconds = time.monotonic() - phase_three_started
|
||||
probe.assert_generation_state(scale.workspaces, 3)
|
||||
probe.assert_generation_collected(2)
|
||||
phase_three = _sample_process()
|
||||
phase_three_state = probe.retained_state()
|
||||
|
||||
cardinality_keys = (
|
||||
'model_providers',
|
||||
'llm_models',
|
||||
'embedding_models',
|
||||
'rerank_models',
|
||||
'model_scopes',
|
||||
'pipelines',
|
||||
'pipeline_scopes',
|
||||
'knowledge_bases',
|
||||
'knowledge_scopes',
|
||||
'mcp_sessions',
|
||||
'mcp_scopes',
|
||||
'bots',
|
||||
'bot_scopes',
|
||||
)
|
||||
if any(
|
||||
phase_two_state[key] != phase_one_state[key] or phase_three_state[key] != phase_one_state[key]
|
||||
for key in cardinality_keys
|
||||
):
|
||||
raise AssertionError(
|
||||
'populated Workspace registries did not plateau: '
|
||||
f'phase_one={phase_one_state}, phase_two={phase_two_state}, '
|
||||
f'phase_three={phase_three_state}'
|
||||
)
|
||||
|
||||
traced_growth = phase_three.traced_current_bytes - phase_two.traced_current_bytes
|
||||
rss_growth = phase_three.rss_bytes - phase_two.rss_bytes
|
||||
max_traced_growth = int(args.max_traced_growth_mib * 1024 * 1024)
|
||||
max_rss_growth = int(args.max_rss_growth_mib * 1024 * 1024)
|
||||
if traced_growth > max_traced_growth:
|
||||
raise AssertionError(f'replacement traced memory grew by {traced_growth} bytes (limit {max_traced_growth})')
|
||||
if rss_growth > max_rss_growth:
|
||||
raise AssertionError(f'replacement RSS grew by {rss_growth} bytes (limit {max_rss_growth})')
|
||||
phase_ratio = max(
|
||||
phase_two_seconds,
|
||||
phase_three_seconds,
|
||||
) / max(phase_one_seconds, 0.000_001)
|
||||
if phase_ratio > args.max_replacement_time_ratio:
|
||||
raise AssertionError(f'replacement phase ratio {phase_ratio:.3f} exceeds {args.max_replacement_time_ratio:.3f}')
|
||||
|
||||
return {
|
||||
'component': 'langbot-populated-workspaces',
|
||||
'scale': args.scale,
|
||||
'workspaces': scale.workspaces,
|
||||
'passed': True,
|
||||
'phase_seconds': {
|
||||
'initial': round(phase_one_seconds, 3),
|
||||
'replacement_one': round(phase_two_seconds, 3),
|
||||
'replacement_two': round(phase_three_seconds, 3),
|
||||
'maximum_replacement_ratio': round(phase_ratio, 3),
|
||||
},
|
||||
'samples': {
|
||||
'baseline': asdict(baseline),
|
||||
'phase_one': asdict(phase_one),
|
||||
'phase_two': asdict(phase_two),
|
||||
'phase_three': asdict(phase_three),
|
||||
},
|
||||
'replacement_growth': {
|
||||
'rss_bytes': rss_growth,
|
||||
'traced_current_bytes': traced_growth,
|
||||
},
|
||||
'retained_state': {
|
||||
'phase_one': phase_one_state,
|
||||
'phase_two': phase_two_state,
|
||||
'phase_three': phase_three_state,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument('--scale', choices=tuple(SCALES), default='quick')
|
||||
parser.add_argument('--max-traced-growth-mib', type=float, default=16.0)
|
||||
parser.add_argument('--max-rss-growth-mib', type=float, default=64.0)
|
||||
parser.add_argument(
|
||||
'--max-replacement-time-ratio',
|
||||
type=float,
|
||||
default=3.0,
|
||||
)
|
||||
parser.add_argument('--json', action='store_true')
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = _parse_args()
|
||||
result = asyncio.run(_run(args))
|
||||
if args.json:
|
||||
print(json.dumps(result, sort_keys=True))
|
||||
else:
|
||||
print(json.dumps(result, indent=2, sort_keys=True))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -6,6 +6,22 @@ from typing import Dict, List, Any, AsyncGenerator
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
from langbot.pkg.utils import httpclient
|
||||
|
||||
_MAX_COZE_RESPONSE_BYTES = 16 * 1024 * 1024
|
||||
_MAX_COZE_EVENT_BYTES = 1024 * 1024
|
||||
_MAX_COZE_MEDIA_BYTES = 10 * 1024 * 1024
|
||||
|
||||
|
||||
def _read_local_media_limited(path: Path) -> bytes:
|
||||
if path.stat().st_size > _MAX_COZE_MEDIA_BYTES:
|
||||
raise ValueError('Coze upload exceeds the size limit')
|
||||
with path.open('rb') as handle:
|
||||
body = handle.read(_MAX_COZE_MEDIA_BYTES + 1)
|
||||
if len(body) > _MAX_COZE_MEDIA_BYTES:
|
||||
raise ValueError('Coze upload exceeds the size limit')
|
||||
return body
|
||||
|
||||
|
||||
class AsyncCozeAPIClient:
|
||||
def __init__(self, api_key: str, api_base: str = 'https://api.coze.cn'):
|
||||
@@ -58,19 +74,24 @@ class AsyncCozeAPIClient:
|
||||
if isinstance(file, Path):
|
||||
if not file.exists():
|
||||
raise ValueError(f'File not found: {file}')
|
||||
with open(file, 'rb') as f:
|
||||
file = f.read()
|
||||
file = await asyncio.to_thread(_read_local_media_limited, file)
|
||||
|
||||
# 处理文件路径字符串
|
||||
elif isinstance(file, str):
|
||||
if not os.path.isfile(file):
|
||||
raise ValueError(f'File not found: {file}')
|
||||
with open(file, 'rb') as f:
|
||||
file = f.read()
|
||||
file = await asyncio.to_thread(
|
||||
_read_local_media_limited,
|
||||
Path(file),
|
||||
)
|
||||
|
||||
# 处理文件对象
|
||||
elif hasattr(file, 'read'):
|
||||
file = file.read()
|
||||
file = await asyncio.to_thread(file.read, _MAX_COZE_MEDIA_BYTES + 1)
|
||||
if not isinstance(file, (bytes, bytearray)):
|
||||
raise ValueError('Unsupported Coze upload type')
|
||||
if len(file) > _MAX_COZE_MEDIA_BYTES:
|
||||
raise ValueError('Coze upload exceeds the size limit')
|
||||
|
||||
session = await self.coze_session()
|
||||
url = f'{self.api_base}/v1/files/upload'
|
||||
@@ -87,13 +108,18 @@ class AsyncCozeAPIClient:
|
||||
if response.status == 401:
|
||||
raise Exception('Coze API 认证失败，请检查 API Key 是否正确')
|
||||
|
||||
response_text = await response.text()
|
||||
response_text = (
|
||||
await httpclient.read_limited(
|
||||
response,
|
||||
max_bytes=_MAX_COZE_EVENT_BYTES,
|
||||
)
|
||||
).decode('utf-8', errors='replace')
|
||||
|
||||
if response.status != 200:
|
||||
raise Exception(f'文件上传失败，状态码: {response.status}, 响应: {response_text}')
|
||||
try:
|
||||
result = await response.json()
|
||||
except json.JSONDecodeError:
|
||||
result = json.loads(response_text)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
raise Exception(f'文件上传响应解析失败: {response_text}')
|
||||
|
||||
if result.get('code') != 0:
|
||||
@@ -158,7 +184,15 @@ class AsyncCozeAPIClient:
|
||||
if response.status != 200:
|
||||
raise Exception(f'Coze API 流式请求失败，状态码: {response.status}')
|
||||
|
||||
total_bytes = 0
|
||||
chunk_type = 'message'
|
||||
chunk_data = ''
|
||||
async for chunk in response.content:
|
||||
total_bytes += len(chunk)
|
||||
if total_bytes > _MAX_COZE_RESPONSE_BYTES:
|
||||
raise Exception('Coze API stream exceeds the runtime limit')
|
||||
if len(chunk) > _MAX_COZE_EVENT_BYTES:
|
||||
raise Exception('Coze API event exceeds the runtime limit')
|
||||
chunk = chunk.decode('utf-8')
|
||||
if chunk != '\n':
|
||||
if chunk.startswith('event:'):
|
||||
|
||||
@@ -12,10 +12,26 @@ from collections.abc import AsyncGenerator
|
||||
|
||||
import httpx
|
||||
|
||||
from langbot.pkg.utils import httpclient
|
||||
|
||||
from .errors import DeerFlowAPIError
|
||||
|
||||
|
||||
SSE_MAX_BUFFER_CHARS = 1_048_576
|
||||
SSE_MAX_TOTAL_BYTES = 16 * 1024 * 1024
|
||||
ERROR_BODY_MAX_BYTES = 1024 * 1024
|
||||
|
||||
|
||||
async def _read_error_body(response: httpx.Response) -> str:
|
||||
body = bytearray()
|
||||
async for chunk in response.aiter_bytes(8192):
|
||||
body.extend(chunk)
|
||||
if len(body) > ERROR_BODY_MAX_BYTES:
|
||||
raise DeerFlowAPIError(
|
||||
operation='read error response',
|
||||
body='response exceeds the runtime limit',
|
||||
)
|
||||
return body.decode('utf-8', errors='replace')
|
||||
|
||||
|
||||
def _normalize_sse_newlines(text: str) -> str:
|
||||
@@ -94,6 +110,7 @@ class AsyncDeerFlowClient:
|
||||
async with httpx.AsyncClient(
|
||||
trust_env=True,
|
||||
timeout=timeout,
|
||||
event_hooks=httpclient.httpx_response_limit_hooks(),
|
||||
) as client:
|
||||
response = await client.post(
|
||||
url,
|
||||
@@ -101,13 +118,14 @@ class AsyncDeerFlowClient:
|
||||
json=payload,
|
||||
)
|
||||
if response.status_code not in (200, 201):
|
||||
body = await httpclient.response_text(response)
|
||||
raise DeerFlowAPIError(
|
||||
operation='create thread',
|
||||
status=response.status_code,
|
||||
body=response.text,
|
||||
body=body,
|
||||
url=url,
|
||||
)
|
||||
return response.json()
|
||||
return await httpclient.parse_json_response(response)
|
||||
|
||||
async def delete_thread(self, thread_id: str, timeout: float = 20) -> None:
|
||||
"""删除指定 thread"""
|
||||
@@ -116,13 +134,15 @@ class AsyncDeerFlowClient:
|
||||
async with httpx.AsyncClient(
|
||||
trust_env=True,
|
||||
timeout=timeout,
|
||||
event_hooks=httpclient.httpx_response_limit_hooks(),
|
||||
) as client:
|
||||
response = await client.delete(url, headers=self.headers)
|
||||
if response.status_code not in (200, 202, 204, 404):
|
||||
body = await httpclient.response_text(response)
|
||||
raise DeerFlowAPIError(
|
||||
operation='delete thread',
|
||||
status=response.status_code,
|
||||
body=response.text,
|
||||
body=body,
|
||||
url=url,
|
||||
thread_id=thread_id,
|
||||
)
|
||||
@@ -163,19 +183,27 @@ class AsyncDeerFlowClient:
|
||||
json=payload,
|
||||
) as resp:
|
||||
if resp.status_code != 200:
|
||||
body = await resp.aread()
|
||||
raise DeerFlowAPIError(
|
||||
operation='runs/stream request',
|
||||
status=resp.status_code,
|
||||
body=body.decode('utf-8', errors='replace'),
|
||||
body=await _read_error_body(resp),
|
||||
url=url,
|
||||
thread_id=thread_id,
|
||||
)
|
||||
|
||||
decoder = codecs.getincrementaldecoder('utf-8')('replace')
|
||||
buffer = ''
|
||||
total_bytes = 0
|
||||
|
||||
async for chunk in resp.aiter_bytes(8192):
|
||||
total_bytes += len(chunk)
|
||||
if total_bytes > SSE_MAX_TOTAL_BYTES:
|
||||
raise DeerFlowAPIError(
|
||||
operation='runs/stream response',
|
||||
body='response exceeds the runtime limit',
|
||||
url=url,
|
||||
thread_id=thread_id,
|
||||
)
|
||||
buffer += _normalize_sse_newlines(decoder.decode(chunk))
|
||||
|
||||
while '\n\n' in buffer:
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import httpx
|
||||
import typing
|
||||
import json
|
||||
@@ -8,6 +9,75 @@ from .errors import DifyAPIError
|
||||
from pathlib import Path
|
||||
import os
|
||||
|
||||
_MAX_DIFY_RESPONSE_BYTES = 1024 * 1024
|
||||
_MAX_DIFY_SSE_LINE_BYTES = 1024 * 1024
|
||||
_MAX_DIFY_STREAM_BYTES = 16 * 1024 * 1024
|
||||
_MAX_DIFY_UPLOAD_BYTES = 10 * 1024 * 1024
|
||||
|
||||
|
||||
async def _read_limited_response(
|
||||
response: httpx.Response,
|
||||
*,
|
||||
max_bytes: int = _MAX_DIFY_RESPONSE_BYTES,
|
||||
) -> bytes:
|
||||
content_length = response.headers.get('Content-Length')
|
||||
if content_length is not None:
|
||||
try:
|
||||
if int(content_length) > max_bytes:
|
||||
raise DifyAPIError(f'Remote response exceeds the {max_bytes}-byte limit')
|
||||
except (TypeError, ValueError):
|
||||
pass
|
||||
|
||||
body = bytearray()
|
||||
async for chunk in response.aiter_bytes(chunk_size=8192):
|
||||
body.extend(chunk)
|
||||
if len(body) > max_bytes:
|
||||
raise DifyAPIError(f'Remote response exceeds the {max_bytes}-byte limit')
|
||||
return bytes(body)
|
||||
|
||||
|
||||
async def _iter_sse_json(
|
||||
response: httpx.Response,
|
||||
) -> typing.AsyncGenerator[dict[str, typing.Any], None]:
|
||||
"""Parse Dify's one-JSON-per-data-line SSE without unbounded line buffering."""
|
||||
|
||||
buffer = bytearray()
|
||||
total = 0
|
||||
async for chunk in response.aiter_bytes(chunk_size=8192):
|
||||
total += len(chunk)
|
||||
if total > _MAX_DIFY_STREAM_BYTES:
|
||||
raise DifyAPIError('Dify SSE stream exceeds the runtime limit')
|
||||
buffer.extend(chunk)
|
||||
while b'\n' in buffer:
|
||||
raw_line, _, remainder = buffer.partition(b'\n')
|
||||
buffer = bytearray(remainder)
|
||||
if len(raw_line) > _MAX_DIFY_SSE_LINE_BYTES:
|
||||
raise DifyAPIError('Dify SSE event exceeds the runtime limit')
|
||||
line = raw_line.rstrip(b'\r').strip()
|
||||
if not line or not line.startswith(b'data:'):
|
||||
continue
|
||||
payload = json.loads(line[5:].decode('utf-8', errors='replace'))
|
||||
if isinstance(payload, dict):
|
||||
yield payload
|
||||
if len(buffer) > _MAX_DIFY_SSE_LINE_BYTES:
|
||||
raise DifyAPIError('Dify SSE event exceeds the runtime limit')
|
||||
|
||||
line = bytes(buffer).rstrip(b'\r').strip()
|
||||
if line.startswith(b'data:'):
|
||||
payload = json.loads(line[5:].decode('utf-8', errors='replace'))
|
||||
if isinstance(payload, dict):
|
||||
yield payload
|
||||
|
||||
|
||||
def _read_local_file_limited(path: Path) -> bytes:
|
||||
if path.stat().st_size > _MAX_DIFY_UPLOAD_BYTES:
|
||||
raise ValueError('Dify upload exceeds the size limit')
|
||||
with path.open('rb') as handle:
|
||||
body = handle.read(_MAX_DIFY_UPLOAD_BYTES + 1)
|
||||
if len(body) > _MAX_DIFY_UPLOAD_BYTES:
|
||||
raise ValueError('Dify upload exceeds the size limit')
|
||||
return body
|
||||
|
||||
|
||||
class AsyncDifyServiceClient:
|
||||
"""Dify Service API 客户端"""
|
||||
@@ -22,6 +92,21 @@ class AsyncDifyServiceClient:
|
||||
) -> None:
|
||||
self.api_key = api_key
|
||||
self.base_url = base_url
|
||||
self._client: httpx.AsyncClient | None = None
|
||||
|
||||
def _get_client(self) -> httpx.AsyncClient:
|
||||
if self._client is None:
|
||||
self._client = httpx.AsyncClient(
|
||||
base_url=self.base_url,
|
||||
trust_env=True,
|
||||
)
|
||||
return self._client
|
||||
|
||||
async def aclose(self) -> None:
|
||||
client = self._client
|
||||
self._client = None
|
||||
if client is not None:
|
||||
await client.aclose()
|
||||
|
||||
async def chat_messages(
|
||||
self,
|
||||
@@ -38,37 +123,32 @@ class AsyncDifyServiceClient:
|
||||
if response_mode != 'streaming':
|
||||
raise DifyAPIError('当前仅支持 streaming 模式')
|
||||
|
||||
async with httpx.AsyncClient(
|
||||
base_url=self.base_url,
|
||||
trust_env=True,
|
||||
timeout=timeout,
|
||||
) as client:
|
||||
payload = {
|
||||
'inputs': inputs,
|
||||
'query': query,
|
||||
'user': user,
|
||||
'response_mode': response_mode,
|
||||
'conversation_id': conversation_id,
|
||||
'files': files,
|
||||
'model_config': model_config or {},
|
||||
}
|
||||
client = self._get_client()
|
||||
payload = {
|
||||
'inputs': inputs,
|
||||
'query': query,
|
||||
'user': user,
|
||||
'response_mode': response_mode,
|
||||
'conversation_id': conversation_id,
|
||||
'files': files,
|
||||
'model_config': model_config or {},
|
||||
}
|
||||
|
||||
async with client.stream(
|
||||
'POST',
|
||||
'/chat-messages',
|
||||
headers={
|
||||
'Authorization': f'Bearer {self.api_key}',
|
||||
'Content-Type': 'application/json',
|
||||
},
|
||||
json=payload,
|
||||
) as r:
|
||||
async for chunk in r.aiter_lines():
|
||||
if r.status_code != 200:
|
||||
raise DifyAPIError(f'{r.status_code} {chunk}')
|
||||
if chunk.strip() == '':
|
||||
continue
|
||||
if chunk.startswith('data:'):
|
||||
yield json.loads(chunk[5:])
|
||||
async with client.stream(
|
||||
'POST',
|
||||
'/chat-messages',
|
||||
headers={
|
||||
'Authorization': f'Bearer {self.api_key}',
|
||||
'Content-Type': 'application/json',
|
||||
},
|
||||
json=payload,
|
||||
timeout=timeout,
|
||||
) as r:
|
||||
if r.status_code != 200:
|
||||
body = await _read_limited_response(r)
|
||||
raise DifyAPIError(f'{r.status_code} {body.decode(errors="replace")}')
|
||||
async for event in _iter_sse_json(r):
|
||||
yield event
|
||||
|
||||
async def workflow_run(
|
||||
self,
|
||||
@@ -82,32 +162,27 @@ class AsyncDifyServiceClient:
|
||||
if response_mode != 'streaming':
|
||||
raise DifyAPIError('当前仅支持 streaming 模式')
|
||||
|
||||
async with httpx.AsyncClient(
|
||||
base_url=self.base_url,
|
||||
trust_env=True,
|
||||
client = self._get_client()
|
||||
async with client.stream(
|
||||
'POST',
|
||||
'/workflows/run',
|
||||
headers={
|
||||
'Authorization': f'Bearer {self.api_key}',
|
||||
'Content-Type': 'application/json',
|
||||
},
|
||||
json={
|
||||
'inputs': inputs,
|
||||
'user': user,
|
||||
'response_mode': response_mode,
|
||||
'files': files,
|
||||
},
|
||||
timeout=timeout,
|
||||
) as client:
|
||||
async with client.stream(
|
||||
'POST',
|
||||
'/workflows/run',
|
||||
headers={
|
||||
'Authorization': f'Bearer {self.api_key}',
|
||||
'Content-Type': 'application/json',
|
||||
},
|
||||
json={
|
||||
'inputs': inputs,
|
||||
'user': user,
|
||||
'response_mode': response_mode,
|
||||
'files': files,
|
||||
},
|
||||
) as r:
|
||||
async for chunk in r.aiter_lines():
|
||||
if r.status_code != 200:
|
||||
raise DifyAPIError(f'{r.status_code} {chunk}')
|
||||
if chunk.strip() == '':
|
||||
continue
|
||||
if chunk.startswith('data:'):
|
||||
yield json.loads(chunk[5:])
|
||||
) as r:
|
||||
if r.status_code != 200:
|
||||
body = await _read_limited_response(r)
|
||||
raise DifyAPIError(f'{r.status_code} {body.decode(errors="replace")}')
|
||||
async for event in _iter_sse_json(r):
|
||||
yield event
|
||||
|
||||
async def workflow_submit(
|
||||
self,
|
||||
@@ -129,41 +204,38 @@ class AsyncDifyServiceClient:
|
||||
'Content-Type': 'application/json',
|
||||
}
|
||||
|
||||
async with httpx.AsyncClient(
|
||||
base_url=self.base_url,
|
||||
trust_env=True,
|
||||
client = self._get_client()
|
||||
# Step 1: Submit the form
|
||||
payload: dict[str, typing.Any] = {
|
||||
'inputs': inputs if isinstance(inputs, dict) else {},
|
||||
'user': user,
|
||||
'action': action,
|
||||
}
|
||||
|
||||
async with client.stream(
|
||||
'POST',
|
||||
f'/form/human_input/{form_token}',
|
||||
headers=headers,
|
||||
json=payload,
|
||||
timeout=timeout,
|
||||
) as client:
|
||||
# Step 1: Submit the form
|
||||
payload: dict[str, typing.Any] = {
|
||||
'inputs': inputs if isinstance(inputs, dict) else {},
|
||||
'user': user,
|
||||
'action': action,
|
||||
}
|
||||
|
||||
submit_resp = await client.post(
|
||||
f'/form/human_input/{form_token}',
|
||||
headers=headers,
|
||||
json=payload,
|
||||
)
|
||||
) as submit_resp:
|
||||
submit_body = await _read_limited_response(submit_resp)
|
||||
if submit_resp.status_code != 200:
|
||||
raise DifyAPIError(f'{submit_resp.status_code} {submit_resp.text}')
|
||||
raise DifyAPIError(f'{submit_resp.status_code} {submit_body.decode(errors="replace")}')
|
||||
|
||||
# Step 2: Stream resumed workflow events
|
||||
async with client.stream(
|
||||
'GET',
|
||||
f'/workflow/{workflow_run_id}/events',
|
||||
headers={'Authorization': f'Bearer {self.api_key}'},
|
||||
params={'user': user},
|
||||
) as r:
|
||||
if r.status_code != 200:
|
||||
body = (await r.aread()).decode(errors='replace')
|
||||
raise DifyAPIError(f'{r.status_code} {body}')
|
||||
async for chunk in r.aiter_lines():
|
||||
if chunk.strip() == '':
|
||||
continue
|
||||
if chunk.startswith('data:'):
|
||||
yield json.loads(chunk[5:])
|
||||
# Step 2: Stream resumed workflow events
|
||||
async with client.stream(
|
||||
'GET',
|
||||
f'/workflow/{workflow_run_id}/events',
|
||||
headers={'Authorization': f'Bearer {self.api_key}'},
|
||||
params={'user': user},
|
||||
timeout=timeout,
|
||||
) as r:
|
||||
if r.status_code != 200:
|
||||
body = await _read_limited_response(r)
|
||||
raise DifyAPIError(f'{r.status_code} {body.decode(errors="replace")}')
|
||||
async for event in _iter_sse_json(r):
|
||||
yield event
|
||||
|
||||
async def upload_file(
|
||||
self,
|
||||
@@ -175,37 +247,30 @@ class AsyncDifyServiceClient:
|
||||
if isinstance(file, Path):
|
||||
if not file.exists():
|
||||
raise ValueError(f'File not found: {file}')
|
||||
with open(file, 'rb') as f:
|
||||
file = f.read()
|
||||
file = await asyncio.to_thread(_read_local_file_limited, file)
|
||||
|
||||
# 处理文件路径字符串
|
||||
elif isinstance(file, str):
|
||||
if not os.path.isfile(file):
|
||||
raise ValueError(f'File not found: {file}')
|
||||
with open(file, 'rb') as f:
|
||||
file = f.read()
|
||||
file = await asyncio.to_thread(_read_local_file_limited, Path(file))
|
||||
|
||||
# 处理文件对象
|
||||
elif hasattr(file, 'read'):
|
||||
file = file.read()
|
||||
async with httpx.AsyncClient(
|
||||
base_url=self.base_url,
|
||||
trust_env=True,
|
||||
file = await asyncio.to_thread(file.read, _MAX_DIFY_UPLOAD_BYTES + 1)
|
||||
if len(file) > _MAX_DIFY_UPLOAD_BYTES:
|
||||
raise ValueError('Dify upload exceeds the size limit')
|
||||
client = self._get_client()
|
||||
# multipart/form-data
|
||||
async with client.stream(
|
||||
'POST',
|
||||
'/files/upload',
|
||||
headers={'Authorization': f'Bearer {self.api_key}'},
|
||||
files={'file': file},
|
||||
data={'user': user},
|
||||
timeout=timeout,
|
||||
) as client:
|
||||
# multipart/form-data
|
||||
response = await client.post(
|
||||
'/files/upload',
|
||||
headers={'Authorization': f'Bearer {self.api_key}'},
|
||||
files={
|
||||
'file': file,
|
||||
},
|
||||
data={
|
||||
'user': user,
|
||||
},
|
||||
)
|
||||
|
||||
) as response:
|
||||
body = await _read_limited_response(response)
|
||||
if response.status_code != 201:
|
||||
raise DifyAPIError(f'{response.status_code} {response.text}')
|
||||
|
||||
return response.json()
|
||||
raise DifyAPIError(f'{response.status_code} {body.decode(errors="replace")}')
|
||||
return json.loads(body)
|
||||
|
||||
@@ -7,6 +7,7 @@ import time
|
||||
import typing
|
||||
import uuid
|
||||
import urllib.parse
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Awaitable, Callable, Optional
|
||||
import dingtalk_stream # type: ignore
|
||||
import websockets
|
||||
@@ -15,12 +16,42 @@ from .card_callback import DingTalkCardActionHandler
|
||||
from .dingtalkevent import DingTalkEvent
|
||||
import httpx
|
||||
import traceback
|
||||
from langbot.pkg.utils import httpclient
|
||||
|
||||
|
||||
_stdout_logger = logging.getLogger('langbot.dingtalk_api')
|
||||
|
||||
|
||||
DINGTALK_OPENAPI_BASE = 'https://api.dingtalk.com'
|
||||
_MAX_MEDIA_BYTES = 10 * 1024 * 1024
|
||||
_MAX_GATEWAY_MESSAGE_BYTES = 1024 * 1024
|
||||
|
||||
|
||||
def _read_local_media_limited(file_path: str) -> bytes:
|
||||
if os.path.getsize(file_path) > _MAX_MEDIA_BYTES:
|
||||
raise ValueError('DingTalk media exceeds the size limit')
|
||||
with open(file_path, 'rb') as file:
|
||||
body = file.read(_MAX_MEDIA_BYTES + 1)
|
||||
if len(body) > _MAX_MEDIA_BYTES:
|
||||
raise ValueError('DingTalk media exceeds the size limit')
|
||||
return body
|
||||
|
||||
|
||||
async def _read_httpx_media_limited(response: httpx.Response) -> bytes:
|
||||
content_length = response.headers.get('Content-Length')
|
||||
if content_length is not None:
|
||||
try:
|
||||
if int(content_length) > _MAX_MEDIA_BYTES:
|
||||
raise ValueError('DingTalk media exceeds the size limit')
|
||||
except (TypeError, ValueError) as exc:
|
||||
if 'exceeds' in str(exc):
|
||||
raise
|
||||
body = bytearray()
|
||||
async for chunk in response.aiter_bytes():
|
||||
body.extend(chunk)
|
||||
if len(body) > _MAX_MEDIA_BYTES:
|
||||
raise ValueError('DingTalk media exceeds the size limit')
|
||||
return bytes(body)
|
||||
|
||||
|
||||
def _stringify_card_param_map(card_param_map: Optional[dict]) -> dict:
|
||||
@@ -44,6 +75,8 @@ def _stringify_card_param_map(card_param_map: Optional[dict]) -> dict:
|
||||
|
||||
|
||||
class DingTalkClient:
|
||||
_MAX_INBOUND_TASKS = 100
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
client_id: str,
|
||||
@@ -86,6 +119,37 @@ class DingTalkClient:
|
||||
self.legacy_access_token = ''
|
||||
self.legacy_access_token_expiry_time: typing.Optional[float] = None
|
||||
self._stopped = False # Flag to control the event loop
|
||||
self._inbound_tasks: set[asyncio.Task] = set()
|
||||
self._http_client: httpx.AsyncClient | None = None
|
||||
|
||||
@asynccontextmanager
|
||||
async def _http_client_context(self):
|
||||
"""Reuse one connection pool while preserving existing call structure."""
|
||||
|
||||
if self._http_client is None or self._http_client.is_closed:
|
||||
self._http_client = httpx.AsyncClient(event_hooks=httpclient.httpx_response_limit_hooks())
|
||||
yield self._http_client
|
||||
|
||||
def _start_inbound_task(self, coro: typing.Coroutine) -> bool:
|
||||
"""Start one bounded inbound callback task."""
|
||||
|
||||
for task in tuple(self._inbound_tasks):
|
||||
if task.done():
|
||||
self._inbound_tasks.discard(task)
|
||||
if len(self._inbound_tasks) >= self._MAX_INBOUND_TASKS:
|
||||
coro.close()
|
||||
return False
|
||||
|
||||
task = asyncio.create_task(coro)
|
||||
self._inbound_tasks.add(task)
|
||||
|
||||
def done(done_task: asyncio.Task) -> None:
|
||||
self._inbound_tasks.discard(done_task)
|
||||
if not done_task.cancelled():
|
||||
done_task.exception()
|
||||
|
||||
task.add_done_callback(done)
|
||||
return True
|
||||
|
||||
async def _on_card_action(self, payload: dict) -> None:
|
||||
"""Dispatch a parsed card-action payload to the adapter callback."""
|
||||
@@ -101,11 +165,11 @@ class DingTalkClient:
|
||||
url = 'https://api.dingtalk.com/v1.0/oauth2/accessToken'
|
||||
headers = {'Content-Type': 'application/json'}
|
||||
data = {'appKey': self.key, 'appSecret': self.secret}
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with self._http_client_context() as client:
|
||||
try:
|
||||
response = await client.post(url, json=data, headers=headers)
|
||||
if response.status_code == 200:
|
||||
response_data = response.json()
|
||||
response_data = await httpclient.parse_json_response(response)
|
||||
self.access_token = response_data.get('accessToken')
|
||||
expires_in = int(response_data.get('expireIn', 7200))
|
||||
self.access_token_expiry_time = time.time() + expires_in - 60
|
||||
@@ -129,28 +193,28 @@ class DingTalkClient:
|
||||
url = 'https://api.dingtalk.com/v1.0/robot/messageFiles/download'
|
||||
params = {'downloadCode': download_code, 'robotCode': self.robot_code}
|
||||
headers = {'x-acs-dingtalk-access-token': self.access_token}
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with self._http_client_context() as client:
|
||||
response = await client.post(url, headers=headers, json=params)
|
||||
if response.status_code == 200:
|
||||
result = response.json()
|
||||
result = await httpclient.parse_json_response(response)
|
||||
download_url = result.get('downloadUrl')
|
||||
else:
|
||||
await self.logger.error(f'failed to get download url: {response.json()}')
|
||||
error_payload = await httpclient.parse_json_response(response)
|
||||
await self.logger.error(f'failed to get download url: {error_payload}')
|
||||
|
||||
if download_url:
|
||||
return await self.download_url_to_base64(download_url)
|
||||
|
||||
async def download_url_to_base64(self, download_url):
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.get(download_url)
|
||||
|
||||
if response.status_code == 200:
|
||||
file_bytes = response.content
|
||||
mime_type = response.headers.get('Content-Type', 'application/octet-stream')
|
||||
base64_str = base64.b64encode(file_bytes).decode('utf-8')
|
||||
return f'data:{mime_type};base64,{base64_str}'
|
||||
else:
|
||||
await self.logger.error(f'failed to get files: {response.json()}')
|
||||
async with self._http_client_context() as client:
|
||||
async with client.stream('GET', download_url) as response:
|
||||
if response.status_code == 200:
|
||||
file_bytes = await _read_httpx_media_limited(response)
|
||||
mime_type = response.headers.get('Content-Type', 'application/octet-stream')
|
||||
base64_str = (await asyncio.to_thread(base64.b64encode, file_bytes)).decode('utf-8')
|
||||
return f'data:{mime_type};base64,{base64_str}'
|
||||
error_body = await _read_httpx_media_limited(response)
|
||||
await self.logger.error(f'failed to get files: {error_body[:300]!r}')
|
||||
|
||||
async def get_audio_url(self, download_code: str):
|
||||
if not await self.check_access_token():
|
||||
@@ -158,17 +222,19 @@ class DingTalkClient:
|
||||
url = 'https://api.dingtalk.com/v1.0/robot/messageFiles/download'
|
||||
params = {'downloadCode': download_code, 'robotCode': self.robot_code}
|
||||
headers = {'x-acs-dingtalk-access-token': self.access_token}
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with self._http_client_context() as client:
|
||||
response = await client.post(url, headers=headers, json=params)
|
||||
if response.status_code == 200:
|
||||
result = response.json()
|
||||
result = await httpclient.parse_json_response(response)
|
||||
download_url = result.get('downloadUrl')
|
||||
if download_url:
|
||||
return await self.download_url_to_base64(download_url)
|
||||
else:
|
||||
await self.logger.error(f'failed to get audio: {response.json()}')
|
||||
error_payload = await httpclient.parse_json_response(response)
|
||||
await self.logger.error(f'failed to get audio: {error_payload}')
|
||||
else:
|
||||
raise Exception(f'Error: {response.status_code}, {response.text}')
|
||||
body = await httpclient.response_text(response)
|
||||
raise Exception(f'Error: {response.status_code}, {body}')
|
||||
|
||||
async def get_file_url(self, download_code: str):
|
||||
if not await self.check_access_token():
|
||||
@@ -176,17 +242,19 @@ class DingTalkClient:
|
||||
url = 'https://api.dingtalk.com/v1.0/robot/messageFiles/download'
|
||||
params = {'downloadCode': download_code, 'robotCode': self.robot_code}
|
||||
headers = {'x-acs-dingtalk-access-token': self.access_token}
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with self._http_client_context() as client:
|
||||
response = await client.post(url, headers=headers, json=params)
|
||||
if response.status_code == 200:
|
||||
result = response.json()
|
||||
result = await httpclient.parse_json_response(response)
|
||||
download_url = result.get('downloadUrl')
|
||||
if download_url:
|
||||
return download_url
|
||||
else:
|
||||
await self.logger.error(f'failed to get file: {response.json()}')
|
||||
error_payload = await httpclient.parse_json_response(response)
|
||||
await self.logger.error(f'failed to get file: {error_payload}')
|
||||
else:
|
||||
raise Exception(f'Error: {response.status_code}, {response.text}')
|
||||
body = await httpclient.response_text(response)
|
||||
raise Exception(f'Error: {response.status_code}, {body}')
|
||||
|
||||
async def update_incoming_message(self, message):
|
||||
"""异步更新 DingTalkClient 中的 incoming_message"""
|
||||
@@ -503,12 +571,13 @@ class DingTalkClient:
|
||||
len(content),
|
||||
)
|
||||
try:
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with self._http_client_context() as client:
|
||||
response = await client.post(url, headers=headers, json=data)
|
||||
response_body = await httpclient.response_text(response, max_chars=500)
|
||||
_stdout_logger.info(
|
||||
'DingTalk send_proactive_message_to_one response: status=%d body=%s',
|
||||
response.status_code,
|
||||
response.text[:500],
|
||||
response_body,
|
||||
)
|
||||
if response.status_code == 200:
|
||||
return
|
||||
@@ -535,7 +604,7 @@ class DingTalkClient:
|
||||
'msgParam': json.dumps({'content': content}),
|
||||
}
|
||||
try:
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with self._http_client_context() as client:
|
||||
response = await client.post(url, headers=headers, json=data)
|
||||
if response.status_code == 200:
|
||||
return
|
||||
@@ -667,22 +736,23 @@ class DingTalkClient:
|
||||
'DingTalk createAndDeliver request body: %s',
|
||||
json.dumps(body, ensure_ascii=False)[:1500],
|
||||
)
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with self._http_client_context() as client:
|
||||
response = await client.post(url, headers=headers, json=body, timeout=30.0)
|
||||
response_body = await httpclient.response_text(response, max_chars=500)
|
||||
if response.status_code == 200:
|
||||
_stdout_logger.info(
|
||||
'DingTalk createAndDeliver response: %s',
|
||||
response.text[:500],
|
||||
response_body,
|
||||
)
|
||||
return True
|
||||
_stdout_logger.error(
|
||||
'DingTalk createAndDeliver failed: status=%s body=%s',
|
||||
response.status_code,
|
||||
response.text,
|
||||
response_body,
|
||||
)
|
||||
if self.logger:
|
||||
await self.logger.error(
|
||||
f'DingTalk createAndDeliver failed: status={response.status_code} body={response.text}'
|
||||
f'DingTalk createAndDeliver failed: status={response.status_code} body={response_body}'
|
||||
)
|
||||
return False
|
||||
except Exception:
|
||||
@@ -725,13 +795,14 @@ class DingTalkClient:
|
||||
'Content-Type': 'application/json',
|
||||
}
|
||||
try:
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with self._http_client_context() as client:
|
||||
response = await client.put(url, headers=headers, json=body, timeout=30.0)
|
||||
if response.status_code == 200:
|
||||
return True
|
||||
if self.logger:
|
||||
response_body = await httpclient.response_text(response)
|
||||
await self.logger.error(
|
||||
f'DingTalk card streaming failed: status={response.status_code} body={response.text}'
|
||||
f'DingTalk card streaming failed: status={response.status_code} body={response_body}'
|
||||
)
|
||||
return False
|
||||
except Exception:
|
||||
@@ -768,18 +839,19 @@ class DingTalkClient:
|
||||
out_track_id,
|
||||
json.dumps(body, ensure_ascii=False)[:1500],
|
||||
)
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with self._http_client_context() as client:
|
||||
response = await client.put(url, headers=headers, json=body, timeout=30.0)
|
||||
response_body = await httpclient.response_text(response, max_chars=300)
|
||||
_stdout_logger.info(
|
||||
'DingTalk update_card_data response: status=%d body=%s',
|
||||
response.status_code,
|
||||
response.text[:300],
|
||||
response_body,
|
||||
)
|
||||
if response.status_code == 200:
|
||||
return True
|
||||
if self.logger:
|
||||
await self.logger.error(
|
||||
f'DingTalk update card failed: status={response.status_code} body={response.text}'
|
||||
f'DingTalk update card failed: status={response.status_code} body={response_body}'
|
||||
)
|
||||
return False
|
||||
except Exception:
|
||||
@@ -808,17 +880,18 @@ class DingTalkClient:
|
||||
|
||||
url = 'https://oapi.dingtalk.com/gettoken'
|
||||
try:
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with self._http_client_context() as client:
|
||||
response = await client.get(url, params={'appkey': self.key, 'appsecret': self.secret}, timeout=15.0)
|
||||
data = response.json() if response.status_code == 200 else {}
|
||||
data = await httpclient.parse_json_response(response) if response.status_code == 200 else {}
|
||||
if data.get('errcode') == 0 and data.get('access_token'):
|
||||
self.legacy_access_token = data['access_token']
|
||||
expires_in = int(data.get('expires_in', 7200))
|
||||
self.legacy_access_token_expiry_time = now + expires_in - 60
|
||||
return self.legacy_access_token
|
||||
if self.logger:
|
||||
response_body = await httpclient.response_text(response, max_chars=200)
|
||||
await self.logger.error(
|
||||
f'DingTalk legacy gettoken failed: status={response.status_code} body={response.text[:200]}'
|
||||
f'DingTalk legacy gettoken failed: status={response.status_code} body={response_body}'
|
||||
)
|
||||
except Exception:
|
||||
_stdout_logger.exception('DingTalk legacy gettoken error')
|
||||
@@ -848,8 +921,7 @@ class DingTalkClient:
|
||||
|
||||
url = 'https://oapi.dingtalk.com/media/upload'
|
||||
try:
|
||||
with open(file_path, 'rb') as f:
|
||||
file_bytes = f.read()
|
||||
file_bytes = await asyncio.to_thread(_read_local_media_limited, file_path)
|
||||
file_name = os.path.basename(file_path)
|
||||
# Best-effort content-type guess; DingTalk accepts the major image
|
||||
# mime types and otherwise infers from the bytes.
|
||||
@@ -857,20 +929,21 @@ class DingTalkClient:
|
||||
mime = {'png': 'image/png', 'jpg': 'image/jpeg', 'jpeg': 'image/jpeg', 'gif': 'image/gif'}.get(
|
||||
ext, 'application/octet-stream'
|
||||
)
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with self._http_client_context() as client:
|
||||
response = await client.post(
|
||||
url,
|
||||
params={'access_token': token, 'type': 'image'},
|
||||
files={'media': (file_name, file_bytes, mime)},
|
||||
timeout=30.0,
|
||||
)
|
||||
data = response.json() if response.status_code == 200 else {}
|
||||
data = await httpclient.parse_json_response(response) if response.status_code == 200 else {}
|
||||
if data.get('errcode') == 0 and data.get('media_id'):
|
||||
_stdout_logger.info('DingTalk upload_image_media OK: media_id=%s', data['media_id'])
|
||||
return data['media_id']
|
||||
if self.logger:
|
||||
response_body = await httpclient.response_text(response, max_chars=300)
|
||||
await self.logger.error(
|
||||
f'DingTalk upload_image_media failed: status={response.status_code} body={response.text[:300]}'
|
||||
f'DingTalk upload_image_media failed: status={response.status_code} body={response_body}'
|
||||
)
|
||||
except Exception:
|
||||
_stdout_logger.exception('DingTalk upload_image_media error')
|
||||
@@ -897,15 +970,19 @@ class DingTalkClient:
|
||||
continue
|
||||
|
||||
uri = '%s?ticket=%s' % (connection['endpoint'], urllib.parse.quote_plus(connection['ticket']))
|
||||
async with websockets.connect(uri) as websocket:
|
||||
async with websockets.connect(uri, max_size=_MAX_GATEWAY_MESSAGE_BYTES) as websocket:
|
||||
self.client.websocket = websocket
|
||||
keepalive_task = asyncio.create_task(self._keepalive(websocket))
|
||||
try:
|
||||
async for raw_message in websocket:
|
||||
if self._stopped:
|
||||
break
|
||||
json_message = json.loads(raw_message)
|
||||
asyncio.create_task(self.client.background_task(json_message))
|
||||
json_message = await asyncio.to_thread(json.loads, raw_message)
|
||||
if not self._start_inbound_task(self.client.background_task(json_message)):
|
||||
if self.logger:
|
||||
await self.logger.warning(
|
||||
'DingTalk inbound task capacity reached; dropping message'
|
||||
)
|
||||
finally:
|
||||
keepalive_task.cancel()
|
||||
try:
|
||||
@@ -948,5 +1025,15 @@ class DingTalkClient:
|
||||
await self.client.websocket.close()
|
||||
except Exception:
|
||||
pass
|
||||
inbound_tasks = list(self._inbound_tasks)
|
||||
for task in inbound_tasks:
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
if inbound_tasks:
|
||||
await asyncio.gather(*inbound_tasks, return_exceptions=True)
|
||||
self._inbound_tasks.clear()
|
||||
# Clear message handlers to prevent stale callbacks
|
||||
self._message_handlers = {'example': []}
|
||||
if self._http_client is not None:
|
||||
await self._http_client.aclose()
|
||||
self._http_client = None
|
||||
|
||||
@@ -21,8 +21,14 @@ xml_template = """
|
||||
</xml>
|
||||
"""
|
||||
|
||||
_MAX_CALLBACK_BODY_BYTES = 1024 * 1024
|
||||
|
||||
|
||||
class OAClient:
|
||||
_STATE_TTL_SECONDS = 600
|
||||
_STATE_MAX = 4096
|
||||
_MAX_CONTENT_CHARS = 200000
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
token: str,
|
||||
@@ -41,6 +47,7 @@ class OAClient:
|
||||
self.access_token = ''
|
||||
self.unified_mode = unified_mode
|
||||
self.app = Quart(__name__)
|
||||
self.app.config['MAX_CONTENT_LENGTH'] = _MAX_CALLBACK_BODY_BYTES
|
||||
|
||||
# 只有在非统一模式下才注册独立路由
|
||||
if not self.unified_mode:
|
||||
@@ -57,8 +64,38 @@ class OAClient:
|
||||
self.access_token_expiry_time = None
|
||||
self.msg_id_map = {}
|
||||
self.generated_content = {}
|
||||
self._msg_seen_at = {}
|
||||
self._generated_at = {}
|
||||
self._last_state_prune = 0.0
|
||||
self.logger = logger
|
||||
|
||||
def _prune_state(self) -> None:
|
||||
now = time.monotonic()
|
||||
if now - self._last_state_prune >= 60:
|
||||
self._last_state_prune = now
|
||||
for message_id, seen_at in tuple(self._msg_seen_at.items()):
|
||||
if now - seen_at > self._STATE_TTL_SECONDS:
|
||||
self._msg_seen_at.pop(message_id, None)
|
||||
self.msg_id_map.pop(message_id, None)
|
||||
for message_id, generated_at in tuple(self._generated_at.items()):
|
||||
if now - generated_at > self._STATE_TTL_SECONDS:
|
||||
self._generated_at.pop(message_id, None)
|
||||
self.generated_content.pop(message_id, None)
|
||||
while len(self.msg_id_map) > self._STATE_MAX:
|
||||
message_id = next(iter(self.msg_id_map))
|
||||
self.msg_id_map.pop(message_id, None)
|
||||
self._msg_seen_at.pop(message_id, None)
|
||||
while len(self.generated_content) > self._STATE_MAX:
|
||||
message_id = next(iter(self.generated_content))
|
||||
self.generated_content.pop(message_id, None)
|
||||
self._generated_at.pop(message_id, None)
|
||||
|
||||
def clear(self) -> None:
|
||||
self.msg_id_map.clear()
|
||||
self.generated_content.clear()
|
||||
self._msg_seen_at.clear()
|
||||
self._generated_at.clear()
|
||||
|
||||
async def handle_callback_request(self):
|
||||
"""处理回调请求（独立端口模式，使用全局 request）。"""
|
||||
return await self._handle_callback_internal(request)
|
||||
@@ -104,8 +141,16 @@ class OAClient:
|
||||
raise Exception('拒绝请求')
|
||||
elif req.method == 'POST':
|
||||
encryt_msg = await req.data
|
||||
if len(encryt_msg) > _MAX_CALLBACK_BODY_BYTES:
|
||||
raise ValueError('Official Account callback body exceeds the size limit')
|
||||
wxcpt = WXBizMsgCrypt(self.token, self.aes, self.appid)
|
||||
ret, xml_msg = wxcpt.DecryptMsg(encryt_msg, msg_signature, timestamp, nonce)
|
||||
ret, xml_msg = await asyncio.to_thread(
|
||||
wxcpt.DecryptMsg,
|
||||
encryt_msg,
|
||||
msg_signature,
|
||||
timestamp,
|
||||
nonce,
|
||||
)
|
||||
xml_msg = xml_msg.decode('utf-8')
|
||||
|
||||
if ret != 0:
|
||||
@@ -118,7 +163,7 @@ class OAClient:
|
||||
if event:
|
||||
await self._handle_message(event)
|
||||
|
||||
root = ET.fromstring(xml_msg)
|
||||
root = await asyncio.to_thread(ET.fromstring, xml_msg)
|
||||
from_user = root.find('FromUserName').text # 发送者
|
||||
to_user = root.find('ToUserName').text # 机器人
|
||||
|
||||
@@ -126,6 +171,7 @@ class OAClient:
|
||||
interval = 0.1
|
||||
while True:
|
||||
content = self.generated_content.pop(message_data['MsgId'], None)
|
||||
self._generated_at.pop(message_data['MsgId'], None)
|
||||
if content:
|
||||
response_xml = xml_template.format(
|
||||
to_user=from_user,
|
||||
@@ -156,7 +202,7 @@ class OAClient:
|
||||
traceback.print_exc()
|
||||
|
||||
async def get_message(self, xml_msg: str):
|
||||
root = ET.fromstring(xml_msg)
|
||||
root = await asyncio.to_thread(ET.fromstring, xml_msg)
|
||||
|
||||
message_data = {
|
||||
'ToUserName': root.find('ToUserName').text,
|
||||
@@ -193,21 +239,30 @@ class OAClient:
|
||||
处理消息事件。
|
||||
"""
|
||||
message_id = event.message_id
|
||||
self._prune_state()
|
||||
if message_id in self.msg_id_map.keys():
|
||||
self.msg_id_map[message_id] += 1
|
||||
self._msg_seen_at[message_id] = time.monotonic()
|
||||
return
|
||||
|
||||
self.msg_id_map[message_id] = 1
|
||||
self._msg_seen_at[message_id] = time.monotonic()
|
||||
msg_type = event.type
|
||||
if msg_type in self._message_handlers:
|
||||
for handler in self._message_handlers[msg_type]:
|
||||
await handler(event)
|
||||
|
||||
async def set_message(self, msg_id: int, content: str):
|
||||
self.generated_content[msg_id] = content
|
||||
self.generated_content[msg_id] = str(content)[: self._MAX_CONTENT_CHARS]
|
||||
self._generated_at[msg_id] = time.monotonic()
|
||||
self._prune_state()
|
||||
|
||||
|
||||
class OAClientForLongerResponse:
|
||||
_MAX_USERS = 4096
|
||||
_MAX_MESSAGES_PER_USER = 20
|
||||
_MAX_CONTENT_CHARS = 200000
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
token: str,
|
||||
@@ -227,6 +282,7 @@ class OAClientForLongerResponse:
|
||||
self.access_token = ''
|
||||
self.unified_mode = unified_mode
|
||||
self.app = Quart(__name__)
|
||||
self.app.config['MAX_CONTENT_LENGTH'] = _MAX_CALLBACK_BODY_BYTES
|
||||
|
||||
# 只有在非统一模式下才注册独立路由
|
||||
if not self.unified_mode:
|
||||
@@ -244,8 +300,28 @@ class OAClientForLongerResponse:
|
||||
self.loading_message = LoadingMessage
|
||||
self.msg_queue = {}
|
||||
self.user_msg_queue = {}
|
||||
self._last_queue_cleanup = 0.0
|
||||
self.logger = logger
|
||||
|
||||
def _prune_queues(self) -> None:
|
||||
now = time.monotonic()
|
||||
if now - self._last_queue_cleanup >= 60:
|
||||
self._last_queue_cleanup = now
|
||||
for user_id, queue in tuple(self.msg_queue.items()):
|
||||
if not queue:
|
||||
self.msg_queue.pop(user_id, None)
|
||||
for user_id, queue in tuple(self.user_msg_queue.items()):
|
||||
if not queue:
|
||||
self.user_msg_queue.pop(user_id, None)
|
||||
while len(self.msg_queue) > self._MAX_USERS:
|
||||
self.msg_queue.pop(next(iter(self.msg_queue)), None)
|
||||
while len(self.user_msg_queue) > self._MAX_USERS:
|
||||
self.user_msg_queue.pop(next(iter(self.user_msg_queue)), None)
|
||||
|
||||
def clear(self) -> None:
|
||||
self.msg_queue.clear()
|
||||
self.user_msg_queue.clear()
|
||||
|
||||
async def handle_callback_request(self):
|
||||
"""处理回调请求（独立端口模式，使用全局 request）。"""
|
||||
return await self._handle_callback_internal(request)
|
||||
@@ -285,8 +361,16 @@ class OAClientForLongerResponse:
|
||||
|
||||
elif req.method == 'POST':
|
||||
encryt_msg = await req.data
|
||||
if len(encryt_msg) > _MAX_CALLBACK_BODY_BYTES:
|
||||
raise ValueError('Official Account callback body exceeds the size limit')
|
||||
wxcpt = WXBizMsgCrypt(self.token, self.aes, self.appid)
|
||||
ret, xml_msg = wxcpt.DecryptMsg(encryt_msg, msg_signature, timestamp, nonce)
|
||||
ret, xml_msg = await asyncio.to_thread(
|
||||
wxcpt.DecryptMsg,
|
||||
encryt_msg,
|
||||
msg_signature,
|
||||
timestamp,
|
||||
nonce,
|
||||
)
|
||||
xml_msg = xml_msg.decode('utf-8')
|
||||
|
||||
if ret != 0:
|
||||
@@ -294,7 +378,7 @@ class OAClientForLongerResponse:
|
||||
raise Exception('消息解密失败')
|
||||
|
||||
# 解析 XML
|
||||
root = ET.fromstring(xml_msg)
|
||||
root = await asyncio.to_thread(ET.fromstring, xml_msg)
|
||||
from_user = root.find('FromUserName').text
|
||||
to_user = root.find('ToUserName').text
|
||||
|
||||
@@ -305,6 +389,7 @@ class OAClientForLongerResponse:
|
||||
# 弹出用户消息
|
||||
if self.user_msg_queue.get(from_user) and self.user_msg_queue[from_user]:
|
||||
self.user_msg_queue[from_user].pop(0)
|
||||
self._prune_queues()
|
||||
|
||||
response_xml = xml_template.format(
|
||||
to_user=from_user,
|
||||
@@ -332,9 +417,13 @@ class OAClientForLongerResponse:
|
||||
if event:
|
||||
self.user_msg_queue.setdefault(from_user, []).append(
|
||||
{
|
||||
'content': event.message,
|
||||
'content': str(event.message)[: self._MAX_CONTENT_CHARS],
|
||||
}
|
||||
)
|
||||
self.user_msg_queue[from_user] = self.user_msg_queue[from_user][
|
||||
-self._MAX_MESSAGES_PER_USER :
|
||||
]
|
||||
self._prune_queues()
|
||||
await self._handle_message(event)
|
||||
|
||||
return response_xml
|
||||
@@ -344,7 +433,7 @@ class OAClientForLongerResponse:
|
||||
traceback.print_exc()
|
||||
|
||||
async def get_message(self, xml_msg: str):
|
||||
root = ET.fromstring(xml_msg)
|
||||
root = await asyncio.to_thread(ET.fromstring, xml_msg)
|
||||
|
||||
message_data = {
|
||||
'ToUserName': root.find('ToUserName').text,
|
||||
@@ -393,6 +482,8 @@ class OAClientForLongerResponse:
|
||||
self.msg_queue[from_user].append(
|
||||
{
|
||||
'msg_id': message_id,
|
||||
'content': content,
|
||||
'content': str(content)[: self._MAX_CONTENT_CHARS],
|
||||
}
|
||||
)
|
||||
self.msg_queue[from_user] = self.msg_queue[from_user][-self._MAX_MESSAGES_PER_USER :]
|
||||
self._prune_queues()
|
||||
|
||||
@@ -10,7 +10,9 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import hashlib
|
||||
import io
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import struct
|
||||
@@ -21,6 +23,8 @@ from urllib.parse import quote
|
||||
|
||||
import aiohttp
|
||||
|
||||
from langbot.pkg.utils import httpclient
|
||||
|
||||
from .types import (
|
||||
ApiError,
|
||||
CDNMedia,
|
||||
@@ -58,6 +62,51 @@ DEFAULT_BOT_TYPE = '3'
|
||||
|
||||
# Maximum text length per message chunk (WeChat limit)
|
||||
MAX_TEXT_CHUNK_SIZE = 2000
|
||||
MAX_CDN_MEDIA_BYTES = 10 * 1024 * 1024
|
||||
|
||||
|
||||
async def _response_text(response: aiohttp.ClientResponse) -> str:
|
||||
body = await httpclient.read_limited(
|
||||
response,
|
||||
max_bytes=MAX_CDN_MEDIA_BYTES,
|
||||
)
|
||||
return body.decode('utf-8', errors='replace')
|
||||
|
||||
|
||||
async def _response_json(response: aiohttp.ClientResponse) -> dict:
|
||||
payload = json.loads(await _response_text(response))
|
||||
if not isinstance(payload, dict):
|
||||
raise ApiError('OpenClaw API returned a non-object response', status=response.status)
|
||||
return payload
|
||||
|
||||
|
||||
def _decrypt_cdn_payload(encrypted: bytes, aes_key: bytes) -> bytes:
|
||||
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
|
||||
from cryptography.hazmat.primitives.padding import PKCS7
|
||||
|
||||
cipher = Cipher(algorithms.AES(aes_key), modes.ECB())
|
||||
decryptor = cipher.decryptor()
|
||||
padded = decryptor.update(encrypted) + decryptor.finalize()
|
||||
unpadder = PKCS7(128).unpadder()
|
||||
return unpadder.update(padded) + unpadder.finalize()
|
||||
|
||||
|
||||
def _encrypt_cdn_payload(
|
||||
file_bytes: bytes,
|
||||
) -> tuple[str, str, bytes, str]:
|
||||
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
|
||||
from cryptography.hazmat.primitives.padding import PKCS7
|
||||
|
||||
raw_key = os.urandom(16)
|
||||
aes_key_hex = raw_key.hex()
|
||||
encoded_key = base64.b64encode(aes_key_hex.encode('utf-8')).decode('utf-8')
|
||||
padder = PKCS7(128).padder()
|
||||
padded = padder.update(file_bytes) + padder.finalize()
|
||||
cipher = Cipher(algorithms.AES(raw_key), modes.ECB())
|
||||
encryptor = cipher.encryptor()
|
||||
encrypted = encryptor.update(padded) + encryptor.finalize()
|
||||
raw_md5 = hashlib.md5(file_bytes).hexdigest()
|
||||
return aes_key_hex, encoded_key, encrypted, raw_md5
|
||||
|
||||
|
||||
def _random_wechat_uin() -> str:
|
||||
@@ -125,12 +174,12 @@ class OpenClawWeixinClient:
|
||||
url, json=payload, headers=headers, timeout=aiohttp.ClientTimeout(total=timeout)
|
||||
) as resp:
|
||||
if resp.status != 200:
|
||||
text = await resp.text()
|
||||
text = await _response_text(resp)
|
||||
raise ApiError(
|
||||
f'OpenClaw API error {resp.status}: {text}',
|
||||
status=resp.status,
|
||||
)
|
||||
data = await resp.json(content_type=None)
|
||||
data = await _response_json(resp)
|
||||
|
||||
# Check for application-level errors in the response body
|
||||
errcode = data.get('errcode') or data.get('ret')
|
||||
@@ -170,12 +219,12 @@ class OpenClawWeixinClient:
|
||||
timeout=aiohttp.ClientTimeout(total=timeout),
|
||||
) as resp:
|
||||
if resp.status != 200:
|
||||
text = await resp.text()
|
||||
text = await _response_text(resp)
|
||||
raise ApiError(
|
||||
f'OpenClaw API error {resp.status}: {text}',
|
||||
status=resp.status,
|
||||
)
|
||||
data = await resp.json(content_type=None)
|
||||
data = await _response_json(resp)
|
||||
|
||||
except (asyncio.TimeoutError, aiohttp.ServerTimeoutError):
|
||||
return GetUpdatesResponse(ret=0, msgs=[], get_updates_buf=get_updates_buf)
|
||||
@@ -258,9 +307,6 @@ class OpenClawWeixinClient:
|
||||
Returns:
|
||||
Decrypted file bytes.
|
||||
"""
|
||||
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
|
||||
from cryptography.hazmat.primitives.padding import PKCS7
|
||||
|
||||
if not media.encrypt_query_param:
|
||||
raise ApiError('CDN media has no encrypt_query_param', status=0)
|
||||
if not media.aes_key:
|
||||
@@ -285,17 +331,14 @@ class OpenClawWeixinClient:
|
||||
|
||||
async with session.get(cdn_url, timeout=aiohttp.ClientTimeout(total=120)) as resp:
|
||||
if resp.status != 200:
|
||||
text = await resp.text()
|
||||
text = await _response_text(resp)
|
||||
raise ApiError(f'CDN download failed: {resp.status} {text}', status=resp.status)
|
||||
encrypted = await resp.read()
|
||||
encrypted = await httpclient.read_limited(
|
||||
resp,
|
||||
max_bytes=MAX_CDN_MEDIA_BYTES,
|
||||
)
|
||||
|
||||
# Decrypt AES-128-ECB with PKCS7 padding
|
||||
cipher = Cipher(algorithms.AES(aes_key), modes.ECB())
|
||||
decryptor = cipher.decryptor()
|
||||
padded = decryptor.update(encrypted) + decryptor.finalize()
|
||||
|
||||
unpadder = PKCS7(128).unpadder()
|
||||
return unpadder.update(padded) + unpadder.finalize()
|
||||
return await asyncio.to_thread(_decrypt_cdn_payload, encrypted, aes_key)
|
||||
|
||||
async def upload_media(
|
||||
self,
|
||||
@@ -313,28 +356,13 @@ class OpenClawWeixinClient:
|
||||
Returns:
|
||||
CDNMedia with encrypt_query_param and aes_key for use in sendMessage.
|
||||
"""
|
||||
import hashlib
|
||||
if len(file_bytes) > MAX_CDN_MEDIA_BYTES:
|
||||
raise ApiError('CDN media exceeds the size limit', status=0)
|
||||
|
||||
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
|
||||
from cryptography.hazmat.primitives.padding import PKCS7
|
||||
|
||||
# 1. Generate random 16-byte AES key
|
||||
raw_key = os.urandom(16)
|
||||
aes_key_hex = raw_key.hex() # 32-char hex string
|
||||
|
||||
# 2. Encode key for CDNMedia: base64(hex_string) — same for all media types
|
||||
# Matches official SDK: Buffer.from(aeskey_hex).toString("base64")
|
||||
encoded_key = base64.b64encode(aes_key_hex.encode('utf-8')).decode('utf-8')
|
||||
|
||||
# 3. Encrypt file with AES-128-ECB + PKCS7
|
||||
padder = PKCS7(128).padder()
|
||||
padded = padder.update(file_bytes) + padder.finalize()
|
||||
cipher = Cipher(algorithms.AES(raw_key), modes.ECB())
|
||||
encryptor = cipher.encryptor()
|
||||
encrypted = encryptor.update(padded) + encryptor.finalize()
|
||||
|
||||
# 4. Get upload URL
|
||||
raw_md5 = hashlib.md5(file_bytes).hexdigest()
|
||||
aes_key_hex, encoded_key, encrypted, raw_md5 = await asyncio.to_thread(
|
||||
_encrypt_cdn_payload,
|
||||
file_bytes,
|
||||
)
|
||||
filekey = os.urandom(16).hex() # 32-char hex, matches official SDK
|
||||
|
||||
upload_resp = await self.get_upload_url(
|
||||
@@ -370,7 +398,7 @@ class OpenClawWeixinClient:
|
||||
timeout=aiohttp.ClientTimeout(total=120),
|
||||
) as resp:
|
||||
if resp.status != 200:
|
||||
text = await resp.text()
|
||||
text = await _response_text(resp)
|
||||
logger.error('CDN upload failed: status=%d url=%s body=%s', resp.status, cdn_url, text[:500])
|
||||
raise ApiError(f'CDN upload failed: {resp.status} {text}', status=resp.status)
|
||||
download_param = resp.headers.get('x-encrypted-param', '')
|
||||
@@ -491,12 +519,12 @@ class OpenClawWeixinClient:
|
||||
|
||||
async with session.get(url, timeout=aiohttp.ClientTimeout(total=DEFAULT_API_TIMEOUT)) as resp:
|
||||
if resp.status != 200:
|
||||
text = await resp.text()
|
||||
text = await _response_text(resp)
|
||||
raise ApiError(
|
||||
f'Failed to fetch QR code: {resp.status} {text}',
|
||||
status=resp.status,
|
||||
)
|
||||
data = await resp.json(content_type=None)
|
||||
data = await _response_json(resp)
|
||||
|
||||
logger.debug(
|
||||
'fetch_qrcode response: qrcode=%s, img=%s', data.get('qrcode'), bool(data.get('qrcode_img_content'))
|
||||
@@ -536,12 +564,12 @@ class OpenClawWeixinClient:
|
||||
url, headers=headers, timeout=aiohttp.ClientTimeout(total=DEFAULT_QR_POLL_TIMEOUT)
|
||||
) as resp:
|
||||
if resp.status != 200:
|
||||
text = await resp.text()
|
||||
text = await _response_text(resp)
|
||||
raise ApiError(
|
||||
f'Failed to poll QR status: {resp.status} {text}',
|
||||
status=resp.status,
|
||||
)
|
||||
data = await resp.json(content_type=None)
|
||||
data = await _response_json(resp)
|
||||
logger.debug('QR status poll response: %s', data)
|
||||
except (asyncio.TimeoutError, aiohttp.ServerTimeoutError):
|
||||
return QRStatusResponse(status='wait')
|
||||
|
||||
@@ -9,10 +9,13 @@ import langbot_plugin.api.entities.builtin.platform.events as platform_events
|
||||
from .qqofficialevent import QQOfficialEvent
|
||||
import json
|
||||
import traceback
|
||||
from contextlib import asynccontextmanager
|
||||
from cryptography.hazmat.primitives.asymmetric import ed25519
|
||||
from langbot.pkg.utils import httpclient
|
||||
|
||||
|
||||
QQ_SELECT_ACTION_PREFIX = '__langbot_select__:'
|
||||
_MAX_CALLBACK_BODY_BYTES = 1024 * 1024
|
||||
|
||||
|
||||
def get_select_field_options(form_data: dict) -> tuple[str, list[str]]:
|
||||
@@ -152,6 +155,7 @@ class QQOfficialClient:
|
||||
def __init__(self, secret: str, token: str, app_id: str, logger: None, unified_mode: bool = False):
|
||||
self.unified_mode = unified_mode
|
||||
self.app = Quart(__name__)
|
||||
self.app.config['MAX_CONTENT_LENGTH'] = _MAX_CALLBACK_BODY_BYTES
|
||||
|
||||
# 只有在非统一模式下才注册独立路由
|
||||
if not self.unified_mode:
|
||||
@@ -176,6 +180,32 @@ class QQOfficialClient:
|
||||
self.logger = logger
|
||||
self._msg_seq_counter = 0
|
||||
self._token_refresh_task: Optional[asyncio.Task] = None
|
||||
self._http_clients: dict[float | None, httpx.AsyncClient] = {}
|
||||
|
||||
@asynccontextmanager
|
||||
async def _http_client_context(self, timeout: float | None = None):
|
||||
client = self._http_clients.get(timeout)
|
||||
if client is None or client.is_closed:
|
||||
response_hooks = httpclient.httpx_response_limit_hooks()
|
||||
client = (
|
||||
httpx.AsyncClient(event_hooks=response_hooks)
|
||||
if timeout is None
|
||||
else httpx.AsyncClient(timeout=timeout, event_hooks=response_hooks)
|
||||
)
|
||||
self._http_clients[timeout] = client
|
||||
yield client
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Stop client-owned background work."""
|
||||
|
||||
if self._token_refresh_task and not self._token_refresh_task.done():
|
||||
self._token_refresh_task.cancel()
|
||||
await asyncio.gather(self._token_refresh_task, return_exceptions=True)
|
||||
self._token_refresh_task = None
|
||||
clients = list(self._http_clients.values())
|
||||
self._http_clients.clear()
|
||||
if clients:
|
||||
await asyncio.gather(*(client.aclose() for client in clients), return_exceptions=True)
|
||||
|
||||
async def check_access_token(self):
|
||||
"""检查access_token是否存在"""
|
||||
@@ -186,7 +216,7 @@ class QQOfficialClient:
|
||||
async def get_access_token(self):
|
||||
"""获取access_token"""
|
||||
url = 'https://bots.qq.com/app/getAppAccessToken'
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with self._http_client_context() as client:
|
||||
params = {
|
||||
'appId': self.app_id,
|
||||
'clientSecret': self.secret,
|
||||
@@ -196,8 +226,9 @@ class QQOfficialClient:
|
||||
}
|
||||
response = await client.post(url, json=params, headers=headers)
|
||||
if response.status_code != 200:
|
||||
raise Exception(f'Failed to get access_token: HTTP {response.status_code} {response.text}')
|
||||
response_data = response.json()
|
||||
body = await httpclient.response_text(response)
|
||||
raise Exception(f'Failed to get access_token: HTTP {response.status_code} {body}')
|
||||
response_data = await httpclient.parse_json_response(response)
|
||||
access_token = response_data.get('access_token')
|
||||
expires_in = int(response_data.get('expires_in', 7200))
|
||||
self.access_token_expiry_time = time.time() + expires_in - 60
|
||||
@@ -236,8 +267,10 @@ class QQOfficialClient:
|
||||
if not body or len(body) == 0:
|
||||
await self.logger.info('Received empty body, might be health check or GET request')
|
||||
return {'code': 0, 'message': 'ok'}, 200
|
||||
if len(body) > _MAX_CALLBACK_BODY_BYTES:
|
||||
return {'error': 'callback body exceeds the size limit'}, 413
|
||||
|
||||
payload = json.loads(body)
|
||||
payload = await asyncio.to_thread(json.loads, body)
|
||||
|
||||
if payload.get('op') == 13:
|
||||
validation_data = payload.get('d')
|
||||
@@ -367,7 +400,7 @@ class QQOfficialClient:
|
||||
await self.get_access_token()
|
||||
|
||||
url = self.base_url + '/v2/users/' + user_openid + '/messages'
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with self._http_client_context() as client:
|
||||
headers = {
|
||||
'Authorization': f'QQBot {self.access_token}',
|
||||
'Content-Type': 'application/json',
|
||||
@@ -382,7 +415,7 @@ class QQOfficialClient:
|
||||
if event_id:
|
||||
data['event_id'] = event_id
|
||||
response = await client.post(url, headers=headers, json=data)
|
||||
response_data = response.json()
|
||||
response_data = await httpclient.parse_json_response(response)
|
||||
if response.status_code == 200:
|
||||
return
|
||||
else:
|
||||
@@ -406,7 +439,7 @@ class QQOfficialClient:
|
||||
await self.get_access_token()
|
||||
|
||||
url = self.base_url + '/v2/groups/' + group_openid + '/messages'
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with self._http_client_context() as client:
|
||||
headers = {
|
||||
'Authorization': f'QQBot {self.access_token}',
|
||||
'Content-Type': 'application/json',
|
||||
@@ -424,8 +457,9 @@ class QQOfficialClient:
|
||||
if response.status_code == 200:
|
||||
return
|
||||
else:
|
||||
await self.logger.error(f'Failed to send group message: {response.json()}')
|
||||
raise Exception(response.read().decode())
|
||||
error_payload = await httpclient.parse_json_response(response)
|
||||
await self.logger.error(f'Failed to send group message: {error_payload}')
|
||||
raise Exception(str(error_payload))
|
||||
|
||||
async def send_channle_group_text_msg(self, channel_id: str, content: str, msg_id: str):
|
||||
"""发送频道群聊消息"""
|
||||
@@ -433,7 +467,7 @@ class QQOfficialClient:
|
||||
await self.get_access_token()
|
||||
|
||||
url = self.base_url + '/channels/' + channel_id + '/messages'
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with self._http_client_context() as client:
|
||||
headers = {
|
||||
'Authorization': f'QQBot {self.access_token}',
|
||||
'Content-Type': 'application/json',
|
||||
@@ -447,7 +481,8 @@ class QQOfficialClient:
|
||||
if response.status_code == 200:
|
||||
return True
|
||||
else:
|
||||
await self.logger.error(f'Failed to send channel group message: {response.json()}')
|
||||
error_payload = await httpclient.parse_json_response(response)
|
||||
await self.logger.error(f'Failed to send channel group message: {error_payload}')
|
||||
raise Exception(response)
|
||||
|
||||
async def send_channle_private_text_msg(self, guild_id: str, content: str, msg_id: str):
|
||||
@@ -456,7 +491,7 @@ class QQOfficialClient:
|
||||
await self.get_access_token()
|
||||
|
||||
url = self.base_url + '/dms/' + guild_id + '/messages'
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with self._http_client_context() as client:
|
||||
headers = {
|
||||
'Authorization': f'QQBot {self.access_token}',
|
||||
'Content-Type': 'application/json',
|
||||
@@ -470,7 +505,8 @@ class QQOfficialClient:
|
||||
if response.status_code == 200:
|
||||
return True
|
||||
else:
|
||||
await self.logger.error(f'Failed to send channel private message: {response.json()}')
|
||||
error_payload = await httpclient.parse_json_response(response)
|
||||
await self.logger.error(f'Failed to send channel private message: {error_payload}')
|
||||
raise Exception(response)
|
||||
|
||||
# ---- 富媒体消息 ----
|
||||
@@ -532,20 +568,21 @@ class QQOfficialClient:
|
||||
if file_type == self.MEDIA_TYPE_FILE and file_name:
|
||||
body['file_name'] = file_name
|
||||
|
||||
async with httpx.AsyncClient(timeout=120) as client:
|
||||
async with self._http_client_context(timeout=120) as client:
|
||||
headers = {
|
||||
'Authorization': f'QQBot {self.access_token}',
|
||||
'Content-Type': 'application/json',
|
||||
}
|
||||
response = await client.post(url, headers=headers, json=body)
|
||||
if response.status_code == 200:
|
||||
data = response.json()
|
||||
data = await httpclient.parse_json_response(response)
|
||||
file_info = data.get('file_info', '')
|
||||
preview = file_info[:80] + '...' if len(file_info) > 80 else file_info
|
||||
await self.logger.info(f'Upload media success, file_info={preview}')
|
||||
return file_info
|
||||
else:
|
||||
raise Exception(f'Failed to upload media: HTTP {response.status_code} {response.text}')
|
||||
body = await httpclient.response_text(response)
|
||||
raise Exception(f'Failed to upload media: HTTP {response.status_code} {body}')
|
||||
|
||||
async def _send_media_msg(
|
||||
self,
|
||||
@@ -578,7 +615,7 @@ class QQOfficialClient:
|
||||
if msg_id:
|
||||
body['msg_id'] = msg_id
|
||||
|
||||
async with httpx.AsyncClient(timeout=120) as client:
|
||||
async with self._http_client_context(timeout=120) as client:
|
||||
headers = {
|
||||
'Authorization': f'QQBot {self.access_token}',
|
||||
'Content-Type': 'application/json',
|
||||
@@ -586,7 +623,8 @@ class QQOfficialClient:
|
||||
await self.logger.info(f'Sending rich media: {json.dumps(body, ensure_ascii=False)[:200]}')
|
||||
response = await client.post(url, headers=headers, json=body)
|
||||
if response.status_code != 200:
|
||||
raise Exception(f'Failed to send rich media message: HTTP {response.status_code} {response.text}')
|
||||
response_body = await httpclient.response_text(response)
|
||||
raise Exception(f'Failed to send rich media message: HTTP {response.status_code} {response_body}')
|
||||
|
||||
async def send_image_msg(
|
||||
self,
|
||||
@@ -678,15 +716,16 @@ class QQOfficialClient:
|
||||
if stream_msg_id:
|
||||
body['stream_msg_id'] = stream_msg_id
|
||||
|
||||
async with httpx.AsyncClient(timeout=120) as client:
|
||||
async with self._http_client_context(timeout=120) as client:
|
||||
headers = {
|
||||
'Authorization': f'QQBot {self.access_token}',
|
||||
'Content-Type': 'application/json',
|
||||
}
|
||||
response = await client.post(url, headers=headers, json=body)
|
||||
if response.status_code != 200:
|
||||
raise Exception(f'Failed to send stream message: HTTP {response.status_code} {response.text}')
|
||||
return response.json()
|
||||
response_body = await httpclient.response_text(response)
|
||||
raise Exception(f'Failed to send stream message: HTTP {response.status_code} {response_body}')
|
||||
return await httpclient.parse_json_response(response)
|
||||
|
||||
async def send_markdown_keyboard(
|
||||
self,
|
||||
@@ -743,18 +782,19 @@ class QQOfficialClient:
|
||||
if event_id:
|
||||
body['event_id'] = event_id
|
||||
|
||||
async with httpx.AsyncClient(timeout=30) as client:
|
||||
async with self._http_client_context(timeout=30) as client:
|
||||
headers = {
|
||||
'Authorization': f'QQBot {self.access_token}',
|
||||
'Content-Type': 'application/json',
|
||||
}
|
||||
response = await client.post(url, headers=headers, json=body)
|
||||
if response.status_code != 200:
|
||||
response_body = await httpclient.response_text(response)
|
||||
await self.logger.error(
|
||||
f'Failed to send markdown+keyboard: HTTP {response.status_code} {response.text}'
|
||||
f'Failed to send markdown+keyboard: HTTP {response.status_code} {response_body}'
|
||||
)
|
||||
raise Exception(f'Failed to send markdown+keyboard: HTTP {response.status_code} {response.text}')
|
||||
return response.json()
|
||||
raise Exception(f'Failed to send markdown+keyboard: HTTP {response.status_code} {response_body}')
|
||||
return await httpclient.parse_json_response(response)
|
||||
|
||||
async def ack_interaction(self, interaction_id: str, code: int = 0) -> None:
|
||||
"""Acknowledge a button-click INTERACTION_CREATE event.
|
||||
@@ -775,7 +815,7 @@ class QQOfficialClient:
|
||||
await self.get_access_token()
|
||||
|
||||
url = f'{self.base_url}/interactions/{interaction_id}'
|
||||
async with httpx.AsyncClient(timeout=10) as client:
|
||||
async with self._http_client_context(timeout=10) as client:
|
||||
headers = {
|
||||
'Authorization': f'QQBot {self.access_token}',
|
||||
'Content-Type': 'application/json',
|
||||
@@ -783,8 +823,9 @@ class QQOfficialClient:
|
||||
try:
|
||||
response = await client.put(url, headers=headers, json={'code': code})
|
||||
if response.status_code >= 400:
|
||||
response_body = await httpclient.response_text(response)
|
||||
await self.logger.warning(
|
||||
f'ack_interaction non-success: HTTP {response.status_code} {response.text}'
|
||||
f'ack_interaction non-success: HTTP {response.status_code} {response_body}'
|
||||
)
|
||||
except Exception as e:
|
||||
await self.logger.warning(f'ack_interaction error (non-fatal): {e}')
|
||||
@@ -796,10 +837,11 @@ class QQOfficialClient:
|
||||
return time.time() > self.access_token_expiry_time
|
||||
|
||||
async def repeat_seed(self, bot_secret: str, target_size: int = 32) -> bytes:
|
||||
seed = bot_secret
|
||||
while len(seed) < target_size:
|
||||
seed *= 2
|
||||
return seed[:target_size].encode('utf-8')
|
||||
if not bot_secret:
|
||||
raise ValueError('QQ bot secret must not be empty')
|
||||
target_size = max(int(target_size), 1)
|
||||
repeats = (target_size + len(bot_secret) - 1) // len(bot_secret)
|
||||
return (bot_secret * repeats)[:target_size].encode('utf-8')
|
||||
|
||||
async def verify(self, validation_payload: dict):
|
||||
seed = await self.repeat_seed(self.secret)
|
||||
@@ -843,19 +885,20 @@ class QQOfficialClient:
|
||||
await self.get_access_token()
|
||||
|
||||
url = f'{self.base_url}/gateway'
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with self._http_client_context() as client:
|
||||
headers = {
|
||||
'Authorization': f'QQBot {self.access_token}',
|
||||
}
|
||||
response = await client.get(url, headers=headers)
|
||||
if response.status_code == 200:
|
||||
data = response.json()
|
||||
data = await httpclient.parse_json_response(response)
|
||||
ws_url = data.get('url', '')
|
||||
if not ws_url:
|
||||
raise Exception('Gateway URL is empty')
|
||||
return ws_url
|
||||
else:
|
||||
raise Exception(f'Failed to get Gateway URL: HTTP {response.status_code} {response.text}')
|
||||
body = await httpclient.response_text(response)
|
||||
raise Exception(f'Failed to get Gateway URL: HTTP {response.status_code} {body}')
|
||||
|
||||
async def _background_token_refresh(self):
|
||||
"""在 token 到期前主动刷新"""
|
||||
@@ -935,7 +978,7 @@ class QQOfficialClient:
|
||||
|
||||
try:
|
||||
await self.logger.info('Connecting to WebSocket gateway...')
|
||||
ws = await websockets.connect(ws_url)
|
||||
ws = await websockets.connect(ws_url, max_size=_MAX_CALLBACK_BODY_BYTES)
|
||||
await self.logger.info('WebSocket connected')
|
||||
except Exception as e:
|
||||
await self.logger.error(f'WebSocket connection failed: {e}')
|
||||
@@ -948,7 +991,7 @@ class QQOfficialClient:
|
||||
try:
|
||||
async for raw_msg in ws:
|
||||
try:
|
||||
payload = json.loads(raw_msg)
|
||||
payload = await asyncio.to_thread(json.loads, raw_msg)
|
||||
except json.JSONDecodeError:
|
||||
await self.logger.error(f'Failed to parse message: {raw_msg}')
|
||||
continue
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
import asyncio
|
||||
import json
|
||||
import traceback
|
||||
from quart import Quart, jsonify, request
|
||||
@@ -6,6 +7,8 @@ from .slackevent import SlackEvent
|
||||
from typing import Callable
|
||||
import langbot_plugin.api.entities.builtin.platform.events as platform_events
|
||||
|
||||
_MAX_CALLBACK_BODY_BYTES = 1024 * 1024
|
||||
|
||||
|
||||
class SlackClient:
|
||||
def __init__(self, bot_token: str, signing_secret: str, logger: None, unified_mode: bool = False):
|
||||
@@ -13,6 +16,7 @@ class SlackClient:
|
||||
self.signing_secret = signing_secret
|
||||
self.unified_mode = unified_mode
|
||||
self.app = Quart(__name__)
|
||||
self.app.config['MAX_CONTENT_LENGTH'] = _MAX_CALLBACK_BODY_BYTES
|
||||
self.client = AsyncWebClient(self.bot_token)
|
||||
|
||||
# 只有在非统一模式下才注册独立路由
|
||||
@@ -50,7 +54,9 @@ class SlackClient:
|
||||
"""
|
||||
try:
|
||||
body = await req.get_data()
|
||||
data = json.loads(body)
|
||||
if len(body) > _MAX_CALLBACK_BODY_BYTES:
|
||||
raise ValueError('Slack callback body exceeds the size limit')
|
||||
data = await asyncio.to_thread(json.loads, body)
|
||||
if 'type' in data:
|
||||
if data['type'] == 'url_verification':
|
||||
return data['challenge']
|
||||
|
||||
@@ -1,7 +1,32 @@
|
||||
from langbot.libs.wechatpad_api.util.http_util import post_json
|
||||
import httpx
|
||||
import asyncio
|
||||
import base64
|
||||
|
||||
import httpx
|
||||
|
||||
from langbot.libs.wechatpad_api.util.http_util import post_json
|
||||
from langbot.pkg.utils import httpclient
|
||||
|
||||
|
||||
_MAX_WECHATPAD_MEDIA_BYTES = 16 * 1024 * 1024
|
||||
|
||||
|
||||
async def _read_media_limited(response: httpx.Response) -> bytes:
|
||||
content_length = response.headers.get('content-length')
|
||||
if content_length is not None:
|
||||
try:
|
||||
declared_size = int(content_length)
|
||||
except ValueError:
|
||||
declared_size = None
|
||||
if declared_size is not None and declared_size > _MAX_WECHATPAD_MEDIA_BYTES:
|
||||
raise RuntimeError('WeChatPad media exceeds the runtime limit')
|
||||
|
||||
body = bytearray()
|
||||
async for chunk in response.aiter_bytes(chunk_size=64 * 1024):
|
||||
body.extend(chunk)
|
||||
if len(body) > _MAX_WECHATPAD_MEDIA_BYTES:
|
||||
raise RuntimeError('WeChatPad media exceeds the runtime limit')
|
||||
return bytes(body)
|
||||
|
||||
|
||||
class DownloadApi:
|
||||
def __init__(self, base_url, token):
|
||||
@@ -19,12 +44,13 @@ class DownloadApi:
|
||||
return post_json(url, token=self.token, data=json_data)
|
||||
|
||||
async def download_url_to_base64(self, download_url):
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.get(download_url)
|
||||
|
||||
if response.status_code == 200:
|
||||
file_bytes = response.content
|
||||
base64_str = base64.b64encode(file_bytes).decode('utf-8') # 返回字符串格式
|
||||
return base64_str
|
||||
else:
|
||||
raise Exception('获取文件失败')
|
||||
async with httpx.AsyncClient(
|
||||
timeout=30,
|
||||
event_hooks=httpclient.httpx_response_limit_hooks(_MAX_WECHATPAD_MEDIA_BYTES),
|
||||
) as client:
|
||||
async with client.stream('GET', download_url) as response:
|
||||
if response.status_code != 200:
|
||||
raise RuntimeError('获取文件失败')
|
||||
file_bytes = await _read_media_limited(response)
|
||||
encoded = await asyncio.to_thread(base64.b64encode, file_bytes)
|
||||
return encoded.decode('utf-8')
|
||||
|
||||
@@ -1,6 +1,29 @@
|
||||
import json as json_module
|
||||
|
||||
import requests
|
||||
from langbot.pkg.utils import httpclient
|
||||
|
||||
_MAX_WECHATPAD_RESPONSE_BYTES = 16 * 1024 * 1024
|
||||
|
||||
|
||||
def _read_requests_response_limited(response: requests.Response) -> dict:
|
||||
content_length = response.headers.get('Content-Length')
|
||||
if content_length is not None:
|
||||
try:
|
||||
if int(content_length) > _MAX_WECHATPAD_RESPONSE_BYTES:
|
||||
raise RuntimeError('WeChatPad response exceeds the runtime limit')
|
||||
except (TypeError, ValueError):
|
||||
pass
|
||||
body = bytearray()
|
||||
for chunk in response.iter_content(chunk_size=64 * 1024):
|
||||
body.extend(chunk)
|
||||
if len(body) > _MAX_WECHATPAD_RESPONSE_BYTES:
|
||||
raise RuntimeError('WeChatPad response exceeds the runtime limit')
|
||||
result = json_module.loads(body)
|
||||
if not isinstance(result, dict):
|
||||
raise RuntimeError('WeChatPad returned a non-object response')
|
||||
return result
|
||||
|
||||
|
||||
def post_json(base_url, token, data=None):
|
||||
headers = {'Content-Type': 'application/json'}
|
||||
@@ -8,16 +31,21 @@ def post_json(base_url, token, data=None):
|
||||
url = base_url + f'?key={token}'
|
||||
|
||||
try:
|
||||
response = requests.post(url, json=data, headers=headers, timeout=60)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
with requests.post(
|
||||
url,
|
||||
json=data,
|
||||
headers=headers,
|
||||
timeout=60,
|
||||
stream=True,
|
||||
) as response:
|
||||
response.raise_for_status()
|
||||
result = _read_requests_response_limited(response)
|
||||
|
||||
if result:
|
||||
return result
|
||||
else:
|
||||
raise RuntimeError(response.text)
|
||||
raise RuntimeError('WeChatPad returned an empty response')
|
||||
except Exception as e:
|
||||
print(f'http请求失败, url={url}, exception={e}')
|
||||
raise RuntimeError(str(e))
|
||||
|
||||
|
||||
@@ -27,16 +55,20 @@ def get_json(base_url, token):
|
||||
url = base_url + f'?key={token}'
|
||||
|
||||
try:
|
||||
response = requests.get(url, headers=headers, timeout=60)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
with requests.get(
|
||||
url,
|
||||
headers=headers,
|
||||
timeout=60,
|
||||
stream=True,
|
||||
) as response:
|
||||
response.raise_for_status()
|
||||
result = _read_requests_response_limited(response)
|
||||
|
||||
if result:
|
||||
return result
|
||||
else:
|
||||
raise RuntimeError(response.text)
|
||||
raise RuntimeError('WeChatPad returned an empty response')
|
||||
except Exception as e:
|
||||
print(f'http请求失败, url={url}, exception={e}')
|
||||
raise RuntimeError(str(e))
|
||||
|
||||
|
||||
@@ -68,7 +100,12 @@ async def async_request(
|
||||
method=method, url=url, params=params, headers=headers, data=data, json=json
|
||||
) as response:
|
||||
response.raise_for_status() # 如果状态码不是200，抛出异常
|
||||
result = await response.json()
|
||||
result = json_module.loads(
|
||||
await httpclient.read_limited(
|
||||
response,
|
||||
max_bytes=_MAX_WECHATPAD_RESPONSE_BYTES,
|
||||
)
|
||||
)
|
||||
# print(result)
|
||||
return result
|
||||
# if result.get('Code') == 200:
|
||||
|
||||
@@ -10,13 +10,16 @@ import re
|
||||
from typing import Any, Callable, Optional, Tuple
|
||||
from urllib.parse import unquote
|
||||
|
||||
import httpx
|
||||
from Crypto.Cipher import AES
|
||||
from quart import Quart, request, Response, jsonify
|
||||
|
||||
from langbot.libs.wecom_ai_bot_api import wecombotevent
|
||||
from langbot.libs.wecom_ai_bot_api.WXBizMsgCrypt3 import WXBizMsgCrypt
|
||||
from langbot.pkg.platform.logger import EventLogger
|
||||
from langbot.pkg.utils import httpclient
|
||||
|
||||
_CLIENT_TRANSIENT_CACHE_MAX = 4096
|
||||
_MAX_STREAM_CONTENT_CHARS = 200000
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -56,7 +59,7 @@ class StreamSession:
|
||||
last_access: float = field(default_factory=time.time)
|
||||
|
||||
# 将流水线增量结果缓存到队列，刷新请求逐条消费
|
||||
queue: asyncio.Queue = field(default_factory=asyncio.Queue)
|
||||
queue: asyncio.Queue = field(default_factory=lambda: asyncio.Queue(maxsize=1))
|
||||
|
||||
# 是否已经完成（收到最终片段）
|
||||
finished: bool = False
|
||||
@@ -85,6 +88,7 @@ class StreamSessionManager:
|
||||
# full like → cancel → dislike feedback flow. Must align with the adapter's
|
||||
# _stream_to_monitoring_msg TTL (wecombot.py).
|
||||
_FEEDBACK_SESSION_TTL = 600 # 10 minutes
|
||||
_MAX_SESSIONS = 4096
|
||||
|
||||
def __init__(self, logger: EventLogger, ttl: int = 60) -> None:
|
||||
self.logger = logger
|
||||
@@ -165,6 +169,26 @@ class StreamSessionManager:
|
||||
if task_id:
|
||||
self._task_index.pop(task_id, None)
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Release every retained stream and reverse index."""
|
||||
|
||||
self._sessions.clear()
|
||||
self._msg_index.clear()
|
||||
self._feedback_index.clear()
|
||||
self._task_index.clear()
|
||||
|
||||
def _drop_session(self, stream_id: str) -> StreamSession | None:
|
||||
session = self._sessions.pop(stream_id, None)
|
||||
if session is None:
|
||||
return None
|
||||
if session.msg_id and self._msg_index.get(session.msg_id) == stream_id:
|
||||
self._msg_index.pop(session.msg_id, None)
|
||||
if session.feedback_id:
|
||||
self._feedback_index.pop(session.feedback_id, None)
|
||||
if session.pending_form_task_id:
|
||||
self._task_index.pop(session.pending_form_task_id, None)
|
||||
return session
|
||||
|
||||
def create_or_get(self, msg_json: dict[str, Any]) -> tuple[StreamSession, bool]:
|
||||
"""根据企业微信回调创建或获取会话。
|
||||
|
||||
@@ -185,6 +209,14 @@ class StreamSessionManager:
|
||||
session.last_access = time.time()
|
||||
return session, False
|
||||
|
||||
self.cleanup()
|
||||
while len(self._sessions) >= self._MAX_SESSIONS:
|
||||
oldest_stream_id = min(
|
||||
self._sessions,
|
||||
key=lambda candidate: self._sessions[candidate].last_access,
|
||||
)
|
||||
self._drop_session(oldest_stream_id)
|
||||
|
||||
stream_id = str(uuid.uuid4())
|
||||
session = StreamSession(
|
||||
stream_id=stream_id,
|
||||
@@ -221,8 +253,13 @@ class StreamSessionManager:
|
||||
try:
|
||||
session.queue.put_nowait(chunk)
|
||||
except asyncio.QueueFull:
|
||||
# 默认无界队列，此处兜底防御
|
||||
await session.queue.put(chunk)
|
||||
# Each chunk is a complete snapshot. Coalesce a slow consumer to
|
||||
# the newest value instead of retaining every intermediate body.
|
||||
try:
|
||||
session.queue.get_nowait()
|
||||
except asyncio.QueueEmpty:
|
||||
pass
|
||||
session.queue.put_nowait(chunk)
|
||||
|
||||
if chunk.is_final:
|
||||
session.finished = True
|
||||
@@ -265,7 +302,7 @@ class StreamSessionManager:
|
||||
session.finished = True
|
||||
session.last_access = time.time()
|
||||
|
||||
def cleanup(self) -> None:
|
||||
def cleanup(self) -> list[str]:
|
||||
"""定期清理过期会话，防止队列与映射无上限累积。
|
||||
|
||||
已注册 feedback_id 的会话使用更长的 TTL，确保用户在点赞/取消/点踩流程中
|
||||
@@ -279,16 +316,14 @@ class StreamSessionManager:
|
||||
if now - session.last_access > effective_ttl:
|
||||
expired.append(stream_id)
|
||||
|
||||
removed_msg_ids: list[str] = []
|
||||
for stream_id in expired:
|
||||
session = self._sessions.pop(stream_id, None)
|
||||
session = self._drop_session(stream_id)
|
||||
if not session:
|
||||
continue
|
||||
msg_id = session.msg_id
|
||||
if msg_id and self._msg_index.get(msg_id) == stream_id:
|
||||
self._msg_index.pop(msg_id, None)
|
||||
# Clean up feedback index for expired sessions
|
||||
if session.feedback_id:
|
||||
self._feedback_index.pop(session.feedback_id, None)
|
||||
if session.msg_id:
|
||||
removed_msg_ids.append(session.msg_id)
|
||||
return removed_msg_ids
|
||||
|
||||
|
||||
def _decrypt_file(encrypted_data: bytes, aes_key_str: str) -> bytes:
|
||||
@@ -405,19 +440,19 @@ async def download_encrypted_file(
|
||||
|
||||
filename: Optional[str] = None
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=30.0) as client:
|
||||
response = await client.get(download_url)
|
||||
if response.status_code != 200:
|
||||
await logger.error(f'Failed to download file (HTTP {response.status_code}): {response.text[:200]}')
|
||||
client = httpclient.get_session()
|
||||
async with client.get(download_url, timeout=30.0) as response:
|
||||
if response.status != 200:
|
||||
await logger.error(f'Failed to download file (HTTP {response.status})')
|
||||
return None, None
|
||||
encrypted_bytes = response.content
|
||||
encrypted_bytes = await httpclient.read_limited(response)
|
||||
filename = _extract_filename(response.headers.get('content-disposition', ''))
|
||||
except Exception:
|
||||
await logger.error(f'Failed to download file: {traceback.format_exc()}')
|
||||
return None, None
|
||||
|
||||
try:
|
||||
decrypted = _decrypt_file(encrypted_bytes, aes_key)
|
||||
decrypted = await asyncio.to_thread(_decrypt_file, encrypted_bytes, aes_key)
|
||||
return decrypted, filename
|
||||
except Exception:
|
||||
await logger.error(f'Failed to decrypt file: {traceback.format_exc()}')
|
||||
@@ -466,7 +501,7 @@ async def parse_wecom_bot_message(
|
||||
"""Download, decrypt, and convert to data URI for backward compatibility."""
|
||||
data, _filename = await _safe_download(url, per_msg_aeskey)
|
||||
if data:
|
||||
return _bytes_to_data_uri(data)
|
||||
return await asyncio.to_thread(_bytes_to_data_uri, data)
|
||||
return None
|
||||
|
||||
if msg_type == 'text':
|
||||
@@ -579,7 +614,10 @@ async def parse_wecom_bot_message(
|
||||
if (file_data.get('filesize') or 0) <= max_inline_file_size:
|
||||
file_bytes, dl_filename = await _safe_download(download_url, item_aeskey)
|
||||
if file_bytes:
|
||||
file_data['base64'] = _bytes_to_data_uri(file_bytes)
|
||||
file_data['base64'] = await asyncio.to_thread(
|
||||
_bytes_to_data_uri,
|
||||
file_bytes,
|
||||
)
|
||||
if dl_filename and not file_data.get('filename'):
|
||||
file_data['filename'] = dl_filename
|
||||
files.append(file_data)
|
||||
@@ -1567,6 +1605,8 @@ def build_multiple_interaction_update_card(
|
||||
|
||||
|
||||
class WecomBotClient:
|
||||
_MAX_DISPATCH_TASKS = 100
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
Token: str,
|
||||
@@ -1613,6 +1653,7 @@ class WecomBotClient:
|
||||
self._feedback_callback: Optional[Callable] = None
|
||||
self._card_action_callback: Optional[Callable] = None
|
||||
self._stream_last_content: dict[str, str] = {}
|
||||
self._dispatch_tasks: set[asyncio.Task] = set()
|
||||
# Optional `source` block injected into every interactive template_card
|
||||
# the client builds. Set via `set_card_source` from the adapter after
|
||||
# reading config. Format: {icon_url, desc, desc_color}.
|
||||
@@ -1695,7 +1736,12 @@ class WecomBotClient:
|
||||
"""
|
||||
reply_plain_str = json.dumps(payload, ensure_ascii=False)
|
||||
reply_timestamp = str(int(time.time()))
|
||||
ret, encrypt_text = self.wxcpt.EncryptMsg(reply_plain_str, nonce, reply_timestamp)
|
||||
ret, encrypt_text = await asyncio.to_thread(
|
||||
self.wxcpt.EncryptMsg,
|
||||
reply_plain_str,
|
||||
nonce,
|
||||
reply_timestamp,
|
||||
)
|
||||
if ret != 0:
|
||||
await self.logger.error(f'加密失败: {ret}')
|
||||
return jsonify({'error': 'encrypt_failed'}), 500
|
||||
@@ -1718,6 +1764,41 @@ class WecomBotClient:
|
||||
except Exception:
|
||||
await self.logger.error(traceback.format_exc())
|
||||
|
||||
def _start_dispatch_task(self, event: wecombotevent.WecomBotEvent) -> bool:
|
||||
"""Start one bounded pipeline dispatch task."""
|
||||
|
||||
for task in tuple(self._dispatch_tasks):
|
||||
if task.done():
|
||||
self._dispatch_tasks.discard(task)
|
||||
if len(self._dispatch_tasks) >= self._MAX_DISPATCH_TASKS:
|
||||
return False
|
||||
|
||||
task = asyncio.create_task(self._dispatch_event(event))
|
||||
self._dispatch_tasks.add(task)
|
||||
|
||||
def done(done_task: asyncio.Task) -> None:
|
||||
self._dispatch_tasks.discard(done_task)
|
||||
if not done_task.cancelled():
|
||||
done_task.exception()
|
||||
|
||||
task.add_done_callback(done)
|
||||
return True
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Cancel callbacks and release retained webhook state."""
|
||||
|
||||
dispatch_tasks = list(self._dispatch_tasks)
|
||||
for task in dispatch_tasks:
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
if dispatch_tasks:
|
||||
await asyncio.gather(*dispatch_tasks, return_exceptions=True)
|
||||
self._dispatch_tasks.clear()
|
||||
self.generated_content.clear()
|
||||
self.msg_id_map.clear()
|
||||
self._stream_last_content.clear()
|
||||
self.stream_sessions.clear()
|
||||
|
||||
async def _handle_post_initial_response(self, msg_json: dict[str, Any], nonce: str) -> tuple[Response, int]:
|
||||
"""处理企业微信首次推送的消息，返回 stream_id 并开启流水线。
|
||||
|
||||
@@ -1747,7 +1828,8 @@ class WecomBotClient:
|
||||
await self.logger.error(traceback.format_exc())
|
||||
else:
|
||||
if is_new:
|
||||
asyncio.create_task(self._dispatch_event(event))
|
||||
if not self._start_dispatch_task(event):
|
||||
await self.logger.warning('WeCom webhook dispatch capacity reached; dropping message')
|
||||
|
||||
payload = self._build_stream_payload(session.stream_id, '', False, feedback_id)
|
||||
return await self._encrypt_and_reply(payload, nonce)
|
||||
@@ -1870,7 +1952,10 @@ class WecomBotClient:
|
||||
async def _handle_post_callback(self, req) -> tuple[Response, int] | Response:
|
||||
"""处理企业微信的 POST 回调请求。"""
|
||||
|
||||
self.stream_sessions.cleanup()
|
||||
for expired_msg_id in self.stream_sessions.cleanup():
|
||||
self.generated_content.pop(expired_msg_id, None)
|
||||
self._stream_last_content.pop(expired_msg_id, None)
|
||||
self.msg_id_map.pop(expired_msg_id, None)
|
||||
|
||||
msg_signature = unquote(req.args.get('msg_signature', ''))
|
||||
timestamp = unquote(req.args.get('timestamp', ''))
|
||||
@@ -1883,12 +1968,18 @@ class WecomBotClient:
|
||||
return Response('Bad Request', status=400)
|
||||
|
||||
xml_post_data = f'<xml><Encrypt><![CDATA[{encrypted_msg}]]></Encrypt></xml>'
|
||||
ret, decrypted_xml = self.wxcpt.DecryptMsg(xml_post_data, msg_signature, timestamp, nonce)
|
||||
ret, decrypted_xml = await asyncio.to_thread(
|
||||
self.wxcpt.DecryptMsg,
|
||||
xml_post_data,
|
||||
msg_signature,
|
||||
timestamp,
|
||||
nonce,
|
||||
)
|
||||
if ret != 0:
|
||||
await self.logger.error('解密失败')
|
||||
return Response('解密失败', status=400)
|
||||
|
||||
msg_json = json.loads(decrypted_xml)
|
||||
msg_json = await asyncio.to_thread(json.loads, decrypted_xml)
|
||||
|
||||
event_type = extract_wecom_event_type(msg_json)
|
||||
|
||||
@@ -2014,6 +2105,8 @@ class WecomBotClient:
|
||||
self.msg_id_map[message_id] += 1
|
||||
return
|
||||
self.msg_id_map[message_id] = 1
|
||||
while len(self.msg_id_map) > _CLIENT_TRANSIENT_CACHE_MAX:
|
||||
self.msg_id_map.pop(next(iter(self.msg_id_map)), None)
|
||||
msg_type = event.type
|
||||
if msg_type in self._message_handlers:
|
||||
for handler in self._message_handlers[msg_type]:
|
||||
@@ -2047,6 +2140,8 @@ class WecomBotClient:
|
||||
next_content = previous_content
|
||||
else:
|
||||
next_content = previous_content + content if previous_content else content
|
||||
if len(next_content) > _MAX_STREAM_CONTENT_CHARS:
|
||||
next_content = next_content[-_MAX_STREAM_CONTENT_CHARS:]
|
||||
|
||||
if not is_final and next_content == previous_content:
|
||||
return True
|
||||
@@ -2096,7 +2191,9 @@ class WecomBotClient:
|
||||
"""
|
||||
handled = await self.push_stream_chunk(msg_id, content, is_final=True)
|
||||
if not handled:
|
||||
self.generated_content[msg_id] = content
|
||||
self.generated_content[msg_id] = content[-_MAX_STREAM_CONTENT_CHARS:]
|
||||
while len(self.generated_content) > _CLIENT_TRANSIENT_CACHE_MAX:
|
||||
self.generated_content.pop(next(iter(self.generated_content)), None)
|
||||
|
||||
def on_message(self, msg_type: str):
|
||||
def decorator(func: Callable[[wecombotevent.WecomBotEvent], None]):
|
||||
@@ -2119,7 +2216,7 @@ class WecomBotClient:
|
||||
async def download_url_to_base64(self, download_url, encoding_aes_key):
|
||||
data, _filename = await download_encrypted_file(download_url, encoding_aes_key, self.logger)
|
||||
if data:
|
||||
return _bytes_to_data_uri(data)
|
||||
return await asyncio.to_thread(_bytes_to_data_uri, data)
|
||||
return None
|
||||
|
||||
async def run_task(self, host: str, port: int, *args, **kwargs):
|
||||
|
||||
@@ -47,6 +47,17 @@ CMD_RESPOND_WELCOME = 'aibot_respond_welcome_msg'
|
||||
CMD_RESPOND_UPDATE = 'aibot_respond_update_msg'
|
||||
CMD_SEND_MSG = 'aibot_send_msg'
|
||||
|
||||
_DEDUP_CACHE_MAX = 4096
|
||||
_STREAM_CACHE_MAX = 1024
|
||||
_FEEDBACK_CACHE_MAX = 4096
|
||||
_PENDING_FORM_MAX = 1024
|
||||
_PENDING_FORM_TTL_SECONDS = 1800
|
||||
_MAX_STREAM_CONTENT_CHARS = 200000
|
||||
_MAX_CALLBACK_TASKS = 100
|
||||
_MAX_REPLY_WORKERS = 100
|
||||
_MAX_REPLY_QUEUE_SIZE = 100
|
||||
_MAX_PENDING_ACKS = 256
|
||||
|
||||
|
||||
def _generate_req_id(prefix: str) -> str:
|
||||
"""Generate a unique request ID in the format: {prefix}_{timestamp}_{random}."""
|
||||
@@ -106,6 +117,7 @@ class WecomBotWsClient:
|
||||
# Per-req_id serial reply queues
|
||||
self._reply_queues: dict[str, asyncio.Queue] = {}
|
||||
self._reply_workers: dict[str, asyncio.Task] = {}
|
||||
self._callback_tasks: set[asyncio.Task] = set()
|
||||
self._reply_ack_timeout = 5.0
|
||||
|
||||
# Stream ID tracking for WebSocket mode
|
||||
@@ -135,6 +147,31 @@ class WecomBotWsClient:
|
||||
# `set_card_source` from the adapter after reading config.
|
||||
self.card_source: Optional[dict] = None
|
||||
|
||||
@staticmethod
|
||||
def _cap_mapping(mapping: dict, max_entries: int) -> None:
|
||||
while len(mapping) > max_entries:
|
||||
mapping.pop(next(iter(mapping)), None)
|
||||
|
||||
def _prune_stream_state(self) -> None:
|
||||
while len(self._stream_sessions) > _STREAM_CACHE_MAX:
|
||||
msg_id = next(iter(self._stream_sessions))
|
||||
self._stream_sessions.pop(msg_id, None)
|
||||
self._stream_ids.pop(msg_id, None)
|
||||
self._stream_last_content.pop(msg_id, None)
|
||||
task_id = self._task_id_by_msg.pop(msg_id, None)
|
||||
if task_id:
|
||||
self._pending_forms_by_task.pop(task_id, None)
|
||||
|
||||
def _prune_pending_forms(self) -> None:
|
||||
cutoff = time.monotonic() - _PENDING_FORM_TTL_SECONDS
|
||||
for task_id, pending in tuple(self._pending_forms_by_task.items()):
|
||||
if float(pending.get('created_at', 0.0)) <= cutoff:
|
||||
self._drop_pending_form_task(task_id, pending)
|
||||
while len(self._pending_forms_by_task) > _PENDING_FORM_MAX:
|
||||
task_id = next(iter(self._pending_forms_by_task))
|
||||
pending = self._pending_forms_by_task.get(task_id, {})
|
||||
self._drop_pending_form_task(task_id, pending)
|
||||
|
||||
# ── Public API ──────────────────────────────────────────────────
|
||||
|
||||
async def connect(self):
|
||||
@@ -173,17 +210,40 @@ class WecomBotWsClient:
|
||||
async def disconnect(self):
|
||||
"""Gracefully disconnect from the WebSocket server."""
|
||||
self._running = False
|
||||
heartbeat_tasks = []
|
||||
if self._heartbeat_task and not self._heartbeat_task.done():
|
||||
self._heartbeat_task.cancel()
|
||||
for task in self._reply_workers.values():
|
||||
heartbeat_tasks.append(self._heartbeat_task)
|
||||
reply_workers = list(self._reply_workers.values())
|
||||
for task in reply_workers:
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
callback_tasks = list(self._callback_tasks)
|
||||
for task in callback_tasks:
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
shutdown_tasks = [*heartbeat_tasks, *reply_workers, *callback_tasks]
|
||||
if shutdown_tasks:
|
||||
await asyncio.gather(*shutdown_tasks, return_exceptions=True)
|
||||
self._clear_pending_acks('Connection closed')
|
||||
if self._ws and not self._ws.closed:
|
||||
await self._ws.close()
|
||||
self._ws = None
|
||||
if self._session and not self._session.closed:
|
||||
await self._session.close()
|
||||
self._session = None
|
||||
self._heartbeat_task = None
|
||||
self._reply_queues.clear()
|
||||
self._reply_workers.clear()
|
||||
self._callback_tasks.clear()
|
||||
self._stream_ids.clear()
|
||||
self._stream_last_content.clear()
|
||||
self._stream_sessions.clear()
|
||||
self._feedback_sessions.clear()
|
||||
self._msg_feedback_ids.clear()
|
||||
self._pending_forms_by_task.clear()
|
||||
self._task_id_by_msg.clear()
|
||||
self._msg_id_map.clear()
|
||||
|
||||
def on_message(self, msg_type: str) -> Callable:
|
||||
"""Decorator to register a message handler.
|
||||
@@ -366,8 +426,10 @@ class WecomBotWsClient:
|
||||
'chat_id': session_info.get('chat_id', ''),
|
||||
'stream_id': stream_id,
|
||||
'req_id': req_id,
|
||||
'created_at': time.monotonic(),
|
||||
}
|
||||
self._task_id_by_msg[msg_id] = task_id
|
||||
self._prune_pending_forms()
|
||||
|
||||
card_payload = build_human_input_template_card_payload(
|
||||
form_data,
|
||||
@@ -458,6 +520,8 @@ class WecomBotWsClient:
|
||||
next_content = previous_content
|
||||
else:
|
||||
next_content = previous_content + content if previous_content else content
|
||||
if len(next_content) > _MAX_STREAM_CONTENT_CHARS:
|
||||
next_content = next_content[-_MAX_STREAM_CONTENT_CHARS:]
|
||||
|
||||
# Skip sending if content hasn't changed (e.g. during tool call argument streaming)
|
||||
if not is_final and next_content == previous_content:
|
||||
@@ -485,6 +549,8 @@ class WecomBotWsClient:
|
||||
session_info = self._stream_sessions.get(msg_id)
|
||||
if session_info:
|
||||
self._feedback_sessions[feedback_id] = session_info
|
||||
self._cap_mapping(self._feedback_sessions, _FEEDBACK_CACHE_MAX)
|
||||
self._cap_mapping(self._msg_feedback_ids, _FEEDBACK_CACHE_MAX)
|
||||
|
||||
# WeCom replaces the displayed stream content on each refresh, so
|
||||
# every frame must contain the complete snapshot, not only a delta.
|
||||
@@ -516,7 +582,7 @@ class WecomBotWsClient:
|
||||
|
||||
self._session = aiohttp.ClientSession()
|
||||
try:
|
||||
self._ws = await self._session.ws_connect(self.ws_url)
|
||||
self._ws = await self._session.ws_connect(self.ws_url, max_msg_size=1024 * 1024)
|
||||
self._missed_pong_count = 0
|
||||
self._reconnect_attempts = 0
|
||||
await self.logger.info('WebSocket connected, sending auth...')
|
||||
@@ -539,6 +605,8 @@ class WecomBotWsClient:
|
||||
finally:
|
||||
if self._heartbeat_task and not self._heartbeat_task.done():
|
||||
self._heartbeat_task.cancel()
|
||||
await asyncio.gather(self._heartbeat_task, return_exceptions=True)
|
||||
self._heartbeat_task = None
|
||||
self._clear_pending_acks('Connection closed')
|
||||
finally:
|
||||
if self._ws and not self._ws.closed:
|
||||
@@ -565,7 +633,7 @@ class WecomBotWsClient:
|
||||
try:
|
||||
msg = await asyncio.wait_for(self._ws.receive(), timeout=10.0)
|
||||
if msg.type in (aiohttp.WSMsgType.TEXT,):
|
||||
frame = json.loads(msg.data)
|
||||
frame = await asyncio.to_thread(json.loads, msg.data)
|
||||
req_id = frame.get('headers', {}).get('req_id', '')
|
||||
if req_id.startswith(CMD_SUBSCRIBE) and frame.get('errcode') == 0:
|
||||
return True
|
||||
@@ -614,7 +682,7 @@ class WecomBotWsClient:
|
||||
break
|
||||
if msg.type == aiohttp.WSMsgType.TEXT:
|
||||
try:
|
||||
frame = json.loads(msg.data)
|
||||
frame = await asyncio.to_thread(json.loads, msg.data)
|
||||
await self._handle_frame(frame)
|
||||
except json.JSONDecodeError:
|
||||
await self.logger.error(f'Failed to parse WebSocket message: {str(msg.data)[:200]}')
|
||||
@@ -622,7 +690,7 @@ class WecomBotWsClient:
|
||||
await self.logger.error(f'Error handling frame: {traceback.format_exc()}')
|
||||
elif msg.type == aiohttp.WSMsgType.BINARY:
|
||||
try:
|
||||
frame = json.loads(msg.data)
|
||||
frame = await asyncio.to_thread(json.loads, msg.data)
|
||||
await self._handle_frame(frame)
|
||||
except Exception:
|
||||
await self.logger.error(f'Error handling binary frame: {traceback.format_exc()}')
|
||||
@@ -638,12 +706,14 @@ class WecomBotWsClient:
|
||||
|
||||
# Message push
|
||||
if cmd == CMD_MSG_CALLBACK:
|
||||
asyncio.create_task(self._handle_message_callback(frame))
|
||||
if not self._start_callback_task(self._handle_message_callback(frame)):
|
||||
await self.logger.warning('WeCom WebSocket callback capacity reached; dropping message')
|
||||
return
|
||||
|
||||
# Event push
|
||||
if cmd == CMD_EVENT_CALLBACK:
|
||||
asyncio.create_task(self._handle_event_callback(frame))
|
||||
if not self._start_callback_task(self._handle_event_callback(frame)):
|
||||
await self.logger.warning('WeCom WebSocket callback capacity reached; dropping event')
|
||||
return
|
||||
|
||||
# No cmd → response/ACK frame, dispatch by req_id prefix
|
||||
@@ -665,6 +735,27 @@ class WecomBotWsClient:
|
||||
# Unknown frame
|
||||
await self.logger.warning(f'Unknown frame: {_frame_snippet(frame)}')
|
||||
|
||||
def _start_callback_task(self, coro) -> bool:
|
||||
"""Start one bounded inbound frame callback."""
|
||||
|
||||
for task in tuple(self._callback_tasks):
|
||||
if task.done():
|
||||
self._callback_tasks.discard(task)
|
||||
if len(self._callback_tasks) >= _MAX_CALLBACK_TASKS:
|
||||
coro.close()
|
||||
return False
|
||||
|
||||
task = asyncio.create_task(coro)
|
||||
self._callback_tasks.add(task)
|
||||
|
||||
def done(done_task: asyncio.Task) -> None:
|
||||
self._callback_tasks.discard(done_task)
|
||||
if not done_task.cancelled():
|
||||
done_task.exception()
|
||||
|
||||
task.add_done_callback(done)
|
||||
return True
|
||||
|
||||
async def _handle_message_callback(self, frame: dict):
|
||||
"""Handle an incoming message callback frame."""
|
||||
try:
|
||||
@@ -697,6 +788,7 @@ class WecomBotWsClient:
|
||||
'chat_id': message_data.get('chatid', ''),
|
||||
'chat_type': message_data.get('type', 'single'),
|
||||
}
|
||||
self._prune_stream_state()
|
||||
message_data['stream_id'] = stream_id
|
||||
message_data['req_id'] = req_id
|
||||
|
||||
@@ -748,7 +840,7 @@ class WecomBotWsClient:
|
||||
)
|
||||
|
||||
# Look up session by feedback_id
|
||||
session_info = self._feedback_sessions.get(feedback_id)
|
||||
session_info = self._feedback_sessions.pop(feedback_id, None)
|
||||
session = None
|
||||
if session_info:
|
||||
session = StreamSession(
|
||||
@@ -806,6 +898,10 @@ class WecomBotWsClient:
|
||||
if pending is None:
|
||||
await self.logger.warning(f'No pending_form found for task_id={task_id} (ws); card event ignored')
|
||||
return
|
||||
if time.monotonic() - float(pending.get('created_at', 0.0)) > _PENDING_FORM_TTL_SECONDS:
|
||||
self._drop_pending_form_task(task_id, pending)
|
||||
await self.logger.warning(f'Pending form expired for task_id={task_id} (ws)')
|
||||
return
|
||||
|
||||
req_id_for_update = frame.get('headers', {}).get('req_id', '')
|
||||
form_data = pending.get('form_data', {}) or {}
|
||||
@@ -868,6 +964,7 @@ class WecomBotWsClient:
|
||||
self._msg_id_map[message_id] += 1
|
||||
return
|
||||
self._msg_id_map[message_id] = 1
|
||||
self._cap_mapping(self._msg_id_map, _DEDUP_CACHE_MAX)
|
||||
|
||||
msg_type = event.type
|
||||
if msg_type in self._message_handlers:
|
||||
@@ -899,40 +996,61 @@ class WecomBotWsClient:
|
||||
|
||||
# Ensure serial delivery per req_id
|
||||
if req_id not in self._reply_queues:
|
||||
self._reply_queues[req_id] = asyncio.Queue()
|
||||
if len(self._reply_queues) >= _MAX_REPLY_WORKERS:
|
||||
await self.logger.warning('WeCom WebSocket reply worker capacity reached; dropping reply')
|
||||
return None
|
||||
self._reply_queues[req_id] = asyncio.Queue(maxsize=_MAX_REPLY_QUEUE_SIZE)
|
||||
self._reply_workers[req_id] = asyncio.create_task(self._reply_queue_worker(req_id))
|
||||
|
||||
future: asyncio.Future = asyncio.get_event_loop().create_future()
|
||||
await self._reply_queues[req_id].put((frame, future))
|
||||
try:
|
||||
self._reply_queues[req_id].put_nowait((frame, future))
|
||||
except asyncio.QueueFull:
|
||||
await self.logger.warning(f'WeCom WebSocket reply queue full for req_id={req_id}; dropping reply')
|
||||
return None
|
||||
return await future
|
||||
|
||||
async def _reply_queue_worker(self, req_id: str):
|
||||
"""Process reply queue items serially for a given req_id."""
|
||||
queue = self._reply_queues[req_id]
|
||||
current_future: asyncio.Future | None = None
|
||||
try:
|
||||
while self._running:
|
||||
try:
|
||||
frame, future = await asyncio.wait_for(queue.get(), timeout=60.0)
|
||||
frame, current_future = await asyncio.wait_for(queue.get(), timeout=60.0)
|
||||
except asyncio.TimeoutError:
|
||||
# Queue idle, clean up worker
|
||||
break
|
||||
|
||||
try:
|
||||
ack = await self._send_and_wait_ack(frame)
|
||||
if not future.done():
|
||||
future.set_result(ack)
|
||||
if not current_future.done():
|
||||
current_future.set_result(ack)
|
||||
except Exception as e:
|
||||
if not future.done():
|
||||
future.set_exception(e)
|
||||
if not current_future.done():
|
||||
current_future.set_exception(e)
|
||||
finally:
|
||||
current_future = None
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
if current_future is not None and not current_future.done():
|
||||
current_future.set_exception(ConnectionError('Connection closed'))
|
||||
finally:
|
||||
while True:
|
||||
try:
|
||||
_, future = queue.get_nowait()
|
||||
except asyncio.QueueEmpty:
|
||||
break
|
||||
if not future.done():
|
||||
future.set_exception(ConnectionError('Reply worker stopped'))
|
||||
self._reply_queues.pop(req_id, None)
|
||||
self._reply_workers.pop(req_id, None)
|
||||
|
||||
async def _send_and_wait_ack(self, frame: dict) -> Optional[dict]:
|
||||
"""Send a frame and wait for the corresponding ACK."""
|
||||
req_id = frame['headers']['req_id']
|
||||
if len(self._pending_acks) >= _MAX_PENDING_ACKS:
|
||||
await self.logger.warning('WeCom WebSocket pending ACK capacity reached; dropping frame')
|
||||
return None
|
||||
ack_future: asyncio.Future = asyncio.get_event_loop().create_future()
|
||||
self._pending_acks[req_id] = ack_future
|
||||
|
||||
|
||||
@@ -1,16 +1,82 @@
|
||||
from quart import request
|
||||
from .WXBizMsgCrypt3 import WXBizMsgCrypt
|
||||
import asyncio
|
||||
import base64
|
||||
import binascii
|
||||
import contextvars
|
||||
import functools
|
||||
import httpx
|
||||
import os
|
||||
import traceback
|
||||
from urllib.parse import quote
|
||||
from quart import Quart
|
||||
import xml.etree.ElementTree as ET
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Callable, Dict, Any
|
||||
from .wecomevent import WecomEvent
|
||||
import langbot_plugin.api.entities.builtin.platform.message as platform_message
|
||||
import aiofiles
|
||||
from langbot.pkg.utils import httpclient
|
||||
|
||||
_MAX_MEDIA_BYTES = 10 * 1024 * 1024
|
||||
_MAX_CALLBACK_BODY_BYTES = 1024 * 1024
|
||||
_EXTENDED_HTTP_TIMEOUT_SECONDS = 120
|
||||
|
||||
|
||||
async def _read_httpx_media_limited(response: httpx.Response) -> bytes:
|
||||
content_length = response.headers.get('Content-Length')
|
||||
if content_length is not None:
|
||||
try:
|
||||
if int(content_length) > _MAX_MEDIA_BYTES:
|
||||
raise ValueError('WeCom media exceeds the size limit')
|
||||
except (TypeError, ValueError) as exc:
|
||||
if 'exceeds' in str(exc):
|
||||
raise
|
||||
content = bytearray()
|
||||
async for chunk in response.aiter_bytes():
|
||||
content.extend(chunk)
|
||||
if len(content) > _MAX_MEDIA_BYTES:
|
||||
raise ValueError('WeCom media exceeds the size limit')
|
||||
return bytes(content)
|
||||
|
||||
|
||||
async def _read_local_media_limited(path: str) -> bytes:
|
||||
if await asyncio.to_thread(os.path.getsize, path) > _MAX_MEDIA_BYTES:
|
||||
raise ValueError('WeCom media exceeds the size limit')
|
||||
async with aiofiles.open(path, 'rb') as file:
|
||||
content = await file.read(_MAX_MEDIA_BYTES + 1)
|
||||
if len(content) > _MAX_MEDIA_BYTES:
|
||||
raise ValueError('WeCom media exceeds the size limit')
|
||||
return content
|
||||
|
||||
|
||||
async def _decode_media_base64_limited(value: str) -> bytes:
|
||||
max_encoded_chars = 4 * ((_MAX_MEDIA_BYTES + 2) // 3) + 4
|
||||
if len(value) > max_encoded_chars:
|
||||
raise ValueError('WeCom media exceeds the size limit')
|
||||
content = await asyncio.to_thread(base64.b64decode, value)
|
||||
if len(content) > _MAX_MEDIA_BYTES:
|
||||
raise ValueError('WeCom media exceeds the size limit')
|
||||
return content
|
||||
|
||||
|
||||
def _bounded_token_retry(method):
|
||||
"""Allow one token-refresh retry without unbounded async recursion."""
|
||||
|
||||
depth = contextvars.ContextVar(f'{method.__name__}_token_retry_depth', default=0)
|
||||
|
||||
@functools.wraps(method)
|
||||
async def wrapped(*args, **kwargs):
|
||||
current_depth = depth.get()
|
||||
if current_depth >= 2:
|
||||
raise RuntimeError(f'{method.__name__} exceeded the token refresh retry limit')
|
||||
token = depth.set(current_depth + 1)
|
||||
try:
|
||||
return await method(*args, **kwargs)
|
||||
finally:
|
||||
depth.reset(token)
|
||||
|
||||
return wrapped
|
||||
|
||||
|
||||
class WecomClient:
|
||||
@@ -36,6 +102,7 @@ class WecomClient:
|
||||
self.logger = logger
|
||||
self.unified_mode = unified_mode
|
||||
self.app = Quart(__name__)
|
||||
self.app.config['MAX_CONTENT_LENGTH'] = _MAX_CALLBACK_BODY_BYTES
|
||||
|
||||
# 只有在非统一模式下才注册独立路由
|
||||
if not self.unified_mode:
|
||||
@@ -49,6 +116,29 @@ class WecomClient:
|
||||
self._message_handlers = {
|
||||
'example': [],
|
||||
}
|
||||
self._http_clients: dict[bool, httpx.AsyncClient] = {}
|
||||
|
||||
@asynccontextmanager
|
||||
async def _http_client_context(self, *, unbounded_timeout: bool = False):
|
||||
client = self._http_clients.get(unbounded_timeout)
|
||||
if client is None or client.is_closed:
|
||||
response_hooks = httpclient.httpx_response_limit_hooks()
|
||||
client = (
|
||||
httpx.AsyncClient(
|
||||
timeout=_EXTENDED_HTTP_TIMEOUT_SECONDS,
|
||||
event_hooks=response_hooks,
|
||||
)
|
||||
if unbounded_timeout
|
||||
else httpx.AsyncClient(event_hooks=response_hooks)
|
||||
)
|
||||
self._http_clients[unbounded_timeout] = client
|
||||
yield client
|
||||
|
||||
async def close(self) -> None:
|
||||
clients = list(self._http_clients.values())
|
||||
self._http_clients.clear()
|
||||
if clients:
|
||||
await asyncio.gather(*(client.aclose() for client in clients), return_exceptions=True)
|
||||
|
||||
# access——token操作
|
||||
async def check_access_token(self):
|
||||
@@ -59,15 +149,16 @@ class WecomClient:
|
||||
|
||||
async def get_access_token(self, secret):
|
||||
url = f'{self.base_url}/gettoken?corpid={self.corpid}&corpsecret={secret}'
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with self._http_client_context() as client:
|
||||
response = await client.get(url)
|
||||
data = response.json()
|
||||
data = await httpclient.parse_json_response(response)
|
||||
if 'access_token' in data:
|
||||
return data['access_token']
|
||||
else:
|
||||
await self.logger.error(f'获取accesstoken失败:{response.json()}')
|
||||
await self.logger.error(f'获取accesstoken失败:{data}')
|
||||
raise Exception(f'未获取access token: {data}')
|
||||
|
||||
@_bounded_token_retry
|
||||
async def get_user_info(self, userid: str) -> dict:
|
||||
"""
|
||||
Get user information by user ID using the application secret.
|
||||
@@ -82,9 +173,9 @@ class WecomClient:
|
||||
self.access_token = await self.get_access_token(self.secret)
|
||||
|
||||
url = self.base_url + '/user/get?access_token=' + self.access_token + '&userid=' + quote(userid)
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with self._http_client_context() as client:
|
||||
response = await client.get(url)
|
||||
data = response.json()
|
||||
data = await httpclient.parse_json_response(response)
|
||||
if data.get('errcode') == 40014 or data.get('errcode') == 42001:
|
||||
self.access_token = await self.get_access_token(self.secret)
|
||||
return await self.get_user_info(userid)
|
||||
@@ -98,13 +189,13 @@ class WecomClient:
|
||||
self.access_token_for_contacts = await self.get_access_token(self.secret_for_contacts)
|
||||
|
||||
url = self.base_url + '/user/list_id?access_token=' + self.access_token_for_contacts
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with self._http_client_context() as client:
|
||||
params = {
|
||||
'cursor': '',
|
||||
'limit': 10000,
|
||||
}
|
||||
response = await client.post(url, json=params)
|
||||
data = response.json()
|
||||
data = await httpclient.parse_json_response(response)
|
||||
if data['errcode'] == 0:
|
||||
dept_users = data['dept_user']
|
||||
userid = []
|
||||
@@ -121,7 +212,7 @@ class WecomClient:
|
||||
url = self.base_url + '/message/send?access_token=' + self.access_token_for_contacts
|
||||
user_ids = await self.get_users()
|
||||
user_ids_string = '|'.join(user_ids)
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with self._http_client_context() as client:
|
||||
params = {
|
||||
'touser': user_ids_string,
|
||||
'msgtype': 'text',
|
||||
@@ -135,16 +226,17 @@ class WecomClient:
|
||||
'duplicate_check_interval': 1800,
|
||||
}
|
||||
response = await client.post(url, json=params)
|
||||
data = response.json()
|
||||
data = await httpclient.parse_json_response(response)
|
||||
if data['errcode'] != 0:
|
||||
raise Exception('Failed to send message: ' + str(data))
|
||||
|
||||
@_bounded_token_retry
|
||||
async def send_image(self, user_id: str, agent_id: int, media_id: str):
|
||||
if not await self.check_access_token():
|
||||
self.access_token = await self.get_access_token(self.secret)
|
||||
|
||||
url = self.base_url + '/message/send?access_token=' + self.access_token
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with self._http_client_context() as client:
|
||||
params = {
|
||||
'touser': user_id,
|
||||
'msgtype': 'image',
|
||||
@@ -158,7 +250,7 @@ class WecomClient:
|
||||
'duplicate_check_interval': 1800,
|
||||
}
|
||||
response = await client.post(url, json=params)
|
||||
data = response.json()
|
||||
data = await httpclient.parse_json_response(response)
|
||||
if data['errcode'] == 40014 or data['errcode'] == 42001:
|
||||
self.access_token = await self.get_access_token(self.secret)
|
||||
return await self.send_image(user_id, agent_id, media_id)
|
||||
@@ -166,11 +258,12 @@ class WecomClient:
|
||||
await self.logger.error(f'发送图片失败:{data}')
|
||||
raise Exception('Failed to send image: ' + str(data))
|
||||
|
||||
@_bounded_token_retry
|
||||
async def send_voice(self, user_id: str, agent_id: int, media_id: str):
|
||||
if not await self.check_access_token():
|
||||
self.access_token = await self.get_access_token(self.secret)
|
||||
url = self.base_url + '/message/send?access_token=' + self.access_token
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with self._http_client_context() as client:
|
||||
params = {
|
||||
'touser': user_id,
|
||||
'msgtype': 'voice',
|
||||
@@ -184,7 +277,7 @@ class WecomClient:
|
||||
'duplicate_check_interval': 1800,
|
||||
}
|
||||
response = await client.post(url, json=params)
|
||||
data = response.json()
|
||||
data = await httpclient.parse_json_response(response)
|
||||
if data['errcode'] == 40014 or data['errcode'] == 42001:
|
||||
self.access_token = await self.get_access_token(self.secret)
|
||||
return await self.send_voice(user_id, agent_id, media_id)
|
||||
@@ -192,11 +285,12 @@ class WecomClient:
|
||||
await self.logger.error(f'发送语音失败:{data}')
|
||||
raise Exception('Failed to send voice: ' + str(data))
|
||||
|
||||
@_bounded_token_retry
|
||||
async def send_file(self, user_id: str, agent_id: int, media_id: str):
|
||||
if not await self.check_access_token():
|
||||
self.access_token = await self.get_access_token(self.secret)
|
||||
url = self.base_url + '/message/send?access_token=' + self.access_token
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with self._http_client_context() as client:
|
||||
params = {
|
||||
'touser': user_id,
|
||||
'msgtype': 'file',
|
||||
@@ -210,7 +304,7 @@ class WecomClient:
|
||||
'duplicate_check_interval': 1800,
|
||||
}
|
||||
response = await client.post(url, json=params)
|
||||
data = response.json()
|
||||
data = await httpclient.parse_json_response(response)
|
||||
if data['errcode'] == 40014 or data['errcode'] == 42001:
|
||||
self.access_token = await self.get_access_token(self.secret)
|
||||
return await self.send_file(user_id, agent_id, media_id)
|
||||
@@ -218,12 +312,13 @@ class WecomClient:
|
||||
await self.logger.error(f'发送文件失败:{data}')
|
||||
raise Exception('Failed to send file: ' + str(data))
|
||||
|
||||
@_bounded_token_retry
|
||||
async def send_private_msg(self, user_id: str, agent_id: int, content: str):
|
||||
if not await self.check_access_token():
|
||||
self.access_token = await self.get_access_token(self.secret)
|
||||
|
||||
url = self.base_url + '/message/send?access_token=' + self.access_token
|
||||
async with httpx.AsyncClient(timeout=None) as client:
|
||||
async with self._http_client_context(unbounded_timeout=True) as client:
|
||||
params = {
|
||||
'touser': user_id,
|
||||
'msgtype': 'text',
|
||||
@@ -237,7 +332,7 @@ class WecomClient:
|
||||
'duplicate_check_interval': 1800,
|
||||
}
|
||||
response = await client.post(url, json=params)
|
||||
data = response.json()
|
||||
data = await httpclient.parse_json_response(response)
|
||||
if data['errcode'] == 40014 or data['errcode'] == 42001:
|
||||
self.access_token = await self.get_access_token(self.secret)
|
||||
return await self.send_private_msg(user_id, agent_id, content)
|
||||
@@ -283,7 +378,15 @@ class WecomClient:
|
||||
|
||||
elif req.method == 'POST':
|
||||
encrypt_msg = await req.data
|
||||
ret, xml_msg = wxcpt.DecryptMsg(encrypt_msg, msg_signature, timestamp, nonce)
|
||||
if len(encrypt_msg) > _MAX_CALLBACK_BODY_BYTES:
|
||||
raise ValueError('WeCom callback body exceeds the size limit')
|
||||
ret, xml_msg = await asyncio.to_thread(
|
||||
wxcpt.DecryptMsg,
|
||||
encrypt_msg,
|
||||
msg_signature,
|
||||
timestamp,
|
||||
nonce,
|
||||
)
|
||||
if ret != 0:
|
||||
await self.logger.error('消息解密失败')
|
||||
raise Exception(f'消息解密失败，错误码: {ret}')
|
||||
@@ -332,7 +435,7 @@ class WecomClient:
|
||||
"""
|
||||
解析微信返回的 XML 消息并转换为字典。
|
||||
"""
|
||||
root = ET.fromstring(xml_msg)
|
||||
root = await asyncio.to_thread(ET.fromstring, xml_msg)
|
||||
message_data = {
|
||||
'ToUserName': root.find('ToUserName').text,
|
||||
'FromUserName': root.find('FromUserName').text,
|
||||
@@ -366,6 +469,7 @@ class WecomClient:
|
||||
return ext
|
||||
return 'jpg' # 默认返回jpg
|
||||
|
||||
@_bounded_token_retry
|
||||
async def upload_image_to_work(self, image: platform_message.Image):
|
||||
"""
|
||||
获取 media_id
|
||||
@@ -379,9 +483,8 @@ class WecomClient:
|
||||
|
||||
# 获取文件的二进制数据
|
||||
if image.path:
|
||||
async with aiofiles.open(image.path, 'rb') as f:
|
||||
file_bytes = await f.read()
|
||||
file_name = image.path.split('/')[-1]
|
||||
file_bytes = await _read_local_media_limited(image.path)
|
||||
file_name = image.path.split('/')[-1]
|
||||
elif image.url:
|
||||
file_bytes = await self.download_media_to_bytes(image.url)
|
||||
file_name = image.url.split('/')[-1]
|
||||
@@ -392,7 +495,7 @@ class WecomClient:
|
||||
base64_data = base64_data.split(',', 1)[1]
|
||||
padding = 4 - (len(base64_data) % 4) if len(base64_data) % 4 else 0
|
||||
padded_base64 = base64_data + '=' * padding
|
||||
file_bytes = base64.b64decode(padded_base64)
|
||||
file_bytes = await _decode_media_base64_limited(padded_base64)
|
||||
except binascii.Error as e:
|
||||
raise ValueError(f'Invalid base64 string: {str(e)}')
|
||||
else:
|
||||
@@ -400,6 +503,8 @@ class WecomClient:
|
||||
raise ValueError('image对象出错')
|
||||
|
||||
# 设置 multipart/form-data 格式的文件
|
||||
if len(file_bytes) > _MAX_MEDIA_BYTES:
|
||||
raise ValueError('WeCom media exceeds the size limit')
|
||||
boundary = '-------------------------acebdf13572468'
|
||||
headers = {'Content-Type': f'multipart/form-data; boundary={boundary}'}
|
||||
body = (
|
||||
@@ -413,9 +518,9 @@ class WecomClient:
|
||||
)
|
||||
|
||||
# 上传文件
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with self._http_client_context() as client:
|
||||
response = await client.post(url, headers=headers, content=body)
|
||||
data = response.json()
|
||||
data = await httpclient.parse_json_response(response)
|
||||
if data['errcode'] == 40014 or data['errcode'] == 42001:
|
||||
self.access_token = await self.get_access_token(self.secret)
|
||||
media_id = await self.upload_image_to_work(image)
|
||||
@@ -426,6 +531,7 @@ class WecomClient:
|
||||
media_id = data.get('media_id')
|
||||
return media_id
|
||||
|
||||
@_bounded_token_retry
|
||||
async def upload_voice_to_work(self, voice: platform_message.Voice):
|
||||
"""
|
||||
上传语音文件到企业微信
|
||||
@@ -437,9 +543,8 @@ class WecomClient:
|
||||
file_name = 'voice.mp3'
|
||||
|
||||
if voice.path:
|
||||
async with aiofiles.open(voice.path, 'rb') as f:
|
||||
file_bytes = await f.read()
|
||||
file_name = voice.path.split('/')[-1]
|
||||
file_bytes = await _read_local_media_limited(voice.path)
|
||||
file_name = voice.path.split('/')[-1]
|
||||
elif voice.url:
|
||||
file_bytes = await self.download_media_to_bytes(voice.url)
|
||||
file_name = voice.url.split('/')[-1]
|
||||
@@ -450,13 +555,15 @@ class WecomClient:
|
||||
base64_data = base64_data.split(',', 1)[1]
|
||||
padding = 4 - (len(base64_data) % 4) if len(base64_data) % 4 else 0
|
||||
padded_base64 = base64_data + '=' * padding
|
||||
file_bytes = base64.b64decode(padded_base64)
|
||||
file_bytes = await _decode_media_base64_limited(padded_base64)
|
||||
except binascii.Error as e:
|
||||
raise ValueError(f'Invalid base64 string: {str(e)}')
|
||||
else:
|
||||
await self.logger.error('Voice对象出错')
|
||||
raise ValueError('voice对象出错')
|
||||
|
||||
if len(file_bytes) > _MAX_MEDIA_BYTES:
|
||||
raise ValueError('WeCom media exceeds the size limit')
|
||||
boundary = '-------------------------acebdf13572468'
|
||||
headers = {'Content-Type': f'multipart/form-data; boundary={boundary}'}
|
||||
body = (
|
||||
@@ -470,9 +577,9 @@ class WecomClient:
|
||||
)
|
||||
|
||||
# print(body)
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with self._http_client_context() as client:
|
||||
response = await client.post(url, headers=headers, content=body)
|
||||
data = response.json()
|
||||
data = await httpclient.parse_json_response(response)
|
||||
if data['errcode'] == 40014 or data['errcode'] == 42001:
|
||||
self.access_token = await self.get_access_token(self.secret)
|
||||
media_id = await self.upload_voice_to_work(voice)
|
||||
@@ -482,6 +589,7 @@ class WecomClient:
|
||||
media_id = data.get('media_id')
|
||||
return media_id
|
||||
|
||||
@_bounded_token_retry
|
||||
async def upload_file_to_work(self, file: platform_message.File):
|
||||
"""
|
||||
上传文件到企业微信
|
||||
@@ -492,9 +600,8 @@ class WecomClient:
|
||||
file_bytes = None
|
||||
file_name = 'file.txt'
|
||||
if file.path:
|
||||
async with aiofiles.open(file.path, 'rb') as f:
|
||||
file_bytes = await f.read()
|
||||
file_name = file.path.split('/')[-1]
|
||||
file_bytes = await _read_local_media_limited(file.path)
|
||||
file_name = file.path.split('/')[-1]
|
||||
elif file.url:
|
||||
file_bytes = await self.download_media_to_bytes(file.url)
|
||||
file_name = file.url.split('/')[-1]
|
||||
@@ -505,12 +612,14 @@ class WecomClient:
|
||||
base64_data = base64_data.split(',', 1)[1]
|
||||
padding = 4 - (len(base64_data) % 4) if len(base64_data) % 4 else 0
|
||||
padded_base64 = base64_data + '=' * padding
|
||||
file_bytes = base64.b64decode(padded_base64)
|
||||
file_bytes = await _decode_media_base64_limited(padded_base64)
|
||||
except binascii.Error as e:
|
||||
raise ValueError(f'Invalid base64 string: {str(e)}')
|
||||
else:
|
||||
await self.logger.error('File对象出错')
|
||||
raise ValueError('file对象出错')
|
||||
if len(file_bytes) > _MAX_MEDIA_BYTES:
|
||||
raise ValueError('WeCom media exceeds the size limit')
|
||||
boundary = '-------------------------acebdf13572468'
|
||||
headers = {'Content-Type': f'multipart/form-data; boundary={boundary}'}
|
||||
body = (
|
||||
@@ -522,9 +631,9 @@ class WecomClient:
|
||||
+ file_bytes
|
||||
+ f'\r\n--{boundary}--\r\n'.encode('utf-8')
|
||||
)
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with self._http_client_context() as client:
|
||||
response = await client.post(url, headers=headers, content=body)
|
||||
data = response.json()
|
||||
data = await httpclient.parse_json_response(response)
|
||||
if data['errcode'] == 40014 or data['errcode'] == 42001:
|
||||
self.access_token = await self.get_access_token(self.secret)
|
||||
media_id = await self.upload_file_to_work(file)
|
||||
@@ -535,10 +644,10 @@ class WecomClient:
|
||||
return media_id
|
||||
|
||||
async def download_media_to_bytes(self, url: str) -> bytes:
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.get(url)
|
||||
response.raise_for_status()
|
||||
return response.content
|
||||
async with self._http_client_context() as client:
|
||||
async with client.stream('GET', url) as response:
|
||||
response.raise_for_status()
|
||||
return await _read_httpx_media_limited(response)
|
||||
|
||||
# 进行media_id的获取
|
||||
async def get_media_id(self, media: platform_message.Image | platform_message.Voice | platform_message.File):
|
||||
|
||||
@@ -1,8 +1,13 @@
|
||||
from quart import request
|
||||
from ..wecom_api.WXBizMsgCrypt3 import WXBizMsgCrypt
|
||||
import asyncio
|
||||
import base64
|
||||
import binascii
|
||||
import contextvars
|
||||
import functools
|
||||
import httpx
|
||||
import json
|
||||
import os
|
||||
import traceback
|
||||
from quart import Quart
|
||||
import xml.etree.ElementTree as ET
|
||||
@@ -11,9 +16,72 @@ from .wecomcsevent import WecomCSEvent
|
||||
import langbot_plugin.api.entities.builtin.platform.message as platform_message
|
||||
import aiofiles
|
||||
import time
|
||||
from contextlib import asynccontextmanager
|
||||
from langbot.pkg.utils import httpclient
|
||||
|
||||
_MAX_MEDIA_BYTES = 10 * 1024 * 1024
|
||||
_MAX_CALLBACK_BODY_BYTES = 1024 * 1024
|
||||
|
||||
|
||||
async def _read_httpx_media_limited(response: httpx.Response) -> bytes:
|
||||
content_length = response.headers.get('Content-Length')
|
||||
if content_length is not None:
|
||||
try:
|
||||
if int(content_length) > _MAX_MEDIA_BYTES:
|
||||
raise ValueError('WeCom customer-service media exceeds the size limit')
|
||||
except (TypeError, ValueError) as exc:
|
||||
if 'exceeds' in str(exc):
|
||||
raise
|
||||
content = bytearray()
|
||||
async for chunk in response.aiter_bytes():
|
||||
content.extend(chunk)
|
||||
if len(content) > _MAX_MEDIA_BYTES:
|
||||
raise ValueError('WeCom customer-service media exceeds the size limit')
|
||||
return bytes(content)
|
||||
|
||||
|
||||
async def _read_local_media_limited(path: str) -> bytes:
|
||||
if await asyncio.to_thread(os.path.getsize, path) > _MAX_MEDIA_BYTES:
|
||||
raise ValueError('WeCom customer-service media exceeds the size limit')
|
||||
async with aiofiles.open(path, 'rb') as file:
|
||||
content = await file.read(_MAX_MEDIA_BYTES + 1)
|
||||
if len(content) > _MAX_MEDIA_BYTES:
|
||||
raise ValueError('WeCom customer-service media exceeds the size limit')
|
||||
return content
|
||||
|
||||
|
||||
async def _decode_media_base64_limited(value: str) -> bytes:
|
||||
max_encoded_chars = 4 * ((_MAX_MEDIA_BYTES + 2) // 3) + 4
|
||||
if len(value) > max_encoded_chars:
|
||||
raise ValueError('WeCom customer-service media exceeds the size limit')
|
||||
content = await asyncio.to_thread(base64.b64decode, value)
|
||||
if len(content) > _MAX_MEDIA_BYTES:
|
||||
raise ValueError('WeCom customer-service media exceeds the size limit')
|
||||
return content
|
||||
|
||||
|
||||
def _bounded_token_retry(method):
|
||||
"""Allow one token-refresh retry without unbounded async recursion."""
|
||||
|
||||
depth = contextvars.ContextVar(f'{method.__name__}_token_retry_depth', default=0)
|
||||
|
||||
@functools.wraps(method)
|
||||
async def wrapped(*args, **kwargs):
|
||||
current_depth = depth.get()
|
||||
if current_depth >= 2:
|
||||
raise RuntimeError(f'{method.__name__} exceeded the token refresh retry limit')
|
||||
token = depth.set(current_depth + 1)
|
||||
try:
|
||||
return await method(*args, **kwargs)
|
||||
finally:
|
||||
depth.reset(token)
|
||||
|
||||
return wrapped
|
||||
|
||||
|
||||
class WecomCSClient:
|
||||
_CUSTOMER_CACHE_MAX = 4096
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
corpid: str,
|
||||
@@ -34,10 +102,12 @@ class WecomCSClient:
|
||||
self.logger = logger
|
||||
self.unified_mode = unified_mode
|
||||
self.app = Quart(__name__)
|
||||
self.app.config['MAX_CONTENT_LENGTH'] = _MAX_CALLBACK_BODY_BYTES
|
||||
|
||||
# Customer info cache: {external_userid: (info_dict, timestamp)}
|
||||
self._customer_cache: dict[str, tuple[dict, float]] = {}
|
||||
self._cache_ttl = 60 # Cache TTL in seconds (1 minute)
|
||||
self._customer_cache_cleanup_at = 0.0
|
||||
|
||||
# 只有在非统一模式下才注册独立路由
|
||||
if not self.unified_mode:
|
||||
@@ -48,29 +118,40 @@ class WecomCSClient:
|
||||
self._message_handlers = {
|
||||
'example': [],
|
||||
}
|
||||
self._http_client: httpx.AsyncClient | None = None
|
||||
|
||||
@asynccontextmanager
|
||||
async def _http_client_context(self):
|
||||
if self._http_client is None or self._http_client.is_closed:
|
||||
self._http_client = httpx.AsyncClient(event_hooks=httpclient.httpx_response_limit_hooks())
|
||||
yield self._http_client
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._http_client is not None:
|
||||
await self._http_client.aclose()
|
||||
self._http_client = None
|
||||
|
||||
@_bounded_token_retry
|
||||
async def get_pic_url(self, media_id: str):
|
||||
if not await self.check_access_token():
|
||||
self.access_token = await self.get_access_token(self.secret)
|
||||
|
||||
url = f'{self.base_url}/media/get?access_token={self.access_token}&media_id={media_id}'
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.get(url)
|
||||
if response.headers.get('Content-Type', '').startswith('application/json'):
|
||||
data = response.json()
|
||||
if data.get('errcode') in [40014, 42001]:
|
||||
self.access_token = await self.get_access_token(self.secret)
|
||||
return await self.get_pic_url(media_id)
|
||||
else:
|
||||
async with self._http_client_context() as client:
|
||||
async with client.stream('GET', url) as response:
|
||||
image_bytes = await _read_httpx_media_limited(response)
|
||||
content_type = response.headers.get('Content-Type', '')
|
||||
if content_type.startswith('application/json'):
|
||||
data = json.loads(image_bytes)
|
||||
if data.get('errcode') in [40014, 42001]:
|
||||
self.access_token = await self.get_access_token(self.secret)
|
||||
return await self.get_pic_url(media_id)
|
||||
raise Exception('Failed to get image: ' + str(data))
|
||||
|
||||
# 否则是图片，转成 base64
|
||||
image_bytes = response.content
|
||||
content_type = response.headers.get('Content-Type', '')
|
||||
base64_str = base64.b64encode(image_bytes).decode('utf-8')
|
||||
base64_str = f'data:{content_type};base64,{base64_str}'
|
||||
return base64_str
|
||||
# 否则是图片，转成 base64
|
||||
base64_str = (await asyncio.to_thread(base64.b64encode, image_bytes)).decode('utf-8')
|
||||
return f'data:{content_type};base64,{base64_str}'
|
||||
|
||||
# access——token操作
|
||||
async def check_access_token(self):
|
||||
@@ -81,19 +162,20 @@ class WecomCSClient:
|
||||
|
||||
async def get_access_token(self, secret):
|
||||
url = f'{self.base_url}/gettoken?corpid={self.corpid}&corpsecret={secret}'
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with self._http_client_context() as client:
|
||||
response = await client.get(url)
|
||||
data = response.json()
|
||||
data = await httpclient.parse_json_response(response)
|
||||
if 'access_token' in data:
|
||||
return data['access_token']
|
||||
else:
|
||||
raise Exception(f'未获取access token: {data}')
|
||||
|
||||
@_bounded_token_retry
|
||||
async def get_detailed_message_list(self, xml_msg: str):
|
||||
# 在本方法中解析消息，并且获得消息的具体内容
|
||||
if isinstance(xml_msg, bytes):
|
||||
xml_msg = xml_msg.decode('utf-8')
|
||||
root = ET.fromstring(xml_msg)
|
||||
root = await asyncio.to_thread(ET.fromstring, xml_msg)
|
||||
token = root.find('Token').text
|
||||
open_kfid = root.find('OpenKfId').text
|
||||
|
||||
@@ -106,14 +188,14 @@ class WecomCSClient:
|
||||
self.access_token = await self.get_access_token(self.secret)
|
||||
|
||||
url = self.base_url + '/kf/sync_msg?access_token=' + self.access_token
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with self._http_client_context() as client:
|
||||
params = {
|
||||
'token': token,
|
||||
'voice_format': 0,
|
||||
'open_kfid': open_kfid,
|
||||
}
|
||||
response = await client.post(url, json=params)
|
||||
data = response.json()
|
||||
data = await httpclient.parse_json_response(response)
|
||||
if data['errcode'] == 40014 or data['errcode'] == 42001:
|
||||
self.access_token = await self.get_access_token(self.secret)
|
||||
return await self.get_detailed_message_list(xml_msg)
|
||||
@@ -130,11 +212,12 @@ class WecomCSClient:
|
||||
# await self.change_service_status(userid=external_userid,openkfid=open_kfid,servicer=servicer)
|
||||
return last_msg_data
|
||||
|
||||
@_bounded_token_retry
|
||||
async def change_service_status(self, userid: str, openkfid: str, servicer: str):
|
||||
if not await self.check_access_token():
|
||||
self.access_token = await self.get_access_token(self.secret)
|
||||
url = self.base_url + '/kf/service_state/get?access_token=' + self.access_token
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with self._http_client_context() as client:
|
||||
params = {
|
||||
'open_kfid': openkfid,
|
||||
'external_userid': userid,
|
||||
@@ -142,18 +225,19 @@ class WecomCSClient:
|
||||
'servicer_userid': servicer,
|
||||
}
|
||||
response = await client.post(url, json=params)
|
||||
data = response.json()
|
||||
data = await httpclient.parse_json_response(response)
|
||||
if data['errcode'] == 40014 or data['errcode'] == 42001:
|
||||
self.access_token = await self.get_access_token(self.secret)
|
||||
return await self.change_service_status(userid, openkfid)
|
||||
return await self.change_service_status(userid, openkfid, servicer)
|
||||
if data['errcode'] != 0:
|
||||
raise Exception('Failed to change service status: ' + str(data))
|
||||
|
||||
@_bounded_token_retry
|
||||
async def send_image(self, user_id: str, agent_id: int, media_id: str):
|
||||
if not await self.check_access_token():
|
||||
self.access_token = await self.get_access_token(self.secret)
|
||||
url = self.base_url + '/media/upload?access_token=' + self.access_token
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with self._http_client_context() as client:
|
||||
params = {
|
||||
'touser': user_id,
|
||||
'toparty': '',
|
||||
@@ -170,7 +254,7 @@ class WecomCSClient:
|
||||
}
|
||||
try:
|
||||
response = await client.post(url, json=params)
|
||||
data = response.json()
|
||||
data = await httpclient.parse_json_response(response)
|
||||
except Exception as e:
|
||||
raise Exception('Failed to send image: ' + str(e))
|
||||
|
||||
@@ -182,6 +266,7 @@ class WecomCSClient:
|
||||
if data['errcode'] != 0:
|
||||
raise Exception('Failed to send image: ' + str(data))
|
||||
|
||||
@_bounded_token_retry
|
||||
async def send_text_msg(self, open_kfid: str, external_userid: str, msgid: str, content: str):
|
||||
if not await self.check_access_token():
|
||||
self.access_token = await self.get_access_token(self.secret)
|
||||
@@ -198,10 +283,10 @@ class WecomCSClient:
|
||||
},
|
||||
}
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with self._http_client_context() as client:
|
||||
response = await client.post(url, json=payload)
|
||||
|
||||
data = response.json()
|
||||
data = await httpclient.parse_json_response(response)
|
||||
if data['errcode'] == 40014 or data['errcode'] == 42001:
|
||||
self.access_token = await self.get_access_token(self.secret)
|
||||
return await self.send_text_msg(open_kfid, external_userid, msgid, content)
|
||||
@@ -250,7 +335,15 @@ class WecomCSClient:
|
||||
|
||||
elif req.method == 'POST':
|
||||
encrypt_msg = await req.data
|
||||
ret, xml_msg = wxcpt.DecryptMsg(encrypt_msg, msg_signature, timestamp, nonce)
|
||||
if len(encrypt_msg) > _MAX_CALLBACK_BODY_BYTES:
|
||||
raise ValueError('WeCom customer-service callback body exceeds the size limit')
|
||||
ret, xml_msg = await asyncio.to_thread(
|
||||
wxcpt.DecryptMsg,
|
||||
encrypt_msg,
|
||||
msg_signature,
|
||||
timestamp,
|
||||
nonce,
|
||||
)
|
||||
if ret != 0:
|
||||
raise Exception(f'消息解密失败，错误码: {ret}')
|
||||
|
||||
@@ -315,6 +408,7 @@ class WecomCSClient:
|
||||
return ext
|
||||
return 'jpg' # 默认返回jpg
|
||||
|
||||
@_bounded_token_retry
|
||||
async def upload_to_work(self, image: platform_message.Image):
|
||||
"""
|
||||
获取 media_id
|
||||
@@ -328,9 +422,8 @@ class WecomCSClient:
|
||||
|
||||
# 获取文件的二进制数据
|
||||
if image.path:
|
||||
async with aiofiles.open(image.path, 'rb') as f:
|
||||
file_bytes = await f.read()
|
||||
file_name = image.path.split('/')[-1]
|
||||
file_bytes = await _read_local_media_limited(image.path)
|
||||
file_name = image.path.split('/')[-1]
|
||||
elif image.url:
|
||||
file_bytes = await self.download_image_to_bytes(image.url)
|
||||
file_name = image.url.split('/')[-1]
|
||||
@@ -341,13 +434,15 @@ class WecomCSClient:
|
||||
base64_data = base64_data.split(',', 1)[1]
|
||||
padding = 4 - (len(base64_data) % 4) if len(base64_data) % 4 else 0
|
||||
padded_base64 = base64_data + '=' * padding
|
||||
file_bytes = base64.b64decode(padded_base64)
|
||||
file_bytes = await _decode_media_base64_limited(padded_base64)
|
||||
except binascii.Error as e:
|
||||
raise ValueError(f'Invalid base64 string: {str(e)}')
|
||||
else:
|
||||
raise ValueError('image对象出错')
|
||||
|
||||
# 设置 multipart/form-data 格式的文件
|
||||
if len(file_bytes) > _MAX_MEDIA_BYTES:
|
||||
raise ValueError('WeCom customer-service media exceeds the size limit')
|
||||
boundary = '-------------------------acebdf13572468'
|
||||
headers = {'Content-Type': f'multipart/form-data; boundary={boundary}'}
|
||||
body = (
|
||||
@@ -361,9 +456,9 @@ class WecomCSClient:
|
||||
)
|
||||
|
||||
# 上传文件
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with self._http_client_context() as client:
|
||||
response = await client.post(url, headers=headers, content=body)
|
||||
data = response.json()
|
||||
data = await httpclient.parse_json_response(response)
|
||||
if data['errcode'] == 40014 or data['errcode'] == 42001:
|
||||
self.access_token = await self.get_access_token(self.secret)
|
||||
media_id = await self.upload_to_work(image)
|
||||
@@ -374,16 +469,17 @@ class WecomCSClient:
|
||||
return media_id
|
||||
|
||||
async def download_image_to_bytes(self, url: str) -> bytes:
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.get(url)
|
||||
response.raise_for_status()
|
||||
return response.content
|
||||
async with self._http_client_context() as client:
|
||||
async with client.stream('GET', url) as response:
|
||||
response.raise_for_status()
|
||||
return await _read_httpx_media_limited(response)
|
||||
|
||||
# 进行media_id的获取
|
||||
async def get_media_id(self, image: platform_message.Image):
|
||||
media_id = await self.upload_to_work(image=image)
|
||||
return media_id
|
||||
|
||||
@_bounded_token_retry
|
||||
async def get_customer_info(self, external_userid: str) -> dict | None:
|
||||
"""
|
||||
Get customer information by external_userid with caching.
|
||||
@@ -398,6 +494,11 @@ class WecomCSClient:
|
||||
"""
|
||||
# Check cache first
|
||||
current_time = time.time()
|
||||
if current_time - self._customer_cache_cleanup_at >= 30:
|
||||
self._customer_cache_cleanup_at = current_time
|
||||
for user_id, (_, cached_time) in tuple(self._customer_cache.items()):
|
||||
if current_time - cached_time >= self._cache_ttl:
|
||||
self._customer_cache.pop(user_id, None)
|
||||
if external_userid in self._customer_cache:
|
||||
cached_info, cached_time = self._customer_cache[external_userid]
|
||||
if current_time - cached_time < self._cache_ttl:
|
||||
@@ -413,9 +514,9 @@ class WecomCSClient:
|
||||
'external_userid_list': [external_userid],
|
||||
}
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with self._http_client_context() as client:
|
||||
response = await client.post(url, json=payload)
|
||||
data = response.json()
|
||||
data = await httpclient.parse_json_response(response)
|
||||
|
||||
if data.get('errcode') in [40014, 42001]:
|
||||
self.access_token = await self.get_access_token(self.secret)
|
||||
@@ -431,5 +532,10 @@ class WecomCSClient:
|
||||
customer_info = customer_list[0]
|
||||
# Store in cache
|
||||
self._customer_cache[external_userid] = (customer_info, current_time)
|
||||
while len(self._customer_cache) > self._CUSTOMER_CACHE_MAX:
|
||||
self._customer_cache.pop(next(iter(self._customer_cache)), None)
|
||||
return customer_info
|
||||
return None
|
||||
|
||||
def clear(self) -> None:
|
||||
self._customer_cache.clear()
|
||||
|
||||
@@ -6,6 +6,56 @@ import json
|
||||
|
||||
from .errors import WeKnoraAPIError
|
||||
|
||||
_MAX_WENKORA_RESPONSE_BYTES = 1024 * 1024
|
||||
_MAX_WENKORA_STREAM_BYTES = 16 * 1024 * 1024
|
||||
_MAX_WENKORA_SSE_LINE_BYTES = 1024 * 1024
|
||||
|
||||
|
||||
async def _read_limited_response(response: httpx.Response) -> bytes:
|
||||
body = bytearray()
|
||||
async for chunk in response.aiter_bytes(chunk_size=8192):
|
||||
body.extend(chunk)
|
||||
if len(body) > _MAX_WENKORA_RESPONSE_BYTES:
|
||||
raise WeKnoraAPIError('WeKnora response exceeds the runtime limit')
|
||||
return bytes(body)
|
||||
|
||||
|
||||
async def _iter_sse_json(
|
||||
response: httpx.Response,
|
||||
) -> typing.AsyncGenerator[dict[str, typing.Any], None]:
|
||||
buffer = bytearray()
|
||||
total = 0
|
||||
async for chunk in response.aiter_bytes(chunk_size=8192):
|
||||
total += len(chunk)
|
||||
if total > _MAX_WENKORA_STREAM_BYTES:
|
||||
raise WeKnoraAPIError('WeKnora stream exceeds the runtime limit')
|
||||
buffer.extend(chunk)
|
||||
while b'\n' in buffer:
|
||||
raw_line, _, remainder = buffer.partition(b'\n')
|
||||
buffer = bytearray(remainder)
|
||||
if len(raw_line) > _MAX_WENKORA_SSE_LINE_BYTES:
|
||||
raise WeKnoraAPIError('WeKnora SSE event exceeds the runtime limit')
|
||||
line = raw_line.rstrip(b'\r').strip()
|
||||
if not line.startswith(b'data:'):
|
||||
continue
|
||||
try:
|
||||
data = json.loads(line[5:].strip())
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
if isinstance(data, dict):
|
||||
yield data
|
||||
if len(buffer) > _MAX_WENKORA_SSE_LINE_BYTES:
|
||||
raise WeKnoraAPIError('WeKnora SSE event exceeds the runtime limit')
|
||||
|
||||
line = bytes(buffer).rstrip(b'\r').strip()
|
||||
if line.startswith(b'data:'):
|
||||
try:
|
||||
data = json.loads(line[5:].strip())
|
||||
except json.JSONDecodeError:
|
||||
return
|
||||
if isinstance(data, dict):
|
||||
yield data
|
||||
|
||||
|
||||
class AsyncWeKnoraClient:
|
||||
"""WeKnora API 客户端"""
|
||||
@@ -39,19 +89,19 @@ class AsyncWeKnoraClient:
|
||||
if description:
|
||||
payload['description'] = description
|
||||
|
||||
response = await client.post(
|
||||
async with client.stream(
|
||||
'POST',
|
||||
'/sessions',
|
||||
headers={
|
||||
'X-API-Key': self.api_key,
|
||||
'Content-Type': 'application/json',
|
||||
},
|
||||
json=payload,
|
||||
)
|
||||
|
||||
if response.status_code not in (200, 201):
|
||||
raise WeKnoraAPIError(f'{response.status_code} {response.text}')
|
||||
|
||||
data = response.json()
|
||||
) as response:
|
||||
body = await _read_limited_response(response)
|
||||
if response.status_code not in (200, 201):
|
||||
raise WeKnoraAPIError(f'{response.status_code} {body.decode("utf-8", errors="replace")}')
|
||||
data = json.loads(body)
|
||||
return data['data']['id']
|
||||
|
||||
async def agent_chat(
|
||||
@@ -107,20 +157,13 @@ class AsyncWeKnoraClient:
|
||||
},
|
||||
json=payload,
|
||||
) as r:
|
||||
async for chunk in r.aiter_lines():
|
||||
if r.status_code != 200:
|
||||
raise WeKnoraAPIError(f'{r.status_code} {chunk}')
|
||||
if chunk.strip() == '':
|
||||
continue
|
||||
if chunk.startswith('data:'):
|
||||
try:
|
||||
data = json.loads(chunk[5:].strip())
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
yield data
|
||||
# 收到 error 事件后主动结束流，避免上层未 raise 时持续等待
|
||||
if data.get('response_type') == 'error':
|
||||
return
|
||||
if r.status_code != 200:
|
||||
body = await _read_limited_response(r)
|
||||
raise WeKnoraAPIError(f'{r.status_code} {body.decode("utf-8", errors="replace")}')
|
||||
async for data in _iter_sse_json(r):
|
||||
yield data
|
||||
if data.get('response_type') == 'error':
|
||||
return
|
||||
|
||||
async def knowledge_chat(
|
||||
self,
|
||||
@@ -164,17 +207,10 @@ class AsyncWeKnoraClient:
|
||||
},
|
||||
json=payload,
|
||||
) as r:
|
||||
async for chunk in r.aiter_lines():
|
||||
if r.status_code != 200:
|
||||
raise WeKnoraAPIError(f'{r.status_code} {chunk}')
|
||||
if chunk.strip() == '':
|
||||
continue
|
||||
if chunk.startswith('data:'):
|
||||
try:
|
||||
data = json.loads(chunk[5:].strip())
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
yield data
|
||||
# 收到 error 事件后主动结束流，避免上层未 raise 时持续等待
|
||||
if data.get('response_type') == 'error':
|
||||
return
|
||||
if r.status_code != 200:
|
||||
body = await _read_limited_response(r)
|
||||
raise WeKnoraAPIError(f'{r.status_code} {body.decode("utf-8", errors="replace")}')
|
||||
async for data in _iter_sse_json(r):
|
||||
yield data
|
||||
if data.get('response_type') == 'error':
|
||||
return
|
||||
|
||||
@@ -10,9 +10,11 @@ import uuid
|
||||
from quart.typing import RouteCallable
|
||||
|
||||
from ....utils import constants
|
||||
from ....utils import bounded_executor
|
||||
from ....workspace.collaboration import MembershipPermissionError, WorkspaceCollaborationError
|
||||
from ....workspace.errors import WorkspaceNotFoundError
|
||||
from ....cloud.entitlements import EntitlementUnavailableError
|
||||
from ....core.errors import TaskCapacityError
|
||||
from ..authz import AuthorizationError, Permission, permissions_for_role, require_permission
|
||||
from ..context import PrincipalContext, PrincipalType, RequestContext, WorkspaceContext
|
||||
|
||||
@@ -185,16 +187,27 @@ class RouterGroup(abc.ABC):
|
||||
|
||||
try:
|
||||
if request_context is not None:
|
||||
persistence_mgr = getattr(self.ap, 'persistence_mgr', None)
|
||||
tenant_scope_descriptor = getattr(type(persistence_mgr), 'tenant_scope', None)
|
||||
if callable(tenant_scope_descriptor):
|
||||
# Authorization discovery is complete. Carry the
|
||||
# trusted Workspace identity across the handler, but
|
||||
# do not reserve a database connection while it waits
|
||||
# on providers, runtimes, uploads, or streamed clients.
|
||||
# Services that need atomic writes open a short UoW.
|
||||
async with persistence_mgr.tenant_scope(request_context.workspace_uuid):
|
||||
return await f(*args, **kwargs)
|
||||
with bounded_executor.blocking_work_scope(request_context.workspace_uuid):
|
||||
persistence_mgr = getattr(
|
||||
self.ap,
|
||||
'persistence_mgr',
|
||||
None,
|
||||
)
|
||||
tenant_scope_descriptor = getattr(
|
||||
type(persistence_mgr),
|
||||
'tenant_scope',
|
||||
None,
|
||||
)
|
||||
if callable(tenant_scope_descriptor):
|
||||
# Authorization discovery is complete. Carry
|
||||
# the trusted Workspace identity across the
|
||||
# handler, but do not reserve a database
|
||||
# connection while it waits on providers,
|
||||
# runtimes, uploads, or streamed clients.
|
||||
# Services that need atomic writes open a UoW.
|
||||
async with persistence_mgr.tenant_scope(request_context.workspace_uuid):
|
||||
return await f(*args, **kwargs)
|
||||
return await f(*args, **kwargs)
|
||||
return await f(*args, **kwargs)
|
||||
|
||||
except Exception as e: # 自动 500
|
||||
@@ -206,6 +219,17 @@ class RouterGroup(abc.ABC):
|
||||
return self.http_status(403, e.code, str(e))
|
||||
if isinstance(e, WorkspaceCollaborationError):
|
||||
return self.http_status(400, e.code, str(e))
|
||||
if isinstance(e, TaskCapacityError):
|
||||
return self.http_status(429, 'task_capacity_exceeded', str(e))
|
||||
if isinstance(
|
||||
e,
|
||||
bounded_executor.BlockingWorkCapacityError,
|
||||
):
|
||||
return self.http_status(
|
||||
429,
|
||||
'blocking_work_capacity_exceeded',
|
||||
str(e),
|
||||
)
|
||||
request_id = self.request_id()
|
||||
logger = getattr(self.ap, 'logger', self.quart_app.logger)
|
||||
logger.error(
|
||||
|
||||
@@ -11,6 +11,7 @@ from ....context import ExecutionContext, RequestContext
|
||||
from ......core import taskmgr
|
||||
from ......entity.persistence import metadata as persistence_metadata
|
||||
from ......workspace.errors import WorkspaceError, WorkspaceNotFoundError
|
||||
from ......utils import httpclient
|
||||
from langbot_plugin.runtime.plugin.mgr import PluginInstallSource
|
||||
|
||||
LANGRAG_PLUGIN_AUTHOR = 'langbot-team'
|
||||
@@ -162,10 +163,15 @@ class KnowledgeMigrationRouterGroup(group.RouterGroup):
|
||||
self.ap.logger.info(f'RAG migration: installing plugin {plugin_id} from marketplace...')
|
||||
task_context.trace(f'Installing plugin {plugin_id} from marketplace...')
|
||||
|
||||
async with httpx.AsyncClient(trust_env=True, timeout=15) as client:
|
||||
async with httpx.AsyncClient(
|
||||
trust_env=True,
|
||||
timeout=15,
|
||||
event_hooks=httpclient.httpx_response_limit_hooks(),
|
||||
) as client:
|
||||
resp = await client.get(f'{space_url}/api/v1/marketplace/plugins/{p_author}/{p_name}')
|
||||
resp.raise_for_status()
|
||||
p_data = resp.json().get('data', {}).get('plugin', {})
|
||||
response_data = await httpclient.parse_json_response(resp)
|
||||
p_data = response_data.get('data', {}).get('plugin', {})
|
||||
p_version = p_data.get('latest_version')
|
||||
if not p_version:
|
||||
raise Exception(f'Could not determine latest version for {plugin_id}')
|
||||
|
||||
@@ -20,8 +20,9 @@ import httpx
|
||||
import quart
|
||||
|
||||
from ... import group
|
||||
from ......utils import paths
|
||||
from ......utils import httpclient, paths
|
||||
from ......platform.sources.websocket_manager import WebSocketScope, is_valid_session_id, ws_connection_manager
|
||||
from .websocket_chat import create_scoped_duplex_tasks, wait_for_duplex_tasks
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
_AUTH_TIMEOUT_SECONDS = 10.0
|
||||
@@ -103,7 +104,7 @@ class EmbedRouterGroup(group.RouterGroup):
|
||||
"""Require the embed session token as the first WebSocket frame."""
|
||||
|
||||
raw_message = await asyncio.wait_for(quart.websocket.receive(), timeout=_AUTH_TIMEOUT_SECONDS)
|
||||
payload = json.loads(raw_message)
|
||||
payload = await asyncio.to_thread(json.loads, raw_message)
|
||||
if not isinstance(payload, dict) or payload.get('type') != 'authenticate':
|
||||
raise ValueError('Authentication is required')
|
||||
token = str(payload.get('token') or '')
|
||||
@@ -160,12 +161,12 @@ class EmbedRouterGroup(group.RouterGroup):
|
||||
ts = time.time()
|
||||
return self.success(data={'token': f'{ts}.dummy'})
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
async with httpx.AsyncClient(event_hooks=httpclient.httpx_response_limit_hooks()) as client:
|
||||
resp = await client.post(
|
||||
'https://challenges.cloudflare.com/turnstile/v0/siteverify',
|
||||
data={'secret': secret, 'response': token},
|
||||
)
|
||||
result = resp.json()
|
||||
result = await httpclient.parse_json_response(resp)
|
||||
|
||||
if not result.get('success'):
|
||||
return self.http_status(403, -1, 'Turnstile verification failed')
|
||||
@@ -371,6 +372,21 @@ class EmbedRouterGroup(group.RouterGroup):
|
||||
session_type=session_type,
|
||||
session_id=session_id,
|
||||
metadata={'user_agent': quart.websocket.headers.get('User-Agent', '')},
|
||||
send_queue_size=(
|
||||
self.ap.instance_config.data.get('system', {})
|
||||
.get('websocket_retention', {})
|
||||
.get('send_queue_size', 100)
|
||||
),
|
||||
max_connections=(
|
||||
self.ap.instance_config.data.get('system', {})
|
||||
.get('websocket_retention', {})
|
||||
.get('max_connections', 1024)
|
||||
),
|
||||
max_connections_per_workspace=(
|
||||
self.ap.instance_config.data.get('system', {})
|
||||
.get('websocket_retention', {})
|
||||
.get('max_connections_per_workspace', 32)
|
||||
),
|
||||
)
|
||||
|
||||
await quart.websocket.send(
|
||||
@@ -390,13 +406,19 @@ class EmbedRouterGroup(group.RouterGroup):
|
||||
f'(bot={bot_uuid}, pipeline={pipeline_uuid}, session_type={session_type})'
|
||||
)
|
||||
|
||||
receive_task = asyncio.create_task(
|
||||
self._handle_receive(connection, websocket_adapter, runtime_bot, pipeline_uuid)
|
||||
receive_task, send_task = create_scoped_duplex_tasks(
|
||||
self._handle_receive(
|
||||
connection,
|
||||
websocket_adapter,
|
||||
runtime_bot,
|
||||
pipeline_uuid,
|
||||
),
|
||||
self._handle_send(connection),
|
||||
runtime_bot.execution_context.workspace_uuid,
|
||||
)
|
||||
send_task = asyncio.create_task(self._handle_send(connection))
|
||||
|
||||
try:
|
||||
await asyncio.gather(receive_task, send_task)
|
||||
await wait_for_duplex_tasks(receive_task, send_task)
|
||||
except Exception as e:
|
||||
logger.error(f'Embed WebSocket task error: {e}')
|
||||
finally:
|
||||
@@ -418,7 +440,7 @@ class EmbedRouterGroup(group.RouterGroup):
|
||||
await ws_connection_manager.update_activity(connection.connection_id)
|
||||
|
||||
try:
|
||||
data = json.loads(message)
|
||||
data = await asyncio.to_thread(json.loads, message)
|
||||
message_type = data.get('type', 'message')
|
||||
|
||||
if message_type == 'ping':
|
||||
@@ -442,13 +464,20 @@ class EmbedRouterGroup(group.RouterGroup):
|
||||
logger.error(f'Embed receive error: {e}', exc_info=True)
|
||||
finally:
|
||||
connection.is_active = False
|
||||
try:
|
||||
connection.send_queue.put_nowait(None)
|
||||
except asyncio.QueueFull:
|
||||
pass
|
||||
|
||||
async def _handle_send(self, connection):
|
||||
try:
|
||||
while connection.is_active or not connection.send_queue.empty():
|
||||
try:
|
||||
message = await asyncio.wait_for(connection.send_queue.get(), timeout=1.0)
|
||||
await quart.websocket.send(json.dumps(message))
|
||||
if message is None:
|
||||
break
|
||||
encoded = await asyncio.to_thread(json.dumps, message)
|
||||
await quart.websocket.send(encoded)
|
||||
except asyncio.TimeoutError:
|
||||
continue
|
||||
except Exception as e:
|
||||
|
||||
@@ -6,6 +6,7 @@ import asyncio
|
||||
import datetime
|
||||
import json
|
||||
import logging
|
||||
import typing
|
||||
import uuid
|
||||
|
||||
import quart
|
||||
@@ -15,9 +16,64 @@ from ....context import PrincipalContext, PrincipalType, RequestContext, Workspa
|
||||
from ... import group
|
||||
from ......core.task_boundary import run_in_workspace_uow
|
||||
from ......platform.sources.websocket_manager import WebSocketScope, ws_connection_manager
|
||||
from ......utils import bounded_executor
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
_AUTH_TIMEOUT_SECONDS = 10.0
|
||||
_DUPLEX_DRAIN_TIMEOUT_SECONDS = 0.25
|
||||
|
||||
|
||||
def create_scoped_duplex_tasks(
|
||||
receive_coro: typing.Coroutine[typing.Any, typing.Any, None],
|
||||
send_coro: typing.Coroutine[typing.Any, typing.Any, None],
|
||||
workspace_uuid: str,
|
||||
) -> tuple[asyncio.Task[None], asyncio.Task[None]]:
|
||||
"""Create both socket directions under one trusted Workspace budget."""
|
||||
|
||||
return (
|
||||
asyncio.create_task(
|
||||
bounded_executor.run_in_blocking_work_scope(
|
||||
receive_coro,
|
||||
workspace_uuid,
|
||||
)
|
||||
),
|
||||
asyncio.create_task(
|
||||
bounded_executor.run_in_blocking_work_scope(
|
||||
send_coro,
|
||||
workspace_uuid,
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
async def wait_for_duplex_tasks(
|
||||
receive_task: asyncio.Task,
|
||||
send_task: asyncio.Task,
|
||||
) -> None:
|
||||
"""Stop the peer direction as soon as either socket task terminates."""
|
||||
|
||||
try:
|
||||
done, _ = await asyncio.wait(
|
||||
{receive_task, send_task},
|
||||
return_when=asyncio.FIRST_COMPLETED,
|
||||
)
|
||||
# A receive task may enqueue a terminal authorization/error frame and
|
||||
# then finish. Give the sender a short deterministic drain window
|
||||
# instead of cancelling it before that frame reaches the client.
|
||||
if receive_task in done and not send_task.done():
|
||||
await asyncio.wait(
|
||||
{send_task},
|
||||
timeout=_DUPLEX_DRAIN_TIMEOUT_SECONDS,
|
||||
)
|
||||
finally:
|
||||
for task in (receive_task, send_task):
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
await asyncio.gather(
|
||||
receive_task,
|
||||
send_task,
|
||||
return_exceptions=True,
|
||||
)
|
||||
|
||||
|
||||
@group.group_class('websocket_chat', '/api/v1/pipelines/<pipeline_uuid>/ws')
|
||||
@@ -32,7 +88,7 @@ class WebSocketChatRouterGroup(group.RouterGroup):
|
||||
"""
|
||||
|
||||
raw_message = await asyncio.wait_for(quart.websocket.receive(), timeout=_AUTH_TIMEOUT_SECONDS)
|
||||
payload = json.loads(raw_message)
|
||||
payload = await asyncio.to_thread(json.loads, raw_message)
|
||||
if not isinstance(payload, dict) or payload.get('type') != 'authenticate':
|
||||
raise ValueError('Authentication is required')
|
||||
|
||||
@@ -155,6 +211,21 @@ class WebSocketChatRouterGroup(group.RouterGroup):
|
||||
pipeline_uuid=pipeline_uuid,
|
||||
session_type=session_type,
|
||||
metadata={'user_agent': quart.websocket.headers.get('User-Agent', '')},
|
||||
send_queue_size=(
|
||||
self.ap.instance_config.data.get('system', {})
|
||||
.get('websocket_retention', {})
|
||||
.get('send_queue_size', 100)
|
||||
),
|
||||
max_connections=(
|
||||
self.ap.instance_config.data.get('system', {})
|
||||
.get('websocket_retention', {})
|
||||
.get('max_connections', 1024)
|
||||
),
|
||||
max_connections_per_workspace=(
|
||||
self.ap.instance_config.data.get('system', {})
|
||||
.get('websocket_retention', {})
|
||||
.get('max_connections_per_workspace', 32)
|
||||
),
|
||||
)
|
||||
|
||||
await quart.websocket.send(
|
||||
@@ -175,17 +246,18 @@ class WebSocketChatRouterGroup(group.RouterGroup):
|
||||
f'session_type={session_type})'
|
||||
)
|
||||
|
||||
receive_task = asyncio.create_task(
|
||||
receive_task, send_task = create_scoped_duplex_tasks(
|
||||
self._handle_receive(
|
||||
connection,
|
||||
websocket_adapter,
|
||||
request_context,
|
||||
token,
|
||||
)
|
||||
),
|
||||
self._handle_send(connection),
|
||||
request_context.workspace_uuid,
|
||||
)
|
||||
send_task = asyncio.create_task(self._handle_send(connection))
|
||||
try:
|
||||
await asyncio.gather(receive_task, send_task)
|
||||
await wait_for_duplex_tasks(receive_task, send_task)
|
||||
except Exception as exc:
|
||||
logger.error(f'WebSocket task execution error: {exc}')
|
||||
finally:
|
||||
@@ -310,7 +382,7 @@ class WebSocketChatRouterGroup(group.RouterGroup):
|
||||
await ws_connection_manager.update_activity(connection.connection_id)
|
||||
|
||||
try:
|
||||
data = json.loads(message)
|
||||
data = await asyncio.to_thread(json.loads, message)
|
||||
message_type = data.get('type', 'message')
|
||||
if message_type == 'ping':
|
||||
await connection.send_queue.put(
|
||||
@@ -334,13 +406,20 @@ class WebSocketChatRouterGroup(group.RouterGroup):
|
||||
logger.error('Dashboard WebSocket receive error', exc_info=True)
|
||||
finally:
|
||||
connection.is_active = False
|
||||
try:
|
||||
connection.send_queue.put_nowait(None)
|
||||
except asyncio.QueueFull:
|
||||
pass
|
||||
|
||||
async def _handle_send(self, connection):
|
||||
try:
|
||||
while connection.is_active or not connection.send_queue.empty():
|
||||
try:
|
||||
message = await asyncio.wait_for(connection.send_queue.get(), timeout=1.0)
|
||||
await quart.websocket.send(json.dumps(message))
|
||||
if message is None:
|
||||
break
|
||||
encoded = await asyncio.to_thread(json.dumps, message)
|
||||
await quart.websocket.send(encoded)
|
||||
except asyncio.TimeoutError:
|
||||
continue
|
||||
except Exception:
|
||||
|
||||
@@ -6,7 +6,8 @@ import quart
|
||||
|
||||
from langbot.pkg.api.http.authz import Permission
|
||||
from langbot.pkg.api.http.context import RequestContext
|
||||
from langbot.pkg.utils import importutil
|
||||
from langbot.pkg.core.errors import TaskCapacityError
|
||||
from langbot.pkg.utils import httpclient, importutil
|
||||
|
||||
from ... import group
|
||||
|
||||
@@ -71,6 +72,64 @@ def _pop_owned_session(
|
||||
return sessions.pop(session_id, None)
|
||||
|
||||
|
||||
_MAX_ADAPTER_SESSIONS = 100
|
||||
_MAX_ADAPTER_SESSIONS_PER_WORKSPACE = 10
|
||||
|
||||
|
||||
def _start_adapter_session_task(
|
||||
ap,
|
||||
coro,
|
||||
*,
|
||||
adapter: str,
|
||||
session_id: str,
|
||||
request_context: RequestContext,
|
||||
) -> asyncio.Task | None:
|
||||
"""Attach one credential exchange to tenant admission and app shutdown."""
|
||||
|
||||
try:
|
||||
wrapper = ap.task_mgr.create_user_task(
|
||||
coro,
|
||||
kind='platform-adapter-credential-exchange',
|
||||
name=f'{adapter}-credential-{session_id}',
|
||||
label=f'{adapter} credential exchange',
|
||||
instance_uuid=request_context.instance_uuid,
|
||||
workspace_uuid=request_context.workspace_uuid,
|
||||
placement_generation=request_context.placement_generation,
|
||||
)
|
||||
except TaskCapacityError:
|
||||
coro.close()
|
||||
return None
|
||||
return wrapper.task
|
||||
|
||||
|
||||
def _make_room_for_session(
|
||||
sessions: dict[str, dict],
|
||||
request_context: RequestContext,
|
||||
) -> None:
|
||||
"""Bound credential-exchange sessions globally and per workspace."""
|
||||
|
||||
workspace_uuid = request_context.workspace_uuid
|
||||
owned = [
|
||||
(session_id, session)
|
||||
for session_id, session in sessions.items()
|
||||
if getattr(session.get('scope'), 'workspace_uuid', None) == workspace_uuid
|
||||
]
|
||||
evict_workspace_session = len(owned) >= _MAX_ADAPTER_SESSIONS_PER_WORKSPACE
|
||||
evict_global_session = len(sessions) >= _MAX_ADAPTER_SESSIONS
|
||||
if not evict_workspace_session and not evict_global_session:
|
||||
return
|
||||
|
||||
candidates = owned if evict_workspace_session else list(sessions.items())
|
||||
session_id, _ = min(
|
||||
candidates,
|
||||
key=lambda item: float(item[1].get('created_at', 0.0)),
|
||||
)
|
||||
session = sessions.pop(session_id, None)
|
||||
task = session.get('task') if session is not None else None
|
||||
if task is not None and not task.done():
|
||||
task.cancel()
|
||||
|
||||
|
||||
def _decrypt_qqofficial_secret(encrypted_b64: str, key: bytes) -> str:
|
||||
"""Decrypt the AppSecret returned by the QQ Official QR binding endpoint.
|
||||
|
||||
@@ -173,6 +232,7 @@ class AdaptersRouterGroup(group.RouterGroup):
|
||||
'created_at': time.time(),
|
||||
}
|
||||
_bind_session_scope(session, request_context)
|
||||
_make_room_for_session(_create_app_sessions, request_context)
|
||||
_create_app_sessions[session_id] = session
|
||||
|
||||
def on_qr_code(info):
|
||||
@@ -204,7 +264,16 @@ class AdaptersRouterGroup(group.RouterGroup):
|
||||
session['status'] = 'error'
|
||||
session['error'] = str(e)
|
||||
|
||||
task = asyncio.create_task(run_registration())
|
||||
task = _start_adapter_session_task(
|
||||
self.ap,
|
||||
run_registration(),
|
||||
adapter='lark',
|
||||
session_id=session_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
if task is None:
|
||||
_create_app_sessions.pop(session_id, None)
|
||||
return self.http_status(429, -1, 'Too many active credential exchanges')
|
||||
session['task'] = task
|
||||
|
||||
# Wait for QR code to be ready (max 10 seconds)
|
||||
@@ -308,6 +377,7 @@ class AdaptersRouterGroup(group.RouterGroup):
|
||||
'created_at': time.time(),
|
||||
}
|
||||
_bind_session_scope(session, request_context)
|
||||
_make_room_for_session(_weixin_login_sessions, request_context)
|
||||
_weixin_login_sessions[session_id] = session
|
||||
|
||||
client = OpenClawWeixinClient(
|
||||
@@ -346,7 +416,16 @@ class AdaptersRouterGroup(group.RouterGroup):
|
||||
finally:
|
||||
await client.close()
|
||||
|
||||
task = asyncio.create_task(run_login())
|
||||
task = _start_adapter_session_task(
|
||||
self.ap,
|
||||
run_login(),
|
||||
adapter='weixin',
|
||||
session_id=session_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
if task is None:
|
||||
_weixin_login_sessions.pop(session_id, None)
|
||||
return self.http_status(429, -1, 'Too many active credential exchanges')
|
||||
session['task'] = task
|
||||
|
||||
# Wait for QR code to be ready (max 10 seconds)
|
||||
@@ -459,6 +538,7 @@ class AdaptersRouterGroup(group.RouterGroup):
|
||||
'interval': 5,
|
||||
}
|
||||
_bind_session_scope(session, request_context)
|
||||
_make_room_for_session(_dingtalk_sessions, request_context)
|
||||
_dingtalk_sessions[session_id] = session
|
||||
|
||||
async def run_device_flow():
|
||||
@@ -471,7 +551,7 @@ class AdaptersRouterGroup(group.RouterGroup):
|
||||
json={'source': 'langbot'},
|
||||
) as resp:
|
||||
try:
|
||||
data = await resp.json()
|
||||
data = await httpclient.read_json_limited(resp)
|
||||
except (aiohttp.ContentTypeError, ValueError):
|
||||
session['status'] = 'error'
|
||||
session['error'] = 'Invalid response from DingTalk service'
|
||||
@@ -488,7 +568,7 @@ class AdaptersRouterGroup(group.RouterGroup):
|
||||
json={'nonce': nonce},
|
||||
) as resp:
|
||||
try:
|
||||
data = await resp.json()
|
||||
data = await httpclient.read_json_limited(resp)
|
||||
except (aiohttp.ContentTypeError, ValueError):
|
||||
session['status'] = 'error'
|
||||
session['error'] = 'Invalid response from DingTalk service'
|
||||
@@ -519,7 +599,7 @@ class AdaptersRouterGroup(group.RouterGroup):
|
||||
json={'device_code': device_code},
|
||||
) as poll_resp:
|
||||
try:
|
||||
poll_data = await poll_resp.json()
|
||||
poll_data = await httpclient.read_json_limited(poll_resp)
|
||||
except (aiohttp.ContentTypeError, ValueError):
|
||||
continue
|
||||
|
||||
@@ -555,7 +635,16 @@ class AdaptersRouterGroup(group.RouterGroup):
|
||||
session['status'] = 'error'
|
||||
session['error'] = str(e)
|
||||
|
||||
task = asyncio.create_task(run_device_flow())
|
||||
task = _start_adapter_session_task(
|
||||
self.ap,
|
||||
run_device_flow(),
|
||||
adapter='dingtalk',
|
||||
session_id=session_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
if task is None:
|
||||
_dingtalk_sessions.pop(session_id, None)
|
||||
return self.http_status(429, -1, 'Too many active credential exchanges')
|
||||
session['task'] = task
|
||||
|
||||
# Wait for QR code to be ready (max 10 seconds)
|
||||
@@ -665,6 +754,7 @@ class AdaptersRouterGroup(group.RouterGroup):
|
||||
'task': None,
|
||||
}
|
||||
_bind_session_scope(session, request_context)
|
||||
_make_room_for_session(_wecombot_sessions, request_context)
|
||||
_wecombot_sessions[session_id] = session
|
||||
|
||||
async def run_qr_flow():
|
||||
@@ -676,7 +766,7 @@ class AdaptersRouterGroup(group.RouterGroup):
|
||||
f'{WECOM_QC_GENERATE_URL}?source=langbot&plat=0',
|
||||
) as resp:
|
||||
try:
|
||||
data = await resp.json()
|
||||
data = await httpclient.read_json_limited(resp)
|
||||
except (aiohttp.ContentTypeError, ValueError):
|
||||
session['status'] = 'error'
|
||||
session['error'] = 'Invalid response from WeCom service'
|
||||
@@ -703,7 +793,7 @@ class AdaptersRouterGroup(group.RouterGroup):
|
||||
f'{WECOM_QC_QUERY_URL}?scode={scode}',
|
||||
) as poll_resp:
|
||||
try:
|
||||
poll_data = await poll_resp.json()
|
||||
poll_data = await httpclient.read_json_limited(poll_resp)
|
||||
except (aiohttp.ContentTypeError, ValueError):
|
||||
continue
|
||||
|
||||
@@ -730,7 +820,16 @@ class AdaptersRouterGroup(group.RouterGroup):
|
||||
session['status'] = 'error'
|
||||
session['error'] = str(e)
|
||||
|
||||
task = asyncio.create_task(run_qr_flow())
|
||||
task = _start_adapter_session_task(
|
||||
self.ap,
|
||||
run_qr_flow(),
|
||||
adapter='wecombot',
|
||||
session_id=session_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
if task is None:
|
||||
_wecombot_sessions.pop(session_id, None)
|
||||
return self.http_status(429, -1, 'Too many active credential exchanges')
|
||||
session['task'] = task
|
||||
|
||||
# Wait for QR code to be ready (max 10 seconds)
|
||||
@@ -852,6 +951,7 @@ class AdaptersRouterGroup(group.RouterGroup):
|
||||
'interval': 2,
|
||||
}
|
||||
_bind_session_scope(session, request_context)
|
||||
_make_room_for_session(_qqofficial_sessions, request_context)
|
||||
_qqofficial_sessions[session_id] = session
|
||||
|
||||
async def run_qr_binding():
|
||||
@@ -865,7 +965,7 @@ class AdaptersRouterGroup(group.RouterGroup):
|
||||
headers={'Accept': 'application/json'},
|
||||
) as resp:
|
||||
try:
|
||||
data = await resp.json(content_type=None)
|
||||
data = await httpclient.read_json_limited(resp)
|
||||
except (aiohttp.ContentTypeError, ValueError):
|
||||
session['status'] = 'error'
|
||||
session['error'] = 'Invalid response from QQ bind service'
|
||||
@@ -903,7 +1003,7 @@ class AdaptersRouterGroup(group.RouterGroup):
|
||||
headers={'Accept': 'application/json'},
|
||||
) as poll_resp:
|
||||
try:
|
||||
poll_data = await poll_resp.json(content_type=None)
|
||||
poll_data = await httpclient.read_json_limited(poll_resp)
|
||||
except (aiohttp.ContentTypeError, ValueError):
|
||||
continue
|
||||
|
||||
@@ -956,7 +1056,16 @@ class AdaptersRouterGroup(group.RouterGroup):
|
||||
session['status'] = 'error'
|
||||
session['error'] = str(e)
|
||||
|
||||
task = asyncio.create_task(run_qr_binding())
|
||||
task = _start_adapter_session_task(
|
||||
self.ap,
|
||||
run_qr_binding(),
|
||||
adapter='qqofficial',
|
||||
session_id=session_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
if task is None:
|
||||
_qqofficial_sessions.pop(session_id, None)
|
||||
return self.http_status(429, -1, 'Too many active credential exchanges')
|
||||
session['task'] = task
|
||||
|
||||
# Wait up to 10s for the QR URL to be ready before responding.
|
||||
|
||||
@@ -1,16 +1,15 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import collections.abc
|
||||
import copy
|
||||
import io
|
||||
import quart
|
||||
import re
|
||||
import httpx
|
||||
import uuid
|
||||
import os
|
||||
import zipfile
|
||||
import yaml
|
||||
from urllib.parse import urlparse
|
||||
import posixpath
|
||||
import sqlalchemy
|
||||
@@ -23,6 +22,8 @@ from ...context import ExecutionContext, RequestContext
|
||||
from .. import group
|
||||
from .....workspace.errors import WorkspaceNotFoundError
|
||||
from .....plugin.github import validate_github_plugin_install_info
|
||||
from .....plugin.archive import inspect_plugin_archive_metadata
|
||||
from .....utils import httpclient
|
||||
from langbot_plugin.runtime.plugin.mgr import PluginInstallSource
|
||||
|
||||
|
||||
@@ -548,7 +549,7 @@ class PluginsRouterGroup(group.RouterGroup):
|
||||
icon_base64 = icon_data['plugin_icon_base64']
|
||||
mime_type = icon_data['mime_type']
|
||||
|
||||
icon_data = base64.b64decode(icon_base64)
|
||||
icon_data = await asyncio.to_thread(base64.b64decode, icon_base64)
|
||||
|
||||
return quart.Response(icon_data, mimetype=mime_type)
|
||||
|
||||
@@ -566,7 +567,10 @@ class PluginsRouterGroup(group.RouterGroup):
|
||||
asset_data = await self.ap.plugin_connector.get_plugin_assets(author, plugin_name, asset_path)
|
||||
if not asset_data.get('asset_base64'):
|
||||
return quart.Response('Asset not found', status=404)
|
||||
asset_bytes = base64.b64decode(asset_data['asset_base64'])
|
||||
asset_bytes = await asyncio.to_thread(
|
||||
base64.b64decode,
|
||||
asset_data['asset_base64'],
|
||||
)
|
||||
mime_type = asset_data['mime_type']
|
||||
resp = quart.Response(asset_bytes, mimetype=mime_type)
|
||||
# CSP for HTML pages served to sandboxed iframes (opaque origin).
|
||||
@@ -662,10 +666,11 @@ class PluginsRouterGroup(group.RouterGroup):
|
||||
trust_env=True,
|
||||
follow_redirects=True,
|
||||
timeout=10,
|
||||
event_hooks=httpclient.httpx_response_limit_hooks(),
|
||||
) as client:
|
||||
response = await client.get(url)
|
||||
response.raise_for_status()
|
||||
releases = response.json()
|
||||
releases = await httpclient.parse_json_response(response)
|
||||
|
||||
# Format releases data for frontend
|
||||
formatted_releases = []
|
||||
@@ -716,12 +721,13 @@ class PluginsRouterGroup(group.RouterGroup):
|
||||
trust_env=True,
|
||||
follow_redirects=True,
|
||||
timeout=10,
|
||||
event_hooks=httpclient.httpx_response_limit_hooks(),
|
||||
) as client:
|
||||
response = await client.get(
|
||||
url,
|
||||
)
|
||||
response.raise_for_status()
|
||||
release = response.json()
|
||||
release = await httpclient.parse_json_response(response)
|
||||
|
||||
# Format assets data for frontend
|
||||
formatted_assets = []
|
||||
@@ -902,51 +908,29 @@ class PluginsRouterGroup(group.RouterGroup):
|
||||
|
||||
file_bytes = file.read()
|
||||
try:
|
||||
with zipfile.ZipFile(io.BytesIO(file_bytes)) as zf:
|
||||
names = [name for name in zf.namelist() if not name.endswith('/')]
|
||||
manifest_name = next(
|
||||
(
|
||||
name
|
||||
for name in names
|
||||
if name.replace('\\', '/').strip('/').lower() in ('manifest.yaml', 'manifest.yml')
|
||||
),
|
||||
None,
|
||||
)
|
||||
if manifest_name is None:
|
||||
return self.http_status(400, -1, 'manifest.yaml is required')
|
||||
manifest, requirements, names = await asyncio.to_thread(
|
||||
inspect_plugin_archive_metadata,
|
||||
file_bytes,
|
||||
)
|
||||
spec = manifest.get('spec') or {}
|
||||
components = spec.get('components') or {}
|
||||
component_counts = self._count_plugin_components(components, names)
|
||||
component_types = list(component_counts.keys())
|
||||
|
||||
manifest = yaml.safe_load(zf.read(manifest_name).decode('utf-8')) or {}
|
||||
requirements: list[str] = []
|
||||
requirements_name = next(
|
||||
(name for name in names if name.replace('\\', '/').strip('/').lower() == 'requirements.txt'),
|
||||
None,
|
||||
)
|
||||
if requirements_name is not None:
|
||||
requirements = [
|
||||
line.strip()
|
||||
for line in zf.read(requirements_name).decode('utf-8', errors='ignore').splitlines()
|
||||
if line.strip() and not line.strip().startswith('#')
|
||||
]
|
||||
|
||||
spec = manifest.get('spec') or {}
|
||||
components = spec.get('components') or {}
|
||||
component_counts = self._count_plugin_components(components, names)
|
||||
component_types = list(component_counts.keys())
|
||||
|
||||
return self.success(
|
||||
data={
|
||||
'filename': file.filename or 'local plugin',
|
||||
'size': len(file_bytes),
|
||||
'manifest': manifest,
|
||||
'metadata': manifest.get('metadata') or {},
|
||||
'component_types': component_types,
|
||||
'component_counts': component_counts,
|
||||
'requirements': requirements,
|
||||
'file_count': len(names),
|
||||
}
|
||||
)
|
||||
except zipfile.BadZipFile:
|
||||
return self.http_status(400, -1, 'invalid .lbpkg file')
|
||||
return self.success(
|
||||
data={
|
||||
'filename': file.filename or 'local plugin',
|
||||
'size': len(file_bytes),
|
||||
'manifest': manifest,
|
||||
'metadata': manifest.get('metadata') or {},
|
||||
'component_types': component_types,
|
||||
'component_counts': component_counts,
|
||||
'requirements': requirements,
|
||||
'file_count': len(names),
|
||||
}
|
||||
)
|
||||
except (zipfile.BadZipFile, ValueError) as exc:
|
||||
return self.http_status(400, -1, str(exc) or 'invalid .lbpkg file')
|
||||
except Exception:
|
||||
raise
|
||||
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
import asyncio
|
||||
import base64
|
||||
|
||||
import quart
|
||||
@@ -59,7 +60,14 @@ class SurveyRouterGroup(group.RouterGroup):
|
||||
continue
|
||||
try:
|
||||
payload = data_url.split(',', 1)[1]
|
||||
if len(base64.b64decode(payload, validate=True)) > 1024 * 1024:
|
||||
if len(payload) > 4 * ((1024 * 1024 + 2) // 3) + 4:
|
||||
return self.fail(5, 'attachment too large')
|
||||
decoded = await asyncio.to_thread(
|
||||
base64.b64decode,
|
||||
payload,
|
||||
validate=True,
|
||||
)
|
||||
if len(decoded) > 1024 * 1024:
|
||||
return self.fail(5, 'attachment too large')
|
||||
except Exception:
|
||||
return self.fail(5, 'attachment too large')
|
||||
|
||||
@@ -4,6 +4,7 @@ import quart
|
||||
import traceback
|
||||
|
||||
from .. import group
|
||||
from .....utils import bounded_executor
|
||||
|
||||
|
||||
@group.group_class('webhooks', '/bots')
|
||||
@@ -55,19 +56,26 @@ class WebhookRouterGroup(group.RouterGroup):
|
||||
request=quart.request,
|
||||
)
|
||||
|
||||
persistence_mgr = self.ap.persistence_mgr
|
||||
cloud_runtime = getattr(getattr(persistence_mgr, 'mode', None), 'value', None) == 'cloud_runtime'
|
||||
if cloud_runtime:
|
||||
tenant_scope = getattr(persistence_mgr, 'tenant_scope', None)
|
||||
if not callable(tenant_scope):
|
||||
raise RuntimeError('Cloud webhook dispatch requires an explicit tenant scope')
|
||||
async with tenant_scope(runtime_bot.workspace_uuid):
|
||||
with bounded_executor.blocking_work_scope(runtime_bot.workspace_uuid):
|
||||
persistence_mgr = self.ap.persistence_mgr
|
||||
cloud_runtime = getattr(getattr(persistence_mgr, 'mode', None), 'value', None) == 'cloud_runtime'
|
||||
if cloud_runtime:
|
||||
tenant_scope = getattr(persistence_mgr, 'tenant_scope', None)
|
||||
if not callable(tenant_scope):
|
||||
raise RuntimeError('Cloud webhook dispatch requires an explicit tenant scope')
|
||||
async with tenant_scope(runtime_bot.workspace_uuid):
|
||||
response = await dispatch()
|
||||
else:
|
||||
response = await dispatch()
|
||||
else:
|
||||
response = await dispatch()
|
||||
|
||||
return response
|
||||
|
||||
except bounded_executor.BlockingWorkCapacityError as exc:
|
||||
return self.http_status(
|
||||
429,
|
||||
'blocking_work_capacity_exceeded',
|
||||
str(exc),
|
||||
)
|
||||
except Exception:
|
||||
request_id = self.request_id()
|
||||
self.ap.logger.error(
|
||||
|
||||
@@ -2,6 +2,7 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import typing
|
||||
|
||||
import quart
|
||||
import quart_cors
|
||||
@@ -27,6 +28,37 @@ importutil.import_modules_in_pkg(groups_knowledge)
|
||||
importutil.import_modules_in_pkg(groups_resources)
|
||||
|
||||
|
||||
class BoundedJSONRequest(quart.Request):
|
||||
"""Parse bounded HTTP JSON bodies outside the shared event loop."""
|
||||
|
||||
async def get_json(
|
||||
self,
|
||||
force: bool = False,
|
||||
silent: bool = False,
|
||||
cache: bool = True,
|
||||
) -> typing.Any:
|
||||
# Keep Quart's cache and error semantics, changing only where the
|
||||
# potentially 10 MiB JSON decoder runs. The RouterGroup establishes a
|
||||
# trusted Workspace blocking-work scope before calling route handlers.
|
||||
if cache and self._cached_json[silent] is not Ellipsis:
|
||||
return self._cached_json[silent]
|
||||
if not (force or self.is_json):
|
||||
return None
|
||||
|
||||
data = await self.get_data(cache=cache, as_text=False)
|
||||
try:
|
||||
result = await asyncio.to_thread(self.json_module.loads, data)
|
||||
except ValueError as error:
|
||||
if silent:
|
||||
result = None
|
||||
else:
|
||||
result = self.on_json_loading_failed(error)
|
||||
|
||||
if cache:
|
||||
self._cached_json[silent] = result
|
||||
return result
|
||||
|
||||
|
||||
class HTTPController:
|
||||
ap: app.Application
|
||||
|
||||
@@ -35,6 +67,7 @@ class HTTPController:
|
||||
def __init__(self, ap: app.Application) -> None:
|
||||
self.ap = ap
|
||||
self.quart_app = quart.Quart(__name__)
|
||||
self.quart_app.request_class = BoundedJSONRequest
|
||||
quart_cors.cors(self.quart_app, allow_origin='*')
|
||||
|
||||
# Set maximum content length to prevent large file uploads
|
||||
@@ -103,6 +136,7 @@ class HTTPController:
|
||||
config.accesslog = '-'
|
||||
config.bind = [f'{host}:{port}']
|
||||
config.errorlog = config.accesslog
|
||||
config.websocket_max_message_size = group.MAX_FILE_SIZE
|
||||
|
||||
asgi_app = self.quart_app
|
||||
if self.mcp_mount is not None:
|
||||
@@ -113,7 +147,16 @@ class HTTPController:
|
||||
async def register_routes(self) -> None:
|
||||
@self.quart_app.route('/healthz')
|
||||
async def healthz():
|
||||
return {'code': 0, 'msg': 'ok'}
|
||||
get_resource_stats = getattr(
|
||||
self.ap,
|
||||
'get_runtime_resource_stats',
|
||||
None,
|
||||
)
|
||||
return {
|
||||
'code': 0,
|
||||
'msg': 'ok',
|
||||
'resources': (get_resource_stats() if callable(get_resource_stats) else {}),
|
||||
}
|
||||
|
||||
for g in group.preregistered_groups:
|
||||
ginst = g(self.ap, self.quart_app)
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import datetime
|
||||
import functools
|
||||
import os
|
||||
@@ -69,7 +70,10 @@ class MaintenanceService:
|
||||
|
||||
return {
|
||||
'uploaded_files': await self._cleanup_expired_uploaded_files(context, upload_retention_days),
|
||||
'log_files': self._cleanup_expired_log_files(log_retention_days)
|
||||
'log_files': await asyncio.to_thread(
|
||||
self._cleanup_expired_log_files,
|
||||
log_retention_days,
|
||||
)
|
||||
if await self._is_oss_singleton(context)
|
||||
else 0,
|
||||
}
|
||||
@@ -108,22 +112,19 @@ class MaintenanceService:
|
||||
scoped_storage_path = Path('data/storage') / self.ap.storage_mgr.scoped_prefix(context)
|
||||
roots = [('storage', scoped_storage_path)]
|
||||
|
||||
sections = []
|
||||
for key, path in roots:
|
||||
sections.append(
|
||||
{
|
||||
'key': key,
|
||||
'path': str(path) if path else '',
|
||||
'exists': path.exists() if path else False,
|
||||
'size_bytes': self._path_size(path) if path else 0,
|
||||
'file_count': self._file_count(path) if path else 0,
|
||||
}
|
||||
)
|
||||
sections = await asyncio.to_thread(self._collect_sections, roots)
|
||||
|
||||
monitoring_counts = await self._monitoring_counts(context)
|
||||
binary_storage = await self._binary_storage_stats(context)
|
||||
upload_candidates = await self._expired_uploaded_candidates(context, upload_retention_days)
|
||||
log_candidates = self._expired_log_candidates(log_retention_days) if is_oss_singleton else []
|
||||
log_candidates = (
|
||||
await asyncio.to_thread(
|
||||
self._expired_log_candidates,
|
||||
log_retention_days,
|
||||
)
|
||||
if is_oss_singleton
|
||||
else []
|
||||
)
|
||||
|
||||
return {
|
||||
'generated_at': datetime.datetime.now(datetime.timezone.utc).isoformat(),
|
||||
@@ -144,6 +145,23 @@ class MaintenanceService:
|
||||
'tasks': self.ap.task_mgr.get_stats() if is_oss_singleton and self.ap.task_mgr else {},
|
||||
}
|
||||
|
||||
def _collect_sections(
|
||||
self,
|
||||
roots: list[tuple[str, Path | None]],
|
||||
) -> list[dict[str, Any]]:
|
||||
sections = []
|
||||
for key, path in roots:
|
||||
sections.append(
|
||||
{
|
||||
'key': key,
|
||||
'path': str(path) if path else '',
|
||||
'exists': path.exists() if path else False,
|
||||
'size_bytes': self._path_size(path) if path else 0,
|
||||
'file_count': self._file_count(path) if path else 0,
|
||||
}
|
||||
)
|
||||
return sections
|
||||
|
||||
async def _is_oss_singleton(self, context: TenantContext) -> bool:
|
||||
try:
|
||||
await self.ap.workspace_service.get_local_execution_binding(
|
||||
@@ -162,21 +180,16 @@ class MaintenanceService:
|
||||
provider = self.ap.storage_mgr.storage_provider
|
||||
provider_name = provider.__class__.__name__
|
||||
if provider_name == 'LocalStorageProvider':
|
||||
candidates = self._expired_local_upload_candidates(
|
||||
candidates = await asyncio.to_thread(
|
||||
self._expired_local_upload_candidates,
|
||||
context,
|
||||
retention_days,
|
||||
include_paths=True,
|
||||
True,
|
||||
)
|
||||
return await asyncio.to_thread(
|
||||
self._delete_local_candidates,
|
||||
candidates,
|
||||
)
|
||||
deleted = 0
|
||||
for item in candidates:
|
||||
try:
|
||||
os.remove(item['path'])
|
||||
deleted += 1
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
except Exception as e:
|
||||
self.ap.logger.warning(f'Failed to delete expired uploaded file {item["key"]}: {e}')
|
||||
return deleted
|
||||
|
||||
if provider_name == 'S3StorageProvider':
|
||||
return await self._cleanup_expired_s3_uploaded_files(context, retention_days)
|
||||
@@ -190,7 +203,11 @@ class MaintenanceService:
|
||||
) -> list[dict[str, Any]]:
|
||||
provider_name = self.ap.storage_mgr.storage_provider.__class__.__name__
|
||||
if provider_name == 'LocalStorageProvider':
|
||||
return self._expired_local_upload_candidates(context, retention_days)
|
||||
return await asyncio.to_thread(
|
||||
self._expired_local_upload_candidates,
|
||||
context,
|
||||
retention_days,
|
||||
)
|
||||
if provider_name == 'S3StorageProvider':
|
||||
return await self._expired_s3_upload_candidates(context, retention_days)
|
||||
return []
|
||||
@@ -212,6 +229,25 @@ class MaintenanceService:
|
||||
self,
|
||||
context: TenantContext,
|
||||
retention_days: int,
|
||||
) -> list[dict[str, Any]]:
|
||||
provider = self.ap.storage_mgr.storage_provider
|
||||
run_io = getattr(provider, '_run_io', None)
|
||||
if callable(run_io):
|
||||
return await run_io(
|
||||
self._expired_s3_upload_candidates_sync,
|
||||
context,
|
||||
retention_days,
|
||||
)
|
||||
return await asyncio.to_thread(
|
||||
self._expired_s3_upload_candidates_sync,
|
||||
context,
|
||||
retention_days,
|
||||
)
|
||||
|
||||
def _expired_s3_upload_candidates_sync(
|
||||
self,
|
||||
context: TenantContext,
|
||||
retention_days: int,
|
||||
) -> list[dict[str, Any]]:
|
||||
provider = self.ap.storage_mgr.storage_provider
|
||||
cutoff = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(days=retention_days)
|
||||
@@ -241,6 +277,18 @@ class MaintenanceService:
|
||||
|
||||
return candidates
|
||||
|
||||
def _delete_local_candidates(self, candidates: list[dict[str, Any]]) -> int:
|
||||
deleted = 0
|
||||
for item in candidates:
|
||||
try:
|
||||
os.remove(item['path'])
|
||||
deleted += 1
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
except Exception as e:
|
||||
self.ap.logger.warning(f'Failed to delete expired uploaded file {item["key"]}: {e}')
|
||||
return deleted
|
||||
|
||||
def _cleanup_expired_log_files(self, retention_days: int) -> int:
|
||||
deleted = 0
|
||||
for item in self._expired_log_candidates(retention_days, include_paths=True):
|
||||
|
||||
@@ -252,8 +252,17 @@ class MCPService:
|
||||
task = create_detached_task(
|
||||
self.ap.tool_mgr.mcp_tool_loader.host_mcp_server(execution_context, created),
|
||||
after_commit_manager=self.ap.persistence_mgr,
|
||||
workspace_uuid=execution_context.workspace_uuid,
|
||||
)
|
||||
self.ap.tool_mgr.mcp_tool_loader._hosted_mcp_tasks.append(task)
|
||||
tracker = getattr(
|
||||
self.ap.tool_mgr.mcp_tool_loader,
|
||||
'track_hosted_task',
|
||||
None,
|
||||
)
|
||||
if callable(tracker):
|
||||
tracker(task, execution_context)
|
||||
else:
|
||||
self.ap.tool_mgr.mcp_tool_loader._hosted_mcp_tasks.append(task)
|
||||
return payload['uuid']
|
||||
|
||||
async def get_mcp_server_by_uuid(self, context: TenantContext, server_uuid: str) -> dict | None:
|
||||
@@ -357,8 +366,13 @@ class MCPService:
|
||||
task = create_detached_task(
|
||||
loader.host_mcp_server(execution_context, updated),
|
||||
after_commit_manager=self.ap.persistence_mgr,
|
||||
workspace_uuid=execution_context.workspace_uuid,
|
||||
)
|
||||
loader._hosted_mcp_tasks.append(task)
|
||||
tracker = getattr(loader, 'track_hosted_task', None)
|
||||
if callable(tracker):
|
||||
tracker(task, execution_context)
|
||||
else:
|
||||
loader._hosted_mcp_tasks.append(task)
|
||||
|
||||
async def delete_mcp_server(self, context: TenantContext, server_uuid: str) -> None:
|
||||
execution_context = await self._execution_context(context)
|
||||
@@ -420,6 +434,7 @@ class MCPService:
|
||||
async def test_mcp_server(self, context: TenantContext, server_name: str, server_data: dict) -> int:
|
||||
execution_context = await self._execution_context(context)
|
||||
runtime_mcp_session: RuntimeMCPSession | None = None
|
||||
test_session: RuntimeMCPSession | None = None
|
||||
ctx = taskmgr.TaskContext.new()
|
||||
|
||||
if server_name != '_':
|
||||
@@ -468,16 +483,27 @@ class MCPService:
|
||||
|
||||
coroutine = _run_and_cleanup()
|
||||
|
||||
wrapper = self.ap.task_mgr.create_user_task(
|
||||
coroutine,
|
||||
kind='mcp-operation',
|
||||
name=f'mcp-test-{execution_context.workspace_uuid}-{server_name}',
|
||||
label=f'Testing MCP server {server_name}',
|
||||
context=ctx,
|
||||
instance_uuid=execution_context.instance_uuid,
|
||||
workspace_uuid=execution_context.workspace_uuid,
|
||||
placement_generation=execution_context.placement_generation,
|
||||
)
|
||||
try:
|
||||
wrapper = self.ap.task_mgr.create_user_task(
|
||||
coroutine,
|
||||
kind='mcp-operation',
|
||||
name=f'mcp-test-{execution_context.workspace_uuid}-{server_name}',
|
||||
label=f'Testing MCP server {server_name}',
|
||||
context=ctx,
|
||||
instance_uuid=execution_context.instance_uuid,
|
||||
workspace_uuid=execution_context.workspace_uuid,
|
||||
placement_generation=execution_context.placement_generation,
|
||||
)
|
||||
except taskmgr.TaskCapacityError:
|
||||
if test_session is not None:
|
||||
try:
|
||||
await test_session.shutdown()
|
||||
except Exception as exc:
|
||||
self.ap.logger.warning(
|
||||
f'Failed to tear down rejected transient MCP test session '
|
||||
f'{test_session.server_name}: {type(exc).__name__}: {exc}'
|
||||
)
|
||||
raise
|
||||
return wrapper.id
|
||||
|
||||
async def get_mcp_server_logs(
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import io
|
||||
import inspect
|
||||
import os
|
||||
@@ -13,6 +14,7 @@ import httpx
|
||||
|
||||
from ....core import app
|
||||
from ....skill.utils import parse_frontmatter
|
||||
from ....utils import httpclient
|
||||
from ..context import ExecutionContext
|
||||
from .tenant import TenantContext, require_workspace_uuid
|
||||
|
||||
@@ -328,7 +330,11 @@ class SkillService:
|
||||
await result
|
||||
|
||||
async def _download_github_asset(self, asset_url: str) -> bytes:
|
||||
async with httpx.AsyncClient(follow_redirects=True, timeout=120) as client:
|
||||
async with httpx.AsyncClient(
|
||||
follow_redirects=True,
|
||||
timeout=120,
|
||||
event_hooks=httpclient.httpx_response_limit_hooks(_MAX_GITHUB_ARCHIVE_BYTES),
|
||||
) as client:
|
||||
async with client.stream('GET', asset_url) as resp:
|
||||
resp.raise_for_status()
|
||||
content_length = resp.headers.get('content-length')
|
||||
@@ -352,7 +358,14 @@ class SkillService:
|
||||
info = self._parse_github_skill_md_url(asset_url, owner=owner, repo=repo)
|
||||
archive_url = f'https://codeload.github.com/{owner}/{repo}/zip/{quote(info["ref"], safe="/")}'
|
||||
archive_bytes = await self._download_github_asset(archive_url)
|
||||
return await asyncio.to_thread(self._build_github_skill_directory_zip, archive_bytes, info)
|
||||
|
||||
def _build_github_skill_directory_zip(
|
||||
self,
|
||||
archive_bytes: bytes,
|
||||
info: dict[str, str],
|
||||
) -> tuple[bytes, str, str]:
|
||||
"""Validate and repack a GitHub skill archive outside the event loop."""
|
||||
try:
|
||||
source_archive = zipfile.ZipFile(io.BytesIO(archive_bytes), 'r')
|
||||
except zipfile.BadZipFile as exc:
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import OrderedDict
|
||||
|
||||
from langbot.pkg.utils import httpclient
|
||||
import typing
|
||||
import datetime
|
||||
@@ -11,6 +13,10 @@ from ....entity.persistence import user
|
||||
from ....entity.dto.space_model import SpaceModel
|
||||
|
||||
|
||||
_CREDITS_CACHE_TTL_SECONDS = 60
|
||||
_CREDITS_CACHE_MAX_ENTRIES = 4096
|
||||
|
||||
|
||||
class SpaceService:
|
||||
"""Service for interacting with LangBot Space API"""
|
||||
|
||||
@@ -19,7 +25,24 @@ class SpaceService:
|
||||
|
||||
def __init__(self, ap: app.Application) -> None:
|
||||
self.ap = ap
|
||||
self._credits_cache = {}
|
||||
self._credits_cache = OrderedDict()
|
||||
|
||||
def _ordered_credits_cache(
|
||||
self,
|
||||
) -> OrderedDict[str, tuple[int, float]]:
|
||||
if not isinstance(self._credits_cache, OrderedDict):
|
||||
# Preserve compatibility with tests and callers that seed the cache.
|
||||
self._credits_cache = OrderedDict(self._credits_cache)
|
||||
return self._credits_cache
|
||||
|
||||
def _prune_credits_cache(self, now: float) -> None:
|
||||
cache = self._ordered_credits_cache()
|
||||
while cache:
|
||||
email = next(iter(cache))
|
||||
_, cached_at = cache[email]
|
||||
if now - cached_at < _CREDITS_CACHE_TTL_SECONDS:
|
||||
break
|
||||
cache.pop(email, None)
|
||||
|
||||
def _get_space_config(self) -> typing.Dict[str, str]:
|
||||
"""Get Space configuration from config file"""
|
||||
@@ -107,8 +130,9 @@ class SpaceService:
|
||||
json={'code': code, 'instance_id': constants.instance_id},
|
||||
) as response:
|
||||
if response.status != 200:
|
||||
raise ValueError(f'Failed to exchange OAuth code: {await response.text()}')
|
||||
data = await response.json()
|
||||
error = await httpclient.read_text_limited(response)
|
||||
raise ValueError(f'Failed to exchange OAuth code: {error}')
|
||||
data = await httpclient.read_json_limited(response)
|
||||
if data.get('code') != 0:
|
||||
raise ValueError(f'Failed to exchange OAuth code: {data.get("msg")}')
|
||||
return data.get('data', {})
|
||||
@@ -123,8 +147,9 @@ class SpaceService:
|
||||
f'{space_url}/api/v1/accounts/token/refresh', json={'refresh_token': refresh_token}
|
||||
) as response:
|
||||
if response.status != 200:
|
||||
raise ValueError(f'Failed to refresh token: {await response.text()}')
|
||||
data = await response.json()
|
||||
error = await httpclient.read_text_limited(response)
|
||||
raise ValueError(f'Failed to refresh token: {error}')
|
||||
data = await httpclient.read_json_limited(response)
|
||||
if data.get('code') != 0:
|
||||
raise ValueError(f'Failed to refresh token: {data.get("msg")}')
|
||||
return data.get('data', {})
|
||||
@@ -139,8 +164,9 @@ class SpaceService:
|
||||
f'{space_url}/api/v1/accounts/me', headers={'Authorization': f'Bearer {access_token}'}
|
||||
) as response:
|
||||
if response.status != 200:
|
||||
raise ValueError(f'Failed to get user info: {await response.text()}')
|
||||
data = await response.json()
|
||||
error = await httpclient.read_text_limited(response)
|
||||
raise ValueError(f'Failed to get user info: {error}')
|
||||
data = await httpclient.read_json_limited(response)
|
||||
if data.get('code') != 0:
|
||||
raise ValueError(f'Failed to get user info: {data.get("msg")}')
|
||||
return data.get('data', {})
|
||||
@@ -156,11 +182,13 @@ class SpaceService:
|
||||
|
||||
async def get_credits(self, user_email: str, force_refresh: bool = False) -> int | None:
|
||||
"""Get Space credits for user with caching (60s TTL)"""
|
||||
cache_ttl = 60
|
||||
now = time.time()
|
||||
cached_fallback = self._credits_cache.get(user_email)
|
||||
self._prune_credits_cache(now)
|
||||
|
||||
if not force_refresh and user_email in self._credits_cache:
|
||||
credits, ts = self._credits_cache[user_email]
|
||||
if time.time() - ts < cache_ttl:
|
||||
if now - ts < _CREDITS_CACHE_TTL_SECONDS:
|
||||
return credits
|
||||
|
||||
try:
|
||||
@@ -169,10 +197,14 @@ class SpaceService:
|
||||
return None
|
||||
credits = info.get('credits')
|
||||
if credits is not None:
|
||||
self._credits_cache[user_email] = (credits, time.time())
|
||||
cache = self._ordered_credits_cache()
|
||||
cache.pop(user_email, None)
|
||||
if len(cache) >= _CREDITS_CACHE_MAX_ENTRIES:
|
||||
cache.popitem(last=False)
|
||||
cache[user_email] = (credits, time.time())
|
||||
return credits
|
||||
except Exception:
|
||||
return self._credits_cache.get(user_email, (None, 0))[0]
|
||||
return cached_fallback[0] if cached_fallback is not None else None
|
||||
|
||||
async def get_models(self) -> typing.List[SpaceModel]:
|
||||
"""Get models from Space"""
|
||||
@@ -183,8 +215,9 @@ class SpaceService:
|
||||
session = httpclient.get_session()
|
||||
async with session.get(f'{space_url}/api/v1/models', params={'page_size': 100}) as response:
|
||||
if response.status != 200:
|
||||
raise ValueError(f'Failed to get models: {await response.text()}')
|
||||
data = await response.json()
|
||||
error = await httpclient.read_text_limited(response)
|
||||
raise ValueError(f'Failed to get models: {error}')
|
||||
data = await httpclient.read_json_limited(response)
|
||||
if data.get('code') != 0:
|
||||
raise ValueError(f'Failed to get models: {data.get("msg")}')
|
||||
models_data = data.get('data', {}).get('models', [])
|
||||
|
||||
@@ -7,6 +7,7 @@ import datetime
|
||||
import typing
|
||||
import asyncio
|
||||
import dataclasses
|
||||
import heapq
|
||||
import hashlib
|
||||
import secrets
|
||||
import time
|
||||
@@ -19,11 +20,17 @@ from ....entity.persistence.workspace import MembershipRole, MembershipStatus, W
|
||||
from ....utils import constants
|
||||
from ....entity.errors import account as account_errors
|
||||
from ....workspace.collaboration import normalize_email
|
||||
from ....utils import bounded_executor
|
||||
|
||||
if typing.TYPE_CHECKING:
|
||||
from ....core.app import Application
|
||||
|
||||
|
||||
_SPACE_OAUTH_STATE_MAX_ENTRIES = 4096
|
||||
_SPACE_OAUTH_STATE_HEAP_COMPACT_FLOOR = 64
|
||||
_SPACE_OAUTH_STATE_HEAP_MAX_MULTIPLIER = 4
|
||||
|
||||
|
||||
class AccountExistsLoginRequiredError(ValueError):
|
||||
code = 'account_exists_login_required'
|
||||
|
||||
@@ -54,14 +61,45 @@ class UserService:
|
||||
def __init__(self, ap: Application) -> None:
|
||||
self.ap = ap
|
||||
self._create_user_lock = asyncio.Lock()
|
||||
self._password_hash_lock = asyncio.Semaphore(1)
|
||||
self._password_hash_lock = asyncio.Lock()
|
||||
self._space_oauth_state_lock = asyncio.Lock()
|
||||
self._space_oauth_states: dict[str, tuple[str, str | None, float, str | None]] = {}
|
||||
self._space_oauth_state_expiry_heap: list[tuple[float, str]] = []
|
||||
|
||||
@staticmethod
|
||||
def _space_oauth_state_digest(state: str) -> str:
|
||||
return hashlib.sha256(state.encode('utf-8')).hexdigest()
|
||||
|
||||
def _prune_space_oauth_states(self, now: float) -> None:
|
||||
while self._space_oauth_state_expiry_heap:
|
||||
expires_at, digest = self._space_oauth_state_expiry_heap[0]
|
||||
entry = self._space_oauth_states.get(digest)
|
||||
if entry is None or entry[2] != expires_at:
|
||||
heapq.heappop(self._space_oauth_state_expiry_heap)
|
||||
continue
|
||||
if expires_at > now:
|
||||
break
|
||||
heapq.heappop(self._space_oauth_state_expiry_heap)
|
||||
self._space_oauth_states.pop(digest, None)
|
||||
|
||||
max_heap_entries = max(
|
||||
_SPACE_OAUTH_STATE_HEAP_COMPACT_FLOOR,
|
||||
len(self._space_oauth_states) * _SPACE_OAUTH_STATE_HEAP_MAX_MULTIPLIER,
|
||||
)
|
||||
if len(self._space_oauth_state_expiry_heap) > max_heap_entries:
|
||||
self._space_oauth_state_expiry_heap[:] = [
|
||||
(entry[2], digest) for digest, entry in self._space_oauth_states.items()
|
||||
]
|
||||
heapq.heapify(self._space_oauth_state_expiry_heap)
|
||||
|
||||
def _evict_earliest_space_oauth_state(self) -> None:
|
||||
while self._space_oauth_state_expiry_heap:
|
||||
expires_at, digest = heapq.heappop(self._space_oauth_state_expiry_heap)
|
||||
entry = self._space_oauth_states.get(digest)
|
||||
if entry is not None and entry[2] == expires_at:
|
||||
self._space_oauth_states.pop(digest, None)
|
||||
return
|
||||
|
||||
async def issue_space_oauth_state(
|
||||
self,
|
||||
purpose: typing.Literal['login', 'bind'],
|
||||
@@ -85,11 +123,14 @@ class UserService:
|
||||
expires_at = time.monotonic() + min(ttl_seconds, 600)
|
||||
async with self._space_oauth_state_lock:
|
||||
now = time.monotonic()
|
||||
self._space_oauth_states = {key: value for key, value in self._space_oauth_states.items() if value[2] > now}
|
||||
if len(self._space_oauth_states) >= 4096:
|
||||
oldest = min(self._space_oauth_states, key=lambda key: self._space_oauth_states[key][2])
|
||||
self._space_oauth_states.pop(oldest, None)
|
||||
self._prune_space_oauth_states(now)
|
||||
if len(self._space_oauth_states) >= _SPACE_OAUTH_STATE_MAX_ENTRIES:
|
||||
self._evict_earliest_space_oauth_state()
|
||||
self._space_oauth_states[digest] = (purpose, account_uuid, expires_at, launch_workspace_uuid)
|
||||
heapq.heappush(
|
||||
self._space_oauth_state_expiry_heap,
|
||||
(expires_at, digest),
|
||||
)
|
||||
return raw_state
|
||||
|
||||
async def consume_space_oauth_state_details(
|
||||
@@ -129,8 +170,14 @@ class UserService:
|
||||
return consumed.account
|
||||
|
||||
async def _hash_password(self, password: str) -> str:
|
||||
if self._password_hash_lock.locked():
|
||||
raise bounded_executor.BlockingWorkCapacityError(
|
||||
'Password hashing capacity reached',
|
||||
scope='system:authentication',
|
||||
)
|
||||
async with self._password_hash_lock:
|
||||
return await asyncio.to_thread(argon2.PasswordHasher().hash, password)
|
||||
with bounded_executor.blocking_work_scope('system:authentication'):
|
||||
return await asyncio.to_thread(argon2.PasswordHasher().hash, password)
|
||||
|
||||
def _require_local_directory(self) -> None:
|
||||
if self._uses_control_plane_directory():
|
||||
@@ -143,8 +190,14 @@ class UserService:
|
||||
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:
|
||||
if self._password_hash_lock.locked():
|
||||
raise bounded_executor.BlockingWorkCapacityError(
|
||||
'Password hashing capacity reached',
|
||||
scope='system:authentication',
|
||||
)
|
||||
async with self._password_hash_lock:
|
||||
await asyncio.to_thread(argon2.PasswordHasher().verify, hashed_password, password)
|
||||
with bounded_executor.blocking_work_scope('system:authentication'):
|
||||
await asyncio.to_thread(argon2.PasswordHasher().verify, hashed_password, password)
|
||||
|
||||
async def _update_space_provider_for_account(self, account: typing.Any, api_key: str) -> None:
|
||||
"""Refresh the OSS Workspace Space provider without guessing a SaaS Workspace.
|
||||
|
||||
@@ -3,6 +3,7 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
import datetime as dt
|
||||
import time
|
||||
import weakref
|
||||
from collections.abc import Callable
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
@@ -52,7 +53,7 @@ class SandboxAdmissionController:
|
||||
self.client = client
|
||||
self.policy = policy
|
||||
self._wall_time = wall_time
|
||||
self._locks: dict[str, asyncio.Lock] = {}
|
||||
self._locks: weakref.WeakValueDictionary[str, asyncio.Lock] = weakref.WeakValueDictionary()
|
||||
self._highest_revisions: dict[str, int] = {}
|
||||
|
||||
def _workspace_lock(self, workspace_uuid: str) -> asyncio.Lock:
|
||||
|
||||
@@ -21,6 +21,7 @@ from .admission import SandboxAdmissionController, require_cloud_admission_polic
|
||||
from .connector import BoxRuntimeConnector, _get_box_config
|
||||
from . import secure_fs
|
||||
from ..telemetry import features as telemetry_features
|
||||
from ..utils import httpclient
|
||||
from ..api.http.context import ExecutionContext
|
||||
from ..api.http.service.tenant import TenantContext, require_workspace_uuid
|
||||
from langbot_plugin.box.errors import BoxAdmissionError, BoxError, BoxValidationError
|
||||
@@ -39,6 +40,53 @@ _MAX_RECENT_ERRORS = 50
|
||||
_MIB = 1024 * 1024
|
||||
|
||||
|
||||
def _create_shared_workspace_probe(root: str, marker_name: str, payload: bytes) -> None:
|
||||
"""Create and durably flush a no-follow probe without blocking the event loop."""
|
||||
|
||||
directory_flags = os.O_RDONLY | getattr(os, 'O_DIRECTORY', 0)
|
||||
nofollow = getattr(os, 'O_NOFOLLOW', 0)
|
||||
root_fd: int | None = None
|
||||
marker_fd: int | None = None
|
||||
marker_created = False
|
||||
try:
|
||||
root_fd = os.open(root, directory_flags | nofollow)
|
||||
marker_fd = os.open(
|
||||
marker_name,
|
||||
os.O_WRONLY | os.O_CREAT | os.O_EXCL | nofollow,
|
||||
0o600,
|
||||
dir_fd=root_fd,
|
||||
)
|
||||
marker_created = True
|
||||
remaining = memoryview(payload)
|
||||
while remaining:
|
||||
written = os.write(marker_fd, remaining)
|
||||
if written <= 0:
|
||||
raise BoxValidationError('Failed to write Cloud Box shared-volume probe')
|
||||
remaining = remaining[written:]
|
||||
os.fsync(marker_fd)
|
||||
except Exception:
|
||||
if root_fd is not None and marker_created:
|
||||
with contextlib.suppress(FileNotFoundError):
|
||||
os.unlink(marker_name, dir_fd=root_fd)
|
||||
raise
|
||||
finally:
|
||||
if marker_fd is not None:
|
||||
os.close(marker_fd)
|
||||
if root_fd is not None:
|
||||
os.close(root_fd)
|
||||
|
||||
|
||||
def _remove_shared_workspace_probe(root: str, marker_name: str) -> None:
|
||||
directory_flags = os.O_RDONLY | getattr(os, 'O_DIRECTORY', 0)
|
||||
nofollow = getattr(os, 'O_NOFOLLOW', 0)
|
||||
root_fd = os.open(root, directory_flags | nofollow)
|
||||
try:
|
||||
with contextlib.suppress(FileNotFoundError):
|
||||
os.unlink(marker_name, dir_fd=root_fd)
|
||||
finally:
|
||||
os.close(root_fd)
|
||||
|
||||
|
||||
def _is_path_under(path: str, root: str) -> bool:
|
||||
"""Check whether *path* equals *root* or is a child of *root*."""
|
||||
return path == root or path.startswith(f'{root}{os.sep}')
|
||||
@@ -267,29 +315,15 @@ class BoxService:
|
||||
marker_name = f'{BOX_SHARED_WORKSPACE_PROBE_PREFIX}{secrets.token_hex(16)}'
|
||||
marker_payload = secrets.token_bytes(64)
|
||||
expected_digest = hashlib.sha256(marker_payload).hexdigest()
|
||||
directory_flags = os.O_RDONLY | getattr(os, 'O_DIRECTORY', 0)
|
||||
nofollow = getattr(os, 'O_NOFOLLOW', 0)
|
||||
root_fd: int | None = None
|
||||
marker_fd: int | None = None
|
||||
marker_created = False
|
||||
probe_created = False
|
||||
try:
|
||||
root_fd = os.open(self.default_workspace, directory_flags | nofollow)
|
||||
marker_fd = os.open(
|
||||
await asyncio.to_thread(
|
||||
_create_shared_workspace_probe,
|
||||
self.default_workspace,
|
||||
marker_name,
|
||||
os.O_WRONLY | os.O_CREAT | os.O_EXCL | nofollow,
|
||||
0o600,
|
||||
dir_fd=root_fd,
|
||||
marker_payload,
|
||||
)
|
||||
marker_created = True
|
||||
remaining = memoryview(marker_payload)
|
||||
while remaining:
|
||||
written = os.write(marker_fd, remaining)
|
||||
if written <= 0:
|
||||
raise BoxValidationError('Failed to write Cloud Box shared-volume probe')
|
||||
remaining = remaining[written:]
|
||||
os.fsync(marker_fd)
|
||||
os.close(marker_fd)
|
||||
marker_fd = None
|
||||
probe_created = True
|
||||
|
||||
result = await self.client.verify_shared_workspace(marker_name)
|
||||
if (
|
||||
@@ -306,15 +340,12 @@ class BoxService:
|
||||
except Exception as exc:
|
||||
raise BoxValidationError('Cloud Box shared durable Workspace volume verification failed') from exc
|
||||
finally:
|
||||
if marker_fd is not None:
|
||||
os.close(marker_fd)
|
||||
if root_fd is not None:
|
||||
if marker_created:
|
||||
try:
|
||||
os.unlink(marker_name, dir_fd=root_fd)
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
os.close(root_fd)
|
||||
if probe_created:
|
||||
await asyncio.to_thread(
|
||||
_remove_shared_workspace_probe,
|
||||
self.default_workspace,
|
||||
marker_name,
|
||||
)
|
||||
|
||||
@property
|
||||
def available(self) -> bool:
|
||||
@@ -886,7 +917,13 @@ class BoxService:
|
||||
mime = data[5:split_index]
|
||||
data = data[split_index + 8 :]
|
||||
try:
|
||||
return _b64.b64decode(data), mime
|
||||
max_encoded_bytes = 4 * ((BoxService._ATTACHMENT_MAX_BYTES + 2) // 3)
|
||||
if not isinstance(data, (str, bytes)) or len(data) > max_encoded_bytes:
|
||||
return None
|
||||
decoded = _b64.b64decode(data)
|
||||
if len(decoded) > BoxService._ATTACHMENT_MAX_BYTES:
|
||||
return None
|
||||
return decoded, mime
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
@@ -895,10 +932,25 @@ class BoxService:
|
||||
try:
|
||||
import httpx
|
||||
|
||||
async with httpx.AsyncClient(timeout=30) as client:
|
||||
resp = await client.get(url)
|
||||
resp.raise_for_status()
|
||||
return resp.content, resp.headers.get('Content-Type', 'application/octet-stream')
|
||||
async with httpx.AsyncClient(
|
||||
timeout=30,
|
||||
event_hooks=httpclient.httpx_response_limit_hooks(BoxService._ATTACHMENT_MAX_BYTES),
|
||||
) as client:
|
||||
async with client.stream('GET', url) as resp:
|
||||
resp.raise_for_status()
|
||||
declared_size = resp.headers.get('content-length')
|
||||
if declared_size is not None:
|
||||
try:
|
||||
if int(declared_size) > BoxService._ATTACHMENT_MAX_BYTES:
|
||||
return None
|
||||
except ValueError:
|
||||
pass
|
||||
body = bytearray()
|
||||
async for chunk in resp.aiter_bytes(chunk_size=64 * 1024):
|
||||
body.extend(chunk)
|
||||
if len(body) > BoxService._ATTACHMENT_MAX_BYTES:
|
||||
return None
|
||||
return bytes(body), resp.headers.get('Content-Type', 'application/octet-stream')
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
@@ -907,8 +959,13 @@ class BoxService:
|
||||
try:
|
||||
import aiofiles
|
||||
|
||||
if await asyncio.to_thread(os.path.getsize, path) > BoxService._ATTACHMENT_MAX_BYTES:
|
||||
return None
|
||||
async with aiofiles.open(path, 'rb') as f:
|
||||
return await f.read(), 'application/octet-stream'
|
||||
data = await f.read(BoxService._ATTACHMENT_MAX_BYTES + 1)
|
||||
if len(data) > BoxService._ATTACHMENT_MAX_BYTES:
|
||||
return None
|
||||
return data, 'application/octet-stream'
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
@@ -34,6 +34,7 @@ from .directory import (
|
||||
DirectorySnapshot,
|
||||
DirectoryWorkspace,
|
||||
)
|
||||
from .entitlements import EntitlementResolver
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -257,6 +258,7 @@ class DirectoryProjectionService:
|
||||
|
||||
await session.flush()
|
||||
|
||||
await self._reconcile_entitlement_snapshot_set(snapshot)
|
||||
self._record_success()
|
||||
self._consumer_cursor = snapshot.cursor
|
||||
|
||||
@@ -343,10 +345,45 @@ class DirectoryProjectionService:
|
||||
await session.flush()
|
||||
projection_caught_up = batch.cursor == batch.high_water_cursor and int(state.cursor) == batch.cursor
|
||||
|
||||
await self._update_entitlement_workspace_activity(
|
||||
returned.values(),
|
||||
requested_workspace_uuids=requested,
|
||||
)
|
||||
if projection_caught_up:
|
||||
self._record_success()
|
||||
self._consumer_cursor = batch.cursor
|
||||
|
||||
async def _reconcile_entitlement_snapshot_set(
|
||||
self,
|
||||
snapshot: DirectorySnapshot,
|
||||
) -> None:
|
||||
resolver = getattr(self.ap, 'entitlement_resolver', None)
|
||||
if not isinstance(resolver, EntitlementResolver):
|
||||
return
|
||||
await resolver.reconcile_active_workspaces(
|
||||
{workspace.uuid for workspace in snapshot.workspaces if workspace.status == WorkspaceStatus.ACTIVE.value}
|
||||
)
|
||||
|
||||
async def _update_entitlement_workspace_activity(
|
||||
self,
|
||||
workspaces: Iterable[DirectoryWorkspace],
|
||||
*,
|
||||
requested_workspace_uuids: set[str],
|
||||
) -> None:
|
||||
resolver = getattr(self.ap, 'entitlement_resolver', None)
|
||||
if not isinstance(resolver, EntitlementResolver):
|
||||
return
|
||||
returned = {workspace.uuid: workspace for workspace in workspaces}
|
||||
active = {
|
||||
workspace_uuid
|
||||
for workspace_uuid, workspace in returned.items()
|
||||
if workspace.status == WorkspaceStatus.ACTIVE.value
|
||||
}
|
||||
await resolver.update_workspace_activity(
|
||||
active_workspace_uuids=active,
|
||||
inactive_workspace_uuids=requested_workspace_uuids - active,
|
||||
)
|
||||
|
||||
async def apply_event_batch(self, batch: DirectoryEventBatch) -> None:
|
||||
"""Advance non-directory events after the adapter refreshes local caches."""
|
||||
|
||||
|
||||
@@ -122,6 +122,7 @@ class EntitlementResolver:
|
||||
self._deployment_admission = deployment_admission
|
||||
self._lock = asyncio.Lock()
|
||||
self._snapshots: dict[str, tuple[int, str, EntitlementSnapshot]] = {}
|
||||
self._active_workspace_uuids: frozenset[str] | None = None
|
||||
|
||||
@staticmethod
|
||||
def _fingerprint(snapshot: EntitlementSnapshot) -> str:
|
||||
@@ -136,6 +137,8 @@ class EntitlementResolver:
|
||||
) -> EntitlementSnapshot:
|
||||
if self._deployment_admission is not None:
|
||||
self._deployment_admission()
|
||||
async with self._lock:
|
||||
self._require_projected_workspace_locked(workspace_uuid)
|
||||
candidate = await self.provider.get_workspace_entitlement(workspace_uuid)
|
||||
if self._deployment_admission is not None:
|
||||
# A provider call may cross the Manifest expiry boundary.
|
||||
@@ -154,6 +157,9 @@ class EntitlementResolver:
|
||||
|
||||
fingerprint = self._fingerprint(candidate)
|
||||
async with self._lock:
|
||||
# The directory may fence a Workspace while the provider call is
|
||||
# in flight. Recheck before retaining or returning its snapshot.
|
||||
self._require_projected_workspace_locked(workspace_uuid)
|
||||
previous = self._snapshots.get(workspace_uuid)
|
||||
if previous is not None:
|
||||
previous_revision, previous_fingerprint, _ = previous
|
||||
@@ -167,3 +173,60 @@ class EntitlementResolver:
|
||||
candidate,
|
||||
)
|
||||
return candidate.model_copy(deep=True)
|
||||
|
||||
def _require_projected_workspace_locked(self, workspace_uuid: str) -> None:
|
||||
active_workspace_uuids = self._active_workspace_uuids
|
||||
if active_workspace_uuids is not None and workspace_uuid not in active_workspace_uuids:
|
||||
raise EntitlementUnavailableError('Workspace is not active in the Cloud directory projection')
|
||||
|
||||
async def reconcile_active_workspaces(
|
||||
self,
|
||||
workspace_uuids: set[str] | frozenset[str],
|
||||
) -> None:
|
||||
"""Drop entitlement history for Workspaces fenced by the directory."""
|
||||
|
||||
active = frozenset(workspace_uuids)
|
||||
async with self._lock:
|
||||
self._active_workspace_uuids = active
|
||||
self._snapshots = {
|
||||
workspace_uuid: cached for workspace_uuid, cached in self._snapshots.items() if workspace_uuid in active
|
||||
}
|
||||
|
||||
async def set_workspace_active(
|
||||
self,
|
||||
workspace_uuid: str,
|
||||
*,
|
||||
active: bool,
|
||||
) -> None:
|
||||
"""Apply one incremental directory activity change."""
|
||||
|
||||
await self.update_workspace_activity(
|
||||
active_workspace_uuids={workspace_uuid} if active else set(),
|
||||
inactive_workspace_uuids=set() if active else {workspace_uuid},
|
||||
)
|
||||
|
||||
async def update_workspace_activity(
|
||||
self,
|
||||
*,
|
||||
active_workspace_uuids: set[str] | frozenset[str],
|
||||
inactive_workspace_uuids: set[str] | frozenset[str],
|
||||
) -> None:
|
||||
"""Apply one directory delta without copying the active set per item."""
|
||||
|
||||
active_updates = set(active_workspace_uuids)
|
||||
inactive_updates = set(inactive_workspace_uuids)
|
||||
if active_updates & inactive_updates:
|
||||
raise ValueError('Workspace activity update contains conflicting entries')
|
||||
async with self._lock:
|
||||
current = set(self._active_workspace_uuids or ())
|
||||
current.update(active_updates)
|
||||
current.difference_update(inactive_updates)
|
||||
for workspace_uuid in inactive_updates:
|
||||
self._snapshots.pop(workspace_uuid, None)
|
||||
self._active_workspace_uuids = frozenset(current)
|
||||
|
||||
def snapshot_counts(self) -> dict[str, int]:
|
||||
return {
|
||||
'active_workspaces': len(self._active_workspace_uuids or ()),
|
||||
'cached_snapshots': len(self._snapshots),
|
||||
}
|
||||
|
||||
@@ -4,6 +4,7 @@ import asyncio
|
||||
import base64
|
||||
import binascii
|
||||
import hashlib
|
||||
import heapq
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
@@ -22,6 +23,9 @@ CONTROL_PLANE_TYP = 'langbot-control-plane+jwt'
|
||||
LAUNCH_KIND = 'workspace.launch'
|
||||
EXPECTED_ISSUER = 'langbot-space'
|
||||
EXPECTED_AUDIENCE = 'langbot-cloud-runtime'
|
||||
_CONSUMED_JTI_MAX_ENTRIES = 4096
|
||||
_CONSUMED_JTI_HEAP_COMPACT_FLOOR = 64
|
||||
_CONSUMED_JTI_HEAP_MAX_MULTIPLIER = 4
|
||||
|
||||
|
||||
class SpaceLaunchError(ValueError):
|
||||
@@ -107,6 +111,7 @@ class SpaceLaunchService:
|
||||
self._wall_time = wall_time
|
||||
self._replay_lock = asyncio.Lock()
|
||||
self._consumed_jtis: dict[str, int] = {}
|
||||
self._consumed_jti_expiry_heap: list[tuple[int, str]] = []
|
||||
|
||||
async def consume_assertion(
|
||||
self,
|
||||
@@ -210,13 +215,38 @@ class SpaceLaunchService:
|
||||
digest = hashlib.sha256(jti.encode('utf-8')).hexdigest()
|
||||
now = int(self._wall_time())
|
||||
async with self._replay_lock:
|
||||
self._consumed_jtis = {existing: expiry for existing, expiry in self._consumed_jtis.items() if expiry > now}
|
||||
self._prune_consumed_jtis(now)
|
||||
if digest in self._consumed_jtis:
|
||||
raise SpaceLaunchError('Launch assertion has already been consumed')
|
||||
if len(self._consumed_jtis) >= 4096:
|
||||
oldest = min(self._consumed_jtis, key=lambda key: self._consumed_jtis[key])
|
||||
self._consumed_jtis.pop(oldest, None)
|
||||
if len(self._consumed_jtis) >= _CONSUMED_JTI_MAX_ENTRIES:
|
||||
# Evicting a still-valid digest would make a signed launch
|
||||
# assertion replayable. Bound memory by failing closed instead.
|
||||
raise SpaceLaunchError('Launch assertion replay cache capacity reached')
|
||||
self._consumed_jtis[digest] = expires_at
|
||||
heapq.heappush(
|
||||
self._consumed_jti_expiry_heap,
|
||||
(expires_at, digest),
|
||||
)
|
||||
|
||||
def _prune_consumed_jtis(self, now: int) -> None:
|
||||
while self._consumed_jti_expiry_heap:
|
||||
expires_at, digest = self._consumed_jti_expiry_heap[0]
|
||||
current_expiry = self._consumed_jtis.get(digest)
|
||||
if current_expiry != expires_at:
|
||||
heapq.heappop(self._consumed_jti_expiry_heap)
|
||||
continue
|
||||
if expires_at > now:
|
||||
break
|
||||
heapq.heappop(self._consumed_jti_expiry_heap)
|
||||
self._consumed_jtis.pop(digest, None)
|
||||
|
||||
max_heap_entries = max(
|
||||
_CONSUMED_JTI_HEAP_COMPACT_FLOOR,
|
||||
len(self._consumed_jtis) * _CONSUMED_JTI_HEAP_MAX_MULTIPLIER,
|
||||
)
|
||||
if len(self._consumed_jti_expiry_heap) > max_heap_entries:
|
||||
self._consumed_jti_expiry_heap[:] = [(expiry, digest) for digest, expiry in self._consumed_jtis.items()]
|
||||
heapq.heapify(self._consumed_jti_expiry_heap)
|
||||
|
||||
@staticmethod
|
||||
def _bounded_float(
|
||||
|
||||
+102
-3
@@ -2,6 +2,7 @@ from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import asyncio
|
||||
import contextlib
|
||||
import traceback
|
||||
import os
|
||||
|
||||
@@ -18,7 +19,7 @@ from ..plugin import connector as plugin_connector
|
||||
from ..pipeline import pool
|
||||
from ..pipeline import controller, pipelinemgr
|
||||
from ..pipeline import aggregator as message_aggregator
|
||||
from ..utils import version as version_mgr, proxy as proxy_mgr
|
||||
from ..utils import version as version_mgr, proxy as proxy_mgr, httpclient
|
||||
from ..persistence import mgr as persistencemgr
|
||||
from ..api.http.controller import main as http_controller
|
||||
from ..api.http.service import user as user_service
|
||||
@@ -36,7 +37,7 @@ from ..api.http.service import skill as skill_service
|
||||
from ..api.http.service import maintenance as maintenance_service
|
||||
from ..discover import engine as discover_engine
|
||||
from ..storage import mgr as storagemgr
|
||||
from ..utils import logcache
|
||||
from ..utils import bounded_executor, event_loop_monitor, logcache
|
||||
from . import taskmgr
|
||||
from . import entities as core_entities
|
||||
from ..rag.knowledge import kbmgr as rag_mgr
|
||||
@@ -191,14 +192,78 @@ class Application:
|
||||
|
||||
maintenance_service: maintenance_service.MaintenanceService = None
|
||||
|
||||
blocking_executor: bounded_executor.BoundedThreadPoolExecutor | None = None
|
||||
event_loop_monitor: event_loop_monitor.EventLoopLagMonitor
|
||||
|
||||
def __init__(self):
|
||||
self._shutdown_lock = asyncio.Lock()
|
||||
self._shutdown_complete = False
|
||||
self._shutdown_task: asyncio.Task | None = None
|
||||
self.event_loop_monitor = event_loop_monitor.EventLoopLagMonitor()
|
||||
|
||||
def get_runtime_resource_stats(self) -> dict[str, object]:
|
||||
"""Return aggregate O(1) counters for liveness and soak validation."""
|
||||
|
||||
try:
|
||||
asyncio_tasks = len(asyncio.all_tasks(self.event_loop))
|
||||
except (RuntimeError, TypeError):
|
||||
asyncio_tasks = 0
|
||||
|
||||
task_stats = self.task_mgr.get_stats() if self.task_mgr is not None else {}
|
||||
query_pool_stats = {}
|
||||
if self.query_pool is not None:
|
||||
query_pool_stats = {
|
||||
'queued': len(self.query_pool.queries),
|
||||
'cached': len(self.query_pool.cached_queries),
|
||||
'active_workspaces': len(self.query_pool.active_query_count_by_workspace),
|
||||
}
|
||||
|
||||
model_stats = {}
|
||||
if self.model_mgr is not None:
|
||||
model_stats = {
|
||||
'providers': len(self.model_mgr.provider_dict),
|
||||
'llms': len(self.model_mgr.llm_model_dict),
|
||||
'embeddings': len(self.model_mgr.embedding_model_dict),
|
||||
'rerankers': len(self.model_mgr.rerank_model_dict),
|
||||
}
|
||||
|
||||
runtime_stats = {
|
||||
'bots': len(getattr(self.platform_mgr, '_bots_by_key', {})),
|
||||
'pipelines': len(getattr(self.pipeline_mgr, '_pipelines_by_key', {})),
|
||||
'knowledge_bases': len(getattr(self.rag_mgr, 'knowledge_bases', {})),
|
||||
'plugin_installations': len(
|
||||
getattr(
|
||||
self.plugin_connector,
|
||||
'_known_desired_states',
|
||||
{},
|
||||
)
|
||||
),
|
||||
}
|
||||
mcp_loader = getattr(self.tool_mgr, 'mcp_tool_loader', None)
|
||||
runtime_stats.update(
|
||||
{
|
||||
'mcp_sessions': len(getattr(mcp_loader, '_sessions', {})),
|
||||
'mcp_host_tasks': len(getattr(mcp_loader, '_hosted_mcp_tasks', ())),
|
||||
'mcp_dispatch_tasks': len(getattr(mcp_loader, '_host_dispatch_tasks', ())),
|
||||
}
|
||||
)
|
||||
|
||||
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,
|
||||
'query_pool': query_pool_stats,
|
||||
'models': model_stats,
|
||||
'runtimes': runtime_stats,
|
||||
'telemetry_tasks': len(getattr(self.telemetry, 'send_tasks', ())),
|
||||
}
|
||||
|
||||
async def initialize(self):
|
||||
pass
|
||||
|
||||
async def run(self):
|
||||
self.event_loop_monitor.start()
|
||||
try:
|
||||
if self.directory_projection_service is not None:
|
||||
self.task_mgr.create_task(
|
||||
@@ -392,18 +457,36 @@ class Application:
|
||||
|
||||
if self.task_mgr is not None:
|
||||
self.task_mgr.cancel_by_scope(core_entities.LifecycleControlScope.APPLICATION)
|
||||
with contextlib.suppress(Exception):
|
||||
await self.event_loop_monitor.stop()
|
||||
mcp_mount = getattr(self.http_ctrl, 'mcp_mount', None)
|
||||
if mcp_mount is not None:
|
||||
with contextlib.suppress(Exception):
|
||||
await mcp_mount.stop_session_manager()
|
||||
if self.platform_mgr is not None:
|
||||
with contextlib.suppress(Exception):
|
||||
await self.platform_mgr.shutdown()
|
||||
if self.tool_mgr is not None:
|
||||
with contextlib.suppress(Exception):
|
||||
await self.tool_mgr.shutdown()
|
||||
if self.model_mgr is not None:
|
||||
with contextlib.suppress(Exception):
|
||||
await self.model_mgr.shutdown()
|
||||
if self.box_service is not None:
|
||||
with contextlib.suppress(Exception):
|
||||
await self.box_service.shutdown()
|
||||
if self.plugin_connector is not None:
|
||||
with contextlib.suppress(Exception):
|
||||
await self.plugin_connector.aclose()
|
||||
if self.telemetry is not None:
|
||||
with contextlib.suppress(Exception):
|
||||
await self.telemetry.shutdown()
|
||||
if self.vector_db_mgr is not None:
|
||||
with contextlib.suppress(Exception):
|
||||
await self.vector_db_mgr.shutdown()
|
||||
if self.storage_mgr is not None:
|
||||
with contextlib.suppress(Exception):
|
||||
await self.storage_mgr.shutdown()
|
||||
manifest_provider = getattr(self.deployment, 'manifest_provider', None)
|
||||
if manifest_provider is not None:
|
||||
with contextlib.suppress(Exception):
|
||||
@@ -413,13 +496,29 @@ class Application:
|
||||
tasks = [wrapper.task for wrapper in self.task_mgr.tasks if not wrapper.task.done()]
|
||||
if tasks:
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
with contextlib.suppress(Exception):
|
||||
await httpclient.close_all()
|
||||
persistence_shutdown = getattr(self.persistence_mgr, 'shutdown', None)
|
||||
if callable(persistence_shutdown):
|
||||
with contextlib.suppress(Exception):
|
||||
await persistence_shutdown()
|
||||
else:
|
||||
# Compatibility for lightweight test/application doubles.
|
||||
persistence_db = getattr(self.persistence_mgr, 'db', None)
|
||||
persistence_engine = getattr(persistence_db, 'engine', None)
|
||||
if persistence_engine is not None:
|
||||
with contextlib.suppress(Exception):
|
||||
await persistence_engine.dispose()
|
||||
self._shutdown_complete = True
|
||||
|
||||
def dispose(self):
|
||||
"""Compatibility wrapper for callers that cannot await shutdown."""
|
||||
if self._shutdown_complete:
|
||||
return
|
||||
loop = self.event_loop
|
||||
if loop is not None and not loop.is_closed():
|
||||
loop.create_task(self.shutdown())
|
||||
if self._shutdown_task is None or self._shutdown_task.done():
|
||||
self._shutdown_task = loop.create_task(self.shutdown())
|
||||
return
|
||||
if self.plugin_connector is not None:
|
||||
self.plugin_connector.dispose()
|
||||
|
||||
@@ -2,6 +2,7 @@ from __future__ import annotations
|
||||
|
||||
import traceback
|
||||
import asyncio
|
||||
import contextlib
|
||||
import os
|
||||
|
||||
from . import app
|
||||
@@ -32,14 +33,22 @@ async def make_app(loop: asyncio.AbstractEventLoop) -> app.Application:
|
||||
|
||||
ap.event_loop = loop
|
||||
|
||||
# Execute startup stage
|
||||
for stage_name in stage_order:
|
||||
stage_cls = stage.preregistered_stages[stage_name]
|
||||
stage_inst = stage_cls()
|
||||
try:
|
||||
# Execute startup stage
|
||||
for stage_name in stage_order:
|
||||
stage_cls = stage.preregistered_stages[stage_name]
|
||||
stage_inst = stage_cls()
|
||||
|
||||
await stage_inst.run(ap)
|
||||
await stage_inst.run(ap)
|
||||
|
||||
await ap.initialize()
|
||||
await ap.initialize()
|
||||
except BaseException:
|
||||
# ``main()`` cannot clean up a partially built application because
|
||||
# ``make_app()`` has not returned it yet. Release managers, pools and
|
||||
# child processes that earlier startup stages already attached.
|
||||
with contextlib.suppress(BaseException):
|
||||
await ap.shutdown()
|
||||
raise
|
||||
|
||||
return ap
|
||||
|
||||
|
||||
@@ -0,0 +1,2 @@
|
||||
class TaskCapacityError(RuntimeError):
|
||||
"""Raised when the configured user-task admission limit is exhausted."""
|
||||
@@ -139,8 +139,8 @@ class BuildAppStage(stage.BootingStage):
|
||||
ap.log_cache = log_cache
|
||||
|
||||
storage_mgr_inst = storagemgr.StorageMgr(ap)
|
||||
await storage_mgr_inst.initialize()
|
||||
ap.storage_mgr = storage_mgr_inst
|
||||
await storage_mgr_inst.initialize()
|
||||
|
||||
persistence_mgr_inst = persistencemgr.PersistenceManager(
|
||||
ap,
|
||||
@@ -166,6 +166,12 @@ class BuildAppStage(stage.BootingStage):
|
||||
if not workspace_policy.multi_workspace_enabled:
|
||||
await workspace_service_inst.ensure_singleton_workspace()
|
||||
ap.workspace_service = workspace_service_inst
|
||||
if workspace_policy.multi_workspace_enabled:
|
||||
# Directory refresh starts in Application.run(), after this serial
|
||||
# build graph. Share one validated immutable binding snapshot
|
||||
# across model/platform/pipeline/RAG/plugin initialization instead
|
||||
# of repeating tenant validation for every manager.
|
||||
await workspace_service_inst.prime_startup_execution_bindings()
|
||||
|
||||
ap.workspace_collaboration_service = workspace_collaboration_module.WorkspaceCollaborationService(
|
||||
ap,
|
||||
@@ -188,14 +194,17 @@ class BuildAppStage(stage.BootingStage):
|
||||
trigger_principal=PrincipalContext(PrincipalType.SYSTEM),
|
||||
)
|
||||
|
||||
concurrency_config = ap.instance_config.data.get('concurrency', {})
|
||||
ap.query_pool = pool.QueryPool(
|
||||
singleton_context_resolver=resolve_singleton_execution_context,
|
||||
max_queries=int(concurrency_config.get('pending_queries', 1000)),
|
||||
max_queries_per_workspace=int(concurrency_config.get('pending_queries_per_workspace', 100)),
|
||||
)
|
||||
|
||||
# Telemetry manager: attach to app so other components can call via self.ap.telemetry
|
||||
telemetry_inst = telemetry_module.TelemetryManager(ap)
|
||||
await telemetry_inst.initialize()
|
||||
ap.telemetry = telemetry_inst
|
||||
await telemetry_inst.initialize()
|
||||
|
||||
# Survey manager
|
||||
survey_inst = survey_module.SurveyManager(ap)
|
||||
@@ -215,16 +224,16 @@ class BuildAppStage(stage.BootingStage):
|
||||
ap.sess_mgr = llm_session_mgr_inst
|
||||
|
||||
box_service_inst = box_service.BoxService(ap)
|
||||
await box_service_inst.initialize()
|
||||
ap.box_service = box_service_inst
|
||||
await box_service_inst.initialize()
|
||||
|
||||
llm_tool_mgr_inst = llm_tool_mgr.ToolManager(ap)
|
||||
await llm_tool_mgr_inst.initialize()
|
||||
ap.tool_mgr = llm_tool_mgr_inst
|
||||
await llm_tool_mgr_inst.initialize()
|
||||
|
||||
im_mgr_inst = im_mgr.PlatformManager(ap=ap)
|
||||
await im_mgr_inst.initialize()
|
||||
ap.platform_mgr = im_mgr_inst
|
||||
await im_mgr_inst.initialize()
|
||||
|
||||
# Initialize webhook pusher
|
||||
webhook_pusher_inst = WebhookPusher(ap)
|
||||
@@ -252,12 +261,12 @@ class BuildAppStage(stage.BootingStage):
|
||||
|
||||
# 初始化向量数据库管理器
|
||||
vectordb_mgr_inst = vectordb_mgr.VectorDBManager(ap)
|
||||
await vectordb_mgr_inst.initialize()
|
||||
ap.vector_db_mgr = vectordb_mgr_inst
|
||||
await vectordb_mgr_inst.initialize()
|
||||
|
||||
http_ctrl = http_controller.HTTPController(ap)
|
||||
await http_ctrl.initialize()
|
||||
ap.http_ctrl = http_ctrl
|
||||
await http_ctrl.initialize()
|
||||
|
||||
monitoring_service_inst = monitoring_service.MonitoringService(ap)
|
||||
ap.monitoring_service = monitoring_service_inst
|
||||
@@ -277,6 +286,7 @@ class BuildAppStage(stage.BootingStage):
|
||||
ap.logger.warning(f'Plugin runtime unavailable during startup; reconnecting in background: {exc}')
|
||||
plugin_connector_inst.schedule_reconnect()
|
||||
ap.plugin_connector = plugin_connector_inst
|
||||
workspace_service_inst.release_startup_execution_bindings()
|
||||
|
||||
ctrl = controller.Controller(ap)
|
||||
ap.ctrl = ctrl
|
||||
|
||||
@@ -3,7 +3,7 @@ from __future__ import annotations
|
||||
import os
|
||||
import copy
|
||||
from typing import Any
|
||||
from langbot.pkg.utils import constants
|
||||
from langbot.pkg.utils import bounded_executor, constants
|
||||
import yaml
|
||||
import importlib.resources as resources
|
||||
import uuid
|
||||
@@ -14,6 +14,13 @@ from ..bootutils import config
|
||||
|
||||
|
||||
_RUNTIME_POLICY_DEFAULTS = {
|
||||
'system': {
|
||||
'blocking_executor': {
|
||||
'max_workers': bounded_executor.DEFAULT_MAX_WORKERS,
|
||||
'max_pending': bounded_executor.DEFAULT_MAX_PENDING,
|
||||
'max_inflight_per_scope': (bounded_executor.DEFAULT_MAX_INFLIGHT_PER_SCOPE),
|
||||
}
|
||||
},
|
||||
'plugin': {
|
||||
'worker': {
|
||||
'max_cpus': 1.0,
|
||||
@@ -21,6 +28,10 @@ _RUNTIME_POLICY_DEFAULTS = {
|
||||
'max_pids': 128,
|
||||
'max_open_files': 256,
|
||||
'max_file_size_mb': 512,
|
||||
'max_workers': 16,
|
||||
'max_total_cpus': 8.0,
|
||||
'max_total_memory_mb': 8192,
|
||||
'max_installations': 10000,
|
||||
'require_hard_limits': False,
|
||||
}
|
||||
},
|
||||
@@ -205,6 +216,14 @@ class LoadConfigStage(stage.BootingStage):
|
||||
# Apply environment variable overrides to data/config.yaml
|
||||
ap.instance_config.data = _apply_env_overrides_to_config(ap.instance_config.data)
|
||||
|
||||
blocking_config = ap.instance_config.data['system']['blocking_executor']
|
||||
ap.blocking_executor = bounded_executor.configure_bounded_default_executor(
|
||||
ap.event_loop,
|
||||
max_workers=blocking_config['max_workers'],
|
||||
max_pending=blocking_config['max_pending'],
|
||||
max_inflight_per_scope=blocking_config['max_inflight_per_scope'],
|
||||
)
|
||||
|
||||
await ap.instance_config.dump_config()
|
||||
|
||||
# load or generate instance id
|
||||
|
||||
@@ -1,8 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
|
||||
from .. import stage, app, note
|
||||
from .. import entities as core_entities
|
||||
from ...utils import importutil
|
||||
|
||||
from .. import notes
|
||||
@@ -31,6 +30,12 @@ class ShowNotesStage(stage.BootingStage):
|
||||
if msg:
|
||||
ap.logger.log(level, msg)
|
||||
|
||||
asyncio.create_task(ayield_note(note_inst))
|
||||
ap.task_mgr.create_task(
|
||||
ayield_note(note_inst),
|
||||
kind='launch-note',
|
||||
name=f'launch-note-{note_cls.__name__}',
|
||||
scopes=[core_entities.LifecycleControlScope.APPLICATION],
|
||||
instance_uuid=ap.workspace_service.instance_uuid,
|
||||
)
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
@@ -4,6 +4,8 @@ import asyncio
|
||||
import contextvars
|
||||
import typing
|
||||
|
||||
from ..utils import bounded_executor
|
||||
|
||||
|
||||
T = typing.TypeVar('T')
|
||||
|
||||
@@ -14,6 +16,7 @@ def create_detached_task(
|
||||
loop: asyncio.AbstractEventLoop | None = None,
|
||||
name: str | None = None,
|
||||
after_commit_manager: typing.Any | None = None,
|
||||
workspace_uuid: str | None = None,
|
||||
) -> asyncio.Task[T]:
|
||||
"""Create a task that inherits no request-local ContextVars.
|
||||
|
||||
@@ -33,6 +36,11 @@ def create_detached_task(
|
||||
if callable(gate_factory):
|
||||
gate = gate_factory(after_commit_manager)
|
||||
task_coro = _wait_for_commit(coro, gate) if gate is not None else coro
|
||||
if workspace_uuid is not None:
|
||||
task_coro = bounded_executor.run_in_blocking_work_scope(
|
||||
task_coro,
|
||||
workspace_uuid,
|
||||
)
|
||||
return task_loop.create_task(task_coro, name=name, context=contextvars.Context())
|
||||
|
||||
|
||||
|
||||
@@ -7,6 +7,7 @@ import time
|
||||
|
||||
from . import app
|
||||
from . import entities as core_entities
|
||||
from .errors import TaskCapacityError
|
||||
from .task_boundary import create_detached_task
|
||||
|
||||
|
||||
@@ -22,13 +23,18 @@ class TaskContext:
|
||||
metadata: dict
|
||||
"""Structured metadata for progress reporting"""
|
||||
|
||||
def __init__(self):
|
||||
def __init__(self, max_log_chars: int = 200000):
|
||||
self.current_action = 'default'
|
||||
self.log = ''
|
||||
self.metadata = {}
|
||||
self.max_log_chars = max(int(max_log_chars), 1)
|
||||
|
||||
def _log(self, msg: str):
|
||||
self.log += msg + '\n'
|
||||
if len(self.log) > self.max_log_chars:
|
||||
marker = '[older task output truncated]\n'
|
||||
keep = max(self.max_log_chars - len(marker), 0)
|
||||
self.log = marker + (self.log[-keep:] if keep else '')
|
||||
|
||||
def set_current_action(self, action: str):
|
||||
self.current_action = action
|
||||
@@ -131,6 +137,7 @@ class TaskWrapper:
|
||||
loop=self.ap.event_loop,
|
||||
name=name or None,
|
||||
after_commit_manager=getattr(self.ap, 'persistence_mgr', None),
|
||||
workspace_uuid=workspace_uuid,
|
||||
)
|
||||
self.task_type = task_type
|
||||
self.kind = kind
|
||||
@@ -207,6 +214,39 @@ class AsyncTaskManager:
|
||||
self.ap = ap
|
||||
self.tasks = []
|
||||
|
||||
def _task_log_limit(self) -> int:
|
||||
value = self.ap.instance_config.data.get('system', {}).get('task_retention', {}).get('max_log_chars', 200000)
|
||||
try:
|
||||
value = int(value)
|
||||
except (TypeError, ValueError):
|
||||
value = 200000
|
||||
return max(value, 1)
|
||||
|
||||
def _user_task_limit(self, name: str, default: int) -> int:
|
||||
value = self.ap.instance_config.data.get('system', {}).get('task_retention', {}).get(name, default)
|
||||
try:
|
||||
value = int(value)
|
||||
except (TypeError, ValueError):
|
||||
value = default
|
||||
return max(value, 1)
|
||||
|
||||
def _admit_user_task(self, coro: typing.Coroutine, workspace_uuid: str | None) -> None:
|
||||
active_user_tasks = [
|
||||
wrapper for wrapper in self.tasks if wrapper.task_type == 'user' and not wrapper.task.done()
|
||||
]
|
||||
global_limit = self._user_task_limit('max_active_user_tasks', 256)
|
||||
if len(active_user_tasks) >= global_limit:
|
||||
coro.close()
|
||||
raise TaskCapacityError('The instance has too many active user operations')
|
||||
|
||||
if workspace_uuid is None:
|
||||
return
|
||||
workspace_limit = self._user_task_limit('max_active_user_tasks_per_workspace', 8)
|
||||
active_workspace_tasks = sum(1 for wrapper in active_user_tasks if wrapper.workspace_uuid == workspace_uuid)
|
||||
if active_workspace_tasks >= workspace_limit:
|
||||
coro.close()
|
||||
raise TaskCapacityError('The Workspace has too many active user operations')
|
||||
|
||||
def create_task(
|
||||
self,
|
||||
coro: typing.Coroutine,
|
||||
@@ -220,6 +260,13 @@ class AsyncTaskManager:
|
||||
workspace_uuid: str | None = None,
|
||||
placement_generation: int | None = None,
|
||||
) -> TaskWrapper:
|
||||
if context is None:
|
||||
context = TaskContext(max_log_chars=self._task_log_limit())
|
||||
else:
|
||||
context.max_log_chars = self._task_log_limit()
|
||||
if len(context.log) > context.max_log_chars:
|
||||
context.log = context.log[-context.max_log_chars :]
|
||||
|
||||
wrapper = TaskWrapper(
|
||||
self.ap,
|
||||
coro,
|
||||
@@ -250,6 +297,7 @@ class AsyncTaskManager:
|
||||
workspace_uuid: str | None = None,
|
||||
placement_generation: int | None = None,
|
||||
) -> TaskWrapper:
|
||||
self._admit_user_task(coro, workspace_uuid)
|
||||
return self.create_task(
|
||||
coro,
|
||||
'user',
|
||||
|
||||
@@ -211,6 +211,7 @@ class ComponentDiscoveryEngine:
|
||||
|
||||
def __init__(self, ap: app.Application):
|
||||
self.ap = ap
|
||||
self.components = {}
|
||||
|
||||
def load_component_manifest(self, path: str, owner: str = 'builtin', no_save: bool = False) -> Component | None:
|
||||
"""加载组件清单"""
|
||||
|
||||
@@ -11,11 +11,27 @@ from ..postgresql_url import normalize_asyncpg_url
|
||||
class PostgreSQLDatabaseManager(database.BaseDatabaseManager):
|
||||
"""PostgreSQL database manager"""
|
||||
|
||||
@staticmethod
|
||||
def _pool_integer(
|
||||
config: dict,
|
||||
name: str,
|
||||
default: int,
|
||||
*,
|
||||
minimum: int,
|
||||
) -> int:
|
||||
value = config.get(name, default)
|
||||
if isinstance(value, bool) or not isinstance(value, int) or value < minimum:
|
||||
comparator = 'non-negative' if minimum == 0 else 'positive'
|
||||
raise ValueError(f'database.postgresql.{name} must be a {comparator} integer')
|
||||
return value
|
||||
|
||||
async def initialize(self) -> None:
|
||||
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')
|
||||
if self.url_override is not None:
|
||||
engine_url = self.url_override
|
||||
else:
|
||||
postgresql_config = self.ap.instance_config.data.get('database', {}).get('postgresql', {})
|
||||
explicit_url = postgresql_config.get('url')
|
||||
if explicit_url:
|
||||
if not isinstance(explicit_url, str):
|
||||
@@ -37,4 +53,31 @@ class PostgreSQLDatabaseManager(database.BaseDatabaseManager):
|
||||
port=postgresql_config.get('port', 5432),
|
||||
database=postgresql_config.get('database', 'postgres'),
|
||||
)
|
||||
self.engine = sqlalchemy_asyncio.create_async_engine(engine_url)
|
||||
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_pre_ping=True,
|
||||
)
|
||||
|
||||
@@ -187,6 +187,14 @@ class PersistenceManager:
|
||||
if self.mode == PersistenceMode.OSS_COMPAT:
|
||||
await self.write_space_model_providers()
|
||||
|
||||
async def shutdown(self) -> None:
|
||||
"""Dispose the owned database engine when initialization or runtime ends."""
|
||||
|
||||
db = getattr(self, 'db', None)
|
||||
engine = getattr(db, 'engine', None)
|
||||
if engine is not None:
|
||||
await engine.dispose()
|
||||
|
||||
@contextlib.asynccontextmanager
|
||||
async def _release_migration_lock(self) -> typing.AsyncIterator[None]:
|
||||
"""Serialize the complete PostgreSQL migration and validation window."""
|
||||
|
||||
@@ -177,10 +177,7 @@ async def run_cloud_release_migration(
|
||||
await manager.initialize()
|
||||
ap.logger.info('Cloud PostgreSQL release migration reached and validated the exact release head.')
|
||||
finally:
|
||||
db = getattr(manager, 'db', None)
|
||||
engine = getattr(db, 'engine', None)
|
||||
if engine is not None:
|
||||
await engine.dispose()
|
||||
await manager.shutdown()
|
||||
|
||||
|
||||
async def run_cloud_release_migration_from_config(loop: asyncio.AbstractEventLoop) -> None:
|
||||
|
||||
@@ -72,6 +72,12 @@ class MessageAggregator:
|
||||
self.ap = ap
|
||||
self.buffers = {}
|
||||
self.lock = asyncio.Lock()
|
||||
concurrency = self.ap.instance_config.data.get('concurrency', {})
|
||||
self.max_buffers = max(int(concurrency.get('pending_queries', 1000)), 1)
|
||||
self.max_buffers_per_workspace = max(
|
||||
int(concurrency.get('pending_queries_per_workspace', 100)),
|
||||
1,
|
||||
)
|
||||
|
||||
def _get_aggregation_key(
|
||||
self,
|
||||
@@ -184,15 +190,26 @@ class MessageAggregator:
|
||||
)
|
||||
|
||||
force_flush = False
|
||||
bypass_aggregation = False
|
||||
async with self.lock:
|
||||
buffer = self.buffers.get(aggregation_key)
|
||||
if buffer is None:
|
||||
buffer = SessionBuffer(
|
||||
aggregation_key=aggregation_key,
|
||||
execution_context=execution_context,
|
||||
messages=[pending_msg],
|
||||
workspace_buffer_count = sum(
|
||||
1
|
||||
for key in self.buffers
|
||||
if key[0] == execution_context.instance_uuid
|
||||
and key[1] == execution_context.workspace_uuid
|
||||
and key[2] == execution_context.placement_generation
|
||||
)
|
||||
self.buffers[aggregation_key] = buffer
|
||||
if len(self.buffers) >= self.max_buffers or workspace_buffer_count >= self.max_buffers_per_workspace:
|
||||
bypass_aggregation = True
|
||||
else:
|
||||
buffer = SessionBuffer(
|
||||
aggregation_key=aggregation_key,
|
||||
execution_context=execution_context,
|
||||
messages=[pending_msg],
|
||||
)
|
||||
self.buffers[aggregation_key] = buffer
|
||||
else:
|
||||
if buffer.execution_context != execution_context:
|
||||
raise ExecutionContextMismatchError('Aggregation buffer ExecutionContext changed for the same key')
|
||||
@@ -200,14 +217,31 @@ class MessageAggregator:
|
||||
buffer.timer_task.cancel()
|
||||
buffer.messages.append(pending_msg)
|
||||
|
||||
buffer.last_message_time = time.time()
|
||||
if len(buffer.messages) >= MAX_BUFFER_MESSAGES:
|
||||
force_flush = True
|
||||
else:
|
||||
buffer.timer_task = create_detached_task(
|
||||
self._delayed_flush(aggregation_key, delay, execution_context),
|
||||
after_commit_manager=getattr(self.ap, 'persistence_mgr', None),
|
||||
)
|
||||
if not bypass_aggregation:
|
||||
buffer.last_message_time = time.time()
|
||||
if len(buffer.messages) >= MAX_BUFFER_MESSAGES:
|
||||
force_flush = True
|
||||
else:
|
||||
buffer.timer_task = create_detached_task(
|
||||
self._delayed_flush(aggregation_key, delay, execution_context),
|
||||
after_commit_manager=getattr(self.ap, 'persistence_mgr', None),
|
||||
workspace_uuid=execution_context.workspace_uuid,
|
||||
)
|
||||
|
||||
if bypass_aggregation:
|
||||
await self.ap.query_pool.add_query(
|
||||
bot_uuid=bot_uuid,
|
||||
launcher_type=launcher_type,
|
||||
launcher_id=launcher_id,
|
||||
sender_id=sender_id,
|
||||
message_event=message_event,
|
||||
message_chain=message_chain,
|
||||
adapter=adapter,
|
||||
pipeline_uuid=pipeline_uuid,
|
||||
routed_by_rule=routed_by_rule,
|
||||
execution_context=execution_context,
|
||||
)
|
||||
return
|
||||
|
||||
if force_flush:
|
||||
await self._flush_buffer(aggregation_key, execution_context)
|
||||
|
||||
@@ -23,7 +23,7 @@ class BaiduCloudExamine(filter_model.ContentFilter):
|
||||
'client_secret': self.ap.pipeline_cfg.data['baidu-cloud-examine']['api-secret'],
|
||||
},
|
||||
) as resp:
|
||||
return (await resp.json())['access_token']
|
||||
return (await httpclient.read_json_limited(resp))['access_token']
|
||||
|
||||
async def process(self, query: pipeline_query.Query, message: str) -> entities.FilterResult:
|
||||
session = httpclient.get_session()
|
||||
@@ -35,7 +35,7 @@ class BaiduCloudExamine(filter_model.ContentFilter):
|
||||
},
|
||||
data=f'text={message}'.encode('utf-8'),
|
||||
) as resp:
|
||||
result = await resp.json()
|
||||
result = await httpclient.read_json_limited(resp)
|
||||
|
||||
if 'error_code' in result:
|
||||
return entities.FilterResult(
|
||||
|
||||
@@ -1,9 +1,9 @@
|
||||
from __future__ import annotations
|
||||
import re
|
||||
|
||||
from .. import filter as filter_model
|
||||
from .. import entities
|
||||
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
|
||||
from ....utils.safe_regex import SafeRegexError, mask_patterns
|
||||
|
||||
|
||||
@filter_model.filter_class('ban-word-filter')
|
||||
@@ -14,22 +14,20 @@ class BanWordFilter(filter_model.ContentFilter):
|
||||
pass
|
||||
|
||||
async def process(self, query: pipeline_query.Query, message: str) -> entities.FilterResult:
|
||||
found = False
|
||||
|
||||
for word in self.ap.sensitive_meta.data['words']:
|
||||
match = re.findall(word, message)
|
||||
|
||||
if len(match) > 0:
|
||||
found = True
|
||||
|
||||
for i in range(len(match)):
|
||||
if self.ap.sensitive_meta.data['mask_word'] == '':
|
||||
message = message.replace(
|
||||
match[i],
|
||||
self.ap.sensitive_meta.data['mask'] * len(match[i]),
|
||||
)
|
||||
else:
|
||||
message = message.replace(match[i], self.ap.sensitive_meta.data['mask_word'])
|
||||
try:
|
||||
found, message = await mask_patterns(
|
||||
self.ap.sensitive_meta.data['words'],
|
||||
message,
|
||||
mask=self.ap.sensitive_meta.data['mask'],
|
||||
mask_word=self.ap.sensitive_meta.data['mask_word'],
|
||||
)
|
||||
except SafeRegexError as exc:
|
||||
return entities.FilterResult(
|
||||
level=entities.ResultLevel.BLOCK,
|
||||
replacement='',
|
||||
user_notice='内容检查规则执行失败，请联系管理员',
|
||||
console_notice=f'Sensitive-word regex rejected: {exc}',
|
||||
)
|
||||
|
||||
return entities.FilterResult(
|
||||
level=entities.ResultLevel.MASKED if found else entities.ResultLevel.PASS,
|
||||
|
||||
@@ -1,9 +1,9 @@
|
||||
from __future__ import annotations
|
||||
import re
|
||||
|
||||
from .. import entities
|
||||
from .. import filter as filter_model
|
||||
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
|
||||
from ....utils.safe_regex import SafeRegexError, matches_any
|
||||
|
||||
|
||||
@filter_model.filter_class('content-ignore')
|
||||
@@ -28,14 +28,25 @@ class ContentIgnore(filter_model.ContentFilter):
|
||||
)
|
||||
|
||||
if 'regexp' in query.pipeline_config['trigger']['ignore-rules']:
|
||||
for rule in query.pipeline_config['trigger']['ignore-rules']['regexp']:
|
||||
if re.search(rule, message):
|
||||
return entities.FilterResult(
|
||||
level=entities.ResultLevel.BLOCK,
|
||||
replacement='',
|
||||
user_notice='',
|
||||
console_notice='Ignore message according to regexp rule in ignore_rules',
|
||||
)
|
||||
try:
|
||||
matches = await matches_any(
|
||||
query.pipeline_config['trigger']['ignore-rules']['regexp'],
|
||||
message,
|
||||
)
|
||||
except SafeRegexError as exc:
|
||||
return entities.FilterResult(
|
||||
level=entities.ResultLevel.BLOCK,
|
||||
replacement='',
|
||||
user_notice='',
|
||||
console_notice=f'Ignore-rule regex rejected: {exc}',
|
||||
)
|
||||
if matches:
|
||||
return entities.FilterResult(
|
||||
level=entities.ResultLevel.BLOCK,
|
||||
replacement='',
|
||||
user_notice='',
|
||||
console_notice='Ignore message according to regexp rule in ignore_rules',
|
||||
)
|
||||
|
||||
return entities.FilterResult(
|
||||
level=entities.ResultLevel.PASS,
|
||||
|
||||
@@ -38,57 +38,80 @@ class Controller:
|
||||
raise WorkspaceInvariantError('Queued query instance does not match the active Workspace binding')
|
||||
return execution_context
|
||||
|
||||
async def _process_query(self, selected_query: pipeline_query.Query) -> None:
|
||||
async def _process_query(
|
||||
self,
|
||||
selected_query: pipeline_query.Query,
|
||||
*,
|
||||
selected_session=None,
|
||||
global_slot_reserved: bool = False,
|
||||
) -> None:
|
||||
"""Run one selected query and always release its scheduling slot."""
|
||||
|
||||
try:
|
||||
async with self.semaphore:
|
||||
queued_context = get_query_execution_context(selected_query)
|
||||
queued_context = get_query_execution_context(selected_query)
|
||||
|
||||
async def run_scoped_query() -> None:
|
||||
execution_context = await self._assert_query_execution_active(selected_query)
|
||||
pipeline_uuid = selected_query.pipeline_uuid
|
||||
async def run_scoped_query() -> None:
|
||||
execution_context = await self._assert_query_execution_active(selected_query)
|
||||
pipeline_uuid = selected_query.pipeline_uuid
|
||||
|
||||
if pipeline_uuid:
|
||||
pipeline = await self.ap.pipeline_mgr.get_pipeline_by_uuid(
|
||||
execution_context,
|
||||
pipeline_uuid,
|
||||
)
|
||||
if pipeline:
|
||||
await pipeline.run(selected_query)
|
||||
else:
|
||||
self.ap.logger.warning(
|
||||
f'Pipeline {pipeline_uuid} not found for query {selected_query.query_id}, query dropped'
|
||||
)
|
||||
if pipeline_uuid:
|
||||
pipeline = await self.ap.pipeline_mgr.get_pipeline_by_uuid(
|
||||
execution_context,
|
||||
pipeline_uuid,
|
||||
)
|
||||
if pipeline:
|
||||
await pipeline.run(selected_query)
|
||||
else:
|
||||
self.ap.logger.warning(f'No pipeline_uuid for query {selected_query.query_id}, query dropped')
|
||||
|
||||
tenant_scope = getattr(self.ap.persistence_mgr, 'tenant_scope', None)
|
||||
cloud_runtime = (
|
||||
getattr(getattr(self.ap.persistence_mgr, 'mode', None), 'value', None) == 'cloud_runtime'
|
||||
)
|
||||
if cloud_runtime:
|
||||
if not callable(tenant_scope):
|
||||
raise RuntimeError('Cloud query processing requires an explicit tenant scope')
|
||||
async with tenant_scope(queued_context.workspace_uuid):
|
||||
await run_scoped_query()
|
||||
self.ap.logger.warning(
|
||||
f'Pipeline {pipeline_uuid} not found for query {selected_query.query_id}, query dropped'
|
||||
)
|
||||
else:
|
||||
self.ap.logger.warning(f'No pipeline_uuid for query {selected_query.query_id}, query dropped')
|
||||
|
||||
tenant_scope = getattr(self.ap.persistence_mgr, 'tenant_scope', None)
|
||||
cloud_runtime = getattr(getattr(self.ap.persistence_mgr, 'mode', None), 'value', None) == 'cloud_runtime'
|
||||
if cloud_runtime:
|
||||
if not callable(tenant_scope):
|
||||
raise RuntimeError('Cloud query processing requires an explicit tenant scope')
|
||||
async with tenant_scope(queued_context.workspace_uuid):
|
||||
await run_scoped_query()
|
||||
else:
|
||||
await run_scoped_query()
|
||||
except WorkspaceError as exc:
|
||||
self.ap.logger.info(
|
||||
f'Dropped query {selected_query.query_id} because its Workspace execution binding is stale: {exc}'
|
||||
)
|
||||
finally:
|
||||
try:
|
||||
try:
|
||||
await self.ap.query_pool.remove_query(selected_query)
|
||||
finally:
|
||||
async with self.ap.query_pool:
|
||||
session = selected_session or await self.ap.sess_mgr.get_session(selected_query)
|
||||
try:
|
||||
session._semaphore.release()
|
||||
finally:
|
||||
self.ap.query_pool.condition.notify_all()
|
||||
finally:
|
||||
if global_slot_reserved:
|
||||
self.semaphore.release()
|
||||
|
||||
async def _drop_selected_query(self, selected_query, selected_session) -> None:
|
||||
"""Undo scheduler ownership when work cannot be handed to a task."""
|
||||
|
||||
try:
|
||||
await self.ap.query_pool.remove_query(selected_query)
|
||||
finally:
|
||||
async with self.ap.query_pool:
|
||||
(await self.ap.sess_mgr.get_session(selected_query))._semaphore.release()
|
||||
selected_session._semaphore.release()
|
||||
self.ap.query_pool.condition.notify_all()
|
||||
|
||||
async def consumer(self):
|
||||
"""事件处理循环"""
|
||||
try:
|
||||
while True:
|
||||
while True:
|
||||
try:
|
||||
selected_query: pipeline_query.Query = None
|
||||
selected_session = None
|
||||
|
||||
# 取请求
|
||||
async with self.ap.query_pool:
|
||||
@@ -101,7 +124,9 @@ class Controller:
|
||||
|
||||
if not session._semaphore.locked():
|
||||
selected_query = query
|
||||
selected_session = session
|
||||
await session._semaphore.acquire()
|
||||
self.ap.query_pool.mark_query_running_locked(query)
|
||||
# Only log when actually selecting a query
|
||||
self.ap.logger.debug(f'Selected query {query.query_id} for processing')
|
||||
|
||||
@@ -114,24 +139,48 @@ class Controller:
|
||||
continue
|
||||
|
||||
if selected_query:
|
||||
execution_context = get_query_execution_context(selected_query)
|
||||
self.ap.task_mgr.create_task(
|
||||
self._process_query(selected_query),
|
||||
kind='query',
|
||||
name=f'query-{selected_query.query_id}',
|
||||
scopes=[
|
||||
core_entities.LifecycleControlScope.APPLICATION,
|
||||
core_entities.LifecycleControlScope.PLATFORM,
|
||||
],
|
||||
instance_uuid=execution_context.instance_uuid,
|
||||
workspace_uuid=execution_context.workspace_uuid,
|
||||
placement_generation=execution_context.placement_generation,
|
||||
)
|
||||
try:
|
||||
# Reserve global capacity before creating the task.
|
||||
# At most one selected query is held by this consumer
|
||||
# while all pipeline slots are busy.
|
||||
await self.semaphore.acquire()
|
||||
except asyncio.CancelledError:
|
||||
await self._drop_selected_query(selected_query, selected_session)
|
||||
raise
|
||||
|
||||
except Exception as e:
|
||||
# traceback.print_exc()
|
||||
self.ap.logger.error(f'控制器循环出错: {e}')
|
||||
self.ap.logger.error(f'Traceback: {traceback.format_exc()}')
|
||||
execution_context = get_query_execution_context(selected_query)
|
||||
process_coro = self._process_query(
|
||||
selected_query,
|
||||
selected_session=selected_session,
|
||||
global_slot_reserved=True,
|
||||
)
|
||||
try:
|
||||
self.ap.task_mgr.create_task(
|
||||
process_coro,
|
||||
kind='query',
|
||||
name=f'query-{selected_query.query_id}',
|
||||
scopes=[
|
||||
core_entities.LifecycleControlScope.APPLICATION,
|
||||
core_entities.LifecycleControlScope.PLATFORM,
|
||||
],
|
||||
instance_uuid=execution_context.instance_uuid,
|
||||
workspace_uuid=execution_context.workspace_uuid,
|
||||
placement_generation=execution_context.placement_generation,
|
||||
)
|
||||
except Exception:
|
||||
process_coro.close()
|
||||
self.semaphore.release()
|
||||
await self._drop_selected_query(selected_query, selected_session)
|
||||
raise
|
||||
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as e:
|
||||
self.ap.logger.error(f'控制器循环出错: {e}')
|
||||
self.ap.logger.error(f'Traceback: {traceback.format_exc()}')
|
||||
# A persistent external failure must not turn this recovery
|
||||
# loop into a CPU spin.
|
||||
await asyncio.sleep(1)
|
||||
|
||||
async def run(self):
|
||||
"""运行控制器"""
|
||||
|
||||
@@ -1,9 +1,11 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import base64
|
||||
import time
|
||||
import re
|
||||
import uuid
|
||||
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
|
||||
@@ -28,28 +30,34 @@ class Text2ImageStrategy(strategy_model.LongTextStrategy):
|
||||
)
|
||||
|
||||
async def process(self, message: str, query: pipeline_query.Query) -> list[platform_message.MessageComponent]:
|
||||
img_path = self.text_to_image(
|
||||
text_str=message,
|
||||
save_as='temp/{}.png'.format(int(time.time())),
|
||||
query=query,
|
||||
)
|
||||
def render() -> str:
|
||||
render_id = f'{int(time.time())}-{uuid.uuid4().hex}'
|
||||
img_path = f'temp/{render_id}.png'
|
||||
compressed_path = f'temp/{render_id}-compressed.png'
|
||||
try:
|
||||
self.text_to_image(
|
||||
text_str=message,
|
||||
save_as=img_path,
|
||||
query=query,
|
||||
)
|
||||
compressed_path, _ = self.compress_image(
|
||||
img_path,
|
||||
outfile=compressed_path,
|
||||
)
|
||||
with open(compressed_path, 'rb') as f:
|
||||
return base64.b64encode(f.read()).decode('utf-8')
|
||||
finally:
|
||||
for path in {img_path, compressed_path}:
|
||||
if os.path.exists(path):
|
||||
os.remove(path)
|
||||
|
||||
compressed_path, size = self.compress_image(img_path, outfile='temp/{}_compressed.png'.format(int(time.time())))
|
||||
|
||||
with open(compressed_path, 'rb') as f:
|
||||
img = f.read()
|
||||
|
||||
b64 = base64.b64encode(img)
|
||||
|
||||
# 删除图片
|
||||
os.remove(img_path)
|
||||
|
||||
if os.path.exists(compressed_path):
|
||||
os.remove(compressed_path)
|
||||
# Font measurement, image rendering and compression are CPU-bound PIL
|
||||
# work and must not block the shared asyncio loop for every tenant.
|
||||
image_base64 = await asyncio.to_thread(render)
|
||||
|
||||
return [
|
||||
platform_message.Image(
|
||||
base64=b64.decode('utf-8'),
|
||||
base64=image_base64,
|
||||
)
|
||||
]
|
||||
|
||||
@@ -126,6 +134,36 @@ class Text2ImageStrategy(strategy_model.LongTextStrategy):
|
||||
o_size = self.get_size(outfile)
|
||||
return outfile, self.get_size(outfile)
|
||||
|
||||
def _split_text_lines(self, text_str: str, text_width: int, font) -> list[str]:
|
||||
"""Split text while guaranteeing that every loop iteration advances."""
|
||||
|
||||
final_lines: list[str] = []
|
||||
text_width = max(int(text_width), 1)
|
||||
for line in text_str.replace('\t', ' ').split('\n'):
|
||||
line_width = font.getlength(line)
|
||||
if not line or line_width < text_width:
|
||||
final_lines.append(line)
|
||||
continue
|
||||
|
||||
rest_text = line
|
||||
while rest_text:
|
||||
line_width = max(font.getlength(rest_text), 1)
|
||||
point = int(len(rest_text) * (text_width / line_width))
|
||||
point = max(1, min(point, len(rest_text)))
|
||||
|
||||
for number, number_index in self.indexNumber(rest_text):
|
||||
if number_index < point < number_index + len(number) and number_index != 0:
|
||||
point = number_index
|
||||
break
|
||||
|
||||
point = max(1, min(point, len(rest_text)))
|
||||
final_lines.append(rest_text[:point])
|
||||
rest_text = rest_text[point:]
|
||||
if rest_text and font.getlength(rest_text) < text_width:
|
||||
final_lines.append(rest_text)
|
||||
break
|
||||
return final_lines
|
||||
|
||||
def text_to_image(
|
||||
self,
|
||||
text_str: str,
|
||||
@@ -133,50 +171,9 @@ class Text2ImageStrategy(strategy_model.LongTextStrategy):
|
||||
width=800,
|
||||
query: pipeline_query.Query = None,
|
||||
):
|
||||
text_str = text_str.replace('\t', ' ')
|
||||
|
||||
# 分行
|
||||
lines = text_str.split('\n')
|
||||
|
||||
# 计算并分割
|
||||
final_lines = []
|
||||
|
||||
text_width = width - 80
|
||||
|
||||
self.ap.logger.debug('lines: {}, text_width: {}'.format(lines, text_width))
|
||||
for line in lines:
|
||||
# 如果长了就分割
|
||||
line_width = self.get_font(query.pipeline_config['output']['long-text-processing']['font-path']).getlength(
|
||||
line
|
||||
)
|
||||
self.ap.logger.debug('line_width: {}'.format(line_width))
|
||||
if line_width < text_width:
|
||||
final_lines.append(line)
|
||||
continue
|
||||
else:
|
||||
rest_text = line
|
||||
while True:
|
||||
# 分割最前面的一行
|
||||
point = int(len(rest_text) * (text_width / line_width))
|
||||
|
||||
# 检查断点是否在数字中间
|
||||
numbers = self.indexNumber(rest_text)
|
||||
|
||||
for number in numbers:
|
||||
if number[1] < point < number[1] + len(number[0]) and number[1] != 0:
|
||||
point = number[1]
|
||||
break
|
||||
|
||||
final_lines.append(rest_text[:point])
|
||||
rest_text = rest_text[point:]
|
||||
line_width = self.get_font(
|
||||
query.pipeline_config['output']['long-text-processing']['font-path']
|
||||
).getlength(rest_text)
|
||||
if line_width < text_width:
|
||||
final_lines.append(rest_text)
|
||||
break
|
||||
else:
|
||||
continue
|
||||
font = self.get_font(query.pipeline_config['output']['long-text-processing']['font-path'])
|
||||
text_width = max(width - 80, 1)
|
||||
final_lines = self._split_text_lines(text_str, text_width, font)
|
||||
# 准备画布
|
||||
img = Image.new('RGBA', (width, max(280, len(final_lines) * 35 + 65)), (255, 255, 255, 255))
|
||||
draw = ImageDraw.Draw(img, mode='RGBA')
|
||||
@@ -191,7 +188,7 @@ class Text2ImageStrategy(strategy_model.LongTextStrategy):
|
||||
(offset_x, offset_y + 35 * line_number),
|
||||
final_line,
|
||||
fill=(0, 0, 0),
|
||||
font=self.get_font(query.pipeline_config['output']['long-text-processing']['font-path']),
|
||||
font=font,
|
||||
)
|
||||
# 遍历此行,检查是否有emoji
|
||||
idx_in_line = 0
|
||||
|
||||
@@ -496,7 +496,50 @@ class PipelineManager:
|
||||
|
||||
def __init__(self, ap: app.Application):
|
||||
self.ap = ap
|
||||
self.pipelines = []
|
||||
self._pipelines_by_key: dict[
|
||||
tuple[str, str, str],
|
||||
RuntimePipeline,
|
||||
] = {}
|
||||
self._pipeline_keys_by_scope: dict[
|
||||
tuple[str, str],
|
||||
set[tuple[str, str, str]],
|
||||
] = {}
|
||||
self._scope_generations: dict[tuple[str, str], int] = {}
|
||||
|
||||
@property
|
||||
def pipelines(self) -> list[RuntimePipeline]:
|
||||
"""Compatibility view over the indexed runtime pipeline registry."""
|
||||
|
||||
return list(self._pipelines_by_key.values())
|
||||
|
||||
@pipelines.setter
|
||||
def pipelines(self, pipelines: list[RuntimePipeline]) -> None:
|
||||
self._pipelines_by_key = {}
|
||||
self._pipeline_keys_by_scope = {}
|
||||
for pipeline in pipelines:
|
||||
context = pipeline.execution_context
|
||||
pipeline_uuid = (
|
||||
getattr(getattr(pipeline, 'pipeline_entity', None), 'uuid', None) or context.pipeline_uuid or ''
|
||||
)
|
||||
key = (
|
||||
context.instance_uuid,
|
||||
pipeline.workspace_uuid,
|
||||
pipeline_uuid,
|
||||
)
|
||||
self._pipelines_by_key[key] = pipeline
|
||||
self._pipeline_keys_by_scope.setdefault(key[:2], set()).add(key)
|
||||
|
||||
def _observe_execution_context(self, context: ExecutionContext) -> None:
|
||||
scope = (context.instance_uuid, context.workspace_uuid)
|
||||
previous_generation = self._scope_generations.get(scope)
|
||||
if previous_generation is not None and context.placement_generation < previous_generation:
|
||||
raise WorkspaceInvariantError('Pipeline runtime placement generation rolled back')
|
||||
if previous_generation == context.placement_generation:
|
||||
return
|
||||
if previous_generation is not None:
|
||||
for key in self._pipeline_keys_by_scope.pop(scope, ()):
|
||||
self._pipelines_by_key.pop(key, None)
|
||||
self._scope_generations[scope] = context.placement_generation
|
||||
|
||||
async def initialize(self):
|
||||
self.stage_dict = {name: cls for name, cls in stage.preregistered_stages.items()}
|
||||
@@ -506,7 +549,9 @@ class PipelineManager:
|
||||
async def load_pipelines_from_db(self):
|
||||
self.ap.logger.info('Loading pipelines from db...')
|
||||
|
||||
self.pipelines = []
|
||||
self._pipelines_by_key = {}
|
||||
self._pipeline_keys_by_scope = {}
|
||||
self._scope_generations = {}
|
||||
list_bindings = getattr(self.ap.workspace_service, 'list_active_execution_bindings', None)
|
||||
tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
|
||||
cloud_runtime = getattr(getattr(self.ap.persistence_mgr, 'mode', None), 'value', None) == 'cloud_runtime'
|
||||
@@ -530,6 +575,7 @@ class PipelineManager:
|
||||
trigger_principal=PrincipalContext(PrincipalType.SYSTEM),
|
||||
),
|
||||
pipeline,
|
||||
_binding_validated=True,
|
||||
)
|
||||
return
|
||||
|
||||
@@ -575,6 +621,8 @@ class PipelineManager:
|
||||
pipeline_entity: persistence_pipeline.LegacyPipeline
|
||||
| sqlalchemy.Row[persistence_pipeline.LegacyPipeline]
|
||||
| dict,
|
||||
*,
|
||||
_binding_validated: bool = False,
|
||||
):
|
||||
if isinstance(pipeline_entity, sqlalchemy.Row):
|
||||
pipeline_entity = persistence_pipeline.LegacyPipeline(**pipeline_entity._mapping)
|
||||
@@ -584,10 +632,12 @@ class PipelineManager:
|
||||
execution_context = self._normalize_execution_context(context, pipeline_entity.uuid)
|
||||
if pipeline_entity.workspace_uuid != execution_context.workspace_uuid:
|
||||
raise WorkspaceRequiredError('Pipeline entity Workspace does not match its runtime context')
|
||||
await self.ap.workspace_service.get_execution_binding(
|
||||
execution_context.workspace_uuid,
|
||||
expected_generation=execution_context.placement_generation,
|
||||
)
|
||||
if not _binding_validated:
|
||||
await self.ap.workspace_service.get_execution_binding(
|
||||
execution_context.workspace_uuid,
|
||||
expected_generation=execution_context.placement_generation,
|
||||
)
|
||||
self._observe_execution_context(execution_context)
|
||||
|
||||
coerce_pipeline_config(
|
||||
pipeline_entity.config,
|
||||
@@ -605,13 +655,27 @@ class PipelineManager:
|
||||
for stage_container in stage_containers:
|
||||
await stage_container.inst.initialize(pipeline_entity.config)
|
||||
|
||||
# Stage initialization can yield while a Workspace is being moved.
|
||||
# Revalidate before publishing the runtime assembled above.
|
||||
if not _binding_validated:
|
||||
await self.ap.workspace_service.get_execution_binding(
|
||||
execution_context.workspace_uuid,
|
||||
expected_generation=execution_context.placement_generation,
|
||||
)
|
||||
self._observe_execution_context(execution_context)
|
||||
runtime_pipeline = RuntimePipeline(
|
||||
self.ap,
|
||||
pipeline_entity,
|
||||
stage_containers,
|
||||
execution_context,
|
||||
)
|
||||
self.pipelines.append(runtime_pipeline)
|
||||
key = (
|
||||
execution_context.instance_uuid,
|
||||
execution_context.workspace_uuid,
|
||||
pipeline_entity.uuid,
|
||||
)
|
||||
self._pipelines_by_key[key] = runtime_pipeline
|
||||
self._pipeline_keys_by_scope.setdefault(key[:2], set()).add(key)
|
||||
|
||||
async def get_pipeline_by_uuid(
|
||||
self,
|
||||
@@ -619,13 +683,24 @@ class PipelineManager:
|
||||
uuid: str,
|
||||
) -> RuntimePipeline | None:
|
||||
execution_context = self._normalize_execution_context(context, uuid)
|
||||
for pipeline in self.pipelines:
|
||||
if (
|
||||
pipeline.workspace_uuid == execution_context.workspace_uuid
|
||||
and pipeline.placement_generation == execution_context.placement_generation
|
||||
and pipeline.pipeline_entity.uuid == uuid
|
||||
):
|
||||
return pipeline
|
||||
await self.ap.workspace_service.get_execution_binding(
|
||||
execution_context.workspace_uuid,
|
||||
expected_generation=execution_context.placement_generation,
|
||||
)
|
||||
self._observe_execution_context(execution_context)
|
||||
key = (
|
||||
execution_context.instance_uuid,
|
||||
execution_context.workspace_uuid,
|
||||
uuid,
|
||||
)
|
||||
pipeline = self._pipelines_by_key.get(key)
|
||||
if pipeline is not None and pipeline.placement_generation == execution_context.placement_generation:
|
||||
return pipeline
|
||||
if not self._pipeline_keys_by_scope.get(key[:2]):
|
||||
self._scope_generations.pop(
|
||||
(execution_context.instance_uuid, execution_context.workspace_uuid),
|
||||
None,
|
||||
)
|
||||
return None
|
||||
|
||||
async def remove_pipeline(
|
||||
@@ -634,7 +709,21 @@ class PipelineManager:
|
||||
uuid: str,
|
||||
) -> None:
|
||||
execution_context = self._normalize_execution_context(context, uuid)
|
||||
for pipeline in self.pipelines:
|
||||
if pipeline.workspace_uuid == execution_context.workspace_uuid and pipeline.pipeline_entity.uuid == uuid:
|
||||
self.pipelines.remove(pipeline)
|
||||
return
|
||||
await self.ap.workspace_service.get_execution_binding(
|
||||
execution_context.workspace_uuid,
|
||||
expected_generation=execution_context.placement_generation,
|
||||
)
|
||||
self._observe_execution_context(execution_context)
|
||||
key = (
|
||||
execution_context.instance_uuid,
|
||||
execution_context.workspace_uuid,
|
||||
uuid,
|
||||
)
|
||||
if self._pipelines_by_key.pop(key, None) is not None:
|
||||
scope_keys = self._pipeline_keys_by_scope.get(key[:2])
|
||||
if scope_keys is not None:
|
||||
scope_keys.discard(key)
|
||||
if not scope_keys:
|
||||
self._pipeline_keys_by_scope.pop(key[:2], None)
|
||||
self._scope_generations.pop(key[:2], None)
|
||||
return
|
||||
|
||||
@@ -159,6 +159,16 @@ def _discard_query_state(query_key: int) -> None:
|
||||
_QUERY_STATES.pop(query_key, None)
|
||||
|
||||
|
||||
def discard_query_state(query: pipeline_query.Query) -> None:
|
||||
"""Release all diagnostics retained for a query leaving the runtime pool."""
|
||||
|
||||
query_key = id(query)
|
||||
state = _QUERY_STATES.get(query_key)
|
||||
if state is not None and state.finalizer is not None:
|
||||
state.finalizer.detach()
|
||||
_discard_query_state(query_key)
|
||||
|
||||
|
||||
def _discard_query_state_if_empty(query: pipeline_query.Query) -> None:
|
||||
query_key = id(query)
|
||||
state = _QUERY_STATES.get(query_key)
|
||||
@@ -166,9 +176,7 @@ def _discard_query_state_if_empty(query: pipeline_query.Query) -> None:
|
||||
return
|
||||
if state.pending_by_chain_id or state.by_response_index:
|
||||
return
|
||||
if state.finalizer is not None:
|
||||
state.finalizer.detach()
|
||||
_discard_query_state(query_key)
|
||||
discard_query_state(query)
|
||||
|
||||
|
||||
def _get_response_sources(
|
||||
|
||||
@@ -13,6 +13,7 @@ import langbot_plugin.api.entities.builtin.platform.message as platform_message
|
||||
import langbot_plugin.api.entities.builtin.provider.session as provider_session
|
||||
|
||||
from ..api.http.context import ExecutionContext
|
||||
from . import plugin_diagnostics
|
||||
|
||||
QueryCacheKey = tuple[str, str]
|
||||
LegacyQueryKey = tuple[str, int]
|
||||
@@ -35,6 +36,10 @@ class QueryNotFoundError(LookupError):
|
||||
"""Raised when a query does not exist inside the requested Workspace."""
|
||||
|
||||
|
||||
class QueryPoolCapacityError(RuntimeError):
|
||||
"""Raised when no queued query can be discarded to admit new work."""
|
||||
|
||||
|
||||
def _validate_execution_context(execution_context: ExecutionContext) -> None:
|
||||
if not isinstance(execution_context, ExecutionContext):
|
||||
raise ExecutionContextRequiredError('A trusted ExecutionContext is required')
|
||||
@@ -120,15 +125,100 @@ class QueryPool:
|
||||
def __init__(
|
||||
self,
|
||||
singleton_context_resolver: SingletonContextResolver | None = None,
|
||||
*,
|
||||
max_queries: int = 1000,
|
||||
max_queries_per_workspace: int = 100,
|
||||
):
|
||||
if max_queries < 1:
|
||||
raise ValueError('max_queries must be positive')
|
||||
if max_queries_per_workspace < 1:
|
||||
raise ValueError('max_queries_per_workspace must be positive')
|
||||
if max_queries_per_workspace > max_queries:
|
||||
raise ValueError('max_queries_per_workspace cannot exceed max_queries')
|
||||
self.query_id_counter = 0
|
||||
self.pool_lock = asyncio.Lock()
|
||||
self.queries = []
|
||||
self.cached_queries = {}
|
||||
self.active_query_count_by_workspace: dict[str, int] = {}
|
||||
self.legacy_query_index = {}
|
||||
self.query_count_by_scope = {}
|
||||
self.dropped_query_count_by_scope: dict[QueryCounterKey, int] = {}
|
||||
self.condition = asyncio.Condition(self.pool_lock)
|
||||
self._singleton_context_resolver = singleton_context_resolver
|
||||
self.max_queries = max_queries
|
||||
self.max_queries_per_workspace = max_queries_per_workspace
|
||||
|
||||
def _discard_queued_query_locked(
|
||||
self,
|
||||
*,
|
||||
workspace_uuid: str | None = None,
|
||||
) -> pipeline_query.Query | None:
|
||||
"""Discard the oldest queued query from one scope and all indexes."""
|
||||
|
||||
for index, query in enumerate(self.queries):
|
||||
execution_context = get_query_execution_context(query)
|
||||
if workspace_uuid is not None and execution_context.workspace_uuid != workspace_uuid:
|
||||
continue
|
||||
self.queries.pop(index)
|
||||
query_uuid = execution_context.query_uuid
|
||||
if query_uuid is not None:
|
||||
self.cached_queries.pop((execution_context.workspace_uuid, query_uuid), None)
|
||||
self.legacy_query_index.pop((execution_context.workspace_uuid, query.query_id), None)
|
||||
query_workspace_uuid = execution_context.workspace_uuid
|
||||
remaining = self.active_query_count_by_workspace.get(query_workspace_uuid, 0) - 1
|
||||
if remaining > 0:
|
||||
self.active_query_count_by_workspace[query_workspace_uuid] = remaining
|
||||
else:
|
||||
self.active_query_count_by_workspace.pop(query_workspace_uuid, None)
|
||||
plugin_diagnostics.discard_query_state(query)
|
||||
counter_key = (
|
||||
execution_context.instance_uuid,
|
||||
execution_context.workspace_uuid,
|
||||
execution_context.placement_generation,
|
||||
)
|
||||
self.dropped_query_count_by_scope[counter_key] = self.dropped_query_count_by_scope.get(counter_key, 0) + 1
|
||||
return query
|
||||
return None
|
||||
|
||||
def _admit_query_locked(self, workspace_uuid: str) -> None:
|
||||
workspace_query_count = self.active_query_count_by_workspace.get(workspace_uuid, 0)
|
||||
if workspace_query_count >= self.max_queries_per_workspace:
|
||||
if self._discard_queued_query_locked(workspace_uuid=workspace_uuid) is None:
|
||||
raise QueryPoolCapacityError(f'Workspace query capacity reached ({self.max_queries_per_workspace})')
|
||||
|
||||
if len(self.cached_queries) >= self.max_queries:
|
||||
if self._discard_queued_query_locked() is None:
|
||||
raise QueryPoolCapacityError(f'Global query capacity reached ({self.max_queries})')
|
||||
|
||||
def mark_query_running_locked(self, query: pipeline_query.Query) -> None:
|
||||
"""Remove a scheduled query from the overload-discardable queue."""
|
||||
|
||||
if not self.pool_lock.locked():
|
||||
raise RuntimeError('Query pool lock is required to schedule a query')
|
||||
for index, queued_query in enumerate(self.queries):
|
||||
if queued_query is query:
|
||||
self.queries.pop(index)
|
||||
return
|
||||
raise QueryNotFoundError('Scheduled query is no longer queued')
|
||||
|
||||
def _make_scope_counter_room_locked(self, counter_key: QueryCounterKey) -> None:
|
||||
"""Retain recent counters without pinning every historical Workspace."""
|
||||
|
||||
if counter_key in self.query_count_by_scope:
|
||||
return
|
||||
while len(self.query_count_by_scope) >= self.max_queries:
|
||||
stale_key = next(
|
||||
(
|
||||
existing_key
|
||||
for existing_key in self.query_count_by_scope
|
||||
if self.active_query_count_by_workspace.get(existing_key[1], 0) <= 0
|
||||
),
|
||||
None,
|
||||
)
|
||||
if stale_key is None:
|
||||
raise QueryPoolCapacityError('Query counter capacity reached while every Workspace is active')
|
||||
self.query_count_by_scope.pop(stale_key, None)
|
||||
self.dropped_query_count_by_scope.pop(stale_key, None)
|
||||
|
||||
async def resolve_execution_context(
|
||||
self,
|
||||
@@ -180,6 +270,7 @@ class QueryPool:
|
||||
)
|
||||
|
||||
async with self.condition:
|
||||
self._admit_query_locked(execution_context.workspace_uuid)
|
||||
query_id = self.query_id_counter
|
||||
initial_variables: dict[str, typing.Any] = {'_routed_by_rule': routed_by_rule}
|
||||
if variables:
|
||||
@@ -217,6 +308,9 @@ class QueryPool:
|
||||
|
||||
self.queries.append(query)
|
||||
self.cached_queries[(execution_context.workspace_uuid, query_uuid)] = query
|
||||
self.active_query_count_by_workspace[execution_context.workspace_uuid] = (
|
||||
self.active_query_count_by_workspace.get(execution_context.workspace_uuid, 0) + 1
|
||||
)
|
||||
self.legacy_query_index[(execution_context.workspace_uuid, query_id)] = query_uuid
|
||||
self.query_id_counter += 1
|
||||
counter_key = (
|
||||
@@ -224,6 +318,13 @@ class QueryPool:
|
||||
execution_context.workspace_uuid,
|
||||
execution_context.placement_generation,
|
||||
)
|
||||
# A Workspace has only one active placement. Drop obsolete
|
||||
# generation counters so deployment churn cannot grow these maps.
|
||||
for existing_key in tuple(self.query_count_by_scope):
|
||||
if existing_key[:2] == counter_key[:2] and existing_key != counter_key:
|
||||
self.query_count_by_scope.pop(existing_key, None)
|
||||
self.dropped_query_count_by_scope.pop(existing_key, None)
|
||||
self._make_scope_counter_room_locked(counter_key)
|
||||
self.query_count_by_scope[counter_key] = self.query_count_by_scope.get(counter_key, 0) + 1
|
||||
self.condition.notify_all()
|
||||
return query
|
||||
@@ -241,6 +342,19 @@ class QueryPool:
|
||||
0,
|
||||
)
|
||||
|
||||
def get_dropped_query_count(self, execution_context: ExecutionContext) -> int:
|
||||
"""Return overload drops for one active placement scope."""
|
||||
|
||||
_validate_execution_context(execution_context)
|
||||
return self.dropped_query_count_by_scope.get(
|
||||
(
|
||||
execution_context.instance_uuid,
|
||||
execution_context.workspace_uuid,
|
||||
execution_context.placement_generation,
|
||||
),
|
||||
0,
|
||||
)
|
||||
|
||||
async def get_query(
|
||||
self,
|
||||
workspace_uuid: str,
|
||||
@@ -290,6 +404,11 @@ class QueryPool:
|
||||
if cached_query is not query:
|
||||
return False
|
||||
del self.cached_queries[cache_key]
|
||||
remaining = self.active_query_count_by_workspace.get(execution_context.workspace_uuid, 0) - 1
|
||||
if remaining > 0:
|
||||
self.active_query_count_by_workspace[execution_context.workspace_uuid] = remaining
|
||||
else:
|
||||
self.active_query_count_by_workspace.pop(execution_context.workspace_uuid, None)
|
||||
self.legacy_query_index.pop(
|
||||
(execution_context.workspace_uuid, query.query_id),
|
||||
None,
|
||||
@@ -298,6 +417,7 @@ class QueryPool:
|
||||
if queued_query is query:
|
||||
self.queries.pop(index)
|
||||
break
|
||||
plugin_diagnostics.discard_query_state(query)
|
||||
return True
|
||||
|
||||
async def __aenter__(self):
|
||||
|
||||
@@ -26,6 +26,25 @@ importutil.import_modules_in_pkg(runners)
|
||||
|
||||
|
||||
class ChatMessageHandler(handler.MessageHandler):
|
||||
def _response_limit(self, name: str, default: int) -> int:
|
||||
instance_config = getattr(self.ap, 'instance_config', None)
|
||||
data = getattr(instance_config, 'data', {})
|
||||
if not isinstance(data, dict):
|
||||
return default
|
||||
value = data.get('system', {}).get('response_limits', {}).get(name, default)
|
||||
try:
|
||||
return max(int(value), 1)
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
|
||||
def _check_response_size(
|
||||
self,
|
||||
result: provider_message.Message | provider_message.MessageChunk,
|
||||
) -> None:
|
||||
content = result.content
|
||||
if isinstance(content, str) and len(content) > self._response_limit('max_generated_chars', 1024 * 1024):
|
||||
raise RuntimeError('Provider response exceeds the configured limit')
|
||||
|
||||
async def handle(
|
||||
self,
|
||||
query: pipeline_query.Query,
|
||||
@@ -87,6 +106,7 @@ class ChatMessageHandler(handler.MessageHandler):
|
||||
query.user_message.content = [event_ctx.event.user_message_alter]
|
||||
|
||||
text_length = 0
|
||||
runner = None
|
||||
try:
|
||||
is_stream = await query.adapter.is_stream_output_supported()
|
||||
except AttributeError:
|
||||
@@ -107,6 +127,7 @@ class ChatMessageHandler(handler.MessageHandler):
|
||||
chunk_count = 0 # Track streaming chunks to reduce excessive logging
|
||||
|
||||
async for result in runner.run(query):
|
||||
self._check_response_size(result)
|
||||
result.resp_message_id = str(resp_message_id)
|
||||
if query.resp_messages:
|
||||
query.resp_messages.pop()
|
||||
@@ -119,6 +140,11 @@ class ChatMessageHandler(handler.MessageHandler):
|
||||
query.resp_messages.append(result)
|
||||
|
||||
chunk_count += 1
|
||||
if chunk_count > self._response_limit(
|
||||
'max_stream_chunks',
|
||||
100_000,
|
||||
):
|
||||
raise RuntimeError('Provider stream exceeds the configured event limit')
|
||||
# Only log every 10th chunk to reduce excessive logging during streaming
|
||||
# This prevents memory overflow from thousands of log entries per conversation
|
||||
# First chunk uses INFO level to confirm connection establishment
|
||||
@@ -145,6 +171,7 @@ class ChatMessageHandler(handler.MessageHandler):
|
||||
|
||||
else:
|
||||
async for result in runner.run(query):
|
||||
self._check_response_size(result)
|
||||
query.resp_messages.append(result)
|
||||
|
||||
summary = self.format_result_log(result)
|
||||
@@ -159,6 +186,10 @@ class ChatMessageHandler(handler.MessageHandler):
|
||||
query.session.using_conversation.messages.append(query.user_message)
|
||||
|
||||
query.session.using_conversation.messages.extend(query.resp_messages)
|
||||
self.ap.sess_mgr.trim_conversation_messages(
|
||||
query.session.using_conversation,
|
||||
max_rounds=query.pipeline_config['ai']['local-agent'].get('max-round', 10),
|
||||
)
|
||||
except Exception as e:
|
||||
error_info = f'{traceback.format_exc()}'
|
||||
self.ap.logger.error(f'Conversation({query.query_id}) Request Failed: {error_info}')
|
||||
@@ -181,6 +212,13 @@ class ChatMessageHandler(handler.MessageHandler):
|
||||
debug_notice=traceback.format_exc(),
|
||||
)
|
||||
finally:
|
||||
if runner is not None:
|
||||
try:
|
||||
close_runner = getattr(runner, 'aclose', None)
|
||||
if close_runner is not None:
|
||||
await close_runner()
|
||||
except Exception as ex:
|
||||
self.ap.logger.warning(f'Failed to close request runner: {ex}')
|
||||
# Telemetry reporting: collect minimal per-query execution info and send asynchronously
|
||||
try:
|
||||
end_ts = time.time()
|
||||
|
||||
@@ -1,9 +1,17 @@
|
||||
from __future__ import annotations
|
||||
import asyncio
|
||||
from collections import OrderedDict
|
||||
import time
|
||||
import typing
|
||||
from .. import algo
|
||||
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
|
||||
from ...pool import get_query_execution_context
|
||||
|
||||
|
||||
_MAX_SESSION_CONTAINERS = 10000
|
||||
_MIN_CONTAINER_TTL_SECONDS = 300
|
||||
_CLEANUP_INTERVAL_SECONDS = 60
|
||||
_MAX_EVICTION_PROBES = 64
|
||||
|
||||
|
||||
# 固定窗口算法
|
||||
@@ -13,9 +21,11 @@ class SessionContainer:
|
||||
records: dict[int, int]
|
||||
"""访问记录，key为每窗口长度的起始时间戳，value为访问次数"""
|
||||
|
||||
def __init__(self):
|
||||
def __init__(self, ttl_seconds: int = _MIN_CONTAINER_TTL_SECONDS):
|
||||
self.wait_lock = asyncio.Lock()
|
||||
self.records = {}
|
||||
self.last_accessed = time.monotonic()
|
||||
self.ttl_seconds = ttl_seconds
|
||||
|
||||
|
||||
@algo.algo_class('fixwin')
|
||||
@@ -28,7 +38,8 @@ class FixedWindowAlgo(algo.ReteLimitAlgo):
|
||||
|
||||
async def initialize(self):
|
||||
self.containers_lock = asyncio.Lock()
|
||||
self.containers = {}
|
||||
self.containers = OrderedDict()
|
||||
self._last_cleanup = time.monotonic()
|
||||
|
||||
async def require_access(
|
||||
self,
|
||||
@@ -39,14 +50,53 @@ class FixedWindowAlgo(algo.ReteLimitAlgo):
|
||||
# 加锁，找容器
|
||||
container: SessionContainer = None
|
||||
|
||||
session_name = f'{launcher_type}_{launcher_id}'
|
||||
execution_context = get_query_execution_context(query)
|
||||
session_name = ':'.join(
|
||||
(
|
||||
execution_context.instance_uuid,
|
||||
execution_context.workspace_uuid,
|
||||
str(execution_context.placement_generation),
|
||||
str(getattr(query, 'bot_uuid', '')),
|
||||
str(getattr(query, 'pipeline_uuid', '')),
|
||||
str(launcher_type),
|
||||
str(launcher_id),
|
||||
)
|
||||
)
|
||||
|
||||
async with self.containers_lock:
|
||||
container = self.containers.get(session_name)
|
||||
|
||||
if container is None:
|
||||
container = SessionContainer()
|
||||
window_size = query.pipeline_config['safety']['rate-limit']['window-length']
|
||||
ttl_seconds = max(int(window_size) * 2, _MIN_CONTAINER_TTL_SECONDS)
|
||||
now_monotonic = time.monotonic()
|
||||
if now_monotonic - self._last_cleanup >= _CLEANUP_INTERVAL_SECONDS:
|
||||
self._last_cleanup = now_monotonic
|
||||
for key, candidate in tuple(self.containers.items()):
|
||||
if (
|
||||
not candidate.wait_lock.locked()
|
||||
and now_monotonic - candidate.last_accessed >= candidate.ttl_seconds
|
||||
):
|
||||
self.containers.pop(key, None)
|
||||
|
||||
if len(self.containers) >= _MAX_SESSION_CONTAINERS:
|
||||
for _ in range(min(_MAX_EVICTION_PROBES, len(self.containers))):
|
||||
oldest_key = next(iter(self.containers))
|
||||
oldest = self.containers[oldest_key]
|
||||
if oldest.wait_lock.locked():
|
||||
self.containers.move_to_end(oldest_key)
|
||||
continue
|
||||
self.containers.pop(oldest_key, None)
|
||||
break
|
||||
if len(self.containers) >= _MAX_SESSION_CONTAINERS:
|
||||
# Every retained session is actively waiting. Reject this
|
||||
# admission instead of growing an attacker-controlled map.
|
||||
return False
|
||||
container = SessionContainer(ttl_seconds=ttl_seconds)
|
||||
self.containers[session_name] = container
|
||||
else:
|
||||
self.containers.move_to_end(session_name)
|
||||
container.last_accessed = time.monotonic()
|
||||
|
||||
# 等待锁
|
||||
async with container.wait_lock:
|
||||
@@ -87,6 +137,7 @@ class FixedWindowAlgo(algo.ReteLimitAlgo):
|
||||
container.records[now] = count + 1
|
||||
|
||||
# 返回True
|
||||
container.last_accessed = time.monotonic()
|
||||
return True
|
||||
|
||||
async def release_access(
|
||||
|
||||
@@ -1,10 +1,8 @@
|
||||
import re
|
||||
|
||||
|
||||
from .. import rule as rule_model
|
||||
from .. import entities
|
||||
import langbot_plugin.api.entities.builtin.platform.message as platform_message
|
||||
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
|
||||
from ....utils.safe_regex import SafeRegexError, matches_any
|
||||
|
||||
|
||||
@rule_model.rule_class('regexp')
|
||||
@@ -16,15 +14,20 @@ class RegExpRule(rule_model.GroupRespondRule):
|
||||
rule_dict: dict,
|
||||
query: pipeline_query.Query,
|
||||
) -> entities.RuleJudgeResult:
|
||||
regexps = rule_dict['regexp']
|
||||
try:
|
||||
matching = await matches_any(
|
||||
rule_dict['regexp'],
|
||||
message_text,
|
||||
mode='match',
|
||||
)
|
||||
except SafeRegexError as exc:
|
||||
self.ap.logger.warning(f'Group response regex rejected: {exc}')
|
||||
matching = False
|
||||
|
||||
for regexp in regexps:
|
||||
match = re.match(regexp, message_text)
|
||||
|
||||
if match:
|
||||
return entities.RuleJudgeResult(
|
||||
matching=True,
|
||||
replacement=message_chain,
|
||||
)
|
||||
if matching:
|
||||
return entities.RuleJudgeResult(
|
||||
matching=True,
|
||||
replacement=message_chain,
|
||||
)
|
||||
|
||||
return entities.RuleJudgeResult(matching=False, replacement=message_chain)
|
||||
|
||||
+373
-106
@@ -1,10 +1,12 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import dataclasses
|
||||
import functools
|
||||
import json
|
||||
import re
|
||||
import time
|
||||
import traceback
|
||||
import uuid
|
||||
import sqlalchemy
|
||||
@@ -20,6 +22,7 @@ from ..entity.persistence import workspace as persistence_workspace
|
||||
from ..entity.errors import platform as platform_errors
|
||||
from ..api.http.context import ExecutionContext, PrincipalContext, PrincipalType, RequestContext
|
||||
from ..api.http.authz import WorkspaceRequiredError
|
||||
from ..workspace.errors import WorkspaceInvariantError
|
||||
|
||||
from .logger import EventLogger
|
||||
|
||||
@@ -40,7 +43,7 @@ class RuntimeBot:
|
||||
|
||||
adapter: abstract_platform_adapter.AbstractMessagePlatformAdapter
|
||||
|
||||
task_wrapper: taskmgr.TaskWrapper
|
||||
task_wrapper: taskmgr.TaskWrapper | None
|
||||
|
||||
task_context: taskmgr.TaskContext
|
||||
|
||||
@@ -80,7 +83,10 @@ class RuntimeBot:
|
||||
self.enable = bot_entity.enable
|
||||
self.adapter = adapter
|
||||
self.task_context = taskmgr.TaskContext()
|
||||
self.task_wrapper = None
|
||||
self.logger = logger
|
||||
self._shutdown_lock = asyncio.Lock()
|
||||
self._shutdown_complete = False
|
||||
|
||||
async def assert_execution_active(self) -> None:
|
||||
"""Fail closed when this long-lived adapter belongs to a stale placement."""
|
||||
@@ -490,12 +496,27 @@ class RuntimeBot:
|
||||
core_entities.LifecycleControlScope.APPLICATION,
|
||||
core_entities.LifecycleControlScope.PLATFORM,
|
||||
],
|
||||
instance_uuid=self.execution_context.instance_uuid,
|
||||
workspace_uuid=self.execution_context.workspace_uuid,
|
||||
placement_generation=(self.execution_context.placement_generation),
|
||||
)
|
||||
|
||||
async def shutdown(self):
|
||||
await self.adapter.kill()
|
||||
async with self._shutdown_lock:
|
||||
if self._shutdown_complete:
|
||||
return
|
||||
|
||||
self.ap.task_mgr.cancel_task(self.task_wrapper.id)
|
||||
wrapper = self.task_wrapper
|
||||
self.task_wrapper = None
|
||||
try:
|
||||
await asyncio.wait_for(self.adapter.kill(), timeout=15)
|
||||
finally:
|
||||
if wrapper is not None:
|
||||
self.ap.task_mgr.cancel_task(wrapper.id)
|
||||
if wrapper.task is not asyncio.current_task():
|
||||
with contextlib.suppress(asyncio.CancelledError, asyncio.TimeoutError):
|
||||
await asyncio.wait_for(wrapper.task, timeout=5)
|
||||
self._shutdown_complete = True
|
||||
|
||||
|
||||
# 控制QQ消息输入输出的类
|
||||
@@ -513,10 +534,140 @@ class PlatformManager:
|
||||
|
||||
def __init__(self, ap: app.Application = None):
|
||||
self.ap = ap
|
||||
self.bots = []
|
||||
self._bots_by_key: dict[tuple[str, str, str], RuntimeBot] = {}
|
||||
self._bot_keys_by_workspace: dict[
|
||||
str,
|
||||
set[tuple[str, str, str]],
|
||||
] = {}
|
||||
self._bot_keys_by_uuid: dict[str, set[tuple[str, str, str]]] = {}
|
||||
self.websocket_proxy_bots = {}
|
||||
self.adapter_components = []
|
||||
self.adapter_dict = {}
|
||||
self._scope_generations: dict[tuple[str, str], int] = {}
|
||||
self._proxy_last_accessed: dict[str, float] = {}
|
||||
self._runtime_mutation_lock = asyncio.Lock()
|
||||
|
||||
@staticmethod
|
||||
def _runtime_bot_key(bot: RuntimeBot) -> tuple[str, str, str]:
|
||||
context = getattr(bot, 'execution_context', None)
|
||||
instance_uuid = str(getattr(context, 'instance_uuid', '__test_instance__'))
|
||||
workspace_uuid = str(
|
||||
getattr(bot, 'workspace_uuid', None) or getattr(context, 'workspace_uuid', '__test_workspace__')
|
||||
)
|
||||
bot_uuid = str(
|
||||
getattr(getattr(bot, 'bot_entity', None), 'uuid', None)
|
||||
or getattr(context, 'bot_uuid', None)
|
||||
or f'__runtime_{id(bot)}'
|
||||
)
|
||||
return instance_uuid, workspace_uuid, bot_uuid
|
||||
|
||||
def _register_runtime_bot(self, bot: RuntimeBot) -> RuntimeBot | None:
|
||||
key = self._runtime_bot_key(bot)
|
||||
previous = self._bots_by_key.get(key)
|
||||
self._bots_by_key[key] = bot
|
||||
self._bot_keys_by_workspace.setdefault(key[1], set()).add(key)
|
||||
self._bot_keys_by_uuid.setdefault(key[2], set()).add(key)
|
||||
return previous
|
||||
|
||||
def _pop_runtime_bot(
|
||||
self,
|
||||
key: tuple[str, str, str],
|
||||
) -> RuntimeBot | None:
|
||||
bot = self._bots_by_key.pop(key, None)
|
||||
if bot is None:
|
||||
return None
|
||||
workspace_keys = self._bot_keys_by_workspace.get(key[1])
|
||||
if workspace_keys is not None:
|
||||
workspace_keys.discard(key)
|
||||
if not workspace_keys:
|
||||
self._bot_keys_by_workspace.pop(key[1], None)
|
||||
uuid_keys = self._bot_keys_by_uuid.get(key[2])
|
||||
if uuid_keys is not None:
|
||||
uuid_keys.discard(key)
|
||||
if not uuid_keys:
|
||||
self._bot_keys_by_uuid.pop(key[2], None)
|
||||
return bot
|
||||
|
||||
@property
|
||||
def bots(self) -> list[RuntimeBot]:
|
||||
"""Compatibility view over the indexed platform runtime registry."""
|
||||
|
||||
return list(self._bots_by_key.values())
|
||||
|
||||
@bots.setter
|
||||
def bots(self, bots: list[RuntimeBot]) -> None:
|
||||
self._bots_by_key = {}
|
||||
self._bot_keys_by_workspace = {}
|
||||
self._bot_keys_by_uuid = {}
|
||||
for bot in bots:
|
||||
self._register_runtime_bot(bot)
|
||||
|
||||
def _max_workspace_proxies(self) -> int:
|
||||
instance_data = getattr(
|
||||
getattr(self.ap, 'instance_config', None),
|
||||
'data',
|
||||
{},
|
||||
)
|
||||
value = instance_data.get('system', {}).get('websocket_retention', {}).get('max_workspace_proxies', 1024)
|
||||
try:
|
||||
return max(int(value), 1)
|
||||
except (TypeError, ValueError):
|
||||
return 1024
|
||||
|
||||
async def _evict_idle_websocket_proxy_unlocked(self) -> None:
|
||||
"""Make room without interrupting a live socket or in-flight query."""
|
||||
|
||||
if len(self.websocket_proxy_bots) < self._max_workspace_proxies():
|
||||
return
|
||||
from .sources.websocket_manager import WebSocketScope, ws_connection_manager
|
||||
|
||||
for workspace_uuid in sorted(
|
||||
self.websocket_proxy_bots,
|
||||
key=lambda item: self._proxy_last_accessed.get(item, 0.0),
|
||||
):
|
||||
proxy_bot = self.websocket_proxy_bots[workspace_uuid]
|
||||
listener_tasks = getattr(proxy_bot.adapter, 'inbound_listener_tasks', ())
|
||||
if any(not task.done() for task in tuple(listener_tasks)):
|
||||
continue
|
||||
scope = WebSocketScope.from_context(proxy_bot.execution_context)
|
||||
if ws_connection_manager.get_stats(scope=scope)['total_connections'] > 0:
|
||||
continue
|
||||
self.websocket_proxy_bots.pop(workspace_uuid, None)
|
||||
self._proxy_last_accessed.pop(workspace_uuid, None)
|
||||
await proxy_bot.shutdown()
|
||||
if not self._bot_keys_by_workspace.get(workspace_uuid):
|
||||
self._scope_generations.pop(
|
||||
(proxy_bot.execution_context.instance_uuid, workspace_uuid),
|
||||
None,
|
||||
)
|
||||
return
|
||||
raise RuntimeError('WebSocket Workspace proxy capacity reached and every proxy is active')
|
||||
|
||||
async def _observe_execution_context(self, context: ExecutionContext) -> None:
|
||||
"""Shutdown superseded Workspace adapters when placement advances."""
|
||||
|
||||
async with self._runtime_mutation_lock:
|
||||
await self._observe_execution_context_unlocked(context)
|
||||
|
||||
async def _observe_execution_context_unlocked(self, context: ExecutionContext) -> None:
|
||||
scope = (context.instance_uuid, context.workspace_uuid)
|
||||
previous_generation = self._scope_generations.get(scope)
|
||||
if previous_generation is not None and context.placement_generation < previous_generation:
|
||||
raise WorkspaceInvariantError('Platform runtime placement generation rolled back')
|
||||
if previous_generation == context.placement_generation:
|
||||
return
|
||||
if previous_generation is not None:
|
||||
proxy_bot = self.websocket_proxy_bots.pop(context.workspace_uuid, None)
|
||||
self._proxy_last_accessed.pop(context.workspace_uuid, None)
|
||||
if proxy_bot is not None and proxy_bot.enable:
|
||||
await proxy_bot.shutdown()
|
||||
for key in tuple(self._bot_keys_by_workspace.get(context.workspace_uuid, ())):
|
||||
bot = self._pop_runtime_bot(key)
|
||||
if bot is None:
|
||||
continue
|
||||
if bot.enable:
|
||||
await bot.shutdown()
|
||||
self._scope_generations[scope] = context.placement_generation
|
||||
|
||||
async def initialize(self):
|
||||
# delete all bot log images
|
||||
@@ -606,47 +757,60 @@ class PlatformManager:
|
||||
context: ExecutionContext | RequestContext,
|
||||
) -> RuntimeBot:
|
||||
execution_context = self._normalize_execution_context(context)
|
||||
existing = self.websocket_proxy_bots.get(execution_context.workspace_uuid)
|
||||
if existing is not None:
|
||||
if existing.placement_generation != execution_context.placement_generation:
|
||||
raise WorkspaceRequiredError('WebSocket proxy placement generation is stale')
|
||||
return existing
|
||||
|
||||
binding = await self.ap.workspace_service.get_execution_binding(
|
||||
await self.ap.workspace_service.get_execution_binding(
|
||||
execution_context.workspace_uuid,
|
||||
expected_generation=execution_context.placement_generation,
|
||||
)
|
||||
websocket_adapter_class = self.adapter_dict['websocket']
|
||||
websocket_logger = EventLogger(
|
||||
name='websocket-adapter',
|
||||
ap=self.ap,
|
||||
execution_context=execution_context,
|
||||
owner='websocket-proxy-bot',
|
||||
)
|
||||
websocket_adapter_inst = websocket_adapter_class({}, websocket_logger, ap=self.ap)
|
||||
proxy_context = dataclasses.replace(
|
||||
execution_context,
|
||||
instance_uuid=binding.instance_uuid,
|
||||
bot_uuid='websocket-proxy-bot',
|
||||
)
|
||||
runtime_bot = RuntimeBot(
|
||||
ap=self.ap,
|
||||
bot_entity=persistence_bot.Bot(
|
||||
uuid='websocket-proxy-bot',
|
||||
workspace_uuid=binding.workspace_uuid,
|
||||
name='WebSocket',
|
||||
description='',
|
||||
adapter='websocket',
|
||||
adapter_config={},
|
||||
enable=True,
|
||||
),
|
||||
adapter=websocket_adapter_inst,
|
||||
logger=websocket_logger,
|
||||
execution_context=proxy_context,
|
||||
)
|
||||
await runtime_bot.initialize()
|
||||
self.websocket_proxy_bots[binding.workspace_uuid] = runtime_bot
|
||||
return runtime_bot
|
||||
async with self._runtime_mutation_lock:
|
||||
await self.ap.workspace_service.get_execution_binding(
|
||||
execution_context.workspace_uuid,
|
||||
expected_generation=execution_context.placement_generation,
|
||||
)
|
||||
await self._observe_execution_context_unlocked(execution_context)
|
||||
existing = self.websocket_proxy_bots.get(execution_context.workspace_uuid)
|
||||
if existing is not None:
|
||||
if existing.placement_generation != execution_context.placement_generation:
|
||||
raise WorkspaceRequiredError('WebSocket proxy placement generation is stale')
|
||||
self._proxy_last_accessed[execution_context.workspace_uuid] = time.monotonic()
|
||||
return existing
|
||||
|
||||
await self._evict_idle_websocket_proxy_unlocked()
|
||||
binding = await self.ap.workspace_service.get_execution_binding(
|
||||
execution_context.workspace_uuid,
|
||||
expected_generation=execution_context.placement_generation,
|
||||
)
|
||||
websocket_adapter_class = self.adapter_dict['websocket']
|
||||
websocket_logger = EventLogger(
|
||||
name='websocket-adapter',
|
||||
ap=self.ap,
|
||||
execution_context=execution_context,
|
||||
owner='websocket-proxy-bot',
|
||||
)
|
||||
websocket_adapter_inst = websocket_adapter_class({}, websocket_logger, ap=self.ap)
|
||||
proxy_context = dataclasses.replace(
|
||||
execution_context,
|
||||
instance_uuid=binding.instance_uuid,
|
||||
bot_uuid='websocket-proxy-bot',
|
||||
)
|
||||
runtime_bot = RuntimeBot(
|
||||
ap=self.ap,
|
||||
bot_entity=persistence_bot.Bot(
|
||||
uuid='websocket-proxy-bot',
|
||||
workspace_uuid=binding.workspace_uuid,
|
||||
name='WebSocket',
|
||||
description='',
|
||||
adapter='websocket',
|
||||
adapter_config={},
|
||||
enable=True,
|
||||
),
|
||||
adapter=websocket_adapter_inst,
|
||||
logger=websocket_logger,
|
||||
execution_context=proxy_context,
|
||||
)
|
||||
await runtime_bot.initialize()
|
||||
self.websocket_proxy_bots[binding.workspace_uuid] = runtime_bot
|
||||
self._proxy_last_accessed[binding.workspace_uuid] = time.monotonic()
|
||||
return runtime_bot
|
||||
|
||||
def get_running_adapters(
|
||||
self,
|
||||
@@ -654,13 +818,52 @@ class PlatformManager:
|
||||
) -> list[abstract_platform_adapter.AbstractMessagePlatformAdapter]:
|
||||
execution_context = self._normalize_execution_context(context)
|
||||
return [
|
||||
bot.adapter for bot in self.bots if bot.enable and bot.workspace_uuid == execution_context.workspace_uuid
|
||||
bot.adapter
|
||||
for bot in self.bots
|
||||
if bot.enable
|
||||
and bot.workspace_uuid == execution_context.workspace_uuid
|
||||
and bot.placement_generation == execution_context.placement_generation
|
||||
]
|
||||
|
||||
async def load_bots_from_db(self):
|
||||
self.ap.logger.info('Loading bots from db...')
|
||||
|
||||
self.bots = []
|
||||
async with self._runtime_mutation_lock:
|
||||
old_bots = [*self.websocket_proxy_bots.values(), *self.bots]
|
||||
self.websocket_proxy_bots = {}
|
||||
self._proxy_last_accessed = {}
|
||||
self.bots = []
|
||||
self._scope_generations = {}
|
||||
for bot in old_bots:
|
||||
if not bot.enable:
|
||||
continue
|
||||
try:
|
||||
await bot.shutdown()
|
||||
except Exception as exc:
|
||||
self.ap.logger.warning(f'Failed to stop old platform runtime during reload: {exc}')
|
||||
|
||||
list_bindings = getattr(
|
||||
self.ap.workspace_service,
|
||||
'list_active_execution_bindings',
|
||||
None,
|
||||
)
|
||||
tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
|
||||
cloud_runtime = getattr(getattr(self.ap.persistence_mgr, 'mode', None), 'value', None) == 'cloud_runtime'
|
||||
if cloud_runtime:
|
||||
if not callable(list_bindings) or not callable(tenant_uow):
|
||||
raise RuntimeError('Cloud platform loading requires explicit instance discovery and tenant UoWs')
|
||||
for binding in await list_bindings():
|
||||
try:
|
||||
async with tenant_uow(binding.workspace_uuid):
|
||||
await self._load_workspace_bots(
|
||||
binding.workspace_uuid,
|
||||
_binding=binding,
|
||||
)
|
||||
except Exception as exc:
|
||||
self.ap.logger.error(
|
||||
f'Failed to load Workspace bots for {binding.workspace_uuid}: {exc}\n{traceback.format_exc()}'
|
||||
)
|
||||
return
|
||||
|
||||
instance_uow = getattr(self.ap.persistence_mgr, 'instance_discovery_uow', None)
|
||||
tenant_scope = getattr(self.ap.persistence_mgr, 'tenant_scope', None)
|
||||
@@ -698,15 +901,22 @@ class PlatformManager:
|
||||
f'Failed to load Workspace bots for {workspace_uuid}: {e}\n{traceback.format_exc()}'
|
||||
)
|
||||
|
||||
async def _load_workspace_bots(self, workspace_uuid: str) -> None:
|
||||
async def _load_workspace_bots(
|
||||
self,
|
||||
workspace_uuid: str,
|
||||
*,
|
||||
_binding=None,
|
||||
) -> None:
|
||||
result = await self.ap.persistence_mgr.execute_async(
|
||||
sqlalchemy.select(persistence_bot.Bot)
|
||||
.where(persistence_bot.Bot.workspace_uuid == workspace_uuid)
|
||||
.order_by(persistence_bot.Bot.uuid)
|
||||
)
|
||||
binding = _binding
|
||||
for bot in result.all():
|
||||
try:
|
||||
binding = await self.ap.workspace_service.get_execution_binding(workspace_uuid)
|
||||
if binding is None:
|
||||
binding = await self.ap.workspace_service.get_execution_binding(workspace_uuid)
|
||||
execution_context = ExecutionContext(
|
||||
instance_uuid=binding.instance_uuid,
|
||||
workspace_uuid=binding.workspace_uuid,
|
||||
@@ -714,7 +924,11 @@ class PlatformManager:
|
||||
bot_uuid=bot.uuid,
|
||||
trigger_principal=PrincipalContext(PrincipalType.SYSTEM),
|
||||
)
|
||||
await self.load_bot(execution_context, bot)
|
||||
await self.load_bot(
|
||||
execution_context,
|
||||
bot,
|
||||
_binding_validated=True,
|
||||
)
|
||||
except platform_errors.AdapterNotFoundError as e:
|
||||
self.ap.logger.warning(f'Adapter {e.adapter_name} not found, skipping bot {bot.uuid}')
|
||||
except Exception as e:
|
||||
@@ -724,6 +938,8 @@ class PlatformManager:
|
||||
self,
|
||||
context: ExecutionContext | RequestContext,
|
||||
bot_entity: persistence_bot.Bot | sqlalchemy.Row[persistence_bot.Bot] | dict,
|
||||
*,
|
||||
_binding_validated: bool = False,
|
||||
) -> RuntimeBot:
|
||||
"""加载机器人"""
|
||||
if isinstance(bot_entity, sqlalchemy.Row):
|
||||
@@ -734,45 +950,63 @@ class PlatformManager:
|
||||
execution_context = self._normalize_execution_context(context, bot_uuid=bot_entity.uuid)
|
||||
if bot_entity.workspace_uuid != execution_context.workspace_uuid:
|
||||
raise WorkspaceRequiredError('Bot entity Workspace does not match its runtime context')
|
||||
await self.ap.workspace_service.get_execution_binding(
|
||||
execution_context.workspace_uuid,
|
||||
expected_generation=execution_context.placement_generation,
|
||||
)
|
||||
if not _binding_validated:
|
||||
await self.ap.workspace_service.get_execution_binding(
|
||||
execution_context.workspace_uuid,
|
||||
expected_generation=execution_context.placement_generation,
|
||||
)
|
||||
async with self._runtime_mutation_lock:
|
||||
if not _binding_validated:
|
||||
await self.ap.workspace_service.get_execution_binding(
|
||||
execution_context.workspace_uuid,
|
||||
expected_generation=execution_context.placement_generation,
|
||||
)
|
||||
await self._observe_execution_context_unlocked(execution_context)
|
||||
|
||||
logger = EventLogger(
|
||||
name=f'platform-adapter-{bot_entity.name}',
|
||||
ap=self.ap,
|
||||
execution_context=execution_context,
|
||||
owner=bot_entity.uuid,
|
||||
)
|
||||
logger = EventLogger(
|
||||
name=f'platform-adapter-{bot_entity.name}',
|
||||
ap=self.ap,
|
||||
execution_context=execution_context,
|
||||
owner=bot_entity.uuid,
|
||||
)
|
||||
|
||||
if bot_entity.adapter not in self.adapter_dict:
|
||||
raise platform_errors.AdapterNotFoundError(bot_entity.adapter)
|
||||
if bot_entity.adapter not in self.adapter_dict:
|
||||
raise platform_errors.AdapterNotFoundError(bot_entity.adapter)
|
||||
|
||||
adapter_inst = self.adapter_dict[bot_entity.adapter](
|
||||
bot_entity.adapter_config,
|
||||
logger,
|
||||
)
|
||||
if hasattr(adapter_inst, 'ap'):
|
||||
adapter_inst.ap = self.ap
|
||||
adapter_inst = self.adapter_dict[bot_entity.adapter](
|
||||
bot_entity.adapter_config,
|
||||
logger,
|
||||
)
|
||||
if hasattr(adapter_inst, 'ap'):
|
||||
adapter_inst.ap = self.ap
|
||||
|
||||
# 如果 adapter 支持 set_bot_uuid 方法，设置 bot_uuid（用于统一 webhook）
|
||||
if hasattr(adapter_inst, 'set_bot_uuid'):
|
||||
adapter_inst.set_bot_uuid(bot_entity.uuid)
|
||||
# 如果 adapter 支持 set_bot_uuid 方法，设置 bot_uuid（用于统一 webhook）
|
||||
if hasattr(adapter_inst, 'set_bot_uuid'):
|
||||
adapter_inst.set_bot_uuid(bot_entity.uuid)
|
||||
|
||||
runtime_bot = RuntimeBot(
|
||||
ap=self.ap,
|
||||
bot_entity=bot_entity,
|
||||
adapter=adapter_inst,
|
||||
logger=logger,
|
||||
execution_context=execution_context,
|
||||
)
|
||||
runtime_bot = RuntimeBot(
|
||||
ap=self.ap,
|
||||
bot_entity=bot_entity,
|
||||
adapter=adapter_inst,
|
||||
logger=logger,
|
||||
execution_context=execution_context,
|
||||
)
|
||||
|
||||
await runtime_bot.initialize()
|
||||
await runtime_bot.initialize()
|
||||
|
||||
self.bots.append(runtime_bot)
|
||||
bot_key = self._runtime_bot_key(runtime_bot)
|
||||
existing_bot = self._bots_by_key.get(bot_key)
|
||||
if existing_bot is not None and existing_bot is not runtime_bot:
|
||||
try:
|
||||
if existing_bot.enable:
|
||||
await existing_bot.shutdown()
|
||||
except BaseException:
|
||||
if runtime_bot.enable:
|
||||
await runtime_bot.shutdown()
|
||||
raise
|
||||
self._register_runtime_bot(runtime_bot)
|
||||
|
||||
return runtime_bot
|
||||
return runtime_bot
|
||||
|
||||
async def get_bot_by_uuid(
|
||||
self,
|
||||
@@ -780,19 +1014,26 @@ class PlatformManager:
|
||||
bot_uuid: str,
|
||||
) -> RuntimeBot | None:
|
||||
execution_context = self._normalize_execution_context(context, bot_uuid=bot_uuid)
|
||||
await self.ap.workspace_service.get_execution_binding(
|
||||
execution_context.workspace_uuid,
|
||||
expected_generation=execution_context.placement_generation,
|
||||
)
|
||||
await self._observe_execution_context(execution_context)
|
||||
proxy_bot = self.websocket_proxy_bots.get(execution_context.workspace_uuid)
|
||||
if proxy_bot and proxy_bot.bot_entity.uuid == bot_uuid:
|
||||
if proxy_bot.placement_generation != execution_context.placement_generation:
|
||||
return None
|
||||
return proxy_bot
|
||||
for bot in self.bots:
|
||||
if (
|
||||
bot.workspace_uuid == execution_context.workspace_uuid
|
||||
and bot.placement_generation == execution_context.placement_generation
|
||||
and bot.bot_entity.uuid == bot_uuid
|
||||
):
|
||||
return bot
|
||||
return None
|
||||
bot = self._bots_by_key.get(
|
||||
(
|
||||
execution_context.instance_uuid,
|
||||
execution_context.workspace_uuid,
|
||||
bot_uuid,
|
||||
)
|
||||
)
|
||||
if bot is None or bot.placement_generation != execution_context.placement_generation:
|
||||
return None
|
||||
return bot
|
||||
|
||||
async def resolve_public_bot(self, route_key: str) -> RuntimeBot | None:
|
||||
"""Resolve an opaque public bot UUID without consulting request headers."""
|
||||
@@ -801,16 +1042,22 @@ class PlatformManager:
|
||||
normalized = str(uuid.UUID(route_key))
|
||||
except (ValueError, AttributeError, TypeError):
|
||||
return None
|
||||
for bot in self.bots:
|
||||
if bot.bot_entity.uuid == normalized:
|
||||
try:
|
||||
await self.ap.workspace_service.get_execution_binding(
|
||||
bot.workspace_uuid,
|
||||
expected_generation=bot.placement_generation,
|
||||
)
|
||||
except Exception:
|
||||
return None
|
||||
return bot
|
||||
keys = tuple(self._bot_keys_by_uuid.get(normalized, ()))
|
||||
if len(keys) != 1:
|
||||
return None
|
||||
key = keys[0]
|
||||
bot = self._bots_by_key.get(key)
|
||||
if bot is None:
|
||||
return None
|
||||
try:
|
||||
await self.ap.workspace_service.get_execution_binding(
|
||||
bot.workspace_uuid,
|
||||
expected_generation=bot.placement_generation,
|
||||
)
|
||||
except Exception:
|
||||
return None
|
||||
if self._bots_by_key.get(key) is bot:
|
||||
return bot
|
||||
return None
|
||||
|
||||
async def remove_bot(
|
||||
@@ -819,12 +1066,26 @@ class PlatformManager:
|
||||
bot_uuid: str,
|
||||
) -> None:
|
||||
execution_context = self._normalize_execution_context(context, bot_uuid=bot_uuid)
|
||||
for bot in self.bots[:]:
|
||||
if bot.workspace_uuid == execution_context.workspace_uuid and bot.bot_entity.uuid == bot_uuid:
|
||||
await self.ap.workspace_service.get_execution_binding(
|
||||
execution_context.workspace_uuid,
|
||||
expected_generation=execution_context.placement_generation,
|
||||
)
|
||||
async with self._runtime_mutation_lock:
|
||||
await self.ap.workspace_service.get_execution_binding(
|
||||
execution_context.workspace_uuid,
|
||||
expected_generation=execution_context.placement_generation,
|
||||
)
|
||||
await self._observe_execution_context_unlocked(execution_context)
|
||||
key = (
|
||||
execution_context.instance_uuid,
|
||||
execution_context.workspace_uuid,
|
||||
bot_uuid,
|
||||
)
|
||||
bot = self._bots_by_key.get(key)
|
||||
if bot is not None and bot.placement_generation == execution_context.placement_generation:
|
||||
if bot.enable:
|
||||
await bot.shutdown()
|
||||
self.bots.remove(bot)
|
||||
return
|
||||
self._pop_runtime_bot(key)
|
||||
|
||||
def get_available_adapters_info(self) -> list[dict]:
|
||||
return [
|
||||
@@ -853,10 +1114,16 @@ class PlatformManager:
|
||||
await bot.run()
|
||||
|
||||
async def shutdown(self):
|
||||
for proxy_bot in self.websocket_proxy_bots.values():
|
||||
if proxy_bot.enable:
|
||||
await proxy_bot.shutdown()
|
||||
for bot in self.bots:
|
||||
if bot.enable:
|
||||
await bot.shutdown()
|
||||
async with self._runtime_mutation_lock:
|
||||
runtime_bots = [*self.websocket_proxy_bots.values(), *self.bots]
|
||||
self.websocket_proxy_bots = {}
|
||||
self._proxy_last_accessed = {}
|
||||
self.bots = []
|
||||
for bot in runtime_bots:
|
||||
if not bot.enable:
|
||||
continue
|
||||
try:
|
||||
await bot.shutdown()
|
||||
except Exception as exc:
|
||||
self.ap.logger.warning(f'Failed to stop platform runtime during shutdown: {exc}')
|
||||
self.ap.task_mgr.cancel_by_scope(core_entities.LifecycleControlScope.PLATFORM)
|
||||
|
||||
@@ -55,6 +55,7 @@ class EventLog(pydantic.BaseModel):
|
||||
|
||||
MAX_LOG_COUNT = 200
|
||||
DELETE_COUNT_PER_TIME = 50
|
||||
MAX_LOG_TEXT_CHARS = 20000
|
||||
|
||||
|
||||
class EventLogger(abstract_platform_event_logger.AbstractEventLogger):
|
||||
@@ -129,7 +130,7 @@ class EventLogger(abstract_platform_event_logger.AbstractEventLogger):
|
||||
async def _truncate_logs(self):
|
||||
if len(self.logs) > MAX_LOG_COUNT:
|
||||
for i in range(DELETE_COUNT_PER_TIME):
|
||||
for image_key in self.logs[i].images: # type: ignore
|
||||
for image_key in self.logs[i].images or []:
|
||||
await self.ap.storage_mgr.delete_scoped_object_key(
|
||||
self.execution_context,
|
||||
image_key,
|
||||
@@ -147,6 +148,10 @@ class EventLogger(abstract_platform_event_logger.AbstractEventLogger):
|
||||
):
|
||||
try:
|
||||
image_keys = []
|
||||
text = str(text)
|
||||
if len(text) > MAX_LOG_TEXT_CHARS:
|
||||
marker = '\n[log truncated]'
|
||||
text = text[: MAX_LOG_TEXT_CHARS - len(marker)] + marker
|
||||
|
||||
if images is None:
|
||||
images = []
|
||||
|
||||
@@ -23,6 +23,7 @@ _GROUP_NAME_LOOKUP_TIMEOUT_SECONDS = 2
|
||||
_GROUP_MEMBER_INFO_CACHE_TTL_SECONDS = 86400
|
||||
_GROUP_MEMBER_INFO_NEGATIVE_CACHE_TTL_SECONDS = 600
|
||||
_GROUP_MEMBER_INFO_LOOKUP_TIMEOUT_SECONDS = 2
|
||||
_LOOKUP_CACHE_MAX = 4096
|
||||
|
||||
|
||||
def _normalize_base64_payload(value: str) -> str:
|
||||
@@ -366,6 +367,31 @@ class AiocqhttpEventConverter(abstract_platform_adapter.AbstractEventConverter):
|
||||
tuple[typing.Union[int, str], typing.Union[int, str]], tuple[dict, float]
|
||||
] = {}
|
||||
self._group_member_info_negative_cache: dict[tuple[typing.Union[int, str], typing.Union[int, str]], float] = {}
|
||||
self._last_cache_cleanup = 0.0
|
||||
|
||||
def _prune_caches(self, now: float) -> None:
|
||||
caches = (
|
||||
self._group_name_cache,
|
||||
self._group_name_negative_cache,
|
||||
self._group_member_info_cache,
|
||||
self._group_member_info_negative_cache,
|
||||
)
|
||||
if now - self._last_cache_cleanup < 60 and all(len(cache) <= _LOOKUP_CACHE_MAX for cache in caches):
|
||||
return
|
||||
self._last_cache_cleanup = now
|
||||
for cache in caches:
|
||||
for key, value in tuple(cache.items()):
|
||||
expires_at = value[1] if isinstance(value, tuple) else value
|
||||
if expires_at <= now:
|
||||
cache.pop(key, None)
|
||||
while len(cache) > _LOOKUP_CACHE_MAX:
|
||||
cache.pop(next(iter(cache)), None)
|
||||
|
||||
def clear(self) -> None:
|
||||
self._group_name_cache.clear()
|
||||
self._group_name_negative_cache.clear()
|
||||
self._group_member_info_cache.clear()
|
||||
self._group_member_info_negative_cache.clear()
|
||||
|
||||
@staticmethod
|
||||
async def yiri2target(event: platform_events.MessageEvent, bot_account_id: int):
|
||||
@@ -373,6 +399,7 @@ class AiocqhttpEventConverter(abstract_platform_adapter.AbstractEventConverter):
|
||||
|
||||
async def _get_group_name(self, group_id: typing.Union[int, str], bot=None) -> str:
|
||||
now = time.monotonic()
|
||||
self._prune_caches(now)
|
||||
if group_id in self._group_name_cache:
|
||||
group_name, expires_at = self._group_name_cache[group_id]
|
||||
if expires_at > now:
|
||||
@@ -408,6 +435,7 @@ class AiocqhttpEventConverter(abstract_platform_adapter.AbstractEventConverter):
|
||||
bot=None,
|
||||
) -> dict:
|
||||
now = time.monotonic()
|
||||
self._prune_caches(now)
|
||||
cache_key = (group_id, user_id)
|
||||
if cache_key in self._group_member_info_cache:
|
||||
member_info, expires_at = self._group_member_info_cache[cache_key]
|
||||
@@ -526,6 +554,8 @@ class AiocqhttpAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
|
||||
return
|
||||
|
||||
self.on_websocket_connection_event_cache.append(event)
|
||||
if len(self.on_websocket_connection_event_cache) > 100:
|
||||
self.on_websocket_connection_event_cache.pop(0)
|
||||
await self.logger.info(f'WebSocket connection established, bot id: {event.self_id}')
|
||||
|
||||
async def send_message(self, target_type: str, target_id: str, message: platform_message.MessageChain):
|
||||
@@ -691,4 +721,6 @@ class AiocqhttpAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
|
||||
async def kill(self) -> bool:
|
||||
# Current issue: existing connection will not be closed
|
||||
# self.should_shutdown = True
|
||||
self.on_websocket_connection_event_cache.clear()
|
||||
self.event_converter.clear()
|
||||
return False
|
||||
|
||||
@@ -5,6 +5,7 @@ import re
|
||||
import traceback
|
||||
import typing
|
||||
import uuid
|
||||
import time
|
||||
|
||||
from langbot.libs.dingtalk_api.dingtalkevent import DingTalkEvent
|
||||
import langbot_plugin.api.entities.builtin.platform.message as platform_message
|
||||
@@ -544,11 +545,59 @@ class DingTalkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
bot=bot,
|
||||
listeners={},
|
||||
)
|
||||
self._background_tasks: set[asyncio.Task] = set()
|
||||
# Wire the card-action callback after super().__init__ so we can reference
|
||||
# self.* — the client's handler stores this as a soft reference and reads
|
||||
# it at fire time.
|
||||
self.bot.card_action_callback = self._on_card_action
|
||||
|
||||
def _start_background_task(self, coro) -> bool:
|
||||
"""Start one bounded adapter-side auxiliary task."""
|
||||
|
||||
background_tasks = getattr(self, '_background_tasks', None)
|
||||
if background_tasks is None:
|
||||
background_tasks = set()
|
||||
object.__setattr__(self, '_background_tasks', background_tasks)
|
||||
for task in tuple(background_tasks):
|
||||
if task.done():
|
||||
background_tasks.discard(task)
|
||||
if len(background_tasks) >= 100:
|
||||
coro.close()
|
||||
return False
|
||||
task = asyncio.create_task(coro)
|
||||
background_tasks.add(task)
|
||||
|
||||
def done(done_task: asyncio.Task) -> None:
|
||||
background_tasks.discard(done_task)
|
||||
if not done_task.cancelled():
|
||||
done_task.exception()
|
||||
|
||||
task.add_done_callback(done)
|
||||
return True
|
||||
|
||||
def _prune_card_state(self) -> None:
|
||||
now = time.monotonic()
|
||||
ttl_seconds = 1800
|
||||
for card_id, state in tuple(self.card_state.items()):
|
||||
if now - float(state.get('created_at', now)) <= ttl_seconds:
|
||||
continue
|
||||
self.card_state.pop(card_id, None)
|
||||
for session_key, active_card_id in tuple(self.active_turn_card.items()):
|
||||
if active_card_id == card_id:
|
||||
self.active_turn_card.pop(session_key, None)
|
||||
self.active_turn_text.pop(session_key, None)
|
||||
while len(self.card_state) > 1000:
|
||||
card_id = next(iter(self.card_state))
|
||||
self.card_state.pop(card_id, None)
|
||||
while len(self.active_turn_card) > 1000:
|
||||
session_key = next(iter(self.active_turn_card))
|
||||
self.active_turn_card.pop(session_key, None)
|
||||
self.active_turn_text.pop(session_key, None)
|
||||
card_instances = getattr(self, 'card_instance_id_dict', None)
|
||||
if isinstance(card_instances, dict):
|
||||
while len(card_instances) > 1000:
|
||||
card_instances.pop(next(iter(card_instances)), None)
|
||||
|
||||
async def reply_message(
|
||||
self,
|
||||
message_source: platform_events.MessageEvent,
|
||||
@@ -674,6 +723,7 @@ class DingTalkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
return is_stream
|
||||
|
||||
async def create_message_card(self, message_id, event):
|
||||
self._prune_card_state()
|
||||
form_template_id = (self.config.get('human_input_card_template_id') or '').strip()
|
||||
legacy_template_id = self.config.get('card_template_id', '')
|
||||
|
||||
@@ -806,6 +856,20 @@ class DingTalkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
return params
|
||||
|
||||
async def kill(self) -> bool:
|
||||
task_set = getattr(self, '_background_tasks', set())
|
||||
background_tasks = list(task_set)
|
||||
for task in background_tasks:
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
if background_tasks:
|
||||
await asyncio.gather(*background_tasks, return_exceptions=True)
|
||||
task_set.clear()
|
||||
card_instances = getattr(self, 'card_instance_id_dict', None)
|
||||
if isinstance(card_instances, dict):
|
||||
card_instances.clear()
|
||||
self.card_state.clear()
|
||||
self.active_turn_card.clear()
|
||||
self.active_turn_text.clear()
|
||||
await self.bot.stop()
|
||||
return True
|
||||
|
||||
@@ -931,6 +995,7 @@ class DingTalkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
)
|
||||
|
||||
# Record form state for the click-handler.
|
||||
self._prune_card_state()
|
||||
launcher_type, launcher_id, sender_user_id = self._derive_session_descriptor(message_source)
|
||||
self.card_state[out_track_id] = {
|
||||
'session_key': session_key,
|
||||
@@ -947,6 +1012,7 @@ class DingTalkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
'current_input_field': str(form_data.get('_current_input_field') or ''),
|
||||
'input_defs': _dingtalk_form_input_defs(form_data),
|
||||
'inputs': form_data.get('inputs') or {},
|
||||
'created_at': time.monotonic(),
|
||||
}
|
||||
|
||||
btns = self._build_btns(actions if should_show_actions else [], out_track_id)
|
||||
@@ -1040,6 +1106,7 @@ class DingTalkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
f'options={len(component_params.get("select_options") or [])}'
|
||||
)
|
||||
|
||||
self._prune_card_state()
|
||||
self.card_state[out_track_id] = {
|
||||
'session_key': session_key,
|
||||
'launcher_type': launcher_type.value,
|
||||
@@ -1057,6 +1124,7 @@ class DingTalkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
'inputs': form_data.get('inputs') or {},
|
||||
'open_space_id': open_space_id,
|
||||
'is_group': is_group,
|
||||
'created_at': time.monotonic(),
|
||||
}
|
||||
|
||||
parts = []
|
||||
@@ -1223,6 +1291,7 @@ class DingTalkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
f'payload_action_id={payload.get("action_id")!r} params={payload.get("params")!r}'
|
||||
)
|
||||
out_track_id = payload.get('out_track_id') or ''
|
||||
self._prune_card_state()
|
||||
params = payload.get('params') or {}
|
||||
# ButtonGroup `sendCardRequest` events surface the click id at the
|
||||
# callback top level as `actionId`; fall back to `params.action_id`
|
||||
@@ -1359,7 +1428,7 @@ class DingTalkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
# output lives on a separate new card (lazy-created in
|
||||
# reply_message_chunk on the synthetic event), so the form card
|
||||
# stays put as a record of the user's selection.
|
||||
asyncio.create_task(
|
||||
self._start_background_task(
|
||||
self._mark_card_resolved(
|
||||
out_track_id,
|
||||
action_title,
|
||||
|
||||
@@ -28,6 +28,21 @@ import langbot_plugin.api.definition.abstract.platform.event_logger as abstract_
|
||||
from ..logger import EventLogger
|
||||
|
||||
|
||||
_MAX_DISCORD_MEDIA_BYTES = 10 * 1024 * 1024
|
||||
|
||||
|
||||
def _decode_discord_base64_limited(value: str) -> bytes:
|
||||
if ',' in value:
|
||||
value = value.split(',', 1)[1]
|
||||
max_encoded_bytes = 4 * ((_MAX_DISCORD_MEDIA_BYTES + 2) // 3)
|
||||
if len(value) > max_encoded_bytes:
|
||||
raise ValueError('Discord media exceeds the size limit')
|
||||
decoded = base64.b64decode(value)
|
||||
if len(decoded) > _MAX_DISCORD_MEDIA_BYTES:
|
||||
raise ValueError('Discord media exceeds the size limit')
|
||||
return decoded
|
||||
|
||||
|
||||
# 语音功能相关异常定义
|
||||
class VoiceConnectionError(Exception):
|
||||
"""语音连接基础异常"""
|
||||
@@ -604,7 +619,6 @@ class DiscordMessageConverter(abstract_platform_adapter.AbstractMessageConverter
|
||||
|
||||
for ele in message_chain:
|
||||
if isinstance(ele, platform_message.Image):
|
||||
image_bytes = None
|
||||
filename = f'{uuid.uuid4()}.png' # 默认文件名
|
||||
|
||||
if ele.base64:
|
||||
@@ -618,60 +632,17 @@ class DiscordMessageConverter(abstract_platform_adapter.AbstractMessageConverter
|
||||
filename = f'{uuid.uuid4()}.gif'
|
||||
elif 'webp' in data_header:
|
||||
filename = f'{uuid.uuid4()}.webp'
|
||||
# 去掉data:image/xxx;base64,前缀
|
||||
base64_data = ele.base64.split(',')[1]
|
||||
else:
|
||||
base64_data = ele.base64
|
||||
image_bytes = base64.b64decode(base64_data)
|
||||
elif ele.url:
|
||||
# 从URL下载图片
|
||||
session = httpclient.get_session()
|
||||
async with session.get(ele.url) as response:
|
||||
image_bytes = await response.read()
|
||||
# 从URL或Content-Type推断文件类型
|
||||
content_type = response.headers.get('Content-Type', '')
|
||||
if 'jpeg' in content_type or 'jpg' in content_type:
|
||||
filename = f'{uuid.uuid4()}.jpg'
|
||||
elif 'gif' in content_type:
|
||||
filename = f'{uuid.uuid4()}.gif'
|
||||
elif 'webp' in content_type:
|
||||
filename = f'{uuid.uuid4()}.webp'
|
||||
elif ele.url.lower().endswith(('.jpg', '.jpeg')):
|
||||
filename = f'{uuid.uuid4()}.jpg'
|
||||
elif ele.url.lower().endswith('.gif'):
|
||||
filename = f'{uuid.uuid4()}.gif'
|
||||
elif ele.url.lower().endswith('.webp'):
|
||||
filename = f'{uuid.uuid4()}.webp'
|
||||
elif ele.path:
|
||||
# 从文件路径读取图片
|
||||
# 确保路径没有空字节
|
||||
clean_path = ele.path.replace('\x00', '')
|
||||
clean_path = os.path.abspath(clean_path)
|
||||
|
||||
if not os.path.exists(clean_path):
|
||||
continue # 跳过不存在的文件
|
||||
|
||||
try:
|
||||
with open(clean_path, 'rb') as f:
|
||||
image_bytes = f.read()
|
||||
# 从文件路径获取文件名，保持原始扩展名
|
||||
original_filename = os.path.basename(clean_path)
|
||||
if original_filename and '.' in original_filename:
|
||||
# 保持原始文件名的扩展名
|
||||
ext = original_filename.split('.')[-1].lower()
|
||||
filename = f'{uuid.uuid4()}.{ext}'
|
||||
else:
|
||||
# 如果没有扩展名，尝试从文件内容检测
|
||||
if image_bytes.startswith(b'\xff\xd8\xff'):
|
||||
filename = f'{uuid.uuid4()}.jpg'
|
||||
elif image_bytes.startswith(b'GIF'):
|
||||
filename = f'{uuid.uuid4()}.gif'
|
||||
elif image_bytes.startswith(b'RIFF') and b'WEBP' in image_bytes[:20]:
|
||||
filename = f'{uuid.uuid4()}.webp'
|
||||
# 默认保持PNG
|
||||
except Exception as e:
|
||||
print(f'Error reading image file {clean_path}: {e}')
|
||||
continue # 跳过读取失败的文件
|
||||
try:
|
||||
image_bytes, mime_type = await ele.get_bytes()
|
||||
except Exception as exc:
|
||||
print(f'Error reading Discord image: {exc}')
|
||||
continue
|
||||
if 'jpeg' in mime_type or 'jpg' in mime_type:
|
||||
filename = f'{uuid.uuid4()}.jpg'
|
||||
elif 'gif' in mime_type:
|
||||
filename = f'{uuid.uuid4()}.gif'
|
||||
elif 'webp' in mime_type:
|
||||
filename = f'{uuid.uuid4()}.webp'
|
||||
|
||||
if image_bytes:
|
||||
files.append(discord.File(fp=io.BytesIO(image_bytes), filename=filename))
|
||||
@@ -702,27 +673,34 @@ class DiscordMessageConverter(abstract_platform_adapter.AbstractMessageConverter
|
||||
elif 'webm' in data_header:
|
||||
filename = f'{uuid.uuid4()}.webm'
|
||||
|
||||
file_base64 = ele.base64.split(',')[-1]
|
||||
file_bytes = base64.b64decode(file_base64)
|
||||
file_bytes = await asyncio.to_thread(
|
||||
_decode_discord_base64_limited,
|
||||
ele.base64,
|
||||
)
|
||||
elif ele.url:
|
||||
session = httpclient.get_session()
|
||||
async with session.get(ele.url) as response:
|
||||
file_bytes = await response.read()
|
||||
file_bytes = await httpclient.read_limited(
|
||||
response,
|
||||
max_bytes=_MAX_DISCORD_MEDIA_BYTES,
|
||||
)
|
||||
if file_bytes:
|
||||
files.append(discord.File(fp=io.BytesIO(file_bytes), filename=filename))
|
||||
elif isinstance(ele, platform_message.File):
|
||||
file_bytes = None
|
||||
filename = f'{uuid.uuid4()}.{ele.name.split(".")[-1]}'
|
||||
if ele.base64:
|
||||
if ele.base64.startswith('data:'):
|
||||
file_base64 = ele.base64.split(',')[1]
|
||||
file_bytes = base64.b64decode(file_base64)
|
||||
else:
|
||||
file_bytes = base64.b64decode(ele.base64)
|
||||
file_bytes = await asyncio.to_thread(
|
||||
_decode_discord_base64_limited,
|
||||
ele.base64,
|
||||
)
|
||||
elif ele.url:
|
||||
session = httpclient.get_session()
|
||||
async with session.get(ele.url) as response:
|
||||
file_bytes = await response.read()
|
||||
file_bytes = await httpclient.read_limited(
|
||||
response,
|
||||
max_bytes=_MAX_DISCORD_MEDIA_BYTES,
|
||||
)
|
||||
if file_bytes:
|
||||
files.append(discord.File(fp=io.BytesIO(file_bytes), filename=filename))
|
||||
elif isinstance(ele, platform_message.Forward):
|
||||
@@ -780,8 +758,11 @@ class DiscordMessageConverter(abstract_platform_adapter.AbstractMessageConverter
|
||||
for attachment in message.attachments:
|
||||
session = httpclient.get_session(trust_env=True)
|
||||
async with session.get(attachment.url) as response:
|
||||
image_data = await response.read()
|
||||
image_base64 = base64.b64encode(image_data).decode('utf-8')
|
||||
image_data = await httpclient.read_limited(
|
||||
response,
|
||||
max_bytes=_MAX_DISCORD_MEDIA_BYTES,
|
||||
)
|
||||
image_base64 = (await asyncio.to_thread(base64.b64encode, image_data)).decode('utf-8')
|
||||
image_format = response.headers['Content-Type']
|
||||
element_list.append(platform_message.Image(base64=f'data:{image_format};base64,{image_base64}'))
|
||||
|
||||
@@ -968,6 +949,20 @@ class DiscordAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
# resume.
|
||||
self._pending_forms: dict[str, dict] = {}
|
||||
|
||||
def _prune_transient_state(self) -> None:
|
||||
now = time.time()
|
||||
ttl_seconds = 1800
|
||||
for message_id, state in tuple(self._stream_buffer.items()):
|
||||
if now - float(state.get('updated_at', now)) > ttl_seconds:
|
||||
self._stream_buffer.pop(message_id, None)
|
||||
for session_key, state in tuple(self._pending_forms.items()):
|
||||
if now - float(state.get('posted_at', now)) > ttl_seconds:
|
||||
self._pending_forms.pop(session_key, None)
|
||||
while len(self._stream_buffer) > 100:
|
||||
self._stream_buffer.pop(next(iter(self._stream_buffer)), None)
|
||||
while len(self._pending_forms) > 1000:
|
||||
self._pending_forms.pop(next(iter(self._pending_forms)), None)
|
||||
|
||||
# Voice functionality methods
|
||||
async def join_voice_channel(self, guild_id: int, channel_id: int, user_id: int = None) -> discord.VoiceClient:
|
||||
"""
|
||||
@@ -1246,11 +1241,13 @@ class DiscordAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
source = event.source_platform_object
|
||||
if not isinstance(source, discord.Message):
|
||||
return False
|
||||
self._prune_transient_state()
|
||||
self._stream_buffer[message_id] = {
|
||||
'channel': source.channel,
|
||||
'sent_message': None, # discord.Message set on first send
|
||||
'last_content': '',
|
||||
'chunk_count': 0,
|
||||
'updated_at': time.time(),
|
||||
}
|
||||
return True
|
||||
|
||||
@@ -1274,6 +1271,8 @@ class DiscordAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
form_data = getattr(bot_message, '_form_data', None) if not isinstance(bot_message, dict) else None
|
||||
|
||||
ctx = self._stream_buffer.get(msg_id) if msg_id else None
|
||||
if ctx is not None:
|
||||
ctx['updated_at'] = time.time()
|
||||
|
||||
# If the stream ctx was not set up (create_message_card wasn't
|
||||
# called, e.g. synthetic event), or the final chunk carries a
|
||||
@@ -1342,6 +1341,7 @@ class DiscordAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
callback synthesizes a ``_dify_form_action`` query so the runner's
|
||||
``_merge_pending_form_action`` resumes the workflow.
|
||||
"""
|
||||
self._prune_transient_state()
|
||||
source = message_source.source_platform_object
|
||||
|
||||
actions = form_data.get('actions') or []
|
||||
@@ -1445,6 +1445,8 @@ class DiscordAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
and disable the View buttons so the choice is visually locked in."""
|
||||
import langbot_plugin.api.entities.builtin.provider.session as provider_session
|
||||
|
||||
self._prune_transient_state()
|
||||
|
||||
# ACK first (3-second deadline before Discord shows "interaction failed").
|
||||
try:
|
||||
await interaction.response.defer()
|
||||
@@ -1653,5 +1655,7 @@ class DiscordAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
if self.voice_manager:
|
||||
await self.voice_manager.disconnect_all()
|
||||
|
||||
self._stream_buffer.clear()
|
||||
self._pending_forms.clear()
|
||||
await self.bot.close()
|
||||
return True
|
||||
|
||||
@@ -54,6 +54,7 @@ _ERR = {
|
||||
'bad_signature': (401, 40101),
|
||||
'duplicate': (409, 40901),
|
||||
'too_large': (413, 41301),
|
||||
'overloaded': (503, 50301),
|
||||
'internal': (500, 50001),
|
||||
}
|
||||
|
||||
@@ -63,16 +64,21 @@ _MAX_BODY = 1 * 1024 * 1024
|
||||
# Idempotency dedup window (seconds) and cap.
|
||||
_IDEMPOTENCY_TTL = 600
|
||||
_IDEMPOTENCY_MAX = 4096
|
||||
_OUTBOUND_QUEUE_MAX = 100
|
||||
_OUTBOUND_IDLE_SECONDS = 60
|
||||
_OUTBOUND_STATE_MAX = 4096
|
||||
_INBOUND_TASK_MAX = 100
|
||||
|
||||
|
||||
class _SessionOutbound:
|
||||
"""Per-session outbound state: ordered delivery queue + sequence counter."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.queue: asyncio.Queue = asyncio.Queue(maxsize=1000)
|
||||
self.queue: asyncio.Queue = asyncio.Queue(maxsize=_OUTBOUND_QUEUE_MAX)
|
||||
self.worker: asyncio.Task | None = None
|
||||
self.sequence: int = 0
|
||||
self.last_was_final: bool = True # so the first reply of a turn starts at seq 1
|
||||
self.last_active: float = time.monotonic()
|
||||
|
||||
|
||||
class _SyncCollector:
|
||||
@@ -99,6 +105,7 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
idempotency_cache: dict[str, float] = pydantic.Field(default_factory=dict, exclude=True)
|
||||
# session_id -> sync collector (set while a /sync request is awaiting a turn)
|
||||
sync_waiters: dict[str, '_SyncCollector'] = pydantic.Field(default_factory=dict, exclude=True)
|
||||
inbound_tasks: set[asyncio.Task] = pydantic.Field(default_factory=set, exclude=True)
|
||||
|
||||
model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
@@ -108,6 +115,7 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
self.outbound_states = {}
|
||||
self.idempotency_cache = {}
|
||||
self.sync_waiters = {}
|
||||
self.inbound_tasks = set()
|
||||
|
||||
# -- framework hooks ------------------------------------------------------
|
||||
|
||||
@@ -156,10 +164,19 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
await asyncio.sleep(3600)
|
||||
|
||||
async def kill(self):
|
||||
# Cancel any outbound workers.
|
||||
tasks = list(self.inbound_tasks)
|
||||
for task in tasks:
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
for state in self.outbound_states.values():
|
||||
if state.worker and not state.worker.done():
|
||||
state.worker.cancel()
|
||||
tasks.append(state.worker)
|
||||
if tasks:
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
self.outbound_states.clear()
|
||||
self.sync_waiters.clear()
|
||||
self.inbound_tasks.clear()
|
||||
return True
|
||||
|
||||
# -- inbound --------------------------------------------------------------
|
||||
@@ -177,6 +194,24 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
for k in expired:
|
||||
self.idempotency_cache.pop(k, None)
|
||||
|
||||
def _start_inbound_task(self, coro: typing.Coroutine) -> asyncio.Task | None:
|
||||
self.inbound_tasks = {task for task in self.inbound_tasks if not task.done()}
|
||||
if len(self.inbound_tasks) >= _INBOUND_TASK_MAX:
|
||||
coro.close()
|
||||
return None
|
||||
task = asyncio.create_task(coro)
|
||||
self.inbound_tasks.add(task)
|
||||
|
||||
def task_done(done_task: asyncio.Task) -> None:
|
||||
self.inbound_tasks.discard(done_task)
|
||||
if not done_task.cancelled():
|
||||
# Retrieve failures so fire-and-forget callbacks never emit
|
||||
# "Task exception was never retrieved" or retain tracebacks.
|
||||
done_task.exception()
|
||||
|
||||
task.add_done_callback(task_done)
|
||||
return task
|
||||
|
||||
async def handle_unified_webhook(self, bot_uuid: str, path: str, request):
|
||||
"""Handle an inbound POST from the unified webhook dispatcher.
|
||||
|
||||
@@ -213,7 +248,7 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
return None, self._err('bad_signature', f'invalid signature: {reason}')
|
||||
|
||||
try:
|
||||
data = json.loads(body)
|
||||
data = await asyncio.to_thread(json.loads, body)
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
return None, self._err('bad_request', 'body is not valid JSON')
|
||||
if not isinstance(data, dict):
|
||||
@@ -282,7 +317,8 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
return await self._run_sync(event, listener, session_id, message_id)
|
||||
|
||||
# Fire-and-collect: kick the pipeline, return 202 immediately.
|
||||
asyncio.create_task(listener(event, self))
|
||||
if self._start_inbound_task(listener(event, self)) is None:
|
||||
return self._err('overloaded', 'too many inbound messages are already being processed')
|
||||
return quart.jsonify(
|
||||
{
|
||||
'code': 0,
|
||||
@@ -361,7 +397,9 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
return ''
|
||||
|
||||
def _next_sequence(self, session_id: str, is_final: bool) -> int:
|
||||
self._prune_outbound_states()
|
||||
state = self.outbound_states.setdefault(session_id, _SessionOutbound())
|
||||
state.last_active = time.monotonic()
|
||||
if state.last_was_final:
|
||||
state.sequence = 1
|
||||
else:
|
||||
@@ -369,8 +407,37 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
state.last_was_final = is_final
|
||||
return state.sequence
|
||||
|
||||
def _prune_outbound_states(self) -> None:
|
||||
now = time.monotonic()
|
||||
removable = [
|
||||
(session_id, state)
|
||||
for session_id, state in self.outbound_states.items()
|
||||
if (state.worker is None or state.worker.done())
|
||||
and state.queue.empty()
|
||||
and now - state.last_active >= _OUTBOUND_IDLE_SECONDS
|
||||
]
|
||||
for session_id, state in removable:
|
||||
if self.outbound_states.get(session_id) is state:
|
||||
self.outbound_states.pop(session_id, None)
|
||||
|
||||
overflow = len(self.outbound_states) - _OUTBOUND_STATE_MAX
|
||||
if overflow <= 0:
|
||||
return
|
||||
idle = sorted(
|
||||
(
|
||||
(session_id, state)
|
||||
for session_id, state in self.outbound_states.items()
|
||||
if (state.worker is None or state.worker.done()) and state.queue.empty()
|
||||
),
|
||||
key=lambda item: item[1].last_active,
|
||||
)
|
||||
for session_id, state in idle[:overflow]:
|
||||
if self.outbound_states.get(session_id) is state:
|
||||
self.outbound_states.pop(session_id, None)
|
||||
|
||||
async def _enqueue_callback(self, session_id: str, payload: dict) -> None:
|
||||
state = self.outbound_states.setdefault(session_id, _SessionOutbound())
|
||||
state.last_active = time.monotonic()
|
||||
if state.worker is None or state.worker.done():
|
||||
state.worker = asyncio.create_task(self._outbound_worker(session_id, state))
|
||||
try:
|
||||
@@ -386,13 +453,23 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
|
||||
async def _outbound_worker(self, session_id: str, state: _SessionOutbound) -> None:
|
||||
while True:
|
||||
payload = await state.queue.get()
|
||||
try:
|
||||
payload = await asyncio.wait_for(
|
||||
state.queue.get(),
|
||||
timeout=_OUTBOUND_IDLE_SECONDS,
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
if self.outbound_states.get(session_id) is state and state.queue.empty():
|
||||
self.outbound_states.pop(session_id, None)
|
||||
return
|
||||
continue
|
||||
try:
|
||||
await self._deliver_callback(payload)
|
||||
except Exception as e: # noqa: BLE001
|
||||
await self.logger.error(f'http_bot callback delivery failed for {session_id}: {e}')
|
||||
finally:
|
||||
state.queue.task_done()
|
||||
state.last_active = time.monotonic()
|
||||
|
||||
async def _deliver_callback(self, payload: dict) -> None:
|
||||
callback_url = self.config.get('callback_url', '')
|
||||
@@ -508,8 +585,11 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
|
||||
collector = _SyncCollector()
|
||||
self.sync_waiters[session_id] = collector
|
||||
listener_task = self._start_inbound_task(listener(event, self))
|
||||
if listener_task is None:
|
||||
self.sync_waiters.pop(session_id, None)
|
||||
return self._err('overloaded', 'too many inbound messages are already being processed')
|
||||
try:
|
||||
asyncio.create_task(listener(event, self))
|
||||
timeout = int(self.config.get('callback_timeout', 15)) * 4
|
||||
try:
|
||||
await asyncio.wait_for(collector.done.wait(), timeout=timeout)
|
||||
@@ -517,6 +597,9 @@ class HttpBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
await self.logger.warning(f'http_bot sync wait timed out for session {session_id}')
|
||||
finally:
|
||||
self.sync_waiters.pop(session_id, None)
|
||||
state = self.outbound_states.get(session_id)
|
||||
if state is not None and state.worker is None and state.queue.empty():
|
||||
self.outbound_states.pop(session_id, None)
|
||||
|
||||
return quart.jsonify(
|
||||
{
|
||||
|
||||
@@ -6,7 +6,6 @@ import json
|
||||
import base64
|
||||
import zlib
|
||||
import traceback
|
||||
import time
|
||||
|
||||
import aiohttp
|
||||
|
||||
@@ -21,6 +20,39 @@ import langbot_plugin.api.entities.builtin.platform.entities as platform_entitie
|
||||
import langbot_plugin.api.definition.abstract.platform.event_logger as abstract_platform_logger
|
||||
|
||||
|
||||
_KOOK_MAX_GATEWAY_MESSAGE_BYTES = 10 * 1024 * 1024
|
||||
|
||||
|
||||
def _bounded_zlib_decompress(payload: bytes) -> bytes:
|
||||
decompressor = zlib.decompressobj()
|
||||
decoded = decompressor.decompress(
|
||||
payload,
|
||||
_KOOK_MAX_GATEWAY_MESSAGE_BYTES + 1,
|
||||
)
|
||||
if len(decoded) > _KOOK_MAX_GATEWAY_MESSAGE_BYTES or decompressor.unconsumed_tail:
|
||||
raise ValueError('KOOK gateway message exceeds the decompressed size limit')
|
||||
decoded += decompressor.flush(_KOOK_MAX_GATEWAY_MESSAGE_BYTES + 1 - len(decoded))
|
||||
if len(decoded) > _KOOK_MAX_GATEWAY_MESSAGE_BYTES or not decompressor.eof:
|
||||
raise ValueError('KOOK gateway message exceeds the decompressed size limit')
|
||||
return decoded
|
||||
|
||||
|
||||
def _decode_gateway_message(message: str | bytes) -> dict:
|
||||
if isinstance(message, bytes):
|
||||
try:
|
||||
message_bytes = _bounded_zlib_decompress(message)
|
||||
except zlib.error:
|
||||
message_bytes = message
|
||||
else:
|
||||
message_bytes = message.encode('utf-8')
|
||||
if len(message_bytes) > _KOOK_MAX_GATEWAY_MESSAGE_BYTES:
|
||||
raise ValueError('KOOK gateway message exceeds the size limit')
|
||||
decoded = json.loads(message_bytes)
|
||||
if not isinstance(decoded, dict):
|
||||
raise ValueError('KOOK gateway message must be a JSON object')
|
||||
return decoded
|
||||
|
||||
|
||||
class KookMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
|
||||
"""Convert between LangBot MessageChain and KOOK message format"""
|
||||
|
||||
@@ -125,8 +157,8 @@ class KookMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
|
||||
session = httpclient.get_session()
|
||||
async with session.get(content) as response:
|
||||
if response.status == 200:
|
||||
image_bytes = await response.read()
|
||||
image_base64 = base64.b64encode(image_bytes).decode('utf-8')
|
||||
image_bytes = await httpclient.read_limited(response)
|
||||
image_base64 = (await asyncio.to_thread(base64.b64encode, image_bytes)).decode('utf-8')
|
||||
# Detect image format
|
||||
content_type = response.headers.get('Content-Type', 'image/png')
|
||||
components.append(
|
||||
@@ -270,10 +302,6 @@ class KookAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
http_session: typing.Optional[aiohttp.ClientSession] = pydantic.Field(exclude=True, default=None)
|
||||
|
||||
def __init__(self, config: dict, logger: abstract_platform_logger.AbstractEventLogger, **kwargs):
|
||||
# Debug: Track init
|
||||
with open('/tmp/kook_adapter_init.txt', 'w') as f:
|
||||
f.write(f'KOOK adapter __init__ called at {time.time()}\n')
|
||||
|
||||
# Validate required config
|
||||
if 'token' not in config:
|
||||
raise Exception('KOOK adapter requires "token" in config')
|
||||
@@ -300,7 +328,7 @@ class KookAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
session = httpclient.get_session()
|
||||
async with session.get(base_url, params=params, headers=headers) as response:
|
||||
if response.status == 200:
|
||||
data = await response.json()
|
||||
data = await httpclient.read_json_limited(response)
|
||||
if data.get('code') == 0:
|
||||
gateway_url = data['data']['url']
|
||||
return gateway_url
|
||||
@@ -320,7 +348,7 @@ class KookAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
session = httpclient.get_session()
|
||||
async with session.get(base_url, headers=headers) as response:
|
||||
if response.status == 200:
|
||||
data = await response.json()
|
||||
data = await httpclient.read_json_limited(response)
|
||||
if data.get('code') == 0:
|
||||
user_info = data['data']
|
||||
return user_info
|
||||
@@ -409,17 +437,10 @@ class KookAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
# Wait for HELLO within 6 seconds
|
||||
try:
|
||||
hello_msg = await asyncio.wait_for(ws.recv(), timeout=6.0)
|
||||
|
||||
# Handle compressed messages (same as main message loop)
|
||||
if isinstance(hello_msg, bytes):
|
||||
# Decompress if compressed
|
||||
try:
|
||||
hello_msg = zlib.decompress(hello_msg).decode('utf-8')
|
||||
except Exception:
|
||||
# Not compressed or decompression failed
|
||||
hello_msg = hello_msg.decode('utf-8')
|
||||
|
||||
hello_data = json.loads(hello_msg)
|
||||
hello_data = await asyncio.to_thread(
|
||||
_decode_gateway_message,
|
||||
hello_msg,
|
||||
)
|
||||
|
||||
if hello_data.get('s') == 1: # HELLO signal
|
||||
await self._handle_hello(hello_data['d'])
|
||||
@@ -433,16 +454,11 @@ class KookAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
|
||||
# Main message loop
|
||||
async for message in ws:
|
||||
if isinstance(message, bytes):
|
||||
# Decompress if compressed
|
||||
try:
|
||||
message = zlib.decompress(message).decode('utf-8')
|
||||
except Exception:
|
||||
# Not compressed or decompression failed
|
||||
message = message.decode('utf-8')
|
||||
|
||||
try:
|
||||
msg_data = json.loads(message)
|
||||
msg_data = await asyncio.to_thread(
|
||||
_decode_gateway_message,
|
||||
message,
|
||||
)
|
||||
signal = msg_data.get('s')
|
||||
|
||||
if signal == 0: # EVENT
|
||||
@@ -516,7 +532,7 @@ class KookAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
|
||||
async with self.http_session.post(url, json=payload, headers=headers) as response:
|
||||
if response.status == 200:
|
||||
result = await response.json()
|
||||
result = await httpclient.read_json_limited(response)
|
||||
if result.get('code') == 0:
|
||||
await self.logger.debug(f'Message sent successfully to {target_id}')
|
||||
else:
|
||||
@@ -582,7 +598,7 @@ class KookAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
|
||||
async with self.http_session.post(url, json=payload, headers=headers) as response:
|
||||
if response.status == 200:
|
||||
result = await response.json()
|
||||
result = await httpclient.read_json_limited(response)
|
||||
if result.get('code') == 0:
|
||||
await self.logger.debug('Reply sent successfully')
|
||||
else:
|
||||
|
||||
@@ -15,7 +15,7 @@ import hashlib
|
||||
from Crypto.Cipher import AES
|
||||
import tempfile
|
||||
import os
|
||||
import mimetypes
|
||||
import threading
|
||||
|
||||
from langbot.pkg.utils import httpclient
|
||||
import lark_oapi.ws.exception
|
||||
@@ -34,6 +34,53 @@ import langbot_plugin.api.definition.abstract.platform.event_logger as abstract_
|
||||
import langbot_plugin.api.entities.builtin.provider.session as provider_session
|
||||
|
||||
|
||||
_MAX_LARK_MEDIA_BYTES = 10 * 1024 * 1024
|
||||
|
||||
|
||||
def _decode_lark_base64_limited(value: str) -> bytes:
|
||||
if ',' in value:
|
||||
value = value.split(',', 1)[1]
|
||||
max_encoded_bytes = 4 * ((_MAX_LARK_MEDIA_BYTES + 2) // 3)
|
||||
if len(value) > max_encoded_bytes:
|
||||
raise ValueError('Lark media exceeds the size limit')
|
||||
decoded = base64.b64decode(value)
|
||||
if len(decoded) > _MAX_LARK_MEDIA_BYTES:
|
||||
raise ValueError('Lark media exceeds the size limit')
|
||||
return decoded
|
||||
|
||||
|
||||
def _read_lark_path_limited(path: str) -> bytes:
|
||||
if os.path.getsize(path) > _MAX_LARK_MEDIA_BYTES:
|
||||
raise ValueError('Lark media exceeds the size limit')
|
||||
with open(path, 'rb') as file:
|
||||
body = file.read(_MAX_LARK_MEDIA_BYTES + 1)
|
||||
if len(body) > _MAX_LARK_MEDIA_BYTES:
|
||||
raise ValueError('Lark media exceeds the size limit')
|
||||
return body
|
||||
|
||||
|
||||
def _write_lark_temp_file(data: bytes) -> str:
|
||||
with tempfile.NamedTemporaryFile(delete=False) as temp_file:
|
||||
temp_file.write(data)
|
||||
temp_file.flush()
|
||||
return temp_file.name
|
||||
|
||||
|
||||
def _read_lark_response_file_limited(response) -> bytes:
|
||||
content_length = response.raw.headers.get('content-length')
|
||||
if content_length is not None:
|
||||
try:
|
||||
if int(content_length) > _MAX_LARK_MEDIA_BYTES:
|
||||
raise ValueError('Lark media exceeds the size limit')
|
||||
except (TypeError, ValueError) as exc:
|
||||
if 'exceeds' in str(exc):
|
||||
raise
|
||||
body = response.file.read(_MAX_LARK_MEDIA_BYTES + 1)
|
||||
if len(body) > _MAX_LARK_MEDIA_BYTES:
|
||||
raise ValueError('Lark media exceeds the size limit')
|
||||
return body
|
||||
|
||||
|
||||
def _lark_form_component_name(prefix: str, field_name: str, index: int) -> str:
|
||||
safe_name = re.sub(r'[^A-Za-z0-9_]', '_', field_name)[:8] or 'field'
|
||||
digest = hashlib.sha1(field_name.encode('utf-8')).hexdigest()[:6]
|
||||
@@ -299,68 +346,33 @@ class LarkMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
|
||||
@staticmethod
|
||||
async def upload_image_to_lark(msg: platform_message.Image, api_client: lark_oapi.Client) -> typing.Optional[str]:
|
||||
"""Upload an image to Lark and return the image_key, or None if upload fails."""
|
||||
image_bytes = None
|
||||
|
||||
if msg.base64:
|
||||
try:
|
||||
# Remove data URL prefix if present
|
||||
base64_data = msg.base64
|
||||
if base64_data.startswith('data:'):
|
||||
base64_data = base64_data.split(',', 1)[1]
|
||||
image_bytes = base64.b64decode(base64_data)
|
||||
except Exception as e:
|
||||
print(f'Failed to decode base64 image: {e}')
|
||||
traceback.print_exc()
|
||||
return None
|
||||
elif msg.url:
|
||||
try:
|
||||
session = httpclient.get_session()
|
||||
async with session.get(msg.url) as response:
|
||||
if response.status == 200:
|
||||
image_bytes = await response.read()
|
||||
else:
|
||||
print(f'Failed to download image from {msg.url}: HTTP {response.status}')
|
||||
return None
|
||||
except Exception as e:
|
||||
print(f'Failed to download image from {msg.url}: {e}')
|
||||
traceback.print_exc()
|
||||
return None
|
||||
elif msg.path:
|
||||
try:
|
||||
with open(msg.path, 'rb') as f:
|
||||
image_bytes = f.read()
|
||||
except Exception as e:
|
||||
print(f'Failed to read image from path {msg.path}: {e}')
|
||||
traceback.print_exc()
|
||||
return None
|
||||
|
||||
if image_bytes is None:
|
||||
try:
|
||||
image_bytes, _mime_type = await msg.get_bytes()
|
||||
except Exception as exc:
|
||||
print(f'Failed to load Lark image: {exc}')
|
||||
traceback.print_exc()
|
||||
return None
|
||||
if not image_bytes:
|
||||
print(
|
||||
f'No image data available for Image message (url={msg.url}, base64={bool(msg.base64)}, path={msg.path})'
|
||||
)
|
||||
return None
|
||||
|
||||
try:
|
||||
# Create a temporary file to store the image bytes
|
||||
import tempfile
|
||||
import os
|
||||
|
||||
with tempfile.NamedTemporaryFile(delete=False) as temp_file:
|
||||
temp_file.write(image_bytes)
|
||||
temp_file.flush()
|
||||
temp_file_path = temp_file.name
|
||||
temp_file_path = await asyncio.to_thread(
|
||||
_write_lark_temp_file,
|
||||
image_bytes,
|
||||
)
|
||||
|
||||
try:
|
||||
# Create image request using the temporary file
|
||||
request = (
|
||||
CreateImageRequest.builder()
|
||||
.request_body(
|
||||
CreateImageRequestBody.builder().image_type('message').image(open(temp_file_path, 'rb')).build()
|
||||
with open(temp_file_path, 'rb') as upload_file:
|
||||
request = (
|
||||
CreateImageRequest.builder()
|
||||
.request_body(CreateImageRequestBody.builder().image_type('message').image(upload_file).build())
|
||||
.build()
|
||||
)
|
||||
.build()
|
||||
)
|
||||
|
||||
response = await api_client.im.v1.image.acreate(request)
|
||||
response = await api_client.im.v1.image.acreate(request)
|
||||
|
||||
if not response.success():
|
||||
print(
|
||||
@@ -395,23 +407,24 @@ class LarkMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
|
||||
duration: Duration in milliseconds (for audio files).
|
||||
"""
|
||||
try:
|
||||
with tempfile.NamedTemporaryFile(delete=False) as temp_file:
|
||||
temp_file.write(file_bytes)
|
||||
temp_file_path = temp_file.name
|
||||
if len(file_bytes) > _MAX_LARK_MEDIA_BYTES:
|
||||
raise ValueError('Lark media exceeds the size limit')
|
||||
temp_file_path = await asyncio.to_thread(
|
||||
_write_lark_temp_file,
|
||||
file_bytes,
|
||||
)
|
||||
|
||||
try:
|
||||
body_builder = (
|
||||
CreateFileRequestBody.builder()
|
||||
.file_type(file_type)
|
||||
.file_name(file_name)
|
||||
.file(open(temp_file_path, 'rb'))
|
||||
)
|
||||
if duration is not None:
|
||||
body_builder = body_builder.duration(duration)
|
||||
with open(temp_file_path, 'rb') as upload_file:
|
||||
body_builder = (
|
||||
CreateFileRequestBody.builder().file_type(file_type).file_name(file_name).file(upload_file)
|
||||
)
|
||||
if duration is not None:
|
||||
body_builder = body_builder.duration(duration)
|
||||
|
||||
request = CreateFileRequest.builder().request_body(body_builder.build()).build()
|
||||
request = CreateFileRequest.builder().request_body(body_builder.build()).build()
|
||||
|
||||
response = await api_client.im.v1.file.acreate(request)
|
||||
response = await api_client.im.v1.file.acreate(request)
|
||||
|
||||
if not response.success():
|
||||
print(
|
||||
@@ -436,10 +449,10 @@ class LarkMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
|
||||
|
||||
if msg.base64:
|
||||
try:
|
||||
base64_str = msg.base64
|
||||
if ',' in base64_str:
|
||||
base64_str = base64_str.split(',', 1)[1]
|
||||
data = base64.b64decode(base64_str)
|
||||
data = await asyncio.to_thread(
|
||||
_decode_lark_base64_limited,
|
||||
msg.base64,
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
elif msg.url:
|
||||
@@ -447,13 +460,18 @@ class LarkMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
|
||||
session = httpclient.get_session()
|
||||
async with session.get(msg.url) as resp:
|
||||
if resp.status == 200:
|
||||
data = await resp.read()
|
||||
data = await httpclient.read_limited(
|
||||
resp,
|
||||
max_bytes=_MAX_LARK_MEDIA_BYTES,
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
elif msg.path:
|
||||
try:
|
||||
with open(msg.path, 'rb') as f:
|
||||
data = f.read()
|
||||
data = await asyncio.to_thread(
|
||||
_read_lark_path_limited,
|
||||
str(msg.path),
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@@ -694,8 +712,11 @@ class LarkMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
|
||||
f'client.im.v1.message_resource.get failed, code: {response.code}, msg: {response.msg}, log_id: {response.get_log_id()}, resp: \n{json.dumps(json.loads(response.raw.content), indent=4, ensure_ascii=False)}'
|
||||
)
|
||||
|
||||
image_bytes = response.file.read()
|
||||
image_base64 = base64.b64encode(image_bytes).decode()
|
||||
image_bytes = await asyncio.to_thread(
|
||||
_read_lark_response_file_limited,
|
||||
response,
|
||||
)
|
||||
image_base64 = (await asyncio.to_thread(base64.b64encode, image_bytes)).decode()
|
||||
|
||||
image_format = response.raw.headers['content-type']
|
||||
|
||||
@@ -721,27 +742,18 @@ class LarkMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
|
||||
lb_msg_list.append(platform_message.Plain(text='[Audio file download failed]'))
|
||||
return platform_message.MessageChain(lb_msg_list)
|
||||
|
||||
# Read audio bytes
|
||||
audio_bytes = response.file.read()
|
||||
audio_base64 = base64.b64encode(audio_bytes).decode()
|
||||
audio_bytes = await asyncio.to_thread(
|
||||
_read_lark_response_file_limited,
|
||||
response,
|
||||
)
|
||||
audio_base64 = (await asyncio.to_thread(base64.b64encode, audio_bytes)).decode()
|
||||
|
||||
# Get content type from response headers
|
||||
content_type = response.raw.headers.get('content-type', 'audio/mpeg')
|
||||
|
||||
mime_main = content_type.split(';')[0].strip()
|
||||
ext = mimetypes.guess_extension(mime_main) or '.bin'
|
||||
temp_dir = tempfile.gettempdir()
|
||||
temp_file_path = os.path.join(temp_dir, f'lark_audio_{file_key}{ext}')
|
||||
|
||||
with open(temp_file_path, 'wb') as f:
|
||||
f.write(audio_bytes)
|
||||
|
||||
# Create Voice message: prefer path/url + length, include base64 as optional data URI
|
||||
lb_msg_list.append(
|
||||
platform_message.Voice(
|
||||
voice_id=file_key,
|
||||
url=f'file://{temp_file_path}',
|
||||
path=temp_file_path,
|
||||
base64=f'data:{content_type};base64,{audio_base64}',
|
||||
length=(duration // 1000) if duration else None,
|
||||
)
|
||||
@@ -770,40 +782,22 @@ class LarkMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
|
||||
f'client.im.v1.message_resource.get failed, code: {response.code}, msg: {response.msg}, log_id: {response.get_log_id()}, resp: \n{json.dumps(json.loads(response.raw.content), indent=4, ensure_ascii=False)}'
|
||||
)
|
||||
|
||||
file_bytes = response.file.read()
|
||||
file_base64 = base64.b64encode(file_bytes).decode()
|
||||
file_bytes = await asyncio.to_thread(
|
||||
_read_lark_response_file_limited,
|
||||
response,
|
||||
)
|
||||
file_base64 = (await asyncio.to_thread(base64.b64encode, file_bytes)).decode()
|
||||
|
||||
file_format = response.raw.headers['content-type']
|
||||
|
||||
file_size = len(file_bytes)
|
||||
|
||||
# Determine extension from content-type if possible
|
||||
content_type = response.raw.headers.get('content-type', '')
|
||||
mime_main = content_type.split(';')[0].strip() if content_type else ''
|
||||
ext = mimetypes.guess_extension(mime_main) or ''
|
||||
|
||||
# Ensure a safe filename (avoid path components)
|
||||
safe_name = os.path.basename(file_name).replace('/', '_').replace('\\', '_')
|
||||
if ext and not safe_name.lower().endswith(ext.lower()):
|
||||
filename_with_ext = f'{safe_name}{ext}'
|
||||
else:
|
||||
filename_with_ext = safe_name
|
||||
|
||||
temp_dir = tempfile.gettempdir()
|
||||
temp_file_path = os.path.join(temp_dir, f'lark_{file_key}_{filename_with_ext}')
|
||||
|
||||
with open(temp_file_path, 'wb') as f:
|
||||
f.write(file_bytes)
|
||||
|
||||
# Create File message with local path and file:// URL
|
||||
lb_msg_list.append(
|
||||
platform_message.File(
|
||||
id=file_key,
|
||||
name=file_name,
|
||||
size=file_size,
|
||||
url=f'file://{temp_file_path}',
|
||||
path=temp_file_path,
|
||||
base64=f'data:{file_format};base64,{file_base64}', # not including base64 by default to save memory; can be added if needed
|
||||
base64=f'data:{file_format};base64,{file_base64}',
|
||||
)
|
||||
)
|
||||
|
||||
@@ -1042,16 +1036,26 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
# card_id → input_defs / inputs captured for the selected-action notice
|
||||
card_form_input_defs: dict[str, list[dict]]
|
||||
card_form_inputs: dict[str, dict]
|
||||
card_last_accessed: dict[str, float]
|
||||
card_cleanup_at: float
|
||||
# set of card_ids that have already transitioned from "buttons visible" to "resume layout"
|
||||
card_resume_transitioned: set[str]
|
||||
inbound_event_tasks: set[asyncio.Task]
|
||||
threadsafe_event_futures: set[typing.Any]
|
||||
threadsafe_event_lock: typing.Any = pydantic.Field(exclude=True)
|
||||
_MONITORING_MAPPING_TTL = 600 # 10 minutes
|
||||
_MAX_INBOUND_EVENTS = 100
|
||||
_MAX_TENANT_ACCESS_TOKENS = 1024
|
||||
|
||||
seq: int # 用于在发送卡片消息中识别消息顺序，直接以seq作为标识
|
||||
bot_uuid: str = None # 机器人UUID
|
||||
app_ticket: str = None # 商店应用用到
|
||||
app_access_token: str = None # 商店应用用到
|
||||
app_access_token_expire_at: int = None
|
||||
tenant_access_tokens: dict[str, dict[str, str]] = {} # 租户access_token映射
|
||||
tenant_access_tokens: dict[str, dict[str, str]] = pydantic.Field(
|
||||
default_factory=dict,
|
||||
exclude=True,
|
||||
)
|
||||
|
||||
def __init__(self, config: dict, logger: abstract_platform_logger.AbstractEventLogger, **kwargs):
|
||||
quart_app = quart.Quart(__name__)
|
||||
@@ -1062,11 +1066,11 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
await self.listeners[type(lb_event)](lb_event, self)
|
||||
|
||||
def sync_on_message(event: lark_oapi.im.v1.P2ImMessageReceiveV1):
|
||||
asyncio.create_task(on_message(event))
|
||||
self._schedule_inbound_event(on_message(event))
|
||||
|
||||
def schedule_on_app_loop(coro):
|
||||
"""Run a coroutine on the application event loop from sync callbacks."""
|
||||
return asyncio.run_coroutine_threadsafe(coro, self.ap.event_loop)
|
||||
return self._schedule_threadsafe_event(coro)
|
||||
|
||||
def sync_on_card_action(event):
|
||||
try:
|
||||
@@ -1289,7 +1293,13 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
card_form_content={},
|
||||
card_form_input_defs={},
|
||||
card_form_inputs={},
|
||||
card_last_accessed={},
|
||||
card_cleanup_at=0.0,
|
||||
card_resume_transitioned=set(),
|
||||
inbound_event_tasks=set(),
|
||||
threadsafe_event_futures=set(),
|
||||
threadsafe_event_lock=threading.Lock(),
|
||||
tenant_access_tokens={},
|
||||
seq=1,
|
||||
listeners={},
|
||||
quart_app=quart_app,
|
||||
@@ -1300,6 +1310,45 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
def _schedule_inbound_event(self, coro) -> None:
|
||||
for task in tuple(self.inbound_event_tasks):
|
||||
if task.done():
|
||||
self.inbound_event_tasks.discard(task)
|
||||
if len(self.inbound_event_tasks) >= self._MAX_INBOUND_EVENTS:
|
||||
coro.close()
|
||||
return
|
||||
task = asyncio.create_task(coro)
|
||||
self.inbound_event_tasks.add(task)
|
||||
|
||||
def done(done_task: asyncio.Task) -> None:
|
||||
self.inbound_event_tasks.discard(done_task)
|
||||
if not done_task.cancelled():
|
||||
done_task.exception()
|
||||
|
||||
task.add_done_callback(done)
|
||||
|
||||
def _schedule_threadsafe_event(self, coro):
|
||||
"""Submit one bounded callback from the Lark SDK's sync boundary."""
|
||||
|
||||
with self.threadsafe_event_lock:
|
||||
for future in tuple(self.threadsafe_event_futures):
|
||||
if future.done():
|
||||
self.threadsafe_event_futures.discard(future)
|
||||
if len(self.threadsafe_event_futures) >= self._MAX_INBOUND_EVENTS:
|
||||
coro.close()
|
||||
return None
|
||||
future = asyncio.run_coroutine_threadsafe(coro, self.ap.event_loop)
|
||||
self.threadsafe_event_futures.add(future)
|
||||
|
||||
def done(done_future) -> None:
|
||||
with self.threadsafe_event_lock:
|
||||
self.threadsafe_event_futures.discard(done_future)
|
||||
if not done_future.cancelled():
|
||||
done_future.exception()
|
||||
|
||||
future.add_done_callback(done)
|
||||
return future
|
||||
|
||||
def request_app_ticket(self, api_client, config):
|
||||
app_id = config['app_id']
|
||||
app_secret = config['app_secret']
|
||||
@@ -1376,6 +1425,12 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
'token': tenant_access_token,
|
||||
'expire_at': int(time.time()) + expire - 300,
|
||||
}
|
||||
now = int(time.time())
|
||||
for cached_key, cached_token in tuple(self.tenant_access_tokens.items()):
|
||||
if int(cached_token.get('expire_at', 0)) <= now:
|
||||
self.tenant_access_tokens.pop(cached_key, None)
|
||||
while len(self.tenant_access_tokens) > self._MAX_TENANT_ACCESS_TOKENS:
|
||||
self.tenant_access_tokens.pop(next(iter(self.tenant_access_tokens)), None)
|
||||
|
||||
def get_tenant_access_token(self, tenant_key: str):
|
||||
if tenant_key is None or 'isv' != self.config.get('app_type', 'self'):
|
||||
@@ -1558,6 +1613,8 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
user_msg_id = query.message_event.message_chain.message_id
|
||||
if user_msg_id:
|
||||
self.pending_monitoring_msg[user_msg_id] = monitoring_message_id
|
||||
while len(self.pending_monitoring_msg) > CARD_ID_CACHE_SIZE:
|
||||
self.pending_monitoring_msg.pop(next(iter(self.pending_monitoring_msg)), None)
|
||||
except Exception as e:
|
||||
await self.logger.debug(f'Failed to map message to monitoring message: {e}')
|
||||
|
||||
@@ -1570,6 +1627,7 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
|
||||
def _next_card_sequence(self, card_id: str, suggested: int = 1) -> int:
|
||||
"""Return the next strictly increasing sequence for a card update."""
|
||||
self._touch_card(card_id)
|
||||
current = self.card_sequence_dict.get(card_id, 0)
|
||||
next_seq = max(current + 1, suggested)
|
||||
self.card_sequence_dict[card_id] = next_seq
|
||||
@@ -1577,6 +1635,7 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
|
||||
def _register_card_for_source(self, card_id: str, *source_ids: str) -> None:
|
||||
"""Register a card_id under one or more source message ids."""
|
||||
self._touch_card(card_id)
|
||||
bucket = self.card_id_to_source_ids.setdefault(card_id, set())
|
||||
for sid in source_ids:
|
||||
if not sid:
|
||||
@@ -1596,8 +1655,24 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
self.card_form_content.pop(card_id, None)
|
||||
self.card_form_input_defs.pop(card_id, None)
|
||||
self.card_form_inputs.pop(card_id, None)
|
||||
self.card_last_accessed.pop(card_id, None)
|
||||
self.card_resume_transitioned.discard(card_id)
|
||||
|
||||
def _touch_card(self, card_id: str) -> None:
|
||||
now = time.monotonic()
|
||||
if now - self.card_cleanup_at >= 60 or len(self.card_last_accessed) >= CARD_ID_CACHE_SIZE:
|
||||
self.card_cleanup_at = now
|
||||
for stale_card_id, last_accessed in tuple(self.card_last_accessed.items()):
|
||||
if now - last_accessed >= CARD_ID_CACHE_MAX_LIFETIME:
|
||||
self._drop_card_state(stale_card_id)
|
||||
while len(self.card_last_accessed) >= CARD_ID_CACHE_SIZE:
|
||||
oldest_card_id = min(
|
||||
self.card_last_accessed,
|
||||
key=self.card_last_accessed.__getitem__,
|
||||
)
|
||||
self._drop_card_state(oldest_card_id)
|
||||
self.card_last_accessed[card_id] = now
|
||||
|
||||
async def create_card_id(self, message_id):
|
||||
try:
|
||||
# self.logger.debug('飞书支持stream输出,创建卡片......')
|
||||
@@ -1793,6 +1868,7 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
self.card_id_dict[message_id] = response.data.card_id
|
||||
|
||||
card_id = response.data.card_id
|
||||
self._touch_card(card_id)
|
||||
self.card_sequence_dict[card_id] = 0
|
||||
return card_id
|
||||
|
||||
@@ -1864,7 +1940,7 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
self.reply_to_monitoring_msg[reply_msg_id] = (monitoring_msg_id, time.time())
|
||||
self._cleanup_monitoring_mapping()
|
||||
except Exception as e:
|
||||
asyncio.create_task(self.logger.debug(f'Failed to transfer monitoring mapping in create_message_card: {e}'))
|
||||
await self.logger.debug(f'Failed to transfer monitoring mapping in create_message_card: {e}')
|
||||
|
||||
return True
|
||||
|
||||
@@ -2872,8 +2948,8 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
data = await request.json
|
||||
|
||||
if 'encrypt' in data:
|
||||
data = self.cipher.decrypt_string(data['encrypt'])
|
||||
data = json.loads(data)
|
||||
encrypted = data['encrypt']
|
||||
data = await asyncio.to_thread(lambda: json.loads(self.cipher.decrypt_string(encrypted)))
|
||||
type = self.get_event_type(data)
|
||||
context = EventContext(data)
|
||||
if 'url_verification' == type:
|
||||
@@ -3143,4 +3219,38 @@ class LarkAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
# 所以要设置_auto_reconnect=False,让其不重连。
|
||||
self.bot._auto_reconnect = False
|
||||
await self.bot._disconnect()
|
||||
inbound_tasks = list(self.inbound_event_tasks)
|
||||
for task in inbound_tasks:
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
if inbound_tasks:
|
||||
await asyncio.gather(*inbound_tasks, return_exceptions=True)
|
||||
self.inbound_event_tasks.clear()
|
||||
with self.threadsafe_event_lock:
|
||||
threadsafe_futures = list(self.threadsafe_event_futures)
|
||||
for future in threadsafe_futures:
|
||||
future.cancel()
|
||||
if threadsafe_futures:
|
||||
await asyncio.gather(
|
||||
*(asyncio.wrap_future(future) for future in threadsafe_futures),
|
||||
return_exceptions=True,
|
||||
)
|
||||
with self.threadsafe_event_lock:
|
||||
self.threadsafe_event_futures.clear()
|
||||
self.tenant_access_tokens.clear()
|
||||
self.card_id_dict.clear()
|
||||
self.pending_monitoring_msg.clear()
|
||||
self.reply_to_monitoring_msg.clear()
|
||||
for card_id in tuple(self.card_last_accessed):
|
||||
self._drop_card_state(card_id)
|
||||
self.card_last_accessed.clear()
|
||||
self.reply_message_card_ids.clear()
|
||||
self.card_sequence_dict.clear()
|
||||
self.card_id_to_source_ids.clear()
|
||||
self.card_streaming_text.clear()
|
||||
self.card_pre_pause_text.clear()
|
||||
self.card_form_content.clear()
|
||||
self.card_form_input_defs.clear()
|
||||
self.card_form_inputs.clear()
|
||||
self.card_resume_transitioned.clear()
|
||||
return False
|
||||
|
||||
@@ -6,7 +6,6 @@ import traceback
|
||||
import time
|
||||
import re
|
||||
import copy
|
||||
import threading
|
||||
|
||||
import quart
|
||||
from langbot.pkg.utils import httpclient
|
||||
@@ -483,6 +482,7 @@ class GeWeChatAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
|
||||
self.message_converter = GewechatMessageConverter(config)
|
||||
self.event_converter = GewechatEventConverter(config)
|
||||
self.listeners = {}
|
||||
|
||||
@self.quart_app.route('/gewechat/callback', methods=['POST'])
|
||||
async def gewechat_callback():
|
||||
@@ -518,9 +518,13 @@ class GeWeChatAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
at_targets = at_targets or []
|
||||
member_info = []
|
||||
if at_targets:
|
||||
member_info = self.bot.get_chatroom_member_detail(self.config['app_id'], target_id, at_targets[::-1])[
|
||||
'data'
|
||||
]
|
||||
member_result = await asyncio.to_thread(
|
||||
self.bot.get_chatroom_member_detail,
|
||||
self.config['app_id'],
|
||||
target_id,
|
||||
at_targets[::-1],
|
||||
)
|
||||
member_info = member_result['data']
|
||||
|
||||
# 处理消息组件
|
||||
for msg in content_list:
|
||||
@@ -596,7 +600,7 @@ class GeWeChatAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
}
|
||||
|
||||
if handler := handler_map.get(msg['type']):
|
||||
handler(msg)
|
||||
await asyncio.to_thread(handler, msg)
|
||||
else:
|
||||
await self.logger.warning(f'未处理的消息类型: {msg["type"]}')
|
||||
continue
|
||||
@@ -645,8 +649,9 @@ class GeWeChatAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
json={'app_id': self.config['app_id']},
|
||||
) as response:
|
||||
if response.status != 200:
|
||||
raise Exception(f'获取gewechat token失败: {await response.text()}')
|
||||
self.config['token'] = (await response.json())['data']
|
||||
error = await httpclient.read_text_limited(response)
|
||||
raise Exception(f'获取gewechat token失败: {error}')
|
||||
self.config['token'] = (await httpclient.read_json_limited(response))['data']
|
||||
|
||||
self.bot = gewechat_client.GewechatClient(f'{self.config["gewechat_url"]}/v2/api', self.config['token'])
|
||||
|
||||
@@ -672,7 +677,7 @@ class GeWeChatAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
except Exception as e:
|
||||
raise Exception(f'设置 Gewechat 回调失败， token失效： {e}')
|
||||
|
||||
threading.Thread(target=gewechat_login_process).start()
|
||||
await asyncio.to_thread(gewechat_login_process)
|
||||
|
||||
async def shutdown_trigger_placeholder():
|
||||
while True:
|
||||
|
||||
@@ -311,15 +311,17 @@ class NakuruAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
try:
|
||||
import requests
|
||||
|
||||
resp = requests.get(
|
||||
url='http://{}:{}/get_login_info'.format(self.cfg['host'], self.cfg['http_port']),
|
||||
resp = await asyncio.to_thread(
|
||||
requests.get,
|
||||
'http://{}:{}/get_login_info'.format(self.cfg['host'], self.cfg['http_port']),
|
||||
headers={'Authorization': 'Bearer ' + self.cfg['token'] if 'token' in self.cfg else ''},
|
||||
timeout=5,
|
||||
proxies=None,
|
||||
)
|
||||
if resp.status_code == 403:
|
||||
raise Exception('go-cqhttp拒绝访问，请检查配置文件中nakuru适配器的配置')
|
||||
self.bot_account_id = int(resp.json()['data']['user_id'])
|
||||
response_data = await httpclient.parse_json_response(resp)
|
||||
self.bot_account_id = int(response_data['data']['user_id'])
|
||||
except Exception:
|
||||
raise Exception('获取go-cqhttp账号信息失败, 请检查是否已启动go-cqhttp并配置正确')
|
||||
await self.bot._run()
|
||||
|
||||
@@ -5,6 +5,7 @@ import typing
|
||||
import datetime
|
||||
import re
|
||||
import traceback
|
||||
from collections import OrderedDict
|
||||
|
||||
import botpy
|
||||
import botpy.message as botpy_message
|
||||
@@ -40,7 +41,8 @@ event_handler_mapping = {
|
||||
}
|
||||
|
||||
|
||||
cached_message_ids = {}
|
||||
_CACHED_MESSAGE_ID_LIMIT = 10000
|
||||
cached_message_ids: OrderedDict[str, str] = OrderedDict()
|
||||
"""由于QQ官方的消息id是字符串，而YiriMirai的消息id是整数，所以需要一个索引来进行转换"""
|
||||
|
||||
id_index = 0
|
||||
@@ -53,6 +55,8 @@ def save_msg_id(message_id: str) -> int:
|
||||
crt_index = id_index
|
||||
id_index += 1
|
||||
cached_message_ids[str(crt_index)] = message_id
|
||||
while len(cached_message_ids) > _CACHED_MESSAGE_ID_LIMIT:
|
||||
cached_message_ids.popitem(last=False)
|
||||
return crt_index
|
||||
|
||||
|
||||
@@ -355,6 +359,7 @@ class OfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
self.cfg = cfg
|
||||
self.ap = ap
|
||||
self.logger = logger
|
||||
self.cached_official_messages = OrderedDict()
|
||||
|
||||
self.group_msg_seq = 1
|
||||
self.c2c_msg_seq = 1
|
||||
@@ -490,6 +495,8 @@ class OfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
],
|
||||
):
|
||||
self.cached_official_messages[str(message.id)] = message
|
||||
while len(self.cached_official_messages) > 1000:
|
||||
self.cached_official_messages.popitem(last=False)
|
||||
await callback(self.event_converter.target2yiri(message), self)
|
||||
|
||||
for event_handler in event_handler_mapping[event_type]:
|
||||
@@ -519,6 +526,8 @@ class OfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
await (await self.bot.start(**self.cfg))
|
||||
|
||||
async def kill(self) -> bool:
|
||||
self.cached_official_messages.clear()
|
||||
if not self.bot.is_closed():
|
||||
await self.bot.close()
|
||||
return True
|
||||
return True
|
||||
|
||||
@@ -13,6 +13,7 @@ import langbot_plugin.api.entities.builtin.platform.message as platform_message
|
||||
import langbot_plugin.api.entities.builtin.platform.events as platform_events
|
||||
import langbot_plugin.api.entities.builtin.platform.entities as platform_entities
|
||||
from ..logger import EventLogger
|
||||
from ...utils import bounded_executor
|
||||
|
||||
|
||||
from linebot.v3 import WebhookHandler
|
||||
@@ -30,6 +31,14 @@ from linebot.v3.webhooks import (
|
||||
from linebot.v3.webhook import WebhookParser
|
||||
from linebot.v3.messaging import MessagingApiBlob
|
||||
|
||||
MAX_LINE_MEDIA_BYTES = 10 * 1024 * 1024
|
||||
|
||||
|
||||
def _validate_line_media_content(content: bytes) -> bytes:
|
||||
if len(content) > MAX_LINE_MEDIA_BYTES:
|
||||
raise ValueError(f'LINE media exceeds the {MAX_LINE_MEDIA_BYTES}-byte limit')
|
||||
return content
|
||||
|
||||
|
||||
class LINEMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
|
||||
@staticmethod
|
||||
@@ -63,9 +72,13 @@ class LINEMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
|
||||
elif isinstance(message.message, VideoMessageContent):
|
||||
pass
|
||||
elif isinstance(message.message, ImageMessageContent):
|
||||
message_content = MessagingApiBlob(bot_client).get_message_content(message.message.id)
|
||||
message_content = await asyncio.to_thread(
|
||||
MessagingApiBlob(bot_client).get_message_content,
|
||||
message.message.id,
|
||||
)
|
||||
_validate_line_media_content(message_content)
|
||||
|
||||
base64_string = base64.b64encode(message_content).decode('utf-8')
|
||||
base64_string = await asyncio.to_thread(lambda: base64.b64encode(message_content).decode('utf-8'))
|
||||
|
||||
# 如果需要Data URI格式（用于直接嵌入HTML等）
|
||||
# 首先需要知道图片类型，LINE图片通常是JPEG
|
||||
@@ -173,20 +186,22 @@ class LINEAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
|
||||
for content in content_list:
|
||||
if content['type'] == 'text':
|
||||
self.bot.reply_message_with_http_info(
|
||||
await asyncio.to_thread(
|
||||
self.bot.reply_message_with_http_info,
|
||||
ReplyMessageRequest(
|
||||
reply_token=message_source.source_platform_object.reply_token,
|
||||
messages=[TextMessage(text=content['content'])],
|
||||
)
|
||||
),
|
||||
)
|
||||
elif content['type'] == 'image':
|
||||
# LINE ImageMessage requires original_content_url and preview_image_url
|
||||
image_url = content['image']
|
||||
self.bot.reply_message_with_http_info(
|
||||
await asyncio.to_thread(
|
||||
self.bot.reply_message_with_http_info,
|
||||
ReplyMessageRequest(
|
||||
reply_token=message_source.source_platform_object.reply_token,
|
||||
messages=[ImageMessage(original_content_url=image_url, preview_image_url=image_url)],
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
async def is_muted(self, group_id: int) -> bool:
|
||||
@@ -266,4 +281,5 @@ class LINEAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
await keep_alive()
|
||||
|
||||
async def kill(self) -> bool:
|
||||
pass
|
||||
await bounded_executor.run_blocking_cleanup(self.api_client.close)
|
||||
return True
|
||||
|
||||
@@ -5,6 +5,8 @@ import asyncio
|
||||
import traceback
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import nio
|
||||
|
||||
@@ -16,6 +18,58 @@ import langbot_plugin.api.entities.builtin.platform.entities as platform_entitie
|
||||
import langbot_plugin.api.definition.abstract.platform.event_logger as abstract_platform_logger
|
||||
|
||||
|
||||
_MAX_MATRIX_MEDIA_BYTES = 10 * 1024 * 1024
|
||||
|
||||
|
||||
def _decode_matrix_base64_limited(value: str) -> bytes:
|
||||
if ';base64,' in value:
|
||||
value = value.split(';base64,', 1)[1]
|
||||
max_encoded_bytes = 4 * ((_MAX_MATRIX_MEDIA_BYTES + 2) // 3)
|
||||
if len(value) > max_encoded_bytes:
|
||||
raise ValueError('Matrix media exceeds the size limit')
|
||||
decoded = base64.b64decode(value)
|
||||
if len(decoded) > _MAX_MATRIX_MEDIA_BYTES:
|
||||
raise ValueError('Matrix media exceeds the size limit')
|
||||
return decoded
|
||||
|
||||
|
||||
def _read_matrix_file_limited(path: str) -> bytes:
|
||||
if os.path.getsize(path) > _MAX_MATRIX_MEDIA_BYTES:
|
||||
raise ValueError('Matrix media exceeds the size limit')
|
||||
with open(path, 'rb') as file:
|
||||
body = file.read(_MAX_MATRIX_MEDIA_BYTES + 1)
|
||||
if len(body) > _MAX_MATRIX_MEDIA_BYTES:
|
||||
raise ValueError('Matrix media exceeds the size limit')
|
||||
return body
|
||||
|
||||
|
||||
async def _download_matrix_media_limited(
|
||||
client: nio.AsyncClient,
|
||||
mxc_url: str,
|
||||
) -> tuple[bytes, str]:
|
||||
parsed = urlparse(mxc_url)
|
||||
if parsed.scheme != 'mxc' or not parsed.netloc or not parsed.path.strip('/'):
|
||||
raise ValueError('Invalid Matrix media URL')
|
||||
method, path = nio.Api.download(
|
||||
parsed.netloc,
|
||||
parsed.path.replace('/', ''),
|
||||
access_token=None,
|
||||
)
|
||||
headers = {}
|
||||
if client.access_token:
|
||||
headers['Authorization'] = f'Bearer {client.access_token}'
|
||||
response = await client.send(method, path, headers=headers, timeout=30)
|
||||
try:
|
||||
response.raise_for_status()
|
||||
body = await httpclient.read_limited(
|
||||
response,
|
||||
max_bytes=_MAX_MATRIX_MEDIA_BYTES,
|
||||
)
|
||||
return body, response.headers.get('Content-Type', 'application/octet-stream')
|
||||
finally:
|
||||
response.release()
|
||||
|
||||
|
||||
class MatrixMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
|
||||
@staticmethod
|
||||
async def yiri2target(message_chain: platform_message.MessageChain, client: nio.AsyncClient) -> list[dict]:
|
||||
@@ -26,17 +80,22 @@ class MatrixMessageConverter(abstract_platform_adapter.AbstractMessageConverter)
|
||||
elif isinstance(component, platform_message.Image):
|
||||
image_bytes = None
|
||||
if component.base64:
|
||||
b64_data = component.base64
|
||||
if ';base64,' in b64_data:
|
||||
b64_data = b64_data.split(';base64,', 1)[1]
|
||||
image_bytes = base64.b64decode(b64_data)
|
||||
image_bytes = await asyncio.to_thread(
|
||||
_decode_matrix_base64_limited,
|
||||
component.base64,
|
||||
)
|
||||
elif component.url:
|
||||
session = httpclient.get_session()
|
||||
async with session.get(component.url) as response:
|
||||
image_bytes = await response.read()
|
||||
image_bytes = await httpclient.read_limited(
|
||||
response,
|
||||
max_bytes=_MAX_MATRIX_MEDIA_BYTES,
|
||||
)
|
||||
elif component.path:
|
||||
with open(component.path, 'rb') as f:
|
||||
image_bytes = f.read()
|
||||
image_bytes = await asyncio.to_thread(
|
||||
_read_matrix_file_limited,
|
||||
str(component.path),
|
||||
)
|
||||
if image_bytes:
|
||||
resp = await client.upload(image_bytes, content_type='image/png')
|
||||
if isinstance(resp, nio.UploadResponse):
|
||||
@@ -44,17 +103,22 @@ class MatrixMessageConverter(abstract_platform_adapter.AbstractMessageConverter)
|
||||
elif isinstance(component, platform_message.File):
|
||||
file_bytes = None
|
||||
if component.base64:
|
||||
b64_data = component.base64
|
||||
if ';base64,' in b64_data:
|
||||
b64_data = b64_data.split(';base64,', 1)[1]
|
||||
file_bytes = base64.b64decode(b64_data)
|
||||
file_bytes = await asyncio.to_thread(
|
||||
_decode_matrix_base64_limited,
|
||||
component.base64,
|
||||
)
|
||||
elif component.url:
|
||||
session = httpclient.get_session()
|
||||
async with session.get(component.url) as response:
|
||||
file_bytes = await response.read()
|
||||
file_bytes = await httpclient.read_limited(
|
||||
response,
|
||||
max_bytes=_MAX_MATRIX_MEDIA_BYTES,
|
||||
)
|
||||
elif component.path:
|
||||
with open(component.path, 'rb') as f:
|
||||
file_bytes = f.read()
|
||||
file_bytes = await asyncio.to_thread(
|
||||
_read_matrix_file_limited,
|
||||
str(component.path),
|
||||
)
|
||||
if file_bytes:
|
||||
file_name = getattr(component, 'name', None) or 'file'
|
||||
resp = await client.upload(file_bytes, content_type='application/octet-stream', filename=file_name)
|
||||
@@ -86,11 +150,12 @@ class MatrixMessageConverter(abstract_platform_adapter.AbstractMessageConverter)
|
||||
elif isinstance(event, nio.RoomMessageImage):
|
||||
mxc_url = event.url
|
||||
if mxc_url:
|
||||
resp = await client.download(mxc_url)
|
||||
if isinstance(resp, nio.DownloadResponse):
|
||||
b64 = base64.b64encode(resp.body).decode('utf-8')
|
||||
content_type = resp.content_type or 'image/png'
|
||||
message_components.append(platform_message.Image(base64=f'data:{content_type};base64,{b64}'))
|
||||
body, content_type = await _download_matrix_media_limited(
|
||||
client,
|
||||
mxc_url,
|
||||
)
|
||||
b64 = (await asyncio.to_thread(base64.b64encode, body)).decode('utf-8')
|
||||
message_components.append(platform_message.Image(base64=f'data:{content_type};base64,{b64}'))
|
||||
if event.body:
|
||||
message_components.append(platform_message.Plain(text=event.body))
|
||||
|
||||
@@ -431,14 +496,15 @@ class MatrixAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
if not mxc_url:
|
||||
return
|
||||
try:
|
||||
resp = await self.client.download(mxc_url)
|
||||
if isinstance(resp, nio.DownloadResponse):
|
||||
b64 = base64.b64encode(resp.body).decode('utf-8')
|
||||
content_type = resp.content_type or 'image/png'
|
||||
await self.logger.info(
|
||||
f'[{_b.user_id}] Bridge 发送了二维码，请扫码登录:',
|
||||
images=[platform_message.Image(base64=f'data:{content_type};base64,{b64}')],
|
||||
)
|
||||
body, content_type = await _download_matrix_media_limited(
|
||||
self.client,
|
||||
mxc_url,
|
||||
)
|
||||
b64 = (await asyncio.to_thread(base64.b64encode, body)).decode('utf-8')
|
||||
await self.logger.info(
|
||||
f'[{_b.user_id}] Bridge 发送了二维码，请扫码登录:',
|
||||
images=[platform_message.Image(base64=f'data:{content_type};base64,{b64}')],
|
||||
)
|
||||
except Exception:
|
||||
await self.logger.error(
|
||||
f'[{_b.user_id}] Failed to download bridge QR image: {traceback.format_exc()}'
|
||||
@@ -672,11 +738,16 @@ class MatrixAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
|
||||
async def kill(self) -> bool:
|
||||
self._running = False
|
||||
bridge_tasks = []
|
||||
for bridge in self._bridges:
|
||||
if bridge.login_task and not bridge.login_task.done():
|
||||
bridge.login_task.cancel()
|
||||
bridge_tasks.append(bridge.login_task)
|
||||
if bridge.check_task and not bridge.check_task.done():
|
||||
bridge.check_task.cancel()
|
||||
bridge_tasks.append(bridge.check_task)
|
||||
if bridge_tasks:
|
||||
await asyncio.gather(*bridge_tasks, return_exceptions=True)
|
||||
if self.client:
|
||||
await self.client.close()
|
||||
await self.logger.debug('Matrix adapter stopped')
|
||||
|
||||
@@ -164,6 +164,7 @@ class OfficialAccountAdapter(abstract_platform_adapter.AbstractMessagePlatformAd
|
||||
await keep_alive()
|
||||
|
||||
async def kill(self) -> bool:
|
||||
self.bot.clear()
|
||||
return False
|
||||
|
||||
async def unregister_listener(
|
||||
|
||||
@@ -10,6 +10,7 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import os
|
||||
import traceback
|
||||
import typing
|
||||
|
||||
@@ -26,6 +27,7 @@ from langbot.libs.openclaw_weixin_api.types import (
|
||||
WeixinMessage,
|
||||
)
|
||||
from langbot.pkg.entity.persistence import bot as persistence_bot
|
||||
from langbot.pkg.utils import httpclient
|
||||
|
||||
import langbot_plugin.api.definition.abstract.platform.adapter as abstract_platform_adapter
|
||||
import langbot_plugin.api.definition.abstract.platform.event_logger as abstract_platform_logger
|
||||
@@ -35,6 +37,8 @@ import langbot_plugin.api.entities.builtin.platform.message as platform_message
|
||||
|
||||
from langbot.pkg.api.http.context import ExecutionContext
|
||||
|
||||
_MAX_OPENCLAW_COMPONENT_BYTES = 10 * 1024 * 1024
|
||||
|
||||
|
||||
class OpenClawWeixinMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
|
||||
"""Converts between LangBot MessageChain and OpenClaw WeChat message items."""
|
||||
@@ -114,7 +118,12 @@ class OpenClawWeixinMessageConverter(abstract_platform_adapter.AbstractMessageCo
|
||||
|
||||
elif item.type == MessageItem.IMAGE and item.image_item:
|
||||
if hasattr(item.image_item, '_downloaded_bytes') and item.image_item._downloaded_bytes:
|
||||
b64 = base64.b64encode(item.image_item._downloaded_bytes).decode('utf-8')
|
||||
b64 = (
|
||||
await asyncio.to_thread(
|
||||
base64.b64encode,
|
||||
item.image_item._downloaded_bytes,
|
||||
)
|
||||
).decode('utf-8')
|
||||
components.append(platform_message.Image(base64=f'data:image/jpeg;base64,{b64}'))
|
||||
else:
|
||||
components.append(platform_message.Unknown(text='[Image]'))
|
||||
@@ -401,19 +410,30 @@ class OpenClawWeixinAdapter(abstract_platform_adapter.AbstractMessagePlatformAda
|
||||
path_val = getattr(component, 'path', None)
|
||||
|
||||
if b64_val:
|
||||
return base64.b64decode(b64_val)
|
||||
max_encoded_chars = 4 * ((_MAX_OPENCLAW_COMPONENT_BYTES + 2) // 3) + 4
|
||||
if len(b64_val) > max_encoded_chars:
|
||||
raise ValueError('OpenClaw media exceeds the size limit')
|
||||
data = await asyncio.to_thread(base64.b64decode, b64_val)
|
||||
if len(data) > _MAX_OPENCLAW_COMPONENT_BYTES:
|
||||
raise ValueError('OpenClaw media exceeds the size limit')
|
||||
return data
|
||||
elif url_val and url_val.startswith(('http://', 'https://')):
|
||||
import aiohttp
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.get(url_val) as resp:
|
||||
if resp.status == 200:
|
||||
return await resp.read()
|
||||
session = httpclient.get_session()
|
||||
async with session.get(url_val) as resp:
|
||||
if resp.status == 200:
|
||||
return await httpclient.read_limited(resp)
|
||||
elif path_val:
|
||||
import asyncio
|
||||
if await asyncio.to_thread(os.path.getsize, path_val) > _MAX_OPENCLAW_COMPONENT_BYTES:
|
||||
raise ValueError('OpenClaw media exceeds the size limit')
|
||||
|
||||
with open(path_val, 'rb') as f:
|
||||
return await asyncio.to_thread(f.read)
|
||||
def read_file() -> bytes:
|
||||
with open(path_val, 'rb') as file:
|
||||
return file.read(_MAX_OPENCLAW_COMPONENT_BYTES + 1)
|
||||
|
||||
data = await asyncio.to_thread(read_file)
|
||||
if len(data) > _MAX_OPENCLAW_COMPONENT_BYTES:
|
||||
raise ValueError('OpenClaw media exceeds the size limit')
|
||||
return data
|
||||
return None
|
||||
|
||||
def register_listener(
|
||||
@@ -544,6 +564,8 @@ class OpenClawWeixinAdapter(abstract_platform_adapter.AbstractMessagePlatformAda
|
||||
"""Process a single inbound message from getUpdates."""
|
||||
if msg.context_token and msg.from_user_id:
|
||||
self._context_tokens[msg.from_user_id] = msg.context_token
|
||||
while len(self._context_tokens) > 4096:
|
||||
self._context_tokens.pop(next(iter(self._context_tokens)), None)
|
||||
|
||||
# Download CDN media (files, images) before converting to LangBot events
|
||||
await self._download_media_items(msg)
|
||||
@@ -599,6 +621,8 @@ class OpenClawWeixinAdapter(abstract_platform_adapter.AbstractMessagePlatformAda
|
||||
await self._poll_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
self._poll_task = None
|
||||
self._context_tokens.clear()
|
||||
await self.client.close()
|
||||
await self.logger.info('OpenClaw WeChat adapter stopped')
|
||||
return True
|
||||
|
||||
@@ -241,6 +241,7 @@ class QQOfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter
|
||||
# per (msg_id|event_id) within 60 min, but each reuse needs a
|
||||
# fresh ``msg_seq`` — re-sending with msg_seq=1 is silently dedup'd.
|
||||
self._anchor_msg_seq: dict[str, int] = {}
|
||||
self._background_tasks: set[asyncio.Task] = set()
|
||||
|
||||
# Wire button-click handler so webhook mode catches INTERACTION_CREATE.
|
||||
# (ws mode is wired separately via on_event in _run_websocket so the
|
||||
@@ -249,6 +250,30 @@ class QQOfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter
|
||||
async def _on_interaction(event_data: dict, interaction_id: typing.Optional[str]):
|
||||
await self._handle_interaction_create(event_data, interaction_id)
|
||||
|
||||
def _start_background_task(self, coro) -> bool:
|
||||
"""Start one bounded adapter-side auxiliary task."""
|
||||
|
||||
background_tasks = getattr(self, '_background_tasks', None)
|
||||
if background_tasks is None:
|
||||
background_tasks = set()
|
||||
object.__setattr__(self, '_background_tasks', background_tasks)
|
||||
for task in tuple(background_tasks):
|
||||
if task.done():
|
||||
background_tasks.discard(task)
|
||||
if len(background_tasks) >= 100:
|
||||
coro.close()
|
||||
return False
|
||||
task = asyncio.create_task(coro)
|
||||
background_tasks.add(task)
|
||||
|
||||
def done(done_task: asyncio.Task) -> None:
|
||||
background_tasks.discard(done_task)
|
||||
if not done_task.cancelled():
|
||||
done_task.exception()
|
||||
|
||||
task.add_done_callback(done)
|
||||
return True
|
||||
|
||||
async def reply_message(
|
||||
self,
|
||||
message_source: platform_events.MessageEvent,
|
||||
@@ -449,6 +474,14 @@ class QQOfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter
|
||||
pass
|
||||
|
||||
async def kill(self) -> bool:
|
||||
task_set = getattr(self, '_background_tasks', set())
|
||||
background_tasks = list(task_set)
|
||||
for task in background_tasks:
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
if background_tasks:
|
||||
await asyncio.gather(*background_tasks, return_exceptions=True)
|
||||
task_set.clear()
|
||||
if self._ws_task:
|
||||
self._ws_task.cancel()
|
||||
try:
|
||||
@@ -456,6 +489,14 @@ class QQOfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
self._ws_task = None
|
||||
await self.bot.close()
|
||||
self._pending_forms.clear()
|
||||
self._session_event_ids.clear()
|
||||
self._anchor_msg_seq.clear()
|
||||
self._stream_ctx.clear()
|
||||
self._stream_ctx_ts.clear()
|
||||
self._fallback_text.clear()
|
||||
self._fallback_text_ts.clear()
|
||||
return True
|
||||
|
||||
# --------------- 流式输出 ---------------
|
||||
@@ -473,6 +514,14 @@ class QQOfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter
|
||||
for mid in stale_fb:
|
||||
self._fallback_text.pop(mid, None)
|
||||
self._fallback_text_ts.pop(mid, None)
|
||||
while len(self._stream_ctx) > 1000:
|
||||
oldest = min(self._stream_ctx_ts, key=self._stream_ctx_ts.__getitem__)
|
||||
self._stream_ctx.pop(oldest, None)
|
||||
self._stream_ctx_ts.pop(oldest, None)
|
||||
while len(self._fallback_text) > 1000:
|
||||
oldest = min(self._fallback_text_ts, key=self._fallback_text_ts.__getitem__)
|
||||
self._fallback_text.pop(oldest, None)
|
||||
self._fallback_text_ts.pop(oldest, None)
|
||||
if stale_ids or stale_fb:
|
||||
await self.logger.debug(f'Cleaned up {len(stale_ids)} stream contexts, {len(stale_fb)} fallback texts')
|
||||
|
||||
@@ -508,6 +557,8 @@ class QQOfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter
|
||||
# msg_seq=2 instead of being deduplicated by QQ as another seq=1 send.
|
||||
if source.d_id:
|
||||
self._anchor_msg_seq[source.d_id] = max(self._anchor_msg_seq.get(source.d_id, 0), 1)
|
||||
while len(self._anchor_msg_seq) > 4096:
|
||||
self._anchor_msg_seq.pop(next(iter(self._anchor_msg_seq)), None)
|
||||
|
||||
ctx = {
|
||||
'user_openid': source.user_openid,
|
||||
@@ -577,7 +628,7 @@ class QQOfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter
|
||||
# 非流式场景（如群聊不支持流式），累积文本后一次性回复
|
||||
if chunk_text:
|
||||
# Chunks carry the latest full snapshot, not a text delta.
|
||||
self._fallback_text[message_id] = chunk_text
|
||||
self._fallback_text[message_id] = chunk_text[:200000]
|
||||
self._fallback_text_ts[message_id] = time.time()
|
||||
if is_final:
|
||||
full_text = self._fallback_text.pop(message_id, '')
|
||||
@@ -590,7 +641,7 @@ class QQOfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter
|
||||
|
||||
# 累积文本
|
||||
if chunk_text:
|
||||
ctx['accumulated_text'] = chunk_text
|
||||
ctx['accumulated_text'] = chunk_text[:200000]
|
||||
|
||||
# 未启动会话时，等第一个有内容的 chunk 来建立会话
|
||||
if not ctx['session_started']:
|
||||
@@ -668,6 +719,8 @@ class QQOfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter
|
||||
if used >= self._MAX_REPLIES_PER_ANCHOR:
|
||||
return None
|
||||
self._anchor_msg_seq[anchor] = used + 1
|
||||
while len(self._anchor_msg_seq) > 4096:
|
||||
self._anchor_msg_seq.pop(next(iter(self._anchor_msg_seq)), None)
|
||||
return used + 1
|
||||
|
||||
async def _reply_synthetic(
|
||||
@@ -791,7 +844,9 @@ class QQOfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter
|
||||
k for k, v in self._session_event_ids.items() if now - v.get('posted_at', 0) > self._PENDING_FORM_TTL
|
||||
]
|
||||
for k in stale_e:
|
||||
self._session_event_ids.pop(k, None)
|
||||
stale_event = self._session_event_ids.pop(k, None)
|
||||
if stale_event:
|
||||
self._anchor_msg_seq.pop(stale_event.get('event_id'), None)
|
||||
|
||||
async def _handle_form_chunk(
|
||||
self,
|
||||
@@ -973,7 +1028,7 @@ class QQOfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter
|
||||
# ACK uses the interaction id, NOT the ws event id.
|
||||
interaction_id = event_data.get('id') or ''
|
||||
if interaction_id:
|
||||
asyncio.create_task(self.bot.ack_interaction(interaction_id, code=0))
|
||||
self._start_background_task(self.bot.ack_interaction(interaction_id, code=0))
|
||||
|
||||
resolved = (event_data.get('data') or {}).get('resolved') or {}
|
||||
action_id = str(resolved.get('button_data') or resolved.get('button_id') or '').strip()
|
||||
@@ -1018,6 +1073,8 @@ class QQOfficialAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter
|
||||
}
|
||||
# New anchor → fresh 5-reply budget.
|
||||
self._anchor_msg_seq[cached_event_id] = 0
|
||||
while len(self._anchor_msg_seq) > 4096:
|
||||
self._anchor_msg_seq.pop(next(iter(self._anchor_msg_seq)), None)
|
||||
if self.ap is not None and not ws_event_id:
|
||||
self.ap.logger.warning(
|
||||
'QQ Official: INTERACTION_CREATE lacked ws_event_id; '
|
||||
|
||||
@@ -18,6 +18,9 @@ import langbot_plugin.api.entities.builtin.platform.message as platform_message
|
||||
import langbot_plugin.api.entities.builtin.platform.events as platform_events
|
||||
import langbot_plugin.api.entities.builtin.platform.entities as platform_entities
|
||||
import langbot_plugin.api.definition.abstract.platform.event_logger as abstract_platform_logger
|
||||
from langbot.pkg.utils import httpclient
|
||||
|
||||
_MAX_GATEWAY_MESSAGE_BYTES = 1024 * 1024
|
||||
|
||||
|
||||
class SatoriMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
|
||||
@@ -63,7 +66,15 @@ class SatoriMessageConverter(abstract_platform_adapter.AbstractMessageConverter)
|
||||
padding = 4 - len(raw_b64) % 4
|
||||
if padding != 4:
|
||||
raw_b64 += '=' * padding
|
||||
image_bytes = base64.b64decode(raw_b64)
|
||||
max_encoded_chars = 4 * ((10 * 1024 * 1024 + 2) // 3) + 4
|
||||
if len(raw_b64) > max_encoded_chars:
|
||||
raise ValueError('Satori image exceeds the 10 MiB limit')
|
||||
image_bytes = await asyncio.to_thread(
|
||||
base64.b64decode,
|
||||
raw_b64,
|
||||
)
|
||||
if len(image_bytes) > 10 * 1024 * 1024:
|
||||
raise ValueError('Satori image exceeds the 10 MiB limit')
|
||||
uploaded_url = await adapter.upload_image(image_bytes, mime_type)
|
||||
if uploaded_url:
|
||||
await adapter.logger.info(f'Satori 图片上传成功: {len(image_bytes)} 字节')
|
||||
@@ -492,7 +503,7 @@ class SatoriAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
raise ValueError(f'WebSocket URL必须以ws://或wss://开头: {self.endpoint}')
|
||||
|
||||
try:
|
||||
self.ws = await websockets.connect(self.endpoint)
|
||||
self.ws = await websockets.connect(self.endpoint, max_size=_MAX_GATEWAY_MESSAGE_BYTES)
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
await self.send_identify()
|
||||
@@ -584,7 +595,7 @@ class SatoriAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
async def handle_message(self, message: str):
|
||||
"""Handle WebSocket message"""
|
||||
try:
|
||||
data = json.loads(message)
|
||||
data = await asyncio.to_thread(json.loads, message)
|
||||
op = data.get('op')
|
||||
body = data.get('body', {})
|
||||
|
||||
@@ -831,9 +842,10 @@ class SatoriAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
try:
|
||||
async with self.session.request(method, url, headers=headers, json=data) as response:
|
||||
if response.status == 200:
|
||||
return await response.json()
|
||||
result = await httpclient.read_json_limited(response)
|
||||
return result if isinstance(result, dict) else None
|
||||
else:
|
||||
text = await response.text()
|
||||
text = await httpclient.read_text_limited(response)
|
||||
await self.logger.error(f'Satori API 请求失败: {response.status} - {text}')
|
||||
return None
|
||||
except Exception as e:
|
||||
@@ -889,7 +901,7 @@ class SatoriAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
|
||||
async with self.session.post(url, headers=headers, data=form_data) as response:
|
||||
if response.status == 200:
|
||||
result = await response.json()
|
||||
result = await httpclient.read_json_limited(response)
|
||||
# The response should contain the URL of the uploaded file
|
||||
if isinstance(result, dict) and 'url' in result:
|
||||
return result['url']
|
||||
@@ -899,7 +911,7 @@ class SatoriAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
await self.logger.warning(f'Satori 图片上传响应格式未知: {result}')
|
||||
return None
|
||||
else:
|
||||
text = await response.text()
|
||||
text = await httpclient.read_text_limited(response)
|
||||
await self.logger.error(f'Satori 图片上传失败: {response.status} - {text}')
|
||||
return None
|
||||
except Exception as e:
|
||||
@@ -911,6 +923,8 @@ class SatoriAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
self.running = False
|
||||
if self.heartbeat_task:
|
||||
self.heartbeat_task.cancel()
|
||||
await asyncio.gather(self.heartbeat_task, return_exceptions=True)
|
||||
self.heartbeat_task = None
|
||||
if self.ws:
|
||||
try:
|
||||
await self.ws.close()
|
||||
|
||||
@@ -10,6 +10,8 @@ import typing
|
||||
import traceback
|
||||
import json
|
||||
import base64
|
||||
import asyncio
|
||||
import os
|
||||
import time
|
||||
import uuid
|
||||
import pydantic
|
||||
@@ -22,6 +24,31 @@ import langbot_plugin.api.entities.builtin.platform.entities as platform_entitie
|
||||
import langbot_plugin.api.definition.abstract.platform.event_logger as abstract_platform_logger
|
||||
|
||||
|
||||
_MAX_TELEGRAM_MEDIA_BYTES = 10 * 1024 * 1024
|
||||
|
||||
|
||||
def _decode_telegram_base64_limited(value: str) -> bytes:
|
||||
if ';base64,' in value:
|
||||
value = value.split(';base64,', 1)[1]
|
||||
max_encoded_bytes = 4 * ((_MAX_TELEGRAM_MEDIA_BYTES + 2) // 3)
|
||||
if len(value) > max_encoded_bytes:
|
||||
raise ValueError('Telegram media exceeds the size limit')
|
||||
decoded = base64.b64decode(value)
|
||||
if len(decoded) > _MAX_TELEGRAM_MEDIA_BYTES:
|
||||
raise ValueError('Telegram media exceeds the size limit')
|
||||
return decoded
|
||||
|
||||
|
||||
def _read_telegram_file_limited(path: str) -> bytes:
|
||||
if os.path.getsize(path) > _MAX_TELEGRAM_MEDIA_BYTES:
|
||||
raise ValueError('Telegram media exceeds the size limit')
|
||||
with open(path, 'rb') as file:
|
||||
body = file.read(_MAX_TELEGRAM_MEDIA_BYTES + 1)
|
||||
if len(body) > _MAX_TELEGRAM_MEDIA_BYTES:
|
||||
raise ValueError('Telegram media exceeds the size limit')
|
||||
return body
|
||||
|
||||
|
||||
def _telegram_select_field_options(form_data: dict) -> tuple[str, list[str]]:
|
||||
"""Return the active select field and its option values."""
|
||||
field_name = str(form_data.get('_current_input_field') or '').strip()
|
||||
@@ -86,35 +113,29 @@ class TelegramMessageConverter(abstract_platform_adapter.AbstractMessageConverte
|
||||
if isinstance(component, platform_message.Plain):
|
||||
components.append({'type': 'text', 'text': component.text})
|
||||
elif isinstance(component, platform_message.Image):
|
||||
photo_bytes = None
|
||||
|
||||
if component.base64:
|
||||
photo_bytes = base64.b64decode(component.base64)
|
||||
elif component.url:
|
||||
session = httpclient.get_session()
|
||||
async with session.get(component.url) as response:
|
||||
photo_bytes = await response.read()
|
||||
elif component.path:
|
||||
with open(component.path, 'rb') as f:
|
||||
photo_bytes = f.read()
|
||||
photo_bytes, _mime_type = await component.get_bytes()
|
||||
|
||||
components.append({'type': 'photo', 'photo': photo_bytes})
|
||||
elif isinstance(component, platform_message.File):
|
||||
file_bytes = None
|
||||
|
||||
if component.base64:
|
||||
# Strip data URI prefix if present (e.g. "data:application/pdf;base64,...")
|
||||
b64_data = component.base64
|
||||
if ';base64,' in b64_data:
|
||||
b64_data = b64_data.split(';base64,', 1)[1]
|
||||
file_bytes = base64.b64decode(b64_data)
|
||||
file_bytes = await asyncio.to_thread(
|
||||
_decode_telegram_base64_limited,
|
||||
component.base64,
|
||||
)
|
||||
elif component.url:
|
||||
session = httpclient.get_session()
|
||||
async with session.get(component.url) as response:
|
||||
file_bytes = await response.read()
|
||||
file_bytes = await httpclient.read_limited(
|
||||
response,
|
||||
max_bytes=_MAX_TELEGRAM_MEDIA_BYTES,
|
||||
)
|
||||
elif component.path:
|
||||
with open(component.path, 'rb') as f:
|
||||
file_bytes = f.read()
|
||||
file_bytes = await asyncio.to_thread(
|
||||
_read_telegram_file_limited,
|
||||
str(component.path),
|
||||
)
|
||||
|
||||
file_name = getattr(component, 'name', None) or 'file'
|
||||
components.append({'type': 'document', 'document': file_bytes, 'filename': file_name})
|
||||
@@ -152,13 +173,15 @@ class TelegramMessageConverter(abstract_platform_adapter.AbstractMessageConverte
|
||||
file_format = ''
|
||||
|
||||
async with httpclient.get_session(trust_env=True).get(file.file_path) as response:
|
||||
file_bytes = await response.read()
|
||||
file_bytes = await httpclient.read_limited(
|
||||
response,
|
||||
max_bytes=_MAX_TELEGRAM_MEDIA_BYTES,
|
||||
)
|
||||
file_format = 'image/jpeg'
|
||||
|
||||
encoded = await asyncio.to_thread(base64.b64encode, file_bytes)
|
||||
message_components.append(
|
||||
platform_message.Image(
|
||||
base64=f'data:{file_format};base64,{base64.b64encode(file_bytes).decode("utf-8")}'
|
||||
)
|
||||
platform_message.Image(base64=f'data:{file_format};base64,{encoded.decode("utf-8")}')
|
||||
)
|
||||
|
||||
if message.voice:
|
||||
@@ -171,11 +194,15 @@ class TelegramMessageConverter(abstract_platform_adapter.AbstractMessageConverte
|
||||
file_format = message.voice.mime_type or 'audio/ogg'
|
||||
|
||||
async with httpclient.get_session(trust_env=True).get(file.file_path) as response:
|
||||
file_bytes = await response.read()
|
||||
file_bytes = await httpclient.read_limited(
|
||||
response,
|
||||
max_bytes=_MAX_TELEGRAM_MEDIA_BYTES,
|
||||
)
|
||||
|
||||
encoded = await asyncio.to_thread(base64.b64encode, file_bytes)
|
||||
message_components.append(
|
||||
platform_message.Voice(
|
||||
base64=f'data:{file_format};base64,{base64.b64encode(file_bytes).decode("utf-8")}',
|
||||
base64=f'data:{file_format};base64,{encoded.decode("utf-8")}',
|
||||
length=message.voice.duration,
|
||||
)
|
||||
)
|
||||
@@ -188,16 +215,22 @@ class TelegramMessageConverter(abstract_platform_adapter.AbstractMessageConverte
|
||||
file_name = message.document.file_name or 'document'
|
||||
file_size = message.document.file_size or 0
|
||||
file_format = message.document.mime_type or 'application/octet-stream'
|
||||
if file_size > _MAX_TELEGRAM_MEDIA_BYTES:
|
||||
raise ValueError('Telegram media exceeds the size limit')
|
||||
|
||||
file_bytes = None
|
||||
async with httpclient.get_session(trust_env=True).get(file.file_path) as response:
|
||||
file_bytes = await response.read()
|
||||
file_bytes = await httpclient.read_limited(
|
||||
response,
|
||||
max_bytes=_MAX_TELEGRAM_MEDIA_BYTES,
|
||||
)
|
||||
|
||||
encoded = await asyncio.to_thread(base64.b64encode, file_bytes)
|
||||
message_components.append(
|
||||
platform_message.File(
|
||||
name=file_name,
|
||||
size=file_size,
|
||||
base64=f'data:{file_format};base64,{base64.b64encode(file_bytes).decode("utf-8")}',
|
||||
base64=f'data:{file_format};base64,{encoded.decode("utf-8")}',
|
||||
)
|
||||
)
|
||||
|
||||
@@ -263,6 +296,8 @@ class TelegramAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
] = {}
|
||||
|
||||
_FORM_ACTION_CACHE_TTL = 30 * 60
|
||||
_MAX_FORM_ACTION_TITLES = 4096
|
||||
_MAX_STREAM_STATES = 1000
|
||||
# callback_data -> (display title, pipeline UUID, expiration time, form group id)
|
||||
_form_action_titles: typing.Dict[str, tuple[str, str, float, str]] = {}
|
||||
|
||||
@@ -285,6 +320,8 @@ class TelegramAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
self._form_action_titles.update(
|
||||
{callback_data: (title, pipeline_uuid, expires_at, group_id) for callback_data, title in mappings.items()}
|
||||
)
|
||||
while len(self._form_action_titles) > self._MAX_FORM_ACTION_TITLES:
|
||||
self._form_action_titles.pop(next(iter(self._form_action_titles)), None)
|
||||
|
||||
def _take_form_action_context(self, callback_data: str, now: float | None = None) -> tuple[str, str] | None:
|
||||
"""Consume a callback and invalidate every button from the same form."""
|
||||
@@ -445,6 +482,11 @@ class TelegramAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
bot_account_id='',
|
||||
listeners={},
|
||||
)
|
||||
self._form_action_titles = {}
|
||||
|
||||
def _cap_stream_states(self) -> None:
|
||||
while len(self.msg_stream_id) > self._MAX_STREAM_STATES:
|
||||
self.msg_stream_id.pop(next(iter(self.msg_stream_id)), None)
|
||||
|
||||
async def send_message(self, target_type: str, target_id: str, message: platform_message.MessageChain):
|
||||
components = await TelegramMessageConverter.yiri2target(message, self.bot)
|
||||
@@ -553,6 +595,7 @@ class TelegramAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
args = self._build_message_args(chat_id, 'Thinking...', message_thread_id)
|
||||
send_msg = await self.bot.send_message(**args)
|
||||
self.msg_stream_id[message_id] = ('message', send_msg.message_id, False)
|
||||
self._cap_stream_states()
|
||||
|
||||
return True
|
||||
|
||||
@@ -844,4 +887,6 @@ class TelegramAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
if self.application.updater:
|
||||
await self.application.updater.stop()
|
||||
await self.logger.info('Telegram adapter stopped')
|
||||
self.msg_stream_id.clear()
|
||||
self._form_action_titles.clear()
|
||||
return True
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
import asyncio
|
||||
import contextvars
|
||||
import logging
|
||||
import time
|
||||
import typing
|
||||
from datetime import datetime
|
||||
|
||||
@@ -14,6 +15,7 @@ import langbot_plugin.api.entities.builtin.platform.events as platform_events
|
||||
import langbot_plugin.api.entities.builtin.platform.entities as platform_entities
|
||||
import langbot_plugin.api.definition.abstract.platform.event_logger as abstract_platform_logger
|
||||
from ...core import app
|
||||
from ...core import entities as core_entities
|
||||
from .websocket_manager import WebSocketConnection, WebSocketScope, is_valid_session_id, ws_connection_manager
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -45,21 +47,82 @@ class WebSocketSession:
|
||||
stream_message_indexes: dict[str, dict[str, int]] = {}
|
||||
"""流式消息索引 {pipeline_uuid: {resp_message_id: message_index}}"""
|
||||
|
||||
def __init__(self, id: str):
|
||||
def __init__(
|
||||
self,
|
||||
id: str = '',
|
||||
*,
|
||||
max_conversations: int = 200,
|
||||
max_messages: int = 100,
|
||||
idle_ttl_seconds: int = 86400,
|
||||
):
|
||||
self.id = id
|
||||
self.message_lists = {}
|
||||
self.stream_message_indexes = {}
|
||||
self.message_counters: dict[str, int] = {}
|
||||
self.last_accessed: dict[str, float] = {}
|
||||
self.max_conversations = max(int(max_conversations), 1)
|
||||
self.max_messages = max(int(max_messages), 1)
|
||||
self.idle_ttl_seconds = max(int(idle_ttl_seconds), 1)
|
||||
|
||||
def _prune(self, now: float) -> None:
|
||||
expired = [
|
||||
key for key, last_accessed in self.last_accessed.items() if now - last_accessed >= self.idle_ttl_seconds
|
||||
]
|
||||
for key in expired:
|
||||
self.reset(key)
|
||||
|
||||
overflow = len(self.message_lists) - self.max_conversations + 1
|
||||
if overflow <= 0:
|
||||
return
|
||||
oldest = sorted(self.last_accessed, key=self.last_accessed.get)
|
||||
for key in oldest[:overflow]:
|
||||
self.reset(key)
|
||||
|
||||
def get_message_list(self, pipeline_uuid: str) -> list[WebSocketMessage]:
|
||||
now = time.monotonic()
|
||||
self._prune(now)
|
||||
if pipeline_uuid not in self.message_lists:
|
||||
self.message_lists[pipeline_uuid] = []
|
||||
self.last_accessed[pipeline_uuid] = now
|
||||
return self.message_lists[pipeline_uuid]
|
||||
|
||||
def get_stream_message_indexes(self, pipeline_uuid: str) -> dict[str, int]:
|
||||
if pipeline_uuid not in self.stream_message_indexes:
|
||||
self.stream_message_indexes[pipeline_uuid] = {}
|
||||
self.last_accessed[pipeline_uuid] = time.monotonic()
|
||||
return self.stream_message_indexes[pipeline_uuid]
|
||||
|
||||
def next_message_id(self, conversation_key: str) -> int:
|
||||
next_id = self.message_counters.get(conversation_key, 0) + 1
|
||||
self.message_counters[conversation_key] = next_id
|
||||
return next_id
|
||||
|
||||
def append_message(self, conversation_key: str, message: WebSocketMessage) -> None:
|
||||
messages = self.get_message_list(conversation_key)
|
||||
messages.append(message)
|
||||
overflow = len(messages) - self.max_messages
|
||||
if overflow <= 0:
|
||||
return
|
||||
del messages[:overflow]
|
||||
indexes = self.stream_message_indexes.get(conversation_key, {})
|
||||
adjusted_indexes = {
|
||||
response_id: index - overflow for response_id, index in indexes.items() if index >= overflow
|
||||
}
|
||||
indexes.clear()
|
||||
indexes.update(adjusted_indexes)
|
||||
|
||||
def reset(self, conversation_key: str) -> None:
|
||||
self.message_lists.pop(conversation_key, None)
|
||||
self.stream_message_indexes.pop(conversation_key, None)
|
||||
self.message_counters.pop(conversation_key, None)
|
||||
self.last_accessed.pop(conversation_key, None)
|
||||
|
||||
def clear(self) -> None:
|
||||
self.message_lists.clear()
|
||||
self.stream_message_indexes.clear()
|
||||
self.message_counters.clear()
|
||||
self.last_accessed.clear()
|
||||
|
||||
|
||||
class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
"""WebSocket适配器 - 支持双向实时通信"""
|
||||
@@ -75,7 +138,14 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
|
||||
ap: app.Application = pydantic.Field(exclude=True)
|
||||
|
||||
# 主动推送消息的队列
|
||||
outbound_message_queue: asyncio.Queue = pydantic.Field(default_factory=asyncio.Queue, exclude=True)
|
||||
outbound_message_queue: asyncio.Queue = pydantic.Field(
|
||||
default_factory=lambda: asyncio.Queue(maxsize=100),
|
||||
exclude=True,
|
||||
)
|
||||
inbound_listener_tasks: set[asyncio.Task] = pydantic.Field(
|
||||
default_factory=set,
|
||||
exclude=True,
|
||||
)
|
||||
"""后端主动推送消息的队列"""
|
||||
|
||||
# 流式输出开关
|
||||
@@ -89,11 +159,26 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
self.websocket_person_session = WebSocketSession(id='websocketperson')
|
||||
self.websocket_group_session = WebSocketSession(id='websocketgroup')
|
||||
application = kwargs.get('ap')
|
||||
instance_data = getattr(getattr(application, 'instance_config', None), 'data', {})
|
||||
retention = (
|
||||
instance_data.get('system', {}).get('websocket_retention', {}) if isinstance(instance_data, dict) else {}
|
||||
)
|
||||
session_options = {
|
||||
'max_conversations': retention.get('max_conversations_per_workspace', 200),
|
||||
'max_messages': retention.get('max_messages_per_conversation', 100),
|
||||
'idle_ttl_seconds': retention.get('conversation_idle_ttl_seconds', 86400),
|
||||
}
|
||||
self.websocket_person_session = WebSocketSession(id='websocketperson', **session_options)
|
||||
self.websocket_group_session = WebSocketSession(id='websocketgroup', **session_options)
|
||||
|
||||
self.bot_account_id = 'websocketbot'
|
||||
self.outbound_message_queue = asyncio.Queue()
|
||||
try:
|
||||
outbound_queue_size = max(int(retention.get('send_queue_size', 100)), 1)
|
||||
except (TypeError, ValueError):
|
||||
outbound_queue_size = 100
|
||||
self.outbound_message_queue = asyncio.Queue(maxsize=outbound_queue_size)
|
||||
self.inbound_listener_tasks = set()
|
||||
self.stream_enabled = True
|
||||
|
||||
@staticmethod
|
||||
@@ -128,6 +213,25 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
|
||||
|
||||
return _current_pipeline_uuid.get()
|
||||
|
||||
def _listener_task_done(self, task: asyncio.Task) -> None:
|
||||
listener_tasks = getattr(self, 'inbound_listener_tasks', None)
|
||||
if listener_tasks is not None:
|
||||
listener_tasks.discard(task)
|
||||
if not task.cancelled():
|
||||
task.exception()
|
||||
|
||||
@staticmethod
|
||||
def _history_message_chain(message_chain: list[dict]) -> list[dict]:
|
||||
"""Remove large transient payloads before retaining browser history."""
|
||||
|
||||
history = []
|
||||
for component in message_chain:
|
||||
copied = dict(component)
|
||||
if copied.get('base64'):
|
||||
copied['base64'] = ''
|
||||
history.append(copied)
|
||||
return history
|
||||
|
||||
async def _get_connection_from_target(self, target_id: str):
|
||||
"""Resolve a person or group WebSocket launcher to its connection."""
|
||||
scope = self._scope()
|
||||
@@ -195,7 +299,7 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
|
||||
|
||||
session = self.websocket_group_session if session_type == 'group' else self.websocket_person_session
|
||||
|
||||
msg_id = len(session.get_message_list(conversation_key)) + 1
|
||||
msg_id = session.next_message_id(conversation_key)
|
||||
|
||||
message_data = WebSocketMessage(
|
||||
id=msg_id,
|
||||
@@ -206,7 +310,7 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
|
||||
is_final=True,
|
||||
)
|
||||
|
||||
session.get_message_list(conversation_key).append(message_data)
|
||||
session.append_message(conversation_key, message_data)
|
||||
|
||||
await ws_connection_manager.broadcast_to_pipeline(
|
||||
pipeline_uuid,
|
||||
@@ -241,7 +345,7 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
|
||||
session_type = 'group' if isinstance(message_source, platform_events.GroupMessage) else 'person'
|
||||
conversation_key = self._conversation_key(pipeline_uuid, session_id)
|
||||
|
||||
msg_id = len(session.get_message_list(conversation_key)) + 1
|
||||
msg_id = session.next_message_id(conversation_key)
|
||||
|
||||
message_data = WebSocketMessage(
|
||||
id=msg_id,
|
||||
@@ -252,7 +356,7 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
|
||||
is_final=True,
|
||||
)
|
||||
|
||||
session.get_message_list(conversation_key).append(message_data)
|
||||
session.append_message(conversation_key, message_data)
|
||||
|
||||
await ws_connection_manager.broadcast_to_pipeline(
|
||||
pipeline_uuid,
|
||||
@@ -300,7 +404,7 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
|
||||
|
||||
if existing_index is None or existing_index >= len(message_list):
|
||||
# 创建新消息
|
||||
msg_id = len(message_list) + 1
|
||||
msg_id = session.next_message_id(conversation_key)
|
||||
message_data = WebSocketMessage(
|
||||
id=msg_id,
|
||||
role='assistant',
|
||||
@@ -311,7 +415,8 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
|
||||
)
|
||||
|
||||
# 立即添加到历史记录（即使is_final=False），以便后续块可以更新它
|
||||
message_list.append(message_data)
|
||||
session.append_message(conversation_key, message_data)
|
||||
message_list = session.get_message_list(conversation_key)
|
||||
if resp_message_id:
|
||||
stream_message_indexes[resp_message_id] = len(message_list) - 1
|
||||
else:
|
||||
@@ -399,7 +504,22 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
|
||||
|
||||
async def kill(self):
|
||||
"""停止适配器"""
|
||||
pass
|
||||
await ws_connection_manager.close_scope(self._scope())
|
||||
listener_tasks = getattr(self, 'inbound_listener_tasks', set())
|
||||
inbound_tasks = list(listener_tasks)
|
||||
for task in inbound_tasks:
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
if inbound_tasks:
|
||||
await asyncio.gather(*inbound_tasks, return_exceptions=True)
|
||||
listener_tasks.clear()
|
||||
self.websocket_person_session.clear()
|
||||
self.websocket_group_session.clear()
|
||||
while not self.outbound_message_queue.empty():
|
||||
try:
|
||||
self.outbound_message_queue.get_nowait()
|
||||
except asyncio.QueueEmpty:
|
||||
break
|
||||
|
||||
async def _process_image_components(
|
||||
self,
|
||||
@@ -445,7 +565,7 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
|
||||
|
||||
try:
|
||||
file_content = await storage_mgr.storage_provider.load(comp_path)
|
||||
base64_str = base64.b64encode(file_content).decode('utf-8')
|
||||
base64_str = (await asyncio.to_thread(base64.b64encode, file_content)).decode('utf-8')
|
||||
|
||||
lowered = comp_path.lower()
|
||||
if comp_type == 'Image':
|
||||
@@ -507,19 +627,19 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
|
||||
|
||||
message_chain = platform_message.MessageChain.model_validate(message_chain_obj)
|
||||
|
||||
message_id = len(use_session.get_message_list(conversation_key)) + 1
|
||||
message_id = use_session.next_message_id(conversation_key)
|
||||
|
||||
# 保存用户消息
|
||||
user_message = WebSocketMessage(
|
||||
id=message_id,
|
||||
role='user',
|
||||
content=str(message_chain),
|
||||
message_chain=message_chain_obj,
|
||||
message_chain=self._history_message_chain(message_chain_obj),
|
||||
timestamp=datetime.now().isoformat(),
|
||||
connection_id=connection.connection_id,
|
||||
is_final=True, # 用户消息始终是完整的，非流式
|
||||
)
|
||||
use_session.get_message_list(conversation_key).append(user_message)
|
||||
use_session.append_message(conversation_key, user_message)
|
||||
|
||||
await ws_connection_manager.broadcast_to_pipeline(
|
||||
pipeline_uuid,
|
||||
@@ -573,9 +693,36 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
|
||||
owner_bot.adapter.set_ws_adapter(self)
|
||||
callback_adapter = owner_bot.adapter if (owner_bot and hasattr(owner_bot, 'adapter')) else self
|
||||
if event.__class__ in listeners:
|
||||
listener_tasks = getattr(self, 'inbound_listener_tasks', None)
|
||||
if listener_tasks is None:
|
||||
listener_tasks = set()
|
||||
object.__setattr__(self, 'inbound_listener_tasks', listener_tasks)
|
||||
for task in tuple(listener_tasks):
|
||||
if task.done():
|
||||
listener_tasks.discard(task)
|
||||
if len(listener_tasks) >= 100:
|
||||
await self.logger.warning('WebSocket inbound listener capacity reached; dropping message')
|
||||
return
|
||||
token = _current_pipeline_uuid.set(pipeline_uuid)
|
||||
try:
|
||||
asyncio.create_task(listeners[event.__class__](event, callback_adapter))
|
||||
task_manager = getattr(self.ap, 'task_mgr', None)
|
||||
if task_manager is None or not isinstance(getattr(task_manager, 'tasks', None), list):
|
||||
listener_task = asyncio.create_task(listeners[event.__class__](event, callback_adapter))
|
||||
else:
|
||||
listener_task = task_manager.create_task(
|
||||
listeners[event.__class__](event, callback_adapter),
|
||||
kind='websocket-message',
|
||||
name=f'websocket-message-{connection.connection_id}',
|
||||
scopes=[
|
||||
core_entities.LifecycleControlScope.APPLICATION,
|
||||
core_entities.LifecycleControlScope.PLATFORM,
|
||||
],
|
||||
instance_uuid=connection.instance_uuid,
|
||||
workspace_uuid=connection.workspace_uuid,
|
||||
placement_generation=connection.placement_generation,
|
||||
).task
|
||||
listener_tasks.add(listener_task)
|
||||
listener_task.add_done_callback(self._listener_task_done)
|
||||
finally:
|
||||
_current_pipeline_uuid.reset(token)
|
||||
|
||||
@@ -599,10 +746,14 @@ class WebSocketAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
|
||||
"""Reset one pipeline/client conversation."""
|
||||
conversation_key = self._conversation_key(pipeline_uuid, session_id)
|
||||
session = self.websocket_person_session if session_type == 'person' else self.websocket_group_session
|
||||
if conversation_key in session.message_lists:
|
||||
session.message_lists[conversation_key] = []
|
||||
if conversation_key in session.stream_message_indexes:
|
||||
session.stream_message_indexes[conversation_key] = {}
|
||||
if isinstance(session, WebSocketSession):
|
||||
session.reset(conversation_key)
|
||||
else:
|
||||
# Compatibility for lightweight adapter doubles.
|
||||
if conversation_key in session.message_lists:
|
||||
session.message_lists[conversation_key] = []
|
||||
if conversation_key in session.stream_message_indexes:
|
||||
session.stream_message_indexes[conversation_key] = {}
|
||||
|
||||
if session_id:
|
||||
launcher_id = (
|
||||
|
||||
@@ -13,6 +13,7 @@ from ...api.http.context import ExecutionContext
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
_SESSION_FILTER_UNSET = object()
|
||||
_DEFAULT_SEND_QUEUE_SIZE = 100
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True, slots=True)
|
||||
@@ -82,7 +83,10 @@ class WebSocketConnection(pydantic.BaseModel):
|
||||
last_active: datetime = pydantic.Field(default_factory=datetime.now)
|
||||
"""最后活跃时间"""
|
||||
|
||||
send_queue: asyncio.Queue = pydantic.Field(default_factory=asyncio.Queue, exclude=True)
|
||||
send_queue: asyncio.Queue = pydantic.Field(
|
||||
default_factory=lambda: asyncio.Queue(maxsize=_DEFAULT_SEND_QUEUE_SIZE),
|
||||
exclude=True,
|
||||
)
|
||||
"""发送消息队列"""
|
||||
|
||||
is_active: bool = True
|
||||
@@ -135,9 +139,32 @@ class WebSocketConnectionManager:
|
||||
session_type: str,
|
||||
metadata: dict | None = None,
|
||||
session_id: str | None = None,
|
||||
send_queue_size: int = _DEFAULT_SEND_QUEUE_SIZE,
|
||||
max_connections: int = 1024,
|
||||
max_connections_per_workspace: int = 32,
|
||||
) -> WebSocketConnection:
|
||||
"""Register a WebSocket connection and its optional embed session."""
|
||||
try:
|
||||
send_queue_size = max(int(send_queue_size), 1)
|
||||
except (TypeError, ValueError):
|
||||
send_queue_size = _DEFAULT_SEND_QUEUE_SIZE
|
||||
max_connections = max(int(max_connections), 1)
|
||||
max_connections_per_workspace = max(
|
||||
min(int(max_connections_per_workspace), max_connections),
|
||||
1,
|
||||
)
|
||||
async with self._lock:
|
||||
if len(self.connections) >= max_connections:
|
||||
raise RuntimeError(f'WebSocket connection capacity reached ({max_connections})')
|
||||
workspace_connection_count = sum(
|
||||
1
|
||||
for connection in self.connections.values()
|
||||
if connection.instance_uuid == scope.instance_uuid
|
||||
and connection.workspace_uuid == scope.workspace_uuid
|
||||
and connection.placement_generation == scope.placement_generation
|
||||
)
|
||||
if workspace_connection_count >= max_connections_per_workspace:
|
||||
raise RuntimeError(f'Workspace WebSocket connection capacity reached ({max_connections_per_workspace})')
|
||||
connection = WebSocketConnection(
|
||||
instance_uuid=scope.instance_uuid,
|
||||
workspace_uuid=scope.workspace_uuid,
|
||||
@@ -147,6 +174,7 @@ class WebSocketConnectionManager:
|
||||
session_id=session_id,
|
||||
websocket=websocket,
|
||||
metadata=metadata or {},
|
||||
send_queue=asyncio.Queue(maxsize=send_queue_size),
|
||||
)
|
||||
|
||||
self.connections[connection.connection_id] = connection
|
||||
@@ -171,6 +199,31 @@ class WebSocketConnectionManager:
|
||||
|
||||
return connection
|
||||
|
||||
async def close_scope(self, scope: WebSocketScope) -> None:
|
||||
"""Close and forget every live connection for one runtime placement."""
|
||||
|
||||
async with self._lock:
|
||||
connection_ids = [
|
||||
connection_id for connection_id, connection in self.connections.items() if connection.scope == scope
|
||||
]
|
||||
for connection_id in connection_ids:
|
||||
connection = self.connections.get(connection_id)
|
||||
if connection is None:
|
||||
continue
|
||||
close = getattr(connection.websocket, 'close', None)
|
||||
if close is not None:
|
||||
try:
|
||||
result = close()
|
||||
if asyncio.iscoroutine(result):
|
||||
await result
|
||||
except Exception:
|
||||
logger.debug(
|
||||
'Failed to close WebSocket connection %s',
|
||||
connection_id,
|
||||
exc_info=True,
|
||||
)
|
||||
await self.remove_connection(connection_id)
|
||||
|
||||
async def remove_connection(self, connection_id: str):
|
||||
"""移除WebSocket连接"""
|
||||
async with self._lock:
|
||||
@@ -237,7 +290,12 @@ class WebSocketConnectionManager:
|
||||
pipeline_uuid: str | None = None,
|
||||
) -> WebSocketConnection | None:
|
||||
"""Get an active embed connection by its stable browser session identifier."""
|
||||
for connection in self.connections.values():
|
||||
candidates: typing.Iterable[WebSocketConnection]
|
||||
if pipeline_uuid is not None:
|
||||
candidates = await self.get_connections_by_pipeline(pipeline_uuid, scope=scope)
|
||||
else:
|
||||
candidates = self.connections.values()
|
||||
for connection in candidates:
|
||||
if (
|
||||
connection.session_id == session_id
|
||||
and connection.is_active
|
||||
@@ -307,7 +365,20 @@ class WebSocketConnectionManager:
|
||||
return
|
||||
|
||||
try:
|
||||
await connection.send_queue.put(message)
|
||||
try:
|
||||
connection.send_queue.put_nowait(message)
|
||||
except asyncio.QueueFull:
|
||||
# A slow or disconnected browser must not backpressure every
|
||||
# other connection or retain an unbounded response stream.
|
||||
try:
|
||||
connection.send_queue.get_nowait()
|
||||
except asyncio.QueueEmpty:
|
||||
pass
|
||||
connection.send_queue.put_nowait(message)
|
||||
logger.warning(
|
||||
'WebSocket send queue full; dropped oldest message for connection %s',
|
||||
connection_id,
|
||||
)
|
||||
connection.last_active = datetime.now()
|
||||
except Exception as e:
|
||||
logger.error(f'Failed to send message to connection {connection_id}: {e}')
|
||||
|
||||
@@ -1,8 +1,6 @@
|
||||
import requests
|
||||
import websocket
|
||||
import json
|
||||
import time
|
||||
import httpx
|
||||
|
||||
from langbot.libs.wechatpad_api.client import WeChatPadClient
|
||||
|
||||
@@ -17,6 +15,7 @@ import threading
|
||||
import quart
|
||||
|
||||
from langbot.pkg.platform.logger import EventLogger
|
||||
from langbot.pkg.utils import bounded_executor, httpclient
|
||||
import xml.etree.ElementTree as ET
|
||||
from typing import Optional, Tuple
|
||||
from functools import partial
|
||||
@@ -27,6 +26,8 @@ import langbot_plugin.api.entities.builtin.platform.entities as platform_entitie
|
||||
import langbot_plugin.api.definition.abstract.platform.adapter as abstract_platform_adapter
|
||||
import langbot_plugin.api.definition.abstract.platform.event_logger as abstract_platform_logger
|
||||
|
||||
_MAX_GATEWAY_MESSAGE_CHARS = 1024 * 1024
|
||||
|
||||
|
||||
class WeChatPadMessageConverter(abstract_platform_adapter.AbstractMessageConverter):
|
||||
def __init__(self, config: dict, logger: abstract_platform_logger.AbstractEventLogger):
|
||||
@@ -53,12 +54,11 @@ class WeChatPadMessageConverter(abstract_platform_adapter.AbstractMessageConvert
|
||||
content_list.append({'type': 'text', 'content': component.text})
|
||||
elif isinstance(component, platform_message.Image):
|
||||
if component.url:
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.get(component.url)
|
||||
|
||||
if response.status_code == 200:
|
||||
file_bytes = response.content
|
||||
base64_str = base64.b64encode(file_bytes).decode('utf-8') # 返回字符串格式
|
||||
session = httpclient.get_session()
|
||||
async with session.get(component.url) as response:
|
||||
if response.status == 200:
|
||||
file_bytes = await httpclient.read_limited(response)
|
||||
base64_str = (await asyncio.to_thread(base64.b64encode, file_bytes)).decode('utf-8')
|
||||
else:
|
||||
raise Exception('获取文件失败')
|
||||
# pass
|
||||
@@ -156,9 +156,19 @@ class WeChatPadMessageConverter(abstract_platform_adapter.AbstractMessageConvert
|
||||
cdnthumburl = img_tag.get('cdnthumburl')
|
||||
# cdnmidimgurl = img_tag.get('cdnmidimgurl')
|
||||
|
||||
image_data = self.bot.cdn_download(aeskey=aeskey, file_type=1, file_url=cdnthumburl)
|
||||
image_data = await asyncio.to_thread(
|
||||
self.bot.cdn_download,
|
||||
aeskey=aeskey,
|
||||
file_type=1,
|
||||
file_url=cdnthumburl,
|
||||
)
|
||||
if image_data['Data']['FileData'] == '':
|
||||
image_data = self.bot.cdn_download(aeskey=aeskey, file_type=2, file_url=cdnthumburl)
|
||||
image_data = await asyncio.to_thread(
|
||||
self.bot.cdn_download,
|
||||
aeskey=aeskey,
|
||||
file_type=2,
|
||||
file_url=cdnthumburl,
|
||||
)
|
||||
base64_str = image_data['Data']['FileData']
|
||||
# self.logger.info(f"data:image/png;base64,{base64_str}")
|
||||
|
||||
@@ -186,7 +196,12 @@ class WeChatPadMessageConverter(abstract_platform_adapter.AbstractMessageConvert
|
||||
if voicemsg is not None:
|
||||
bufid = voicemsg.get('bufid')
|
||||
length = voicemsg.get('voicelength')
|
||||
voice_data = self.bot.get_msg_voice(buf_id=str(bufid), length=int(length), msgid=str(new_msg_id))
|
||||
voice_data = await asyncio.to_thread(
|
||||
self.bot.get_msg_voice,
|
||||
buf_id=str(bufid),
|
||||
length=int(length),
|
||||
msgid=str(new_msg_id),
|
||||
)
|
||||
audio_base64 = voice_data['Data']['Base64']
|
||||
|
||||
# 验证语音数据有效性
|
||||
@@ -319,7 +334,12 @@ class WeChatPadMessageConverter(abstract_platform_adapter.AbstractMessageConvert
|
||||
|
||||
# print(aeskey,cdnthumburl)
|
||||
|
||||
file_data = self.bot.cdn_download(aeskey=aeskey, file_type=5, file_url=cdnthumburl)
|
||||
file_data = await asyncio.to_thread(
|
||||
self.bot.cdn_download,
|
||||
aeskey=aeskey,
|
||||
file_type=5,
|
||||
file_url=cdnthumburl,
|
||||
)
|
||||
|
||||
file_base64 = file_data['Data']['FileData']
|
||||
# print(file_data)
|
||||
@@ -538,6 +558,7 @@ class WeChatPadAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
|
||||
typing.Type[platform_events.Event],
|
||||
typing.Callable[[platform_events.Event, abstract_platform_adapter.AbstractMessagePlatformAdapter], None],
|
||||
] = {}
|
||||
_MAX_CALLBACK_FUTURES = 100
|
||||
|
||||
def __init__(self, config: dict, logger: EventLogger):
|
||||
quart_app = quart.Quart(__name__)
|
||||
@@ -556,6 +577,12 @@ class WeChatPadAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
|
||||
name='WeChatPad',
|
||||
bot=bot,
|
||||
)
|
||||
self._event_loop: asyncio.AbstractEventLoop | None = None
|
||||
self._ws_app: websocket.WebSocketApp | None = None
|
||||
self._ws_thread: threading.Thread | None = None
|
||||
self._stop_event = threading.Event()
|
||||
self._callback_futures: set = set()
|
||||
self._callback_futures_lock = threading.Lock()
|
||||
|
||||
async def ws_message(self, data):
|
||||
"""处理接收到的消息"""
|
||||
@@ -565,7 +592,7 @@ class WeChatPadAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
|
||||
except Exception:
|
||||
await self.logger.error(f'Error in wechatpad callback: {traceback.format_exc()}')
|
||||
|
||||
if event.__class__ in self.listeners:
|
||||
if event is not None and event.__class__ in self.listeners:
|
||||
await self.listeners[event.__class__](event, self)
|
||||
|
||||
return 'ok'
|
||||
@@ -580,9 +607,8 @@ class WeChatPadAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
|
||||
at_targets = at_targets or []
|
||||
member_info = []
|
||||
if at_targets:
|
||||
member_info = self.bot.get_chatroom_member_detail(
|
||||
target_id,
|
||||
)['Data']['member_data']['chatroom_member_list']
|
||||
member_result = await asyncio.to_thread(self.bot.get_chatroom_member_detail, target_id)
|
||||
member_info = member_result['Data']['member_data']['chatroom_member_list']
|
||||
|
||||
# 处理消息组件
|
||||
for msg in content_list:
|
||||
@@ -623,11 +649,35 @@ class WeChatPadAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
|
||||
}
|
||||
|
||||
if handler := handler_map.get(msg['type']):
|
||||
handler(msg)
|
||||
await asyncio.to_thread(handler, msg)
|
||||
else:
|
||||
self.logger.warning(f'未处理的消息类型: {msg["type"]}')
|
||||
await self.logger.warning(f'未处理的消息类型: {msg["type"]}')
|
||||
continue
|
||||
|
||||
def _schedule_ws_message(self, data: dict) -> None:
|
||||
loop = self._event_loop
|
||||
if loop is None or loop.is_closed() or self._stop_event.is_set():
|
||||
return
|
||||
with self._callback_futures_lock:
|
||||
if len(self._callback_futures) >= self._MAX_CALLBACK_FUTURES:
|
||||
return
|
||||
future = asyncio.run_coroutine_threadsafe(self.ws_message(data), loop)
|
||||
self._callback_futures.add(future)
|
||||
|
||||
def done(completed) -> None:
|
||||
with self._callback_futures_lock:
|
||||
self._callback_futures.discard(completed)
|
||||
if completed.cancelled():
|
||||
return
|
||||
try:
|
||||
completed.result()
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
except Exception:
|
||||
logging.getLogger(__name__).exception('WeChatPad callback failed')
|
||||
|
||||
future.add_done_callback(done)
|
||||
|
||||
async def send_message(self, target_type: str, target_id: str, message: platform_message.MessageChain):
|
||||
"""主动发送消息"""
|
||||
return await self._handle_message(message, target_id)
|
||||
@@ -665,86 +715,113 @@ class WeChatPadAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter)
|
||||
pass
|
||||
|
||||
async def run_async(self):
|
||||
self._event_loop = asyncio.get_running_loop()
|
||||
self._stop_event.clear()
|
||||
if not self.config['admin_key'] and not self.config['token']:
|
||||
raise RuntimeError('无wechatpad管理密匙，请填入配置文件后重启')
|
||||
else:
|
||||
if self.config['token']:
|
||||
self.bot = WeChatPadClient(self.config['wechatpad_url'], self.config['token'])
|
||||
data = self.bot.get_login_status()
|
||||
data = await asyncio.to_thread(self.bot.get_login_status)
|
||||
if data['Code'] == 300 and data['Text'] == '你已退出微信':
|
||||
response = requests.post(
|
||||
response = await asyncio.to_thread(
|
||||
requests.post,
|
||||
f'{self.config["wechatpad_url"]}/admin/GenAuthKey1?key={self.config["admin_key"]}',
|
||||
json={'Count': 1, 'Days': 365},
|
||||
timeout=10,
|
||||
)
|
||||
if response.status_code != 200:
|
||||
raise Exception(f'获取token失败: {response.text}')
|
||||
self.config['token'] = response.json()['Data'][0]
|
||||
body = await httpclient.response_text(response)
|
||||
raise Exception(f'获取token失败: {body}')
|
||||
response_data = await httpclient.parse_json_response(response)
|
||||
self.config['token'] = response_data['Data'][0]
|
||||
|
||||
elif not self.config['token']:
|
||||
response = requests.post(
|
||||
response = await asyncio.to_thread(
|
||||
requests.post,
|
||||
f'{self.config["wechatpad_url"]}/admin/GenAuthKey1?key={self.config["admin_key"]}',
|
||||
json={'Count': 1, 'Days': 365},
|
||||
timeout=10,
|
||||
)
|
||||
if response.status_code != 200:
|
||||
raise Exception(f'获取token失败: {response.text}')
|
||||
self.config['token'] = response.json()['Data'][0]
|
||||
body = await httpclient.response_text(response)
|
||||
raise Exception(f'获取token失败: {body}')
|
||||
response_data = await httpclient.parse_json_response(response)
|
||||
self.config['token'] = response_data['Data'][0]
|
||||
|
||||
self.bot = WeChatPadClient(self.config['wechatpad_url'], self.config['token'], logger=self.logger)
|
||||
await self.logger.info(self.config['token'])
|
||||
thread_1 = threading.Event()
|
||||
|
||||
def wechat_login_process():
|
||||
# 不登录，这些先注释掉，避免登陆态尝试拉qrcode。
|
||||
# login_data =self.bot.get_login_qr()
|
||||
|
||||
# url = login_data['Data']["QrCodeUrl"]
|
||||
|
||||
profile = self.bot.get_profile()
|
||||
# self.logger.info(profile)
|
||||
|
||||
self.bot_account_id = profile['Data']['userInfo']['nickName']['str']
|
||||
self.config['wxid'] = profile['Data']['userInfo']['userName']['str']
|
||||
thread_1.set()
|
||||
|
||||
# asyncio.create_task(wechat_login_process)
|
||||
threading.Thread(target=wechat_login_process).start()
|
||||
profile = await asyncio.to_thread(self.bot.get_profile)
|
||||
self.bot_account_id = profile['Data']['userInfo']['nickName']['str']
|
||||
self.config['wxid'] = profile['Data']['userInfo']['userName']['str']
|
||||
|
||||
def connect_websocket_sync() -> None:
|
||||
thread_1.wait()
|
||||
uri = f'{self.config["wechatpad_ws"]}/GetSyncMsg?key={self.config["token"]}'
|
||||
print(f'Connecting to WebSocket: {uri}')
|
||||
|
||||
def on_message(ws, message):
|
||||
try:
|
||||
if len(message) > _MAX_GATEWAY_MESSAGE_CHARS:
|
||||
logging.getLogger(__name__).warning('WeChatPad WebSocket message exceeds the size limit')
|
||||
return
|
||||
data = json.loads(message)
|
||||
# 这里需要确保ws_message是同步的，或者使用asyncio.run调用异步方法
|
||||
asyncio.run(self.ws_message(data))
|
||||
self._schedule_ws_message(data)
|
||||
except json.JSONDecodeError:
|
||||
self.logger.error(f'Non-JSON message: {message[:100]}...')
|
||||
logging.getLogger(__name__).warning('WeChatPad received a non-JSON message')
|
||||
|
||||
def on_error(ws, error):
|
||||
self.logger.error(f'WebSocket error: {str(error)[:200]}')
|
||||
logging.getLogger(__name__).warning('WeChatPad WebSocket error: %s', str(error)[:200])
|
||||
|
||||
def on_close(ws, close_status_code, close_msg):
|
||||
self.logger.info('WebSocket closed, reconnecting...')
|
||||
time.sleep(5)
|
||||
connect_websocket_sync() # 自动重连
|
||||
logging.getLogger(__name__).info('WeChatPad WebSocket closed')
|
||||
|
||||
def on_open(ws):
|
||||
self.logger.info('WebSocket connected successfully!')
|
||||
logging.getLogger(__name__).info('WeChatPad WebSocket connected')
|
||||
|
||||
ws = websocket.WebSocketApp(
|
||||
uri, on_message=on_message, on_error=on_error, on_close=on_close, on_open=on_open
|
||||
)
|
||||
ws.run_forever(ping_interval=60, ping_timeout=20)
|
||||
while not self._stop_event.is_set():
|
||||
ws = websocket.WebSocketApp(
|
||||
uri,
|
||||
on_message=on_message,
|
||||
on_error=on_error,
|
||||
on_close=on_close,
|
||||
on_open=on_open,
|
||||
)
|
||||
self._ws_app = ws
|
||||
ws.run_forever(ping_interval=60, ping_timeout=20)
|
||||
self._ws_app = None
|
||||
if not self._stop_event.wait(5):
|
||||
logging.getLogger(__name__).info('Reconnecting WeChatPad WebSocket')
|
||||
|
||||
# 直接调用同步版本（会阻塞）
|
||||
# connect_websocket_sync()
|
||||
|
||||
# 这行代码会在WebSocket连接断开后才会执行
|
||||
thread = threading.Thread(target=connect_websocket_sync, name='WebSocketClientThread', daemon=True)
|
||||
thread.start()
|
||||
self.logger.info('WebSocket client thread started')
|
||||
self._ws_thread = threading.Thread(
|
||||
target=connect_websocket_sync,
|
||||
name='WebSocketClientThread',
|
||||
daemon=True,
|
||||
)
|
||||
self._ws_thread.start()
|
||||
await self.logger.info('WebSocket client thread started')
|
||||
while not self._stop_event.is_set() and self._ws_thread.is_alive():
|
||||
await asyncio.sleep(1)
|
||||
if not self._stop_event.is_set():
|
||||
raise RuntimeError('WeChatPad WebSocket client thread exited unexpectedly')
|
||||
|
||||
async def kill(self) -> bool:
|
||||
pass
|
||||
self._stop_event.set()
|
||||
ws = self._ws_app
|
||||
if ws is not None:
|
||||
await bounded_executor.run_blocking_cleanup(ws.close)
|
||||
|
||||
with self._callback_futures_lock:
|
||||
futures = list(self._callback_futures)
|
||||
for future in futures:
|
||||
future.cancel()
|
||||
if futures:
|
||||
await asyncio.gather(
|
||||
*(asyncio.wrap_future(future) for future in futures),
|
||||
return_exceptions=True,
|
||||
)
|
||||
|
||||
thread = self._ws_thread
|
||||
if thread is not None and thread.is_alive():
|
||||
await bounded_executor.run_blocking_cleanup(thread.join, 5)
|
||||
self._ws_thread = None
|
||||
self._ws_app = None
|
||||
self._event_loop = None
|
||||
return True
|
||||
|
||||
@@ -319,6 +319,7 @@ class WecomAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
await keep_alive()
|
||||
|
||||
async def kill(self) -> bool:
|
||||
await self.bot.close()
|
||||
return False
|
||||
|
||||
async def unregister_listener(
|
||||
|
||||
@@ -516,6 +516,8 @@ class WecomBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
# do work. Lazy-create on first call.
|
||||
object.__setattr__(self, '_synthetic_buffers', {})
|
||||
buffers: dict[str, str] = self._synthetic_buffers
|
||||
if buf_key not in buffers and len(buffers) >= 100:
|
||||
buffers.pop(next(iter(buffers)), None)
|
||||
if content and not form_data:
|
||||
previous = buffers.get(buf_key, '')
|
||||
if previous and content.startswith(previous):
|
||||
@@ -524,6 +526,8 @@ class WecomBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
buffers[buf_key] = previous
|
||||
else:
|
||||
buffers[buf_key] = previous + content
|
||||
if len(buffers[buf_key]) > 200000:
|
||||
buffers[buf_key] = buffers[buf_key][-200000:]
|
||||
|
||||
if not is_final:
|
||||
return {'stream': True, 'synthetic': True, 'buffered': True}
|
||||
@@ -613,7 +617,11 @@ class WecomBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
'chat_id': chat_id,
|
||||
'stream_id': '',
|
||||
'req_id': '',
|
||||
'created_at': time.monotonic(),
|
||||
}
|
||||
prune = getattr(self.bot, '_prune_pending_forms', None)
|
||||
if callable(prune):
|
||||
prune()
|
||||
return payload
|
||||
|
||||
async def send_message(self, target_type, target_id, message):
|
||||
@@ -745,10 +753,13 @@ class WecomBotAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
await keep_alive()
|
||||
|
||||
async def kill(self) -> bool:
|
||||
if hasattr(self, '_synthetic_buffers'):
|
||||
self._synthetic_buffers.clear()
|
||||
_ws_mode = not self.config.get('enable-webhook', False)
|
||||
if _ws_mode:
|
||||
await self.bot.disconnect()
|
||||
return True
|
||||
await self.bot.close()
|
||||
return False
|
||||
|
||||
async def unregister_listener(
|
||||
|
||||
@@ -254,6 +254,8 @@ class WecomCSAdapter(abstract_platform_adapter.AbstractMessagePlatformAdapter):
|
||||
await keep_alive()
|
||||
|
||||
async def kill(self) -> bool:
|
||||
self.bot.clear()
|
||||
await self.bot.close()
|
||||
return False
|
||||
|
||||
async def is_muted(self, group_id: int) -> bool:
|
||||
|
||||
@@ -147,7 +147,8 @@ class WebhookPusher:
|
||||
else:
|
||||
self.logger.debug(f'Successfully pushed to webhook {url}')
|
||||
try:
|
||||
return await response.json()
|
||||
result = await httpclient.read_json_limited(response)
|
||||
return result if isinstance(result, dict) else None
|
||||
except Exception as json_error:
|
||||
self.logger.debug(f'Failed to parse JSON response from webhook {url}: {json_error}')
|
||||
return None
|
||||
|
||||
@@ -0,0 +1,98 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import zipfile
|
||||
|
||||
import yaml
|
||||
|
||||
|
||||
_PLUGIN_ARCHIVE_MAX_ENTRIES = 512
|
||||
_PLUGIN_ARCHIVE_MAX_ENTRY_BYTES = 16 * 1024 * 1024
|
||||
_PLUGIN_ARCHIVE_MAX_TOTAL_BYTES = 64 * 1024 * 1024
|
||||
_PLUGIN_ARCHIVE_MAX_COMPRESSION_RATIO = 100
|
||||
_PLUGIN_METADATA_MAX_BYTES = 1024 * 1024
|
||||
_PLUGIN_REQUIREMENTS_MAX_ENTRIES = 1000
|
||||
|
||||
|
||||
def _read_plugin_archive_member(
|
||||
archive: zipfile.ZipFile,
|
||||
member: zipfile.ZipInfo,
|
||||
*,
|
||||
max_bytes: int = _PLUGIN_METADATA_MAX_BYTES,
|
||||
) -> bytes:
|
||||
if member.file_size > max_bytes:
|
||||
raise ValueError(f'Plugin metadata file exceeds the {max_bytes}-byte limit: {member.filename}')
|
||||
with archive.open(member, 'r') as source:
|
||||
content = source.read(max_bytes + 1)
|
||||
if len(content) > max_bytes or len(content) != member.file_size:
|
||||
raise ValueError(f'Plugin metadata file has an invalid size: {member.filename}')
|
||||
return content
|
||||
|
||||
|
||||
def inspect_plugin_archive_metadata(
|
||||
file_bytes: bytes,
|
||||
*,
|
||||
require_manifest: bool = True,
|
||||
) -> tuple[dict, list[str], list[str]]:
|
||||
"""Validate archive size metadata and read only bounded preview fields."""
|
||||
|
||||
with zipfile.ZipFile(io.BytesIO(file_bytes)) as archive:
|
||||
members = archive.infolist()
|
||||
if len(members) > _PLUGIN_ARCHIVE_MAX_ENTRIES:
|
||||
raise ValueError('Plugin archive contains too many entries')
|
||||
|
||||
total_uncompressed = 0
|
||||
files: dict[str, zipfile.ZipInfo] = {}
|
||||
names: list[str] = []
|
||||
for member in members:
|
||||
if member.is_dir():
|
||||
continue
|
||||
if member.flag_bits & 0x1:
|
||||
raise ValueError('Encrypted plugin archives are not supported')
|
||||
if member.file_size > _PLUGIN_ARCHIVE_MAX_ENTRY_BYTES:
|
||||
raise ValueError(f'Plugin archive entry exceeds the size limit: {member.filename}')
|
||||
if (
|
||||
member.file_size
|
||||
and member.file_size > max(member.compress_size, 1) * _PLUGIN_ARCHIVE_MAX_COMPRESSION_RATIO
|
||||
):
|
||||
raise ValueError(f'Plugin archive entry exceeds the compression-ratio limit: {member.filename}')
|
||||
total_uncompressed += member.file_size
|
||||
if total_uncompressed > _PLUGIN_ARCHIVE_MAX_TOTAL_BYTES:
|
||||
raise ValueError('Plugin archive exceeds the uncompressed size limit')
|
||||
normalized = member.filename.replace('\\', '/').strip('/')
|
||||
names.append(member.filename)
|
||||
files.setdefault(normalized.lower(), member)
|
||||
|
||||
manifest_member = files.get('manifest.yaml') or files.get('manifest.yml')
|
||||
if manifest_member is None:
|
||||
if require_manifest:
|
||||
raise ValueError('manifest.yaml is required')
|
||||
manifest = {}
|
||||
else:
|
||||
manifest = yaml.safe_load(_read_plugin_archive_member(archive, manifest_member).decode('utf-8')) or {}
|
||||
if not isinstance(manifest, dict):
|
||||
raise ValueError('Plugin manifest must be an object')
|
||||
|
||||
requirements: list[str] = []
|
||||
requirements_member = next(
|
||||
(
|
||||
member
|
||||
for normalized, member in files.items()
|
||||
if normalized == 'requirements.txt' or normalized.endswith('/requirements.txt')
|
||||
),
|
||||
None,
|
||||
)
|
||||
if requirements_member is not None:
|
||||
content = _read_plugin_archive_member(
|
||||
archive,
|
||||
requirements_member,
|
||||
).decode(
|
||||
'utf-8',
|
||||
errors='ignore',
|
||||
)
|
||||
requirements = [
|
||||
line.strip()[:1000]
|
||||
for line in content.splitlines()
|
||||
if line.strip() and not line.strip().startswith('#')
|
||||
][:_PLUGIN_REQUIREMENTS_MAX_ENTRIES]
|
||||
return manifest, requirements, names
|
||||
@@ -5,10 +5,9 @@ import asyncio
|
||||
import contextlib
|
||||
import contextvars
|
||||
import hashlib
|
||||
import io
|
||||
import json
|
||||
import time
|
||||
import uuid
|
||||
import zipfile
|
||||
from typing import Any
|
||||
import typing
|
||||
import os
|
||||
@@ -16,17 +15,17 @@ import secrets
|
||||
import sys
|
||||
import httpx
|
||||
import sqlalchemy
|
||||
import yaml
|
||||
from urllib.parse import urljoin, urlparse
|
||||
from langbot_plugin.api.entities.builtin.pipeline.query import provider_session
|
||||
|
||||
from ..core import app
|
||||
from . import handler
|
||||
from .archive import inspect_plugin_archive_metadata
|
||||
from .github import (
|
||||
validate_github_plugin_install_info,
|
||||
validate_github_release_asset_url,
|
||||
)
|
||||
from ..utils import constants, platform
|
||||
from ..utils import constants, httpclient, platform
|
||||
from ..utils.managed_runtime import ManagedRuntimeConnector
|
||||
from langbot_plugin.runtime.io.controllers.stdio import (
|
||||
client as stdio_client_controller,
|
||||
@@ -64,6 +63,9 @@ _PLUGIN_ARTIFACT_OWNER_TYPE = 'plugin_artifact'
|
||||
_PLUGIN_ARTIFACT_KEY = 'package.lbpkg'
|
||||
_PLUGIN_ARTIFACT_STORAGE_MARKER = 'tenant_binary_storage_v1'
|
||||
_GITHUB_PLUGIN_DOWNLOAD_MAX_BYTES = 10 * 1024 * 1024
|
||||
_MARKETPLACE_METADATA_MAX_BYTES = 1024 * 1024
|
||||
_MARKETPLACE_PLUGIN_DOWNLOAD_MAX_BYTES = 64 * 1024 * 1024
|
||||
_MARKETPLACE_SKILL_DOWNLOAD_MAX_BYTES = 10 * 1024 * 1024
|
||||
_GITHUB_PLUGIN_DOWNLOAD_MAX_REDIRECTS = 5
|
||||
_GITHUB_ASSET_HOSTS = frozenset(
|
||||
{
|
||||
@@ -80,6 +82,55 @@ _HEARTBEAT_FAILURE_THRESHOLD = 3
|
||||
_RECONNECT_MAX_DELAY_SEC = 60.0
|
||||
|
||||
|
||||
async def _read_httpx_response_limited(
|
||||
response: httpx.Response,
|
||||
*,
|
||||
max_bytes: int,
|
||||
) -> bytes:
|
||||
content_length = response.headers.get('content-length')
|
||||
if content_length is not None:
|
||||
try:
|
||||
declared_size = int(content_length)
|
||||
except ValueError:
|
||||
declared_size = None
|
||||
if declared_size is not None and declared_size > max_bytes:
|
||||
raise ValueError(f'Remote response exceeds the {max_bytes}-byte limit')
|
||||
|
||||
body = bytearray()
|
||||
async for chunk in response.aiter_bytes(chunk_size=64 * 1024):
|
||||
body.extend(chunk)
|
||||
if len(body) > max_bytes:
|
||||
raise ValueError(f'Remote response exceeds the {max_bytes}-byte limit')
|
||||
return bytes(body)
|
||||
|
||||
|
||||
async def _marketplace_get(
|
||||
client: httpx.AsyncClient,
|
||||
url: str,
|
||||
*,
|
||||
max_bytes: int,
|
||||
allow_not_found: bool = False,
|
||||
) -> tuple[int, bytes]:
|
||||
async with client.stream('GET', url) as response:
|
||||
if allow_not_found and response.status_code == 404:
|
||||
return response.status_code, b''
|
||||
response.raise_for_status()
|
||||
return response.status_code, await _read_httpx_response_limited(
|
||||
response,
|
||||
max_bytes=max_bytes,
|
||||
)
|
||||
|
||||
|
||||
def _decode_json_object(body: bytes, *, subject: str) -> dict[str, Any]:
|
||||
try:
|
||||
payload = json.loads(body)
|
||||
except (UnicodeDecodeError, json.JSONDecodeError) as exc:
|
||||
raise ValueError(f'{subject} returned invalid JSON') from exc
|
||||
if not isinstance(payload, dict):
|
||||
raise ValueError(f'{subject} returned a non-object response')
|
||||
return payload
|
||||
|
||||
|
||||
class PluginRuntimeNotConnectedError(RuntimeError):
|
||||
"""Raised when plugin runtime operations are requested before connection."""
|
||||
|
||||
@@ -175,16 +226,23 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
|
||||
worker = self.ap.instance_config.data.get('plugin', {}).get('worker')
|
||||
if not isinstance(worker, dict):
|
||||
raise ValueError('plugin.worker must be configured')
|
||||
return PluginWorkerPolicy.model_validate(
|
||||
{
|
||||
'max_cpus': worker.get('max_cpus'),
|
||||
'max_memory_mb': worker.get('max_memory_mb'),
|
||||
'max_pids': worker.get('max_pids'),
|
||||
'max_open_files': worker.get('max_open_files'),
|
||||
'max_file_size_mb': worker.get('max_file_size_mb'),
|
||||
'require_hard_limits': worker.get('require_hard_limits', False),
|
||||
}
|
||||
)
|
||||
policy_data = {
|
||||
'max_cpus': worker.get('max_cpus'),
|
||||
'max_memory_mb': worker.get('max_memory_mb'),
|
||||
'max_pids': worker.get('max_pids'),
|
||||
'max_open_files': worker.get('max_open_files'),
|
||||
'max_file_size_mb': worker.get('max_file_size_mb'),
|
||||
'require_hard_limits': worker.get('require_hard_limits', False),
|
||||
}
|
||||
for field_name in (
|
||||
'max_workers',
|
||||
'max_total_cpus',
|
||||
'max_total_memory_mb',
|
||||
'max_installations',
|
||||
):
|
||||
if field_name in PluginWorkerPolicy.model_fields and field_name in worker:
|
||||
policy_data[field_name] = worker.get(field_name)
|
||||
return PluginWorkerPolicy.model_validate(policy_data)
|
||||
|
||||
def _control_headers(self, *, allow_generate: bool) -> dict[str, str]:
|
||||
if not self._control_token and allow_generate:
|
||||
@@ -643,9 +701,11 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
|
||||
for context in contexts:
|
||||
execution_context = await self._validate_execution_context(context)
|
||||
states = await self._load_workspace_desired_states(execution_context)
|
||||
workspace_installations[execution_context.workspace_uuid] = {
|
||||
state.binding.installation_uuid for state in states
|
||||
}
|
||||
installation_ids = {state.binding.installation_uuid for state in states}
|
||||
if installation_ids:
|
||||
# A newly registered Cloud Workspace normally has no
|
||||
# plugins. Avoid retaining an empty set for every account.
|
||||
workspace_installations[execution_context.workspace_uuid] = installation_ids
|
||||
for state in states:
|
||||
if state.binding.installation_uuid in all_states:
|
||||
raise ValueError('Duplicate plugin installation UUID across projected Workspaces')
|
||||
@@ -706,7 +766,13 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
|
||||
# restoring the remaining desired state in this Workspace.
|
||||
pass
|
||||
self._known_desired_states[installation_uuid] = desired
|
||||
self._workspace_installations[execution_context.workspace_uuid] = set(desired_by_uuid)
|
||||
if desired_by_uuid:
|
||||
self._workspace_installations[execution_context.workspace_uuid] = set(desired_by_uuid)
|
||||
else:
|
||||
self._workspace_installations.pop(
|
||||
execution_context.workspace_uuid,
|
||||
None,
|
||||
)
|
||||
|
||||
async def _current_execution_context(self) -> ExecutionContext:
|
||||
current = self._execution_context.get()
|
||||
@@ -991,27 +1057,20 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
|
||||
plugin_name = None
|
||||
|
||||
try:
|
||||
with zipfile.ZipFile(io.BytesIO(file_bytes)) as zf:
|
||||
try:
|
||||
manifest = yaml.safe_load(zf.read('manifest.yaml').decode('utf-8', errors='ignore')) or {}
|
||||
metadata = manifest.get('metadata', {})
|
||||
plugin_author = metadata.get('author')
|
||||
plugin_name = metadata.get('name')
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if task_context is not None:
|
||||
for name in zf.namelist():
|
||||
if name.endswith('requirements.txt'):
|
||||
content = zf.read(name).decode('utf-8', errors='ignore')
|
||||
deps = [
|
||||
line.strip()
|
||||
for line in content.splitlines()
|
||||
if line.strip() and not line.strip().startswith('#')
|
||||
]
|
||||
task_context.metadata['deps_total'] = len(deps)
|
||||
task_context.metadata['deps_list'] = deps
|
||||
break
|
||||
manifest, dependencies, archive_names = inspect_plugin_archive_metadata(
|
||||
file_bytes,
|
||||
require_manifest=False,
|
||||
)
|
||||
metadata = manifest.get('metadata', {})
|
||||
if isinstance(metadata, dict):
|
||||
plugin_author = metadata.get('author')
|
||||
plugin_name = metadata.get('name')
|
||||
has_requirements = any(
|
||||
name.replace('\\', '/').lower().endswith('requirements.txt') for name in archive_names
|
||||
)
|
||||
if task_context is not None and has_requirements:
|
||||
task_context.metadata['deps_total'] = len(dependencies)
|
||||
task_context.metadata['deps_list'] = dependencies
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@@ -1329,6 +1388,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
|
||||
trust_env=False,
|
||||
follow_redirects=False,
|
||||
timeout=httpx.Timeout(60, connect=10),
|
||||
event_hooks=httpclient.httpx_response_limit_hooks(_GITHUB_PLUGIN_DOWNLOAD_MAX_BYTES),
|
||||
) as client:
|
||||
asset_id: int | None = None
|
||||
if 'asset_id' in normalized:
|
||||
@@ -1346,7 +1406,7 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
|
||||
if response.status_code in _HTTP_REDIRECT_STATUSES:
|
||||
raise ValueError('GitHub release metadata unexpectedly redirected')
|
||||
response.raise_for_status()
|
||||
release = response.json()
|
||||
release = await httpclient.parse_json_response(response)
|
||||
if not isinstance(release, dict):
|
||||
raise ValueError('GitHub release metadata is invalid')
|
||||
if release.get('id') != release_id or str(release.get('tag_name') or '') != release_tag:
|
||||
@@ -1500,50 +1560,79 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
|
||||
"""Return a plugin package, or install an MCP/skill and return none."""
|
||||
|
||||
space_url = self.ap.instance_config.data.get('space', {}).get('url', 'https://space.langbot.app').rstrip('/')
|
||||
async with httpx.AsyncClient(trust_env=True, timeout=15) as client:
|
||||
mcp_resp = await client.get(f'{space_url}/api/v1/marketplace/mcps/{plugin_author}/{plugin_name}')
|
||||
if mcp_resp.status_code == 200:
|
||||
mcp_data = mcp_resp.json().get('data', {}).get('mcp', {})
|
||||
async with httpx.AsyncClient(
|
||||
trust_env=True,
|
||||
timeout=15,
|
||||
event_hooks=httpclient.httpx_response_limit_hooks(_MARKETPLACE_PLUGIN_DOWNLOAD_MAX_BYTES),
|
||||
) as client:
|
||||
mcp_status, mcp_body = await _marketplace_get(
|
||||
client,
|
||||
f'{space_url}/api/v1/marketplace/mcps/{plugin_author}/{plugin_name}',
|
||||
max_bytes=_MARKETPLACE_METADATA_MAX_BYTES,
|
||||
allow_not_found=True,
|
||||
)
|
||||
if mcp_status == 200:
|
||||
mcp_payload = _decode_json_object(mcp_body, subject='Marketplace MCP metadata')
|
||||
mcp_data = mcp_payload.get('data', {}).get('mcp', {})
|
||||
if not isinstance(mcp_data, dict):
|
||||
raise ValueError(f'MCP {plugin_author}/{plugin_name} metadata is invalid')
|
||||
if not mcp_data.get('mode'):
|
||||
raise ValueError(f'MCP {plugin_author}/{plugin_name} has no mode')
|
||||
await self._install_mcp_from_marketplace(execution_context, mcp_data, task_context)
|
||||
try:
|
||||
await client.post(f'{space_url}/api/v1/marketplace/mcps/{plugin_author}/{plugin_name}/install')
|
||||
async with client.stream(
|
||||
'POST',
|
||||
f'{space_url}/api/v1/marketplace/mcps/{plugin_author}/{plugin_name}/install',
|
||||
):
|
||||
pass
|
||||
except Exception as report_err:
|
||||
self.ap.logger.debug(f'Failed to report MCP install: {report_err}')
|
||||
return None, None
|
||||
if mcp_resp.status_code != 404:
|
||||
mcp_resp.raise_for_status()
|
||||
|
||||
skill_resp = await client.get(f'{space_url}/api/v1/marketplace/skills/{plugin_author}/{plugin_name}')
|
||||
if skill_resp.status_code == 200:
|
||||
download_resp = await client.get(
|
||||
f'{space_url}/api/v1/marketplace/skills/download/{plugin_author}/{plugin_name}'
|
||||
skill_status, _skill_body = await _marketplace_get(
|
||||
client,
|
||||
f'{space_url}/api/v1/marketplace/skills/{plugin_author}/{plugin_name}',
|
||||
max_bytes=_MARKETPLACE_METADATA_MAX_BYTES,
|
||||
allow_not_found=True,
|
||||
)
|
||||
if skill_status == 200:
|
||||
_download_status, skill_package = await _marketplace_get(
|
||||
client,
|
||||
f'{space_url}/api/v1/marketplace/skills/download/{plugin_author}/{plugin_name}',
|
||||
max_bytes=_MARKETPLACE_SKILL_DOWNLOAD_MAX_BYTES,
|
||||
)
|
||||
download_resp.raise_for_status()
|
||||
await self._install_skill_from_zip(
|
||||
execution_context,
|
||||
download_resp.content,
|
||||
skill_package,
|
||||
f'{plugin_author}-{plugin_name}',
|
||||
task_context,
|
||||
)
|
||||
return None, None
|
||||
if skill_resp.status_code != 404:
|
||||
skill_resp.raise_for_status()
|
||||
|
||||
versions_resp = await client.get(
|
||||
f'{space_url}/api/v1/marketplace/plugins/{plugin_author}/{plugin_name}/versions'
|
||||
_versions_status, versions_body = await _marketplace_get(
|
||||
client,
|
||||
f'{space_url}/api/v1/marketplace/plugins/{plugin_author}/{plugin_name}/versions',
|
||||
max_bytes=_MARKETPLACE_METADATA_MAX_BYTES,
|
||||
)
|
||||
versions_resp.raise_for_status()
|
||||
versions = versions_resp.json().get('data', {}).get('versions', [])
|
||||
if not versions or not versions[0].get('version'):
|
||||
versions_payload = _decode_json_object(
|
||||
versions_body,
|
||||
subject='Marketplace plugin versions',
|
||||
)
|
||||
versions = versions_payload.get('data', {}).get('versions', [])
|
||||
if (
|
||||
not isinstance(versions, list)
|
||||
or not versions
|
||||
or not isinstance(versions[0], dict)
|
||||
or not versions[0].get('version')
|
||||
):
|
||||
raise ValueError(f'Plugin {plugin_author}/{plugin_name} has no versions')
|
||||
latest_version = str(versions[0]['version'])
|
||||
download_resp = await client.get(
|
||||
f'{space_url}/api/v1/marketplace/plugins/download/{plugin_author}/{plugin_name}/{latest_version}'
|
||||
_download_status, plugin_package = await _marketplace_get(
|
||||
client,
|
||||
f'{space_url}/api/v1/marketplace/plugins/download/{plugin_author}/{plugin_name}/{latest_version}',
|
||||
max_bytes=_MARKETPLACE_PLUGIN_DOWNLOAD_MAX_BYTES,
|
||||
)
|
||||
download_resp.raise_for_status()
|
||||
return download_resp.content, latest_version
|
||||
return plugin_package, latest_version
|
||||
|
||||
async def install_plugin(
|
||||
self,
|
||||
@@ -1698,7 +1787,11 @@ class PluginRuntimeConnector(ManagedRuntimeConnector):
|
||||
await delete(self.ap.persistence_mgr.execute_async)
|
||||
|
||||
self._known_desired_states.pop(binding.installation_uuid, None)
|
||||
self._workspace_installations.setdefault(binding.workspace_uuid, set()).discard(binding.installation_uuid)
|
||||
workspace_installations = self._workspace_installations.get(binding.workspace_uuid)
|
||||
if workspace_installations is not None:
|
||||
workspace_installations.discard(binding.installation_uuid)
|
||||
if not workspace_installations:
|
||||
self._workspace_installations.pop(binding.workspace_uuid, None)
|
||||
if task_context is not None:
|
||||
task_context.set_current_action('plugin removed')
|
||||
return {}
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import inspect
|
||||
import typing
|
||||
from typing import Any
|
||||
@@ -43,6 +44,24 @@ from ..entity.persistence import model as persistence_model
|
||||
from ..core import app
|
||||
from ..utils import constants
|
||||
|
||||
_DEFAULT_BINARY_STORAGE_VALUE_BYTES = 10 * 1024 * 1024
|
||||
_HARD_MAX_BINARY_STORAGE_VALUE_BYTES = 64 * 1024 * 1024
|
||||
|
||||
|
||||
def _binary_storage_value_limit(ap: Any) -> int:
|
||||
configured = (
|
||||
ap.instance_config.data.get('plugin', {})
|
||||
.get('binary_storage', {})
|
||||
.get('max_value_bytes', _DEFAULT_BINARY_STORAGE_VALUE_BYTES)
|
||||
)
|
||||
try:
|
||||
configured = int(configured)
|
||||
except (TypeError, ValueError):
|
||||
configured = _DEFAULT_BINARY_STORAGE_VALUE_BYTES
|
||||
if configured < 0:
|
||||
configured = _DEFAULT_BINARY_STORAGE_VALUE_BYTES
|
||||
return min(configured, _HARD_MAX_BINARY_STORAGE_VALUE_BYTES)
|
||||
|
||||
|
||||
class _RawAction:
|
||||
def __init__(self, value: str):
|
||||
@@ -859,20 +878,15 @@ class RuntimeConnectionHandler(handler.Handler):
|
||||
return handler.ActionResponse.error(
|
||||
message=str(e),
|
||||
)
|
||||
value = base64.b64decode(data['value_base64'])
|
||||
max_value_bytes = (
|
||||
self.ap.instance_config.data.get('plugin', {})
|
||||
.get('binary_storage', {})
|
||||
.get(
|
||||
'max_value_bytes',
|
||||
10 * 1024 * 1024,
|
||||
max_value_bytes = _binary_storage_value_limit(self.ap)
|
||||
encoded_value = data['value_base64']
|
||||
max_encoded_chars = 4 * ((max_value_bytes + 2) // 3) + 4
|
||||
if len(encoded_value) > max_encoded_chars:
|
||||
return handler.ActionResponse.error(
|
||||
message=f'Binary storage value exceeds the {max_value_bytes}-byte limit',
|
||||
)
|
||||
)
|
||||
try:
|
||||
max_value_bytes = int(max_value_bytes)
|
||||
except (TypeError, ValueError):
|
||||
max_value_bytes = 10 * 1024 * 1024
|
||||
if max_value_bytes >= 0 and len(value) > max_value_bytes:
|
||||
value = await asyncio.to_thread(base64.b64decode, encoded_value)
|
||||
if len(value) > max_value_bytes:
|
||||
return handler.ActionResponse.error(
|
||||
message=f'Binary storage value exceeds limit ({len(value)} > {max_value_bytes} bytes)',
|
||||
)
|
||||
@@ -936,10 +950,15 @@ class RuntimeConnectionHandler(handler.Handler):
|
||||
return handler.ActionResponse.error(
|
||||
message=f'Storage with key {key} not found',
|
||||
)
|
||||
max_value_bytes = _binary_storage_value_limit(self.ap)
|
||||
if len(storage.value) > max_value_bytes:
|
||||
return handler.ActionResponse.error(
|
||||
message=f'Binary storage value exceeds the {max_value_bytes}-byte limit',
|
||||
)
|
||||
|
||||
return handler.ActionResponse.success(
|
||||
data={
|
||||
'value_base64': base64.b64encode(storage.value).decode('utf-8'),
|
||||
'value_base64': (await asyncio.to_thread(base64.b64encode, storage.value)).decode('utf-8'),
|
||||
},
|
||||
)
|
||||
|
||||
@@ -1019,7 +1038,7 @@ class RuntimeConnectionHandler(handler.Handler):
|
||||
|
||||
return handler.ActionResponse.success(
|
||||
data={
|
||||
'file_base64': base64.b64encode(file_bytes).decode('utf-8'),
|
||||
'file_base64': (await asyncio.to_thread(base64.b64encode, file_bytes)).decode('utf-8'),
|
||||
},
|
||||
)
|
||||
except Exception as e:
|
||||
@@ -1744,7 +1763,7 @@ class RuntimeConnectionHandler(handler.Handler):
|
||||
await self.delete_local_file(plugin_icon_file_key)
|
||||
|
||||
return {
|
||||
'plugin_icon_base64': base64.b64encode(plugin_icon_bytes).decode('utf-8'),
|
||||
'plugin_icon_base64': (await asyncio.to_thread(base64.b64encode, plugin_icon_bytes)).decode('utf-8'),
|
||||
'mime_type': mime_type,
|
||||
}
|
||||
|
||||
@@ -1819,7 +1838,7 @@ class RuntimeConnectionHandler(handler.Handler):
|
||||
asset_bytes = await self.read_local_file(asset_file_key)
|
||||
await self.delete_local_file(asset_file_key)
|
||||
return {
|
||||
'asset_base64': base64.b64encode(asset_bytes).decode('utf-8'),
|
||||
'asset_base64': (await asyncio.to_thread(base64.b64encode, asset_bytes)).decode('utf-8'),
|
||||
'mime_type': mime_type,
|
||||
}
|
||||
|
||||
|
||||
@@ -52,6 +52,128 @@ class ModelManager:
|
||||
self.rerank_model_dict = {}
|
||||
self.requester_components = []
|
||||
self.requester_dict = {}
|
||||
self._scope_generations: dict[tuple[str, str], int] = {}
|
||||
self._provider_keys_by_scope: dict[tuple[str, str], set[_CacheKey]] = {}
|
||||
self._llm_keys_by_scope: dict[tuple[str, str], set[_CacheKey]] = {}
|
||||
self._embedding_keys_by_scope: dict[tuple[str, str], set[_CacheKey]] = {}
|
||||
self._rerank_keys_by_scope: dict[tuple[str, str], set[_CacheKey]] = {}
|
||||
|
||||
def _cache_index(self, cache: dict) -> dict[tuple[str, str], set[_CacheKey]]:
|
||||
if cache is self.provider_dict:
|
||||
return self._provider_keys_by_scope
|
||||
if cache is self.llm_model_dict:
|
||||
return self._llm_keys_by_scope
|
||||
if cache is self.embedding_model_dict:
|
||||
return self._embedding_keys_by_scope
|
||||
if cache is self.rerank_model_dict:
|
||||
return self._rerank_keys_by_scope
|
||||
raise ValueError('Unknown model runtime cache')
|
||||
|
||||
def _cache_set(self, cache: dict, key: _CacheKey, value: object) -> None:
|
||||
cache[key] = value
|
||||
self._cache_index(cache).setdefault(key[:2], set()).add(key)
|
||||
|
||||
def _cache_pop(self, cache: dict, key: _CacheKey) -> object | None:
|
||||
removed = cache.pop(key, None)
|
||||
scope = key[:2]
|
||||
index = self._cache_index(cache)
|
||||
keys = index.get(scope)
|
||||
if keys is not None:
|
||||
keys.discard(key)
|
||||
if not keys:
|
||||
index.pop(scope, None)
|
||||
if not any(
|
||||
scope in candidate
|
||||
for candidate in (
|
||||
self._provider_keys_by_scope,
|
||||
self._llm_keys_by_scope,
|
||||
self._embedding_keys_by_scope,
|
||||
self._rerank_keys_by_scope,
|
||||
)
|
||||
):
|
||||
self._scope_generations.pop(scope, None)
|
||||
return removed
|
||||
|
||||
def _observe_execution_context(
|
||||
self,
|
||||
context: ExecutionContext,
|
||||
) -> tuple[requester.RuntimeProvider, ...]:
|
||||
"""Prune superseded runtime objects when a Workspace generation advances."""
|
||||
|
||||
scope = (context.instance_uuid, context.workspace_uuid)
|
||||
previous_generation = self._scope_generations.get(scope)
|
||||
if previous_generation is not None and context.placement_generation < previous_generation:
|
||||
raise WorkspaceInvariantError('Model runtime placement generation rolled back')
|
||||
if previous_generation == context.placement_generation:
|
||||
return ()
|
||||
retired_providers: list[requester.RuntimeProvider] = []
|
||||
if previous_generation is not None:
|
||||
for cache, index in (
|
||||
(self.provider_dict, self._provider_keys_by_scope),
|
||||
(self.llm_model_dict, self._llm_keys_by_scope),
|
||||
(self.embedding_model_dict, self._embedding_keys_by_scope),
|
||||
(self.rerank_model_dict, self._rerank_keys_by_scope),
|
||||
):
|
||||
for key in index.pop(scope, ()):
|
||||
removed = cache.pop(key, None)
|
||||
if cache is self.provider_dict and removed is not None:
|
||||
retired_providers.append(removed)
|
||||
self._scope_generations[scope] = context.placement_generation
|
||||
return tuple(retired_providers)
|
||||
|
||||
async def _close_runtime_providers(
|
||||
self,
|
||||
providers: tuple[requester.RuntimeProvider, ...] | list[requester.RuntimeProvider],
|
||||
) -> None:
|
||||
"""Close each retired requester once without blocking other cleanup."""
|
||||
|
||||
seen: set[int] = set()
|
||||
for provider in providers:
|
||||
provider_id = id(provider)
|
||||
if provider_id in seen:
|
||||
continue
|
||||
seen.add(provider_id)
|
||||
try:
|
||||
await provider.requester.aclose()
|
||||
except Exception as exc:
|
||||
self.ap.logger.warning(
|
||||
f'Failed to close model requester for provider {provider.provider_entity.uuid}: {exc}'
|
||||
)
|
||||
|
||||
async def _observe_and_close_execution_context(
|
||||
self,
|
||||
context: ExecutionContext,
|
||||
*,
|
||||
retain_empty: bool = True,
|
||||
) -> None:
|
||||
await self._close_runtime_providers(self._observe_execution_context(context))
|
||||
if not retain_empty:
|
||||
scope = (context.instance_uuid, context.workspace_uuid)
|
||||
if not any(
|
||||
scope in candidate
|
||||
for candidate in (
|
||||
self._provider_keys_by_scope,
|
||||
self._llm_keys_by_scope,
|
||||
self._embedding_keys_by_scope,
|
||||
self._rerank_keys_by_scope,
|
||||
)
|
||||
):
|
||||
self._scope_generations.pop(scope, None)
|
||||
|
||||
async def shutdown(self) -> None:
|
||||
"""Release every requester owned by the model runtime cache."""
|
||||
|
||||
providers = list(self.provider_dict.values())
|
||||
self.provider_dict = {}
|
||||
self.llm_model_dict = {}
|
||||
self.embedding_model_dict = {}
|
||||
self.rerank_model_dict = {}
|
||||
self._scope_generations = {}
|
||||
self._provider_keys_by_scope = {}
|
||||
self._llm_keys_by_scope = {}
|
||||
self._embedding_keys_by_scope = {}
|
||||
self._rerank_keys_by_scope = {}
|
||||
await self._close_runtime_providers(providers)
|
||||
|
||||
@staticmethod
|
||||
def _get_litellm_provider_from_manifest(component: engine.Component | None) -> str | None:
|
||||
@@ -137,7 +259,17 @@ class ModelManager:
|
||||
if supplied_instance_uuid is not None and supplied_instance_uuid != binding.instance_uuid:
|
||||
raise WorkspaceInvariantError('Runtime context belongs to another LangBot instance')
|
||||
|
||||
return self._context_from_binding(binding, trigger_principal=trigger_principal)
|
||||
execution_context = self._context_from_binding(binding, trigger_principal=trigger_principal)
|
||||
scope = (
|
||||
execution_context.instance_uuid,
|
||||
execution_context.workspace_uuid,
|
||||
)
|
||||
if scope in self._scope_generations:
|
||||
await self._observe_and_close_execution_context(
|
||||
execution_context,
|
||||
retain_empty=False,
|
||||
)
|
||||
return execution_context
|
||||
|
||||
async def initialize(self) -> None:
|
||||
self.requester_components = self.ap.discover.get_components_by_kind('LLMAPIRequester')
|
||||
@@ -199,10 +331,16 @@ class ModelManager:
|
||||
"""Load every active projected Workspace into isolated runtime caches."""
|
||||
|
||||
self.ap.logger.info('Loading models from db...')
|
||||
await self._close_runtime_providers(list(self.provider_dict.values()))
|
||||
self.provider_dict = {}
|
||||
self.llm_model_dict = {}
|
||||
self.embedding_model_dict = {}
|
||||
self.rerank_model_dict = {}
|
||||
self._scope_generations = {}
|
||||
self._provider_keys_by_scope = {}
|
||||
self._llm_keys_by_scope = {}
|
||||
self._embedding_keys_by_scope = {}
|
||||
self._rerank_keys_by_scope = {}
|
||||
|
||||
list_bindings = getattr(self.ap.workspace_service, 'list_active_execution_bindings', None)
|
||||
tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
|
||||
@@ -233,6 +371,7 @@ class ModelManager:
|
||||
binding,
|
||||
trigger_principal=PrincipalContext(principal_type=PrincipalType.SYSTEM),
|
||||
)
|
||||
await self._observe_and_close_execution_context(resolved)
|
||||
contexts[workspace_uuid] = resolved
|
||||
return resolved
|
||||
|
||||
@@ -243,7 +382,11 @@ class ModelManager:
|
||||
try:
|
||||
context = await context_for(provider_entity.workspace_uuid)
|
||||
runtime_provider = await self._build_provider(context, provider_entity)
|
||||
self.provider_dict[self._cache_key(context, provider_entity.uuid)] = runtime_provider
|
||||
self._cache_set(
|
||||
self.provider_dict,
|
||||
self._cache_key(context, provider_entity.uuid),
|
||||
runtime_provider,
|
||||
)
|
||||
except provider_errors.RequesterNotFoundError as exc:
|
||||
self.ap.logger.warning(
|
||||
f'Requester {exc.requester_name} not found, skipping provider {provider_entity.uuid}'
|
||||
@@ -282,7 +425,11 @@ class ModelManager:
|
||||
)
|
||||
continue
|
||||
runtime_model = builder(context, model_entity, provider)
|
||||
cache[self._cache_key(context, model_entity.uuid)] = runtime_model
|
||||
self._cache_set(
|
||||
cache,
|
||||
self._cache_key(context, model_entity.uuid),
|
||||
runtime_model,
|
||||
)
|
||||
except Exception as exc:
|
||||
self.ap.logger.error(f'Failed to load model {model_entity.uuid}: {exc}\n{traceback.format_exc()}')
|
||||
|
||||
@@ -294,10 +441,20 @@ class ModelManager:
|
||||
persistence_model.ModelProvider.workspace_uuid == context.workspace_uuid
|
||||
)
|
||||
)
|
||||
for provider_entity in providers_result.all():
|
||||
provider_entities = providers_result.all()
|
||||
if provider_entities:
|
||||
# Empty Workspaces are the dominant SaaS registration case. Do
|
||||
# not retain one generation record per account until the
|
||||
# Workspace owns an actual runtime model resource.
|
||||
await self._observe_and_close_execution_context(context)
|
||||
for provider_entity in provider_entities:
|
||||
try:
|
||||
runtime_provider = await self._build_provider(context, provider_entity)
|
||||
self.provider_dict[self._cache_key(context, provider_entity.uuid)] = runtime_provider
|
||||
self._cache_set(
|
||||
self.provider_dict,
|
||||
self._cache_key(context, provider_entity.uuid),
|
||||
runtime_provider,
|
||||
)
|
||||
except provider_errors.RequesterNotFoundError as exc:
|
||||
self.ap.logger.warning(
|
||||
f'Requester {exc.requester_name} not found, skipping provider {provider_entity.uuid}'
|
||||
@@ -337,7 +494,11 @@ class ModelManager:
|
||||
)
|
||||
continue
|
||||
runtime_model = builder(context, model_entity, provider)
|
||||
cache[self._cache_key(context, model_entity.uuid)] = runtime_model
|
||||
self._cache_set(
|
||||
cache,
|
||||
self._cache_key(context, model_entity.uuid),
|
||||
runtime_model,
|
||||
)
|
||||
except Exception as exc:
|
||||
self.ap.logger.error(f'Failed to load model {model_entity.uuid}: {exc}\n{traceback.format_exc()}')
|
||||
|
||||
@@ -562,7 +723,12 @@ class ModelManager:
|
||||
execution_context = await self.resolve_execution_context(context)
|
||||
self._ensure_same_scope(execution_context, provider.execution_context, resource='Provider')
|
||||
self._ensure_entity_workspace(provider.provider_entity, execution_context, resource='Provider')
|
||||
self.provider_dict[self._cache_key(execution_context, provider.provider_entity.uuid)] = provider
|
||||
self._observe_execution_context(execution_context)
|
||||
self._cache_set(
|
||||
self.provider_dict,
|
||||
self._cache_key(execution_context, provider.provider_entity.uuid),
|
||||
provider,
|
||||
)
|
||||
|
||||
async def get_provider_by_uuid(
|
||||
self,
|
||||
@@ -578,7 +744,12 @@ class ModelManager:
|
||||
|
||||
async def remove_provider(self, context: TenantContext, provider_uuid: str) -> None:
|
||||
execution_context = await self.resolve_execution_context(context)
|
||||
self.provider_dict.pop(self._cache_key(execution_context, provider_uuid), None)
|
||||
removed = self._cache_pop(
|
||||
self.provider_dict,
|
||||
self._cache_key(execution_context, provider_uuid),
|
||||
)
|
||||
if removed is not None:
|
||||
await self._close_runtime_providers([removed])
|
||||
|
||||
async def reload_provider(self, context: TenantContext, provider_uuid: str) -> None:
|
||||
execution_context = await self.resolve_execution_context(context)
|
||||
@@ -593,12 +764,26 @@ class ModelManager:
|
||||
raise provider_errors.ProviderNotFoundError(provider_uuid)
|
||||
|
||||
new_provider = await self._build_provider(execution_context, provider_entity)
|
||||
cache_prefix = self._cache_key(execution_context, '')[:3]
|
||||
for cache in (self.llm_model_dict, self.embedding_model_dict, self.rerank_model_dict):
|
||||
for key, model in cache.items():
|
||||
if key[:3] == cache_prefix and model.provider.provider_entity.uuid == provider_uuid:
|
||||
scope = (execution_context.instance_uuid, execution_context.workspace_uuid)
|
||||
for cache, index in (
|
||||
(self.llm_model_dict, self._llm_keys_by_scope),
|
||||
(self.embedding_model_dict, self._embedding_keys_by_scope),
|
||||
(self.rerank_model_dict, self._rerank_keys_by_scope),
|
||||
):
|
||||
for key in tuple(index.get(scope, ())):
|
||||
model = cache.get(key)
|
||||
if model is not None and model.provider.provider_entity.uuid == provider_uuid:
|
||||
model.provider = new_provider
|
||||
self.provider_dict[self._cache_key(execution_context, provider_uuid)] = new_provider
|
||||
self._observe_execution_context(execution_context)
|
||||
provider_key = self._cache_key(execution_context, provider_uuid)
|
||||
old_provider = self.provider_dict.get(provider_key)
|
||||
self._cache_set(
|
||||
self.provider_dict,
|
||||
provider_key,
|
||||
new_provider,
|
||||
)
|
||||
if old_provider is not None and old_provider is not new_provider:
|
||||
await self._close_runtime_providers([old_provider])
|
||||
|
||||
@staticmethod
|
||||
def _coerce_model(model_info: _ModelEntity | sqlalchemy.Row, entity_type: type[_ModelEntity]) -> _ModelEntity:
|
||||
@@ -689,7 +874,12 @@ class ModelManager:
|
||||
async def cache_llm_model(self, context: TenantContext, model: requester.RuntimeLLMModel) -> None:
|
||||
execution_context = await self.resolve_execution_context(context)
|
||||
self._ensure_same_scope(execution_context, model.execution_context, resource='LLM model')
|
||||
self.llm_model_dict[self._cache_key(execution_context, model.model_entity.uuid)] = model
|
||||
self._observe_execution_context(execution_context)
|
||||
self._cache_set(
|
||||
self.llm_model_dict,
|
||||
self._cache_key(execution_context, model.model_entity.uuid),
|
||||
model,
|
||||
)
|
||||
|
||||
async def cache_embedding_model(
|
||||
self,
|
||||
@@ -698,12 +888,22 @@ class ModelManager:
|
||||
) -> None:
|
||||
execution_context = await self.resolve_execution_context(context)
|
||||
self._ensure_same_scope(execution_context, model.execution_context, resource='Embedding model')
|
||||
self.embedding_model_dict[self._cache_key(execution_context, model.model_entity.uuid)] = model
|
||||
self._observe_execution_context(execution_context)
|
||||
self._cache_set(
|
||||
self.embedding_model_dict,
|
||||
self._cache_key(execution_context, model.model_entity.uuid),
|
||||
model,
|
||||
)
|
||||
|
||||
async def cache_rerank_model(self, context: TenantContext, model: requester.RuntimeRerankModel) -> None:
|
||||
execution_context = await self.resolve_execution_context(context)
|
||||
self._ensure_same_scope(execution_context, model.execution_context, resource='Rerank model')
|
||||
self.rerank_model_dict[self._cache_key(execution_context, model.model_entity.uuid)] = model
|
||||
self._observe_execution_context(execution_context)
|
||||
self._cache_set(
|
||||
self.rerank_model_dict,
|
||||
self._cache_key(execution_context, model.model_entity.uuid),
|
||||
model,
|
||||
)
|
||||
|
||||
async def get_model_by_uuid(self, context: TenantContext, model_uuid: str) -> requester.RuntimeLLMModel:
|
||||
execution_context = await self.resolve_execution_context(context)
|
||||
@@ -739,15 +939,24 @@ class ModelManager:
|
||||
|
||||
async def remove_llm_model(self, context: TenantContext, model_uuid: str) -> None:
|
||||
execution_context = await self.resolve_execution_context(context)
|
||||
self.llm_model_dict.pop(self._cache_key(execution_context, model_uuid), None)
|
||||
self._cache_pop(
|
||||
self.llm_model_dict,
|
||||
self._cache_key(execution_context, model_uuid),
|
||||
)
|
||||
|
||||
async def remove_embedding_model(self, context: TenantContext, model_uuid: str) -> None:
|
||||
execution_context = await self.resolve_execution_context(context)
|
||||
self.embedding_model_dict.pop(self._cache_key(execution_context, model_uuid), None)
|
||||
self._cache_pop(
|
||||
self.embedding_model_dict,
|
||||
self._cache_key(execution_context, model_uuid),
|
||||
)
|
||||
|
||||
async def remove_rerank_model(self, context: TenantContext, model_uuid: str) -> None:
|
||||
execution_context = await self.resolve_execution_context(context)
|
||||
self.rerank_model_dict.pop(self._cache_key(execution_context, model_uuid), None)
|
||||
self._cache_pop(
|
||||
self.rerank_model_dict,
|
||||
self._cache_key(execution_context, model_uuid),
|
||||
)
|
||||
|
||||
def get_available_requesters_info(self, model_type: str) -> list[dict]:
|
||||
if model_type:
|
||||
|
||||
@@ -463,6 +463,17 @@ class ProviderAPIRequester(metaclass=abc.ABCMeta):
|
||||
async def initialize(self):
|
||||
pass
|
||||
|
||||
async def aclose(self) -> None:
|
||||
"""Release requester-owned clients when its runtime provider retires.
|
||||
|
||||
Most built-in requesters are currently stateless, but provider
|
||||
extensions may own connection pools or background resources. Keeping
|
||||
the lifecycle hook on the base class lets Workspace generation changes
|
||||
and application shutdown retire them deterministically.
|
||||
"""
|
||||
|
||||
return None
|
||||
|
||||
async def scan_models(self, api_key: str | None = None) -> dict[str, typing.Any] | list[dict[str, typing.Any]]:
|
||||
"""Scan models supported by the provider.
|
||||
|
||||
|
||||
@@ -8,6 +8,7 @@ import litellm
|
||||
from litellm import acompletion, aembedding, arerank
|
||||
|
||||
from .. import errors, requester
|
||||
from ....utils import httpclient
|
||||
import langbot_plugin.api.entities.builtin.resource.tool as resource_tool
|
||||
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
|
||||
import langbot_plugin.api.entities.builtin.provider.message as provider_message
|
||||
@@ -955,14 +956,17 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
|
||||
rerank_url = f'{base_url}/rerank'
|
||||
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=timeout) as client:
|
||||
async with httpx.AsyncClient(
|
||||
timeout=timeout,
|
||||
event_hooks=httpclient.httpx_response_limit_hooks(),
|
||||
) as client:
|
||||
resp = await client.post(rerank_url, headers=headers, json=payload)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
data = await httpclient.parse_json_response(resp)
|
||||
except httpx.HTTPStatusError as e:
|
||||
body = ''
|
||||
try:
|
||||
body = e.response.text
|
||||
body = await httpclient.response_text(e.response)
|
||||
except Exception:
|
||||
pass
|
||||
raise errors.RequesterError(f'rerank 请求失败 (HTTP {e.response.status_code}): {body or str(e)}')
|
||||
@@ -998,10 +1002,14 @@ class LiteLLMRequester(requester.ProviderAPIRequester):
|
||||
models_url = f'{base_url}/models'
|
||||
|
||||
try:
|
||||
async with httpx.AsyncClient(trust_env=True, timeout=timeout) as client:
|
||||
async with httpx.AsyncClient(
|
||||
trust_env=True,
|
||||
timeout=timeout,
|
||||
event_hooks=httpclient.httpx_response_limit_hooks(),
|
||||
) as client:
|
||||
response = await client.get(models_url, headers=headers)
|
||||
response.raise_for_status()
|
||||
payload = response.json()
|
||||
payload = await httpclient.parse_json_response(response)
|
||||
|
||||
models = []
|
||||
for item in payload.get('data', []):
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import abc
|
||||
import asyncio
|
||||
import typing
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
@@ -11,6 +12,32 @@ if TYPE_CHECKING:
|
||||
|
||||
|
||||
preregistered_runners: list[typing.Type[RequestRunner]] = []
|
||||
_DEFAULT_SYNC_ITERATION_LIMIT = 100_000
|
||||
|
||||
_T = typing.TypeVar('_T')
|
||||
|
||||
|
||||
def _next_sync(iterator: typing.Iterator[_T]) -> tuple[bool, _T | None]:
|
||||
try:
|
||||
return True, next(iterator)
|
||||
except StopIteration:
|
||||
return False, None
|
||||
|
||||
|
||||
async def iterate_sync(
|
||||
iterable: typing.Iterable[_T],
|
||||
*,
|
||||
max_items: int = _DEFAULT_SYNC_ITERATION_LIMIT,
|
||||
) -> typing.AsyncGenerator[_T, None]:
|
||||
"""Consume a blocking SDK iterator without stalling the event loop."""
|
||||
|
||||
iterator = iter(iterable)
|
||||
for _ in range(max(max_items, 1)):
|
||||
has_item, item = await asyncio.to_thread(_next_sync, iterator)
|
||||
if not has_item:
|
||||
return
|
||||
yield typing.cast(_T, item)
|
||||
raise RuntimeError('Synchronous provider stream exceeded the event limit')
|
||||
|
||||
|
||||
def runner_class(name: str):
|
||||
@@ -43,3 +70,6 @@ class RequestRunner(abc.ABC):
|
||||
) -> typing.AsyncGenerator[provider_message.Message | provider_message.MessageChunk, None]:
|
||||
"""运行请求"""
|
||||
pass
|
||||
|
||||
async def aclose(self) -> None:
|
||||
"""Release request-scoped resources after one runner invocation."""
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user