From ae85ac2b1689cd39c48bcbbbcb1fe02fa123510a Mon Sep 17 00:00:00 2001 From: Junyan Qin Date: Wed, 29 Jul 2026 11:32:26 +0800 Subject: [PATCH] feat(cloud): harden multi-tenant runtime resources --- docker/docker-compose.yaml | 16 +- docker/kubernetes.yaml | 60 +- docs/multi-tenant/cloud-runtime-soak-gate.md | 106 ++ .../cloud-v2-pending-verification.md | 212 +++ .../runtime-resource-audit-2026-07-28.md | 254 ++++ pyproject.toml | 3 +- pytest.ini | 6 + scripts/cloud_runtime_soak.py | 1170 +++++++++++++++++ scripts/runtime_resource_probe.py | 466 +++++++ scripts/workspace_runtime_capacity_probe.py | 561 ++++++++ src/langbot/libs/coze_server_api/client.py | 50 +- src/langbot/libs/deerflow_api/client.py | 38 +- .../libs/dify_service_api/v1/client.py | 287 ++-- src/langbot/libs/dingtalk_api/api.py | 179 ++- src/langbot/libs/official_account_api/api.py | 109 +- .../libs/openclaw_weixin_api/client.py | 112 +- src/langbot/libs/qq_official_api/api.py | 115 +- src/langbot/libs/slack_api/api.py | 8 +- .../libs/wechatpad_api/api/downloadpai.py | 48 +- .../libs/wechatpad_api/util/http_util.py | 59 +- src/langbot/libs/wecom_ai_bot_api/api.py | 151 ++- .../libs/wecom_ai_bot_api/ws_client.py | 150 ++- src/langbot/libs/wecom_api/api.py | 191 ++- .../libs/wecom_customer_service_api/api.py | 184 ++- src/langbot/libs/weknora_api/client.py | 106 +- src/langbot/pkg/api/http/controller/group.py | 44 +- .../controller/groups/knowledge/migration.py | 10 +- .../http/controller/groups/pipelines/embed.py | 49 +- .../groups/pipelines/websocket_chat.py | 93 +- .../controller/groups/platform/adapters.py | 135 +- .../pkg/api/http/controller/groups/plugins.py | 84 +- .../pkg/api/http/controller/groups/survey.py | 10 +- .../api/http/controller/groups/webhooks.py | 26 +- src/langbot/pkg/api/http/controller/main.py | 45 +- .../pkg/api/http/service/maintenance.py | 100 +- src/langbot/pkg/api/http/service/mcp.py | 50 +- src/langbot/pkg/api/http/service/skill.py | 15 +- src/langbot/pkg/api/http/service/space.py | 59 +- src/langbot/pkg/api/http/service/user.py | 67 +- src/langbot/pkg/box/admission.py | 3 +- src/langbot/pkg/box/service.py | 127 +- src/langbot/pkg/cloud/directory_projection.py | 37 + src/langbot/pkg/cloud/entitlements.py | 63 + src/langbot/pkg/cloud/launch.py | 38 +- src/langbot/pkg/core/app.py | 105 +- src/langbot/pkg/core/boot.py | 21 +- src/langbot/pkg/core/errors.py | 2 + src/langbot/pkg/core/stages/build_app.py | 24 +- src/langbot/pkg/core/stages/load_config.py | 21 +- src/langbot/pkg/core/stages/show_notes.py | 11 +- src/langbot/pkg/core/task_boundary.py | 8 + src/langbot/pkg/core/taskmgr.py | 50 +- src/langbot/pkg/discover/engine.py | 1 + .../pkg/persistence/databases/postgresql.py | 47 +- src/langbot/pkg/persistence/mgr.py | 8 + .../pkg/persistence/release_migration.py | 5 +- src/langbot/pkg/pipeline/aggregator.py | 60 +- .../cntfilter/filters/baiduexamine.py | 4 +- .../pipeline/cntfilter/filters/banwords.py | 32 +- .../pipeline/cntfilter/filters/cntignore.py | 29 +- src/langbot/pkg/pipeline/controller.py | 145 +- .../pkg/pipeline/longtext/strategies/image.py | 123 +- src/langbot/pkg/pipeline/pipelinemgr.py | 125 +- .../pkg/pipeline/plugin_diagnostics.py | 14 +- src/langbot/pkg/pipeline/pool.py | 120 ++ .../pkg/pipeline/process/handlers/chat.py | 38 + .../pkg/pipeline/ratelimit/algos/fixedwin.py | 59 +- .../pkg/pipeline/resprule/rules/regexp.py | 27 +- src/langbot/pkg/platform/botmgr.py | 479 +++++-- src/langbot/pkg/platform/logger.py | 7 +- src/langbot/pkg/platform/sources/aiocqhttp.py | 32 + src/langbot/pkg/platform/sources/dingtalk.py | 71 +- src/langbot/pkg/platform/sources/discord.py | 136 +- src/langbot/pkg/platform/sources/http_bot.py | 95 +- src/langbot/pkg/platform/sources/kook.py | 78 +- src/langbot/pkg/platform/sources/lark.py | 348 +++-- .../pkg/platform/sources/legacy/gewechat.py | 21 +- .../pkg/platform/sources/legacy/nakuru.py | 8 +- .../pkg/platform/sources/legacy/qqbotpy.py | 11 +- src/langbot/pkg/platform/sources/line.py | 30 +- src/langbot/pkg/platform/sources/matrix.py | 125 +- .../pkg/platform/sources/officialaccount.py | 1 + .../pkg/platform/sources/openclaw_weixin.py | 46 +- .../pkg/platform/sources/qqofficial.py | 65 +- src/langbot/pkg/platform/sources/satori.py | 28 +- src/langbot/pkg/platform/sources/telegram.py | 99 +- .../pkg/platform/sources/websocket_adapter.py | 193 ++- .../pkg/platform/sources/websocket_manager.py | 77 +- src/langbot/pkg/platform/sources/wechatpad.py | 207 ++- src/langbot/pkg/platform/sources/wecom.py | 1 + src/langbot/pkg/platform/sources/wecombot.py | 11 + src/langbot/pkg/platform/sources/wecomcs.py | 2 + src/langbot/pkg/platform/webhook_pusher.py | 3 +- src/langbot/pkg/plugin/archive.py | 98 ++ src/langbot/pkg/plugin/connector.py | 223 +++- src/langbot/pkg/plugin/handler.py | 53 +- src/langbot/pkg/provider/modelmgr/modelmgr.py | 247 +++- .../pkg/provider/modelmgr/requester.py | 11 + .../modelmgr/requesters/litellmchat.py | 18 +- src/langbot/pkg/provider/runner.py | 30 + src/langbot/pkg/provider/runners/cozeapi.py | 37 +- .../pkg/provider/runners/dashscopeapi.py | 129 +- src/langbot/pkg/provider/runners/difysvapi.py | 193 ++- .../pkg/provider/runners/langflowapi.py | 55 +- src/langbot/pkg/provider/runners/n8nsvapi.py | 17 +- src/langbot/pkg/provider/runners/tboxapi.py | 102 +- .../pkg/provider/runners/weknoraapi.py | 25 +- .../pkg/provider/session/sessionmgr.py | 287 +++- src/langbot/pkg/provider/tools/loaders/mcp.py | 420 +++++- .../pkg/provider/tools/loaders/mcp_stdio.py | 10 +- .../pkg/provider/tools/loaders/native.py | 287 +++- src/langbot/pkg/rag/knowledge/kbmgr.py | 168 ++- src/langbot/pkg/skill/manager.py | 3 + src/langbot/pkg/storage/mgr.py | 27 +- src/langbot/pkg/storage/provider.py | 5 + .../pkg/storage/providers/localstorage.py | 31 +- .../pkg/storage/providers/s3storage.py | 92 +- src/langbot/pkg/survey/manager.py | 37 +- src/langbot/pkg/telemetry/heartbeat.py | 48 +- src/langbot/pkg/telemetry/telemetry.py | 50 +- src/langbot/pkg/utils/bounded_executor.py | 312 +++++ src/langbot/pkg/utils/event_loop_monitor.py | 97 ++ src/langbot/pkg/utils/httpclient.py | 121 ++ src/langbot/pkg/utils/image.py | 89 +- src/langbot/pkg/utils/logcache.py | 5 + src/langbot/pkg/utils/safe_regex.py | 177 +++ src/langbot/pkg/utils/version.py | 10 +- src/langbot/pkg/vector/mgr.py | 8 + src/langbot/pkg/vector/vdb.py | 33 +- src/langbot/pkg/vector/vdbs/chroma.py | 22 +- src/langbot/pkg/vector/vdbs/milvus.py | 12 +- src/langbot/pkg/vector/vdbs/qdrant.py | 11 +- src/langbot/pkg/vector/vdbs/seekdb.py | 65 +- src/langbot/pkg/vector/vdbs/valkey_search.py | 14 +- .../pkg/workspace/invitation_delivery.py | 8 +- src/langbot/pkg/workspace/service.py | 39 + src/langbot/templates/config.yaml | 69 + tests/e2e/test_startup.py | 15 + tests/integration/api/test_smoke.py | 6 +- .../persistence/test_migrations_postgres.py | 422 ++++++ .../test_release_migration_postgres.py | 11 + .../api/http/test_bounded_json_request.py | 39 + .../api/http/test_internal_error_responses.py | 64 + .../api/service/test_mcp_service.py | 21 + .../api/service/test_space_service.py | 62 + .../api/service/test_user_service.py | 40 + .../api/test_adapter_session_scoping.py | 61 +- .../unit_tests/api/test_websocket_chat_uow.py | 75 ++ tests/unit_tests/box/test_box_service.py | 31 +- tests/unit_tests/cloud/test_entitlements.py | 52 + tests/unit_tests/cloud/test_space_launch.py | 44 + tests/unit_tests/core/test_app_shutdown.py | 122 ++ tests/unit_tests/core/test_boot.py | 26 +- tests/unit_tests/core/test_load_config.py | 12 + tests/unit_tests/core/test_taskmgr.py | 73 + .../persistence/test_postgresql_database.py | 69 +- .../persistence/test_release_migration.py | 2 + .../unit_tests/pipeline/test_chat_handler.py | 10 + .../pipeline/test_controller_tenancy.py | 2 + .../pipeline/test_longtext_image.py | 17 + tests/unit_tests/pipeline/test_n8nsvapi.py | 17 +- tests/unit_tests/pipeline/test_pipelinemgr.py | 93 +- tests/unit_tests/pipeline/test_pool.py | 53 + tests/unit_tests/pipeline/test_ratelimit.py | 28 +- .../test_aiocqhttp_message_converter.py | 19 + .../platform/test_botmgr_tenancy.py | 219 ++- .../platform/test_dingtalk_adapter.py | 12 + .../platform/test_discord_limits.py | 12 + .../platform/test_http_bot_tenancy.py | 28 + tests/unit_tests/platform/test_kook_limits.py | 24 + .../unit_tests/platform/test_lark_adapter.py | 27 + tests/unit_tests/platform/test_line_limits.py | 30 + .../unit_tests/platform/test_matrix_limits.py | 21 + .../platform/test_openclaw_weixin_client.py | 31 + .../platform/test_openclaw_weixin_tenancy.py | 11 + .../platform/test_qqofficial_api.py | 23 + .../platform/test_telegram_adapter.py | 22 + .../test_websocket_session_isolation.py | 65 + .../platform/test_wechatpad_limits.py | 39 + .../platform/test_wecom_api_limits.py | 31 + .../platform/test_wecombot_template_card.py | 33 +- .../unit_tests/plugin/test_connector_pure.py | 22 + .../plugin/test_connector_reconcile.py | 17 + .../unit_tests/plugin/test_handler_actions.py | 16 +- .../provider/runners/test_difysvapi_runner.py | 163 ++- .../runners/test_remote_stream_limits.py | 49 + .../runners/test_runner_resource_limits.py | 62 + .../unit_tests/provider/test_mcp_resources.py | 168 ++- .../unit_tests/provider/test_model_manager.py | 125 +- .../provider/test_session_manager.py | 130 ++ .../provider/test_tool_manager_native.py | 65 + tests/unit_tests/rag/test_file_storage.py | 56 +- tests/unit_tests/rag/test_kbmgr.py | 58 +- tests/unit_tests/storage/test_s3storage.py | 13 + .../storage/test_storage_manager.py | 9 + .../storage/test_workspace_scoping.py | 8 + .../unit_tests/survey/test_survey_manager.py | 7 + tests/unit_tests/telemetry/test_heartbeat.py | 30 + tests/unit_tests/test_telemetry.py | 7 +- .../unit_tests/utils/test_bounded_executor.py | 278 ++++ .../utils/test_cloud_runtime_soak.py | 446 +++++++ .../utils/test_event_loop_monitor.py | 75 ++ tests/unit_tests/utils/test_httpclient.py | 88 ++ tests/unit_tests/utils/test_image.py | 17 + tests/unit_tests/utils/test_logcache.py | 16 +- tests/unit_tests/utils/test_safe_regex.py | 73 + tests/unit_tests/vector/test_mgr.py | 22 +- .../vector/test_valkey_search_filter.py | 1 + tests/unit_tests/vector/test_vdb_base.py | 19 +- .../workspace/test_workspace_service.py | 27 + uv.lock | 6 +- 211 files changed, 14963 insertions(+), 1968 deletions(-) create mode 100644 docs/multi-tenant/cloud-runtime-soak-gate.md create mode 100644 docs/multi-tenant/cloud-v2-pending-verification.md create mode 100644 docs/multi-tenant/runtime-resource-audit-2026-07-28.md create mode 100644 scripts/cloud_runtime_soak.py create mode 100644 scripts/runtime_resource_probe.py create mode 100644 scripts/workspace_runtime_capacity_probe.py create mode 100644 src/langbot/pkg/core/errors.py create mode 100644 src/langbot/pkg/plugin/archive.py create mode 100644 src/langbot/pkg/utils/bounded_executor.py create mode 100644 src/langbot/pkg/utils/event_loop_monitor.py create mode 100644 src/langbot/pkg/utils/safe_regex.py create mode 100644 tests/unit_tests/api/http/test_bounded_json_request.py create mode 100644 tests/unit_tests/core/test_app_shutdown.py create mode 100644 tests/unit_tests/pipeline/test_longtext_image.py create mode 100644 tests/unit_tests/platform/test_discord_limits.py create mode 100644 tests/unit_tests/platform/test_kook_limits.py create mode 100644 tests/unit_tests/platform/test_line_limits.py create mode 100644 tests/unit_tests/platform/test_matrix_limits.py create mode 100644 tests/unit_tests/platform/test_openclaw_weixin_client.py create mode 100644 tests/unit_tests/platform/test_wechatpad_limits.py create mode 100644 tests/unit_tests/platform/test_wecom_api_limits.py create mode 100644 tests/unit_tests/provider/runners/test_remote_stream_limits.py create mode 100644 tests/unit_tests/provider/runners/test_runner_resource_limits.py create mode 100644 tests/unit_tests/utils/test_bounded_executor.py create mode 100644 tests/unit_tests/utils/test_cloud_runtime_soak.py create mode 100644 tests/unit_tests/utils/test_event_loop_monitor.py create mode 100644 tests/unit_tests/utils/test_safe_regex.py diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml index dd2444c49..b86f79d43 100644 --- a/docker/docker-compose.yaml +++ b/docker/docker-compose.yaml @@ -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__ `). 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. diff --git a/docker/kubernetes.yaml b/docker/kubernetes.yaml index b88ba6980..5504e5219 100644 --- a/docker/kubernetes.yaml +++ b/docker/kubernetes.yaml @@ -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 diff --git a/docs/multi-tenant/cloud-runtime-soak-gate.md b/docs/multi-tenant/cloud-runtime-soak-gate.md new file mode 100644 index 000000000..27c4ccec8 --- /dev/null +++ b/docs/multi-tenant/cloud-runtime-soak-gate.md @@ -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、生产配置摘要和工作负载版本。滚动更新、节点迁移或镜像变化后，旧报告不能继续作为新候选版本的批准证据。 diff --git a/docs/multi-tenant/cloud-v2-pending-verification.md b/docs/multi-tenant/cloud-v2-pending-verification.md new file mode 100644 index 000000000..511e7acea --- /dev/null +++ b/docs/multi-tenant/cloud-v2-pending-verification.md @@ -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`。 diff --git a/docs/multi-tenant/runtime-resource-audit-2026-07-28.md b/docs/multi-tenant/runtime-resource-audit-2026-07-28.md new file mode 100644 index 000000000..87933306d --- /dev/null +++ b/docs/multi-tenant/runtime-resource-audit-2026-07-28.md @@ -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、文件和子进程数量回到稳定基线。 diff --git a/pyproject.toml b/pyproject.toml index 7d3449e16..075655ae6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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", diff --git a/pytest.ini b/pytest.ini index a430a96e5..d86cbba9f 100644 --- a/pytest.ini +++ b/pytest.ini @@ -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 diff --git a/scripts/cloud_runtime_soak.py b/scripts/cloud_runtime_soak.py new file mode 100644 index 000000000..54ed5519d --- /dev/null +++ b/scripts/cloud_runtime_soak.py @@ -0,0 +1,1170 @@ +#!/usr/bin/env python3 +"""Run the final LangBot Cloud resource-stability acceptance gate. + +The short synthetic probes in this repository prove that selected registries +plateau. This tool is for the production-candidate topology: it samples Core, +Plugin Runtime and Box health endpoints together with their Linux process trees +and cgroup v2 accounting, streams raw evidence to JSONL, and fails when the +post-load tail shows a material leak or sustained CPU pressure. + +Examples: + + uv run python scripts/cloud_runtime_soak.py \ + --duration 24h --startup-grace 5m --cooldown 30m \ + --endpoint core=http://langbot:5300/healthz \ + --endpoint plugin=http://langbot-plugin-runtime:5400/healthz \ + --endpoint box=http://langbot-box:5410/readyz \ + --cgroup core=/sys/fs/cgroup/langbot \ + --cgroup plugin=/sys/fs/cgroup/langbot-plugin-runtime \ + --cgroup box=/sys/fs/cgroup/langbot-box \ + --samples-file artifacts/cloud-soak-samples.jsonl \ + --report-file artifacts/cloud-soak-report.json \ + --workload uv run python tests/load/cloud_candidate_workload.py + +Run this from a node/sidecar that can read the target cgroups. A process target +(``--pid name=PID``) is useful when cgroup paths are not directly mounted, but +cgroup evidence is still required for the production gate because only cgroups +report OOM, PID-limit and CPU-throttling events. +""" + +from __future__ import annotations + +import argparse +import concurrent.futures +import contextlib +import dataclasses +import datetime as dt +import json +import math +import os +import re +import signal +import statistics +import subprocess +import sys +import time +import urllib.error +import urllib.parse +import urllib.request +from collections import deque +from collections.abc import Callable, Iterable, Mapping, Sequence +from pathlib import Path +from typing import Any, TextIO + + +BYTES_PER_MIB = 1024 * 1024 +MAX_HTTP_BODY_BYTES = BYTES_PER_MIB +MAX_FLATTENED_HTTP_METRICS = 256 +MAX_PROCESS_TREE_SIZE = 4096 +MAX_TARGETS = 32 +MAX_SAMPLER_THREADS = 8 +_DIRECT_HTTP_OPENER = urllib.request.build_opener(urllib.request.ProxyHandler({})) +TARGET_NAME_RE = re.compile(r'^[A-Za-z0-9][A-Za-z0-9_.-]{0,63}$') +COUNTER_SUFFIXES = ( + '.memory.events.high', + '.memory.events.max', + '.memory.events.oom', + '.memory.events.oom_kill', + '.memory.events.oom_group_kill', + '.pids.events.max', +) +REJECTION_SUFFIXES = ( + '.blocking_executor.global_rejected_total', + '.blocking_executor.scope_rejected_total', +) +PENDING_SUFFIX = '.blocking_executor.pending' +TRANSIENT_GAUGE_SUFFIXES = ( + '.resources.telemetry_tasks', + '.resources.query_pool.queued', + '.resources.runtimes.mcp_host_tasks', + '.resources.runtimes.mcp_dispatch_tasks', + '.resources.creating_session_tasks', + '.resources.closing_session_tasks', + '.resources.background_tasks', +) + + +@dataclasses.dataclass(frozen=True, slots=True) +class Target: + name: str + kind: str + location: str + + +@dataclasses.dataclass(frozen=True, slots=True) +class MetricSample: + monotonic_seconds: float + wall_time: str + metrics: dict[str, float] + + +@dataclasses.dataclass(slots=True) +class TargetState: + target: Target + samples: deque[MetricSample] + baseline_metrics: dict[str, float] | None = None + last_metrics: dict[str, float] | None = None + observed_max_metrics: dict[str, float] = dataclasses.field(default_factory=dict) + attempted_samples: int = 0 + successful_samples: int = 0 + failed_samples: int = 0 + consecutive_failures: int = 0 + max_consecutive_failures: int = 0 + first_error: str | None = None + + +@dataclasses.dataclass(frozen=True, slots=True) +class Thresholds: + max_memory_growth_bytes: int + max_memory_slope_bytes_per_hour: float + max_tail_cpu_cores: float + max_throttled_period_ratio: float + allow_rejections: bool + max_transient_gauge_growth: float + require_hard_limits: bool + max_event_loop_lag_ms: float + max_event_loop_p95_lag_ms: float + require_event_loop_metrics: bool + + +@dataclasses.dataclass(frozen=True, slots=True) +class WorkloadResult: + executable: str | None + argument_count: int + started_at_seconds: float | None + completed_at_seconds: float | None + return_code: int | None + timed_out: bool + + +@dataclasses.dataclass(frozen=True, slots=True) +class GateResult: + failures: tuple[str, ...] + warnings: tuple[str, ...] + targets: dict[str, dict[str, Any]] + + @property + def passed(self) -> bool: + return not self.failures + + +def parse_duration(value: str) -> float: + """Parse a positive duration such as 30s, 5m, 2h or 1d.""" + + match = re.fullmatch(r'\s*(\d+(?:\.\d+)?)\s*([smhd]?)\s*', value, re.IGNORECASE) + if match is None: + raise argparse.ArgumentTypeError(f'invalid duration: {value!r}') + amount = float(match.group(1)) + if amount <= 0: + raise argparse.ArgumentTypeError('duration must be greater than zero') + multiplier = { + '': 1.0, + 's': 1.0, + 'm': 60.0, + 'h': 3600.0, + 'd': 86400.0, + }[match.group(2).lower()] + return amount * multiplier + + +def _parse_named_values(values: Iterable[str], *, option: str) -> dict[str, str]: + parsed: dict[str, str] = {} + for value in values: + name, separator, location = value.partition('=') + name = name.strip() + location = location.strip() + if not separator or not TARGET_NAME_RE.fullmatch(name) or not location: + raise ValueError(f'{option} must use NAME=VALUE with a safe, non-empty name: {value!r}') + if name in parsed: + raise ValueError(f'duplicate target name {name!r} in {option}') + parsed[name] = location + return parsed + + +def _safe_endpoint_url(value: str) -> str: + parsed = urllib.parse.urlsplit(value) + if parsed.scheme not in {'http', 'https'} or not parsed.hostname: + raise ValueError(f'health endpoint must be an http(s) URL: {value!r}') + if parsed.username or parsed.password or parsed.query or parsed.fragment: + raise ValueError('health endpoint URLs must not contain credentials, query parameters or fragments') + return urllib.parse.urlunsplit((parsed.scheme, parsed.netloc, parsed.path or '/', '', '')) + + +def build_targets( + *, + endpoints: Iterable[str], + cgroups: Iterable[str], + pids: Iterable[str], +) -> list[Target]: + targets: list[Target] = [] + seen_names: set[str] = set() + for kind, option, values in ( + ('endpoint', '--endpoint', endpoints), + ('cgroup', '--cgroup', cgroups), + ('process', '--pid', pids), + ): + for name, location in _parse_named_values(values, option=option).items(): + qualified_name = f'{kind}:{name}' + if qualified_name in seen_names: + raise ValueError(f'duplicate target {qualified_name!r}') + seen_names.add(qualified_name) + if kind == 'endpoint': + location = _safe_endpoint_url(location) + elif kind == 'cgroup': + location = str(Path(location).resolve()) + else: + try: + pid = int(location) + except ValueError as exc: + raise ValueError(f'{option} PID must be an integer: {location!r}') from exc + if pid <= 0: + raise ValueError(f'{option} PID must be greater than zero') + location = str(pid) + targets.append(Target(name=name, kind=kind, location=location)) + return targets + + +def _read_text(path: Path) -> str | None: + try: + return path.read_text(encoding='utf-8').strip() + except (FileNotFoundError, PermissionError, ProcessLookupError, OSError): + return None + + +def _read_integer(path: Path) -> int | None: + value = _read_text(path) + if value is None or value == 'max': + return None + try: + return int(value) + except ValueError: + return None + + +def _read_key_value_file(path: Path) -> dict[str, int]: + values: dict[str, int] = {} + text = _read_text(path) + if text is None: + return values + for line in text.splitlines(): + fields = line.split() + if len(fields) != 2: + continue + try: + values[fields[0]] = int(fields[1]) + except ValueError: + continue + return values + + +def read_cgroup_snapshot(path: Path) -> dict[str, float]: + """Read bounded cgroup v2 resource accounting from *path*.""" + + if not path.is_dir(): + raise RuntimeError(f'cgroup directory is unavailable: {path}') + metrics: dict[str, float] = {} + scalar_files = { + 'memory.current': 'memory.current_bytes', + 'memory.peak': 'memory.peak_bytes', + 'memory.swap.current': 'memory.swap.current_bytes', + 'memory.max': 'memory.max_bytes', + 'memory.swap.max': 'memory.swap.max_bytes', + 'pids.current': 'pids.current', + 'pids.max': 'pids.max', + } + for filename, metric_name in scalar_files.items(): + value = _read_integer(path / filename) + if value is not None: + metrics[metric_name] = float(value) + + for filename, prefix in ( + ('cpu.stat', 'cpu'), + ('memory.events', 'memory.events'), + ('pids.events', 'pids.events'), + ): + for key, value in _read_key_value_file(path / filename).items(): + metrics[f'{prefix}.{key}'] = float(value) + + cpu_max = _read_text(path / 'cpu.max') + if cpu_max: + fields = cpu_max.split() + if len(fields) == 2: + if fields[0] != 'max': + with contextlib.suppress(ValueError): + metrics['cpu.quota_usec'] = float(int(fields[0])) + with contextlib.suppress(ValueError): + metrics['cpu.period_usec'] = float(int(fields[1])) + required_metrics = { + 'memory.current_bytes', + 'memory.events.oom', + 'cpu.usage_usec', + 'pids.current', + 'pids.events.max', + } + missing_metrics = sorted(required_metrics - metrics.keys()) + if missing_metrics: + raise RuntimeError(f'cgroup v2 metrics missing from {path}: {", ".join(missing_metrics)}') + return metrics + + +def _read_process_ids(proc_root: Path, root_pid: int) -> list[int]: + pending = [root_pid] + discovered: list[int] = [] + seen: set[int] = set() + while pending and len(discovered) < MAX_PROCESS_TREE_SIZE: + pid = pending.pop() + if pid in seen: + continue + seen.add(pid) + if not (proc_root / str(pid)).is_dir(): + continue + discovered.append(pid) + task_root = proc_root / str(pid) / 'task' + try: + task_ids = tuple(task_root.iterdir()) + except (FileNotFoundError, PermissionError, ProcessLookupError, OSError): + task_ids = () + for task_path in task_ids: + children = _read_text(task_path / 'children') + if not children: + continue + for child in children.split(): + with contextlib.suppress(ValueError): + child_pid = int(child) + if child_pid not in seen: + pending.append(child_pid) + if pending: + raise RuntimeError(f'process tree exceeded {MAX_PROCESS_TREE_SIZE} processes') + return discovered + + +def _read_process_status(path: Path) -> tuple[int, int]: + rss_bytes = 0 + threads = 0 + text = _read_text(path) + if text is None: + return rss_bytes, threads + for line in text.splitlines(): + key, separator, raw_value = line.partition(':') + if not separator: + continue + fields = raw_value.split() + if not fields: + continue + if key == 'VmRSS': + with contextlib.suppress(ValueError): + rss_bytes = int(fields[0]) * 1024 + elif key == 'Threads': + with contextlib.suppress(ValueError): + threads = int(fields[0]) + return rss_bytes, threads + + +def _read_process_cpu_seconds(path: Path, clock_ticks: int) -> float: + text = _read_text(path) + if text is None: + return 0.0 + _, separator, remainder = text.rpartition(')') + if not separator: + return 0.0 + fields = remainder.split() + if len(fields) <= 12: + return 0.0 + try: + return (int(fields[11]) + int(fields[12])) / clock_ticks + except ValueError: + return 0.0 + + +def _count_open_fds(path: Path) -> int: + try: + with os.scandir(path) as entries: + return sum(1 for _ in entries) + except (FileNotFoundError, PermissionError, ProcessLookupError, OSError): + return 0 + + +def read_process_snapshot( + pid: int, + *, + proc_root: Path = Path('/proc'), + clock_ticks: int | None = None, +) -> dict[str, float]: + """Aggregate RSS, CPU, threads and file descriptors for a process tree.""" + + process_ids = _read_process_ids(proc_root, pid) + if not process_ids: + raise RuntimeError(f'process {pid} is unavailable') + ticks = clock_ticks or int(os.sysconf('SC_CLK_TCK')) + rss_bytes = 0 + threads = 0 + open_fds = 0 + cpu_seconds = 0.0 + for process_id in process_ids: + process_root = proc_root / str(process_id) + process_rss, process_threads = _read_process_status(process_root / 'status') + rss_bytes += process_rss + threads += process_threads + open_fds += _count_open_fds(process_root / 'fd') + cpu_seconds += _read_process_cpu_seconds(process_root / 'stat', ticks) + return { + 'rss_bytes': float(rss_bytes), + 'cpu_seconds': cpu_seconds, + 'threads': float(threads), + 'open_fds': float(open_fds), + 'processes': float(len(process_ids)), + } + + +def _flatten_numeric_json( + value: Any, + *, + prefix: str = 'body', + output: dict[str, float] | None = None, + depth: int = 0, +) -> dict[str, float]: + if output is None: + output = {} + if depth > 8 or len(output) >= MAX_FLATTENED_HTTP_METRICS: + return output + if isinstance(value, bool): + output[prefix] = float(value) + elif isinstance(value, (int, float)) and not isinstance(value, bool) and math.isfinite(float(value)): + output[prefix] = float(value) + elif isinstance(value, Mapping): + for key, nested_value in value.items(): + if len(output) >= MAX_FLATTENED_HTTP_METRICS: + break + safe_key = str(key).replace('.', '_')[:128] + _flatten_numeric_json( + nested_value, + prefix=f'{prefix}.{safe_key}', + output=output, + depth=depth + 1, + ) + return output + + +def read_endpoint_snapshot( + url: str, + *, + timeout_seconds: float, + opener: Callable[..., Any] | None = None, +) -> dict[str, float]: + """Fetch a bounded health response and flatten numeric resource gauges.""" + + request = urllib.request.Request( + url, + headers={ + 'Accept': 'application/json, text/plain', + 'User-Agent': 'langbot-cloud-runtime-soak/1', + }, + method='GET', + ) + started_at = time.monotonic() + direct_opener = opener or _DIRECT_HTTP_OPENER.open + try: + response_context = direct_opener(request, timeout=timeout_seconds) + with response_context as response: + status = int(getattr(response, 'status', response.getcode())) + content_type = str(response.headers.get('Content-Type', '')) + body = response.read(MAX_HTTP_BODY_BYTES + 1) + except urllib.error.HTTPError as exc: + raise RuntimeError(f'health endpoint returned HTTP {exc.code}') from exc + except (urllib.error.URLError, TimeoutError, OSError) as exc: + raise RuntimeError(f'health endpoint request failed: {type(exc).__name__}') from exc + if len(body) > MAX_HTTP_BODY_BYTES: + raise RuntimeError(f'health endpoint response exceeded {MAX_HTTP_BODY_BYTES} bytes') + if not 200 <= status < 300: + raise RuntimeError(f'health endpoint returned HTTP {status}') + + metrics = { + 'http.status': float(status), + 'http.latency_ms': (time.monotonic() - started_at) * 1000, + 'http.ok': 1.0, + } + if 'json' in content_type.lower() or body.lstrip().startswith((b'{', b'[')): + try: + payload = json.loads(body) + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + raise RuntimeError('health endpoint returned invalid JSON') from exc + if isinstance(payload, Mapping): + if 'code' in payload and payload['code'] != 0: + raise RuntimeError(f'health endpoint reported code={payload["code"]!r}') + if payload.get('live') is False or payload.get('ready') is False: + raise RuntimeError('health endpoint reported not ready') + metrics.update(_flatten_numeric_json(payload)) + return metrics + + +def _sample_target(target: Target, *, http_timeout_seconds: float) -> dict[str, float]: + if target.kind == 'endpoint': + return read_endpoint_snapshot(target.location, timeout_seconds=http_timeout_seconds) + if target.kind == 'cgroup': + return read_cgroup_snapshot(Path(target.location)) + if target.kind == 'process': + return read_process_snapshot(int(target.location)) + raise AssertionError(f'unsupported target kind: {target.kind}') + + +def _tail_metric_points( + samples: Sequence[MetricSample], + *, + analysis_start_seconds: float, + metric: str, +) -> list[tuple[float, float]]: + return [ + (sample.monotonic_seconds, sample.metrics[metric]) + for sample in samples + if sample.monotonic_seconds >= analysis_start_seconds and metric in sample.metrics + ] + + +def _robust_growth(points: Sequence[tuple[float, float]]) -> float: + width = max(1, min(len(points) // 5, 12)) + initial = statistics.median(value for _, value in points[:width]) + final = statistics.median(value for _, value in points[-width:]) + return final - initial + + +def _slope_per_hour(points: Sequence[tuple[float, float]]) -> float: + if len(points) < 2: + return 0.0 + origin = points[0][0] + xs = [timestamp - origin for timestamp, _ in points] + ys = [value for _, value in points] + x_mean = statistics.fmean(xs) + y_mean = statistics.fmean(ys) + denominator = sum((value - x_mean) ** 2 for value in xs) + if denominator <= 0: + return 0.0 + slope_per_second = ( + sum((x_value - x_mean) * (y_value - y_mean) for x_value, y_value in zip(xs, ys, strict=True)) / denominator + ) + return slope_per_second * 3600 + + +def _counter_delta( + first: MetricSample | Mapping[str, float], + last: MetricSample | Mapping[str, float], + metric: str, +) -> float | None: + first_metrics = first.metrics if isinstance(first, MetricSample) else first + last_metrics = last.metrics if isinstance(last, MetricSample) else last + if metric not in first_metrics or metric not in last_metrics: + return None + return last_metrics[metric] - first_metrics[metric] + + +def evaluate_gate( + states: Sequence[TargetState], + *, + analysis_start_seconds: float, + thresholds: Thresholds, +) -> GateResult: + failures: list[str] = [] + warnings: list[str] = [] + summaries: dict[str, dict[str, Any]] = {} + + for state in states: + target_id = f'{state.target.kind}:{state.target.name}' + target_summary: dict[str, Any] = { + 'attempted_samples': state.attempted_samples, + 'successful_samples': state.successful_samples, + 'failed_samples': state.failed_samples, + 'max_consecutive_failures': state.max_consecutive_failures, + 'first_error': state.first_error, + } + summaries[target_id] = target_summary + if state.failed_samples: + failures.append( + f'{target_id} had {state.failed_samples} failed samples ' + f'(max consecutive {state.max_consecutive_failures}): {state.first_error}' + ) + samples = list(state.samples) + tail_samples = [sample for sample in samples if sample.monotonic_seconds >= analysis_start_seconds] + target_summary['tail_samples'] = len(tail_samples) + if len(tail_samples) < 2: + failures.append(f'{target_id} has fewer than two successful tail samples') + continue + + first = tail_samples[0] + last = tail_samples[-1] + tail_elapsed = last.monotonic_seconds - first.monotonic_seconds + target_summary['tail_elapsed_seconds'] = tail_elapsed + if tail_elapsed <= 0: + failures.append(f'{target_id} has no measurable tail interval') + continue + + if state.target.kind == 'cgroup': + baseline_metrics = state.baseline_metrics or first.metrics + resource_limits = { + 'cpu_quota_usec': baseline_metrics.get('cpu.quota_usec'), + 'cpu_period_usec': baseline_metrics.get('cpu.period_usec'), + 'memory_max_bytes': baseline_metrics.get('memory.max_bytes'), + 'memory_swap_max_bytes': baseline_metrics.get('memory.swap.max_bytes'), + 'pids_max': baseline_metrics.get('pids.max'), + } + hard_limits = { + 'cpu': resource_limits['cpu_quota_usec'] is not None, + 'memory': resource_limits['memory_max_bytes'] is not None, + 'swap': resource_limits['memory_swap_max_bytes'] is not None, + 'pids': resource_limits['pids_max'] is not None, + } + target_summary['resource_limits'] = resource_limits + target_summary['hard_limits'] = hard_limits + if state.last_metrics is not None: + target_summary['memory.peak_bytes'] = state.last_metrics.get('memory.peak_bytes') + if thresholds.require_hard_limits: + missing_limits = sorted(name for name, enabled in hard_limits.items() if not enabled) + if missing_limits: + failures.append(f'{target_id} is missing hard cgroup limits: {", ".join(missing_limits)}') + for suffix in COUNTER_SUFFIXES: + metric = suffix.lstrip('.') + delta = _counter_delta( + state.baseline_metrics or first, + state.last_metrics or last, + metric, + ) + if delta is not None: + target_summary[f'delta.{metric}'] = delta + if delta < 0: + failures.append(f'{target_id} monotonic counter {metric} reset; the cgroup may have restarted') + elif delta > 0: + failures.append(f'{target_id} increased {metric} by {delta:g}') + + period_delta = _counter_delta( + state.baseline_metrics or first, + state.last_metrics or last, + 'cpu.nr_periods', + ) + throttled_delta = _counter_delta( + state.baseline_metrics or first, + state.last_metrics or last, + 'cpu.nr_throttled', + ) + if period_delta is not None and throttled_delta is not None and period_delta > 0: + throttled_ratio = max(throttled_delta, 0) / period_delta + target_summary['cpu.throttled_period_ratio'] = throttled_ratio + if throttled_ratio > thresholds.max_throttled_period_ratio: + failures.append( + f'{target_id} CPU throttled-period ratio {throttled_ratio:.3f} ' + f'exceeded {thresholds.max_throttled_period_ratio:.3f}' + ) + elif period_delta is not None and period_delta < 0: + failures.append(f'{target_id} monotonic CPU counters reset; the cgroup may have restarted') + + memory_metric = { + 'cgroup': 'memory.current_bytes', + 'process': 'rss_bytes', + }.get(state.target.kind) + if memory_metric is not None: + memory_points = _tail_metric_points( + tail_samples, + analysis_start_seconds=analysis_start_seconds, + metric=memory_metric, + ) + if len(memory_points) >= 2: + growth = _robust_growth(memory_points) + slope = _slope_per_hour(memory_points) + target_summary['memory.metric'] = memory_metric + target_summary['memory.robust_growth_bytes'] = growth + target_summary['memory.slope_bytes_per_hour'] = slope + if growth > thresholds.max_memory_growth_bytes and slope > thresholds.max_memory_slope_bytes_per_hour: + failures.append( + f'{target_id} {memory_metric} grew {growth / BYTES_PER_MIB:.2f} MiB ' + f'at {slope / BYTES_PER_MIB:.2f} MiB/hour' + ) + else: + warnings.append(f'{target_id} did not expose {memory_metric} throughout the tail') + + cpu_metric = { + 'cgroup': 'cpu.usage_usec', + 'process': 'cpu_seconds', + }.get(state.target.kind) + if cpu_metric is not None: + cpu_delta = _counter_delta(first, last, cpu_metric) + if cpu_delta is not None: + cpu_seconds = cpu_delta / 1_000_000 if cpu_metric.endswith('_usec') else cpu_delta + average_cores = max(cpu_seconds, 0) / tail_elapsed + target_summary['cpu.average_cores'] = average_cores + if average_cores > thresholds.max_tail_cpu_cores: + failures.append( + f'{target_id} tail CPU averaged {average_cores:.3f} cores, ' + f'above {thresholds.max_tail_cpu_cores:.3f}' + ) + + if state.target.kind == 'endpoint': + metric_keys = ( + set(first.metrics) + | set(last.metrics) + | set(state.baseline_metrics or ()) + | set(state.last_metrics or ()) + ) + for suffix in REJECTION_SUFFIXES: + if thresholds.allow_rejections: + break + for metric in sorted(key for key in metric_keys if key.endswith(suffix)): + delta = _counter_delta( + state.baseline_metrics or first, + state.last_metrics or last, + metric, + ) + if delta is not None and delta > 0: + failures.append(f'{target_id} increased {metric} by {delta:g}') + elif delta is not None and delta < 0: + failures.append(f'{target_id} monotonic counter {metric} reset; the runtime may have restarted') + for metric in sorted(key for key in metric_keys if key.endswith(PENDING_SUFFIX)): + values = [sample.metrics[metric] for sample in tail_samples if metric in sample.metrics] + if values and min(values) > 0: + failures.append(f'{target_id} kept {metric} above zero for the entire tail') + for suffix in TRANSIENT_GAUGE_SUFFIXES: + for metric in sorted(key for key in metric_keys if key.endswith(suffix)): + growth = _counter_delta(first, last, metric) + if growth is not None and growth > thresholds.max_transient_gauge_growth: + failures.append(f'{target_id} transient gauge {metric} grew by {growth:g} during the idle tail') + + loop_running_metrics = sorted(key for key in metric_keys if key.endswith('.event_loop.running')) + if not loop_running_metrics: + if thresholds.require_event_loop_metrics: + failures.append(f'{target_id} did not expose event-loop health metrics') + else: + for metric in loop_running_metrics: + running_values = [sample.metrics[metric] for sample in tail_samples if metric in sample.metrics] + if not running_values or min(running_values) < 1: + failures.append(f'{target_id} event-loop monitor was not running throughout the tail') + + recent_max_metrics = sorted( + key for key in state.observed_max_metrics if key.endswith('.event_loop.recent_max_lag_ms') + ) + for metric in recent_max_metrics: + observed_max = state.observed_max_metrics[metric] + target_summary['event_loop.max_observed_recent_lag_ms'] = observed_max + if observed_max > thresholds.max_event_loop_lag_ms: + failures.append( + f'{target_id} event-loop lag reached ' + f'{observed_max:.2f} ms, above ' + f'{thresholds.max_event_loop_lag_ms:.2f} ms' + ) + + recent_p95_metrics = sorted(key for key in metric_keys if key.endswith('.event_loop.recent_p95_lag_ms')) + for metric in recent_p95_metrics: + p95_values = [sample.metrics[metric] for sample in tail_samples if metric in sample.metrics] + if not p95_values: + continue + maximum_p95 = max(p95_values) + target_summary['event_loop.max_tail_recent_p95_lag_ms'] = maximum_p95 + if maximum_p95 > thresholds.max_event_loop_p95_lag_ms: + failures.append( + f'{target_id} event-loop recent p95 reached ' + f'{maximum_p95:.2f} ms, above ' + f'{thresholds.max_event_loop_p95_lag_ms:.2f} ms' + ) + + sample_total_metrics = sorted(key for key in metric_keys if key.endswith('.event_loop.samples_total')) + for metric in sample_total_metrics: + delta = _counter_delta( + state.baseline_metrics or first, + state.last_metrics or last, + metric, + ) + if delta is not None and delta < 0: + failures.append(f'{target_id} event-loop sample counter reset; the runtime may have restarted') + + return GateResult( + failures=tuple(dict.fromkeys(failures)), + warnings=tuple(dict.fromkeys(warnings)), + targets=summaries, + ) + + +def _utc_now() -> str: + return dt.datetime.now(dt.UTC).isoformat() + + +def _open_output(path: Path | None) -> TextIO | None: + if path is None: + return None + path.parent.mkdir(parents=True, exist_ok=True) + return path.open('w', encoding='utf-8') + + +def _write_json_line(stream: TextIO | None, payload: Mapping[str, Any]) -> None: + if stream is None: + return + stream.write(json.dumps(payload, sort_keys=True, separators=(',', ':')) + '\n') + stream.flush() + + +def _terminate_workload(process: subprocess.Popen[Any], *, grace_seconds: float = 10.0) -> None: + if process.poll() is not None: + return + if os.name == 'posix': + with contextlib.suppress(ProcessLookupError): + os.killpg(process.pid, signal.SIGTERM) + else: + process.terminate() + try: + process.wait(timeout=grace_seconds) + except subprocess.TimeoutExpired: + if os.name == 'posix': + with contextlib.suppress(ProcessLookupError): + os.killpg(process.pid, signal.SIGKILL) + else: + process.kill() + process.wait(timeout=grace_seconds) + + +def run_soak( + *, + targets: Sequence[Target], + duration_seconds: float, + sample_interval_seconds: float, + startup_grace_seconds: float, + cooldown_seconds: float, + analysis_window_seconds: float, + http_timeout_seconds: float, + thresholds: Thresholds, + workload_command: Sequence[str] | None, + samples_stream: TextIO | None, +) -> tuple[dict[str, Any], int]: + started_at_wall = _utc_now() + started_at = time.monotonic() + deadline = started_at + duration_seconds + max_tail_seconds = max(analysis_window_seconds, cooldown_seconds, sample_interval_seconds) + retained_samples = min( + max(4, math.ceil(max_tail_seconds / sample_interval_seconds) + 4), + 100_000, + ) + states = [TargetState(target=target, samples=deque(maxlen=retained_samples)) for target in targets] + workload: subprocess.Popen[Any] | None = None + workload_started_at: float | None = None + workload_completed_at: float | None = None + workload_return_code: int | None = None + workload_timed_out = False + workload_start_error: str | None = None + interrupted = False + next_sample_at = started_at + previous_sigterm_handler: Any = None + + def interrupt_on_sigterm(_signum: int, _frame: Any) -> None: + raise KeyboardInterrupt + + with contextlib.suppress(ValueError): + previous_sigterm_handler = signal.signal(signal.SIGTERM, interrupt_on_sigterm) + sampler_pool = concurrent.futures.ThreadPoolExecutor( + max_workers=max(1, min(len(states), MAX_SAMPLER_THREADS)), + thread_name_prefix='langbot-soak-sampler', + ) + + try: + while True: + now = time.monotonic() + if workload is None and workload_command and now - started_at >= startup_grace_seconds: + workload_started_at = now + try: + workload = subprocess.Popen( + list(workload_command), + stdin=subprocess.DEVNULL, + stdout=sys.stderr, + stderr=sys.stderr, + start_new_session=True, + ) + except OSError as exc: + workload_start_error = f'{type(exc).__name__}: {exc}' + break + + if workload is not None and workload_return_code is None: + polled = workload.poll() + if polled is not None: + workload_return_code = polled + workload_completed_at = now + + if now >= next_sample_at: + sample_wall_time = _utc_now() + sample_record: dict[str, Any] = { + 'schema_version': 1, + 'wall_time': sample_wall_time, + 'elapsed_seconds': now - started_at, + 'targets': {}, + } + in_gate_window = now - started_at >= startup_grace_seconds + sample_futures = [ + ( + state, + sampler_pool.submit( + _sample_target, + state.target, + http_timeout_seconds=http_timeout_seconds, + ), + ) + for state in states + ] + for state, sample_future in sample_futures: + state.attempted_samples += 1 + target_id = f'{state.target.kind}:{state.target.name}' + try: + metrics = sample_future.result() + except Exception as exc: + error = f'{type(exc).__name__}: {exc}' + sample_record['targets'][target_id] = {'error': error} + if in_gate_window: + state.failed_samples += 1 + state.consecutive_failures += 1 + state.max_consecutive_failures = max( + state.max_consecutive_failures, + state.consecutive_failures, + ) + if state.first_error is None: + state.first_error = error + else: + sample = MetricSample( + monotonic_seconds=now, + wall_time=sample_wall_time, + metrics=metrics, + ) + state.samples.append(sample) + state.successful_samples += 1 + state.consecutive_failures = 0 + state.last_metrics = metrics + if in_gate_window and state.baseline_metrics is None: + state.baseline_metrics = dict(metrics) + if in_gate_window: + for metric, value in metrics.items(): + previous_max = state.observed_max_metrics.get(metric) + if previous_max is None or value > previous_max: + state.observed_max_metrics[metric] = value + sample_record['targets'][target_id] = {'metrics': metrics} + _write_json_line(samples_stream, sample_record) + next_sample_at = max( + next_sample_at + sample_interval_seconds, + time.monotonic(), + ) + + now = time.monotonic() + if workload_command and workload_completed_at is not None: + if now >= min(deadline, workload_completed_at + cooldown_seconds): + break + elif now >= deadline: + if workload is not None and workload.poll() is None: + workload_timed_out = True + _terminate_workload(workload) + workload_return_code = workload.returncode + workload_completed_at = time.monotonic() + break + + sleep_seconds = min(max(next_sample_at - now, 0.01), 0.25) + time.sleep(sleep_seconds) + except KeyboardInterrupt: + interrupted = True + finally: + if workload is not None and workload.poll() is None: + _terminate_workload(workload) + workload_return_code = workload.returncode + workload_completed_at = time.monotonic() + if previous_sigterm_handler is not None: + signal.signal(signal.SIGTERM, previous_sigterm_handler) + sampler_pool.shutdown(wait=True, cancel_futures=True) + + completed_at = time.monotonic() + if workload_completed_at is not None: + analysis_start = workload_completed_at + else: + analysis_start = max( + started_at + startup_grace_seconds, + completed_at - analysis_window_seconds, + ) + gate = evaluate_gate( + states, + analysis_start_seconds=analysis_start, + thresholds=thresholds, + ) + failures = list(gate.failures) + warnings = list(gate.warnings) + if interrupted: + failures.append('soak was interrupted') + if workload_timed_out: + failures.append('workload exceeded the soak duration and was terminated') + if workload_command and workload is None: + if workload_start_error is not None: + failures.append(f'workload failed to start: {workload_start_error}') + else: + failures.append('workload did not start before the soak ended') + if workload_return_code not in (None, 0): + failures.append(f'workload exited with status {workload_return_code}') + + workload_result = WorkloadResult( + executable=workload_command[0] if workload_command else None, + argument_count=max(len(workload_command) - 1, 0) if workload_command else 0, + started_at_seconds=(workload_started_at - started_at if workload_started_at is not None else None), + completed_at_seconds=(workload_completed_at - started_at if workload_completed_at is not None else None), + return_code=workload_return_code, + timed_out=workload_timed_out, + ) + report = { + 'schema_version': 1, + 'verdict': 'pass' if not failures else 'fail', + 'started_at': started_at_wall, + 'completed_at': _utc_now(), + 'elapsed_seconds': completed_at - started_at, + 'analysis_start_seconds': analysis_start - started_at, + 'config': { + 'duration_seconds': duration_seconds, + 'sample_interval_seconds': sample_interval_seconds, + 'startup_grace_seconds': startup_grace_seconds, + 'cooldown_seconds': cooldown_seconds, + 'analysis_window_seconds': analysis_window_seconds, + 'http_timeout_seconds': http_timeout_seconds, + 'thresholds': dataclasses.asdict(thresholds), + 'targets': [ + { + 'name': target.name, + 'kind': target.kind, + 'location': target.location, + } + for target in targets + ], + }, + 'workload': dataclasses.asdict(workload_result), + 'failures': list(dict.fromkeys(failures)), + 'warnings': list(dict.fromkeys(warnings)), + 'targets': gate.targets, + } + return report, 0 if report['verdict'] == 'pass' else 1 + + +def _positive_float(value: str) -> float: + parsed = float(value) + if not math.isfinite(parsed) or parsed <= 0: + raise argparse.ArgumentTypeError('value must be a finite number greater than zero') + return parsed + + +def _nonnegative_float(value: str) -> float: + parsed = float(value) + if not math.isfinite(parsed) or parsed < 0: + raise argparse.ArgumentTypeError('value must be a finite number greater than or equal to zero') + return parsed + + +def _ratio(value: str) -> float: + parsed = float(value) + if not math.isfinite(parsed) or not 0 <= parsed <= 1: + raise argparse.ArgumentTypeError('ratio must be between zero and one') + return parsed + + +def create_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description='Monitor LangBot Cloud process/cgroup resources and enforce a post-load stability gate.', + ) + parser.add_argument('--endpoint', action='append', default=[], metavar='NAME=URL') + parser.add_argument('--cgroup', action='append', default=[], metavar='NAME=PATH') + parser.add_argument('--pid', action='append', default=[], metavar='NAME=PID') + parser.add_argument('--duration', type=parse_duration, default=parse_duration('24h')) + parser.add_argument('--sample-interval', type=parse_duration, default=parse_duration('15s')) + parser.add_argument('--startup-grace', type=parse_duration, default=parse_duration('5m')) + parser.add_argument('--cooldown', type=parse_duration, default=parse_duration('30m')) + parser.add_argument('--analysis-window', type=parse_duration, default=parse_duration('30m')) + parser.add_argument('--http-timeout', type=parse_duration, default=parse_duration('5s')) + parser.add_argument('--max-memory-growth-mib', type=_positive_float, default=64.0) + parser.add_argument('--max-memory-slope-mib-per-hour', type=_positive_float, default=32.0) + parser.add_argument('--max-tail-cpu-cores', type=_positive_float, default=0.5) + parser.add_argument('--max-throttled-period-ratio', type=_ratio, default=0.25) + parser.add_argument('--max-transient-gauge-growth', type=_nonnegative_float, default=0.0) + parser.add_argument( + '--allow-rejections', + action='store_true', + help='Do not fail when blocking-executor capacity rejection counters increase.', + ) + parser.add_argument( + '--require-hard-limits', + action='store_true', + help='Require finite CPU, memory, swap and PID limits on every cgroup target.', + ) + parser.add_argument( + '--max-event-loop-lag-ms', + type=_positive_float, + default=1000.0, + ) + parser.add_argument( + '--max-event-loop-p95-lag-ms', + type=_positive_float, + default=250.0, + ) + parser.add_argument( + '--allow-missing-event-loop-metrics', + action='store_true', + help='Permit endpoint targets that do not expose resources.event_loop.', + ) + parser.add_argument('--samples-file', type=Path) + parser.add_argument('--report-file', type=Path) + parser.add_argument( + '--workload', + nargs=argparse.REMAINDER, + help='Command to start after startup grace; its completion begins the cooldown tail.', + ) + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + parser = create_parser() + args = parser.parse_args(argv) + try: + targets = build_targets( + endpoints=args.endpoint, + cgroups=args.cgroup, + pids=args.pid, + ) + except ValueError as exc: + parser.error(str(exc)) + if not targets: + parser.error('at least one --endpoint, --cgroup or --pid target is required') + if len(targets) > MAX_TARGETS: + parser.error(f'at most {MAX_TARGETS} targets may be monitored') + if args.startup_grace >= args.duration: + parser.error('--startup-grace must be shorter than --duration') + if args.sample_interval >= args.duration: + parser.error('--sample-interval must be shorter than --duration') + workload_command = tuple(args.workload or ()) + if args.workload is not None and not workload_command: + parser.error('--workload requires a command') + thresholds = Thresholds( + max_memory_growth_bytes=int(args.max_memory_growth_mib * BYTES_PER_MIB), + max_memory_slope_bytes_per_hour=args.max_memory_slope_mib_per_hour * BYTES_PER_MIB, + max_tail_cpu_cores=args.max_tail_cpu_cores, + max_throttled_period_ratio=args.max_throttled_period_ratio, + allow_rejections=args.allow_rejections, + max_transient_gauge_growth=args.max_transient_gauge_growth, + require_hard_limits=args.require_hard_limits, + max_event_loop_lag_ms=args.max_event_loop_lag_ms, + max_event_loop_p95_lag_ms=args.max_event_loop_p95_lag_ms, + require_event_loop_metrics=not args.allow_missing_event_loop_metrics, + ) + samples_stream = _open_output(args.samples_file) + try: + report, exit_code = run_soak( + targets=targets, + duration_seconds=args.duration, + sample_interval_seconds=args.sample_interval, + startup_grace_seconds=args.startup_grace, + cooldown_seconds=args.cooldown, + analysis_window_seconds=args.analysis_window, + http_timeout_seconds=args.http_timeout, + thresholds=thresholds, + workload_command=workload_command or None, + samples_stream=samples_stream, + ) + finally: + if samples_stream is not None: + samples_stream.close() + rendered_report = json.dumps(report, indent=2, sort_keys=True) + if args.report_file is not None: + args.report_file.parent.mkdir(parents=True, exist_ok=True) + args.report_file.write_text(rendered_report + '\n', encoding='utf-8') + print(rendered_report) + return exit_code + + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/scripts/runtime_resource_probe.py b/scripts/runtime_resource_probe.py new file mode 100644 index 000000000..b7b61896e --- /dev/null +++ b/scripts/runtime_resource_probe.py @@ -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() diff --git a/scripts/workspace_runtime_capacity_probe.py b/scripts/workspace_runtime_capacity_probe.py new file mode 100644 index 000000000..4b41cc4ad --- /dev/null +++ b/scripts/workspace_runtime_capacity_probe.py @@ -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() diff --git a/src/langbot/libs/coze_server_api/client.py b/src/langbot/libs/coze_server_api/client.py index 54fb47495..6ce035c94 100644 --- a/src/langbot/libs/coze_server_api/client.py +++ b/src/langbot/libs/coze_server_api/client.py @@ -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:'): diff --git a/src/langbot/libs/deerflow_api/client.py b/src/langbot/libs/deerflow_api/client.py index b66bf7e2e..dbf25ead7 100644 --- a/src/langbot/libs/deerflow_api/client.py +++ b/src/langbot/libs/deerflow_api/client.py @@ -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: diff --git a/src/langbot/libs/dify_service_api/v1/client.py b/src/langbot/libs/dify_service_api/v1/client.py index a00eedbc1..d1317e00e 100644 --- a/src/langbot/libs/dify_service_api/v1/client.py +++ b/src/langbot/libs/dify_service_api/v1/client.py @@ -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) diff --git a/src/langbot/libs/dingtalk_api/api.py b/src/langbot/libs/dingtalk_api/api.py index 8d93ee27f..c0a7492ad 100644 --- a/src/langbot/libs/dingtalk_api/api.py +++ b/src/langbot/libs/dingtalk_api/api.py @@ -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 diff --git a/src/langbot/libs/official_account_api/api.py b/src/langbot/libs/official_account_api/api.py index b474205dc..38935f8a4 100644 --- a/src/langbot/libs/official_account_api/api.py +++ b/src/langbot/libs/official_account_api/api.py @@ -21,8 +21,14 @@ xml_template = """ """ +_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() diff --git a/src/langbot/libs/openclaw_weixin_api/client.py b/src/langbot/libs/openclaw_weixin_api/client.py index d713f2996..fa89c0427 100644 --- a/src/langbot/libs/openclaw_weixin_api/client.py +++ b/src/langbot/libs/openclaw_weixin_api/client.py @@ -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') diff --git a/src/langbot/libs/qq_official_api/api.py b/src/langbot/libs/qq_official_api/api.py index 0ba1917d9..dfca4378c 100644 --- a/src/langbot/libs/qq_official_api/api.py +++ b/src/langbot/libs/qq_official_api/api.py @@ -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 diff --git a/src/langbot/libs/slack_api/api.py b/src/langbot/libs/slack_api/api.py index 6f869b932..69d70b402 100644 --- a/src/langbot/libs/slack_api/api.py +++ b/src/langbot/libs/slack_api/api.py @@ -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'] diff --git a/src/langbot/libs/wechatpad_api/api/downloadpai.py b/src/langbot/libs/wechatpad_api/api/downloadpai.py index 3fbdb624a..b1c527e10 100644 --- a/src/langbot/libs/wechatpad_api/api/downloadpai.py +++ b/src/langbot/libs/wechatpad_api/api/downloadpai.py @@ -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') diff --git a/src/langbot/libs/wechatpad_api/util/http_util.py b/src/langbot/libs/wechatpad_api/util/http_util.py index 7390f43ec..312afb416 100644 --- a/src/langbot/libs/wechatpad_api/util/http_util.py +++ b/src/langbot/libs/wechatpad_api/util/http_util.py @@ -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: diff --git a/src/langbot/libs/wecom_ai_bot_api/api.py b/src/langbot/libs/wecom_ai_bot_api/api.py index 4e2c6a28a..63687943d 100644 --- a/src/langbot/libs/wecom_ai_bot_api/api.py +++ b/src/langbot/libs/wecom_ai_bot_api/api.py @@ -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'' - 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): diff --git a/src/langbot/libs/wecom_ai_bot_api/ws_client.py b/src/langbot/libs/wecom_ai_bot_api/ws_client.py index 32c11bd50..3fe985255 100644 --- a/src/langbot/libs/wecom_ai_bot_api/ws_client.py +++ b/src/langbot/libs/wecom_ai_bot_api/ws_client.py @@ -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 diff --git a/src/langbot/libs/wecom_api/api.py b/src/langbot/libs/wecom_api/api.py index b6c7ef0b7..72bd7ade7 100644 --- a/src/langbot/libs/wecom_api/api.py +++ b/src/langbot/libs/wecom_api/api.py @@ -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): diff --git a/src/langbot/libs/wecom_customer_service_api/api.py b/src/langbot/libs/wecom_customer_service_api/api.py index 70270b727..18561da0e 100644 --- a/src/langbot/libs/wecom_customer_service_api/api.py +++ b/src/langbot/libs/wecom_customer_service_api/api.py @@ -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() diff --git a/src/langbot/libs/weknora_api/client.py b/src/langbot/libs/weknora_api/client.py index f753136d8..c0a7724cb 100644 --- a/src/langbot/libs/weknora_api/client.py +++ b/src/langbot/libs/weknora_api/client.py @@ -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 diff --git a/src/langbot/pkg/api/http/controller/group.py b/src/langbot/pkg/api/http/controller/group.py index 33c7d0d06..1fc09a0b2 100644 --- a/src/langbot/pkg/api/http/controller/group.py +++ b/src/langbot/pkg/api/http/controller/group.py @@ -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( diff --git a/src/langbot/pkg/api/http/controller/groups/knowledge/migration.py b/src/langbot/pkg/api/http/controller/groups/knowledge/migration.py index 85e949603..0006d37e7 100644 --- a/src/langbot/pkg/api/http/controller/groups/knowledge/migration.py +++ b/src/langbot/pkg/api/http/controller/groups/knowledge/migration.py @@ -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}') diff --git a/src/langbot/pkg/api/http/controller/groups/pipelines/embed.py b/src/langbot/pkg/api/http/controller/groups/pipelines/embed.py index 599a09ba3..26f255e37 100644 --- a/src/langbot/pkg/api/http/controller/groups/pipelines/embed.py +++ b/src/langbot/pkg/api/http/controller/groups/pipelines/embed.py @@ -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: diff --git a/src/langbot/pkg/api/http/controller/groups/pipelines/websocket_chat.py b/src/langbot/pkg/api/http/controller/groups/pipelines/websocket_chat.py index b47aabe4c..300b9f81f 100644 --- a/src/langbot/pkg/api/http/controller/groups/pipelines/websocket_chat.py +++ b/src/langbot/pkg/api/http/controller/groups/pipelines/websocket_chat.py @@ -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//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: diff --git a/src/langbot/pkg/api/http/controller/groups/platform/adapters.py b/src/langbot/pkg/api/http/controller/groups/platform/adapters.py index fe8d8275d..edfbf9615 100644 --- a/src/langbot/pkg/api/http/controller/groups/platform/adapters.py +++ b/src/langbot/pkg/api/http/controller/groups/platform/adapters.py @@ -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. diff --git a/src/langbot/pkg/api/http/controller/groups/plugins.py b/src/langbot/pkg/api/http/controller/groups/plugins.py index eeccf3228..bbcda6c6b 100644 --- a/src/langbot/pkg/api/http/controller/groups/plugins.py +++ b/src/langbot/pkg/api/http/controller/groups/plugins.py @@ -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 diff --git a/src/langbot/pkg/api/http/controller/groups/survey.py b/src/langbot/pkg/api/http/controller/groups/survey.py index a65d51a85..3e73f1f83 100644 --- a/src/langbot/pkg/api/http/controller/groups/survey.py +++ b/src/langbot/pkg/api/http/controller/groups/survey.py @@ -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') diff --git a/src/langbot/pkg/api/http/controller/groups/webhooks.py b/src/langbot/pkg/api/http/controller/groups/webhooks.py index ad8e5782a..dbb9626cf 100644 --- a/src/langbot/pkg/api/http/controller/groups/webhooks.py +++ b/src/langbot/pkg/api/http/controller/groups/webhooks.py @@ -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( diff --git a/src/langbot/pkg/api/http/controller/main.py b/src/langbot/pkg/api/http/controller/main.py index 835617e16..3d20c4e71 100644 --- a/src/langbot/pkg/api/http/controller/main.py +++ b/src/langbot/pkg/api/http/controller/main.py @@ -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) diff --git a/src/langbot/pkg/api/http/service/maintenance.py b/src/langbot/pkg/api/http/service/maintenance.py index cb680ecf8..5d64daa6c 100644 --- a/src/langbot/pkg/api/http/service/maintenance.py +++ b/src/langbot/pkg/api/http/service/maintenance.py @@ -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): diff --git a/src/langbot/pkg/api/http/service/mcp.py b/src/langbot/pkg/api/http/service/mcp.py index 1527c40ba..09bc185f7 100644 --- a/src/langbot/pkg/api/http/service/mcp.py +++ b/src/langbot/pkg/api/http/service/mcp.py @@ -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( diff --git a/src/langbot/pkg/api/http/service/skill.py b/src/langbot/pkg/api/http/service/skill.py index 59dff4e50..9073ccd4b 100644 --- a/src/langbot/pkg/api/http/service/skill.py +++ b/src/langbot/pkg/api/http/service/skill.py @@ -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: diff --git a/src/langbot/pkg/api/http/service/space.py b/src/langbot/pkg/api/http/service/space.py index d280d5ff1..855b009c1 100644 --- a/src/langbot/pkg/api/http/service/space.py +++ b/src/langbot/pkg/api/http/service/space.py @@ -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', []) diff --git a/src/langbot/pkg/api/http/service/user.py b/src/langbot/pkg/api/http/service/user.py index 47bfe8c6c..87c9b8891 100644 --- a/src/langbot/pkg/api/http/service/user.py +++ b/src/langbot/pkg/api/http/service/user.py @@ -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. diff --git a/src/langbot/pkg/box/admission.py b/src/langbot/pkg/box/admission.py index dc46c34a1..8eab5b16c 100644 --- a/src/langbot/pkg/box/admission.py +++ b/src/langbot/pkg/box/admission.py @@ -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: diff --git a/src/langbot/pkg/box/service.py b/src/langbot/pkg/box/service.py index f4224fb0c..65022482e 100644 --- a/src/langbot/pkg/box/service.py +++ b/src/langbot/pkg/box/service.py @@ -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 diff --git a/src/langbot/pkg/cloud/directory_projection.py b/src/langbot/pkg/cloud/directory_projection.py index dcd8e9d7e..0430ef07e 100644 --- a/src/langbot/pkg/cloud/directory_projection.py +++ b/src/langbot/pkg/cloud/directory_projection.py @@ -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.""" diff --git a/src/langbot/pkg/cloud/entitlements.py b/src/langbot/pkg/cloud/entitlements.py index 3e9e7d134..b9476e6ba 100644 --- a/src/langbot/pkg/cloud/entitlements.py +++ b/src/langbot/pkg/cloud/entitlements.py @@ -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), + } diff --git a/src/langbot/pkg/cloud/launch.py b/src/langbot/pkg/cloud/launch.py index d37329724..ac5006870 100644 --- a/src/langbot/pkg/cloud/launch.py +++ b/src/langbot/pkg/cloud/launch.py @@ -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( diff --git a/src/langbot/pkg/core/app.py b/src/langbot/pkg/core/app.py index 333c657cb..fb461b056 100644 --- a/src/langbot/pkg/core/app.py +++ b/src/langbot/pkg/core/app.py @@ -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() diff --git a/src/langbot/pkg/core/boot.py b/src/langbot/pkg/core/boot.py index fb6919630..1b15e9922 100644 --- a/src/langbot/pkg/core/boot.py +++ b/src/langbot/pkg/core/boot.py @@ -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 diff --git a/src/langbot/pkg/core/errors.py b/src/langbot/pkg/core/errors.py new file mode 100644 index 000000000..24bf39bf7 --- /dev/null +++ b/src/langbot/pkg/core/errors.py @@ -0,0 +1,2 @@ +class TaskCapacityError(RuntimeError): + """Raised when the configured user-task admission limit is exhausted.""" diff --git a/src/langbot/pkg/core/stages/build_app.py b/src/langbot/pkg/core/stages/build_app.py index 21b543409..16dbd9462 100644 --- a/src/langbot/pkg/core/stages/build_app.py +++ b/src/langbot/pkg/core/stages/build_app.py @@ -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 diff --git a/src/langbot/pkg/core/stages/load_config.py b/src/langbot/pkg/core/stages/load_config.py index 37ba31ac9..5db354cc8 100644 --- a/src/langbot/pkg/core/stages/load_config.py +++ b/src/langbot/pkg/core/stages/load_config.py @@ -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 diff --git a/src/langbot/pkg/core/stages/show_notes.py b/src/langbot/pkg/core/stages/show_notes.py index d0f861ba3..912c0fb4b 100644 --- a/src/langbot/pkg/core/stages/show_notes.py +++ b/src/langbot/pkg/core/stages/show_notes.py @@ -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 diff --git a/src/langbot/pkg/core/task_boundary.py b/src/langbot/pkg/core/task_boundary.py index 412e47bc4..c9e652be2 100644 --- a/src/langbot/pkg/core/task_boundary.py +++ b/src/langbot/pkg/core/task_boundary.py @@ -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()) diff --git a/src/langbot/pkg/core/taskmgr.py b/src/langbot/pkg/core/taskmgr.py index 54cb4dc11..25cc38e0d 100644 --- a/src/langbot/pkg/core/taskmgr.py +++ b/src/langbot/pkg/core/taskmgr.py @@ -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', diff --git a/src/langbot/pkg/discover/engine.py b/src/langbot/pkg/discover/engine.py index 713420d14..273dcbf34 100644 --- a/src/langbot/pkg/discover/engine.py +++ b/src/langbot/pkg/discover/engine.py @@ -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: """加载组件清单""" diff --git a/src/langbot/pkg/persistence/databases/postgresql.py b/src/langbot/pkg/persistence/databases/postgresql.py index 0dae7899a..3e99e2129 100644 --- a/src/langbot/pkg/persistence/databases/postgresql.py +++ b/src/langbot/pkg/persistence/databases/postgresql.py @@ -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, + ) diff --git a/src/langbot/pkg/persistence/mgr.py b/src/langbot/pkg/persistence/mgr.py index e35cd6704..749a3d04e 100644 --- a/src/langbot/pkg/persistence/mgr.py +++ b/src/langbot/pkg/persistence/mgr.py @@ -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.""" diff --git a/src/langbot/pkg/persistence/release_migration.py b/src/langbot/pkg/persistence/release_migration.py index 4e169ad8d..0834c0d28 100644 --- a/src/langbot/pkg/persistence/release_migration.py +++ b/src/langbot/pkg/persistence/release_migration.py @@ -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: diff --git a/src/langbot/pkg/pipeline/aggregator.py b/src/langbot/pkg/pipeline/aggregator.py index 34101c136..de9b206ca 100644 --- a/src/langbot/pkg/pipeline/aggregator.py +++ b/src/langbot/pkg/pipeline/aggregator.py @@ -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) diff --git a/src/langbot/pkg/pipeline/cntfilter/filters/baiduexamine.py b/src/langbot/pkg/pipeline/cntfilter/filters/baiduexamine.py index a376310f6..8c74150b0 100644 --- a/src/langbot/pkg/pipeline/cntfilter/filters/baiduexamine.py +++ b/src/langbot/pkg/pipeline/cntfilter/filters/baiduexamine.py @@ -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( diff --git a/src/langbot/pkg/pipeline/cntfilter/filters/banwords.py b/src/langbot/pkg/pipeline/cntfilter/filters/banwords.py index 05b25013d..39796b33c 100644 --- a/src/langbot/pkg/pipeline/cntfilter/filters/banwords.py +++ b/src/langbot/pkg/pipeline/cntfilter/filters/banwords.py @@ -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, diff --git a/src/langbot/pkg/pipeline/cntfilter/filters/cntignore.py b/src/langbot/pkg/pipeline/cntfilter/filters/cntignore.py index 731ab3924..1b6515891 100644 --- a/src/langbot/pkg/pipeline/cntfilter/filters/cntignore.py +++ b/src/langbot/pkg/pipeline/cntfilter/filters/cntignore.py @@ -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, diff --git a/src/langbot/pkg/pipeline/controller.py b/src/langbot/pkg/pipeline/controller.py index d897ab15e..d14288c29 100644 --- a/src/langbot/pkg/pipeline/controller.py +++ b/src/langbot/pkg/pipeline/controller.py @@ -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): """运行控制器""" diff --git a/src/langbot/pkg/pipeline/longtext/strategies/image.py b/src/langbot/pkg/pipeline/longtext/strategies/image.py index 110f1f81c..10be8998c 100644 --- a/src/langbot/pkg/pipeline/longtext/strategies/image.py +++ b/src/langbot/pkg/pipeline/longtext/strategies/image.py @@ -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 diff --git a/src/langbot/pkg/pipeline/pipelinemgr.py b/src/langbot/pkg/pipeline/pipelinemgr.py index bde3b5228..911235d6a 100644 --- a/src/langbot/pkg/pipeline/pipelinemgr.py +++ b/src/langbot/pkg/pipeline/pipelinemgr.py @@ -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 diff --git a/src/langbot/pkg/pipeline/plugin_diagnostics.py b/src/langbot/pkg/pipeline/plugin_diagnostics.py index 3e1195cf2..c019d9548 100644 --- a/src/langbot/pkg/pipeline/plugin_diagnostics.py +++ b/src/langbot/pkg/pipeline/plugin_diagnostics.py @@ -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( diff --git a/src/langbot/pkg/pipeline/pool.py b/src/langbot/pkg/pipeline/pool.py index fecdf84d3..bc96b054f 100644 --- a/src/langbot/pkg/pipeline/pool.py +++ b/src/langbot/pkg/pipeline/pool.py @@ -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): diff --git a/src/langbot/pkg/pipeline/process/handlers/chat.py b/src/langbot/pkg/pipeline/process/handlers/chat.py index a15226a9a..ef6dd874d 100644 --- a/src/langbot/pkg/pipeline/process/handlers/chat.py +++ b/src/langbot/pkg/pipeline/process/handlers/chat.py @@ -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() diff --git a/src/langbot/pkg/pipeline/ratelimit/algos/fixedwin.py b/src/langbot/pkg/pipeline/ratelimit/algos/fixedwin.py index 6a2a8e97d..278550290 100644 --- a/src/langbot/pkg/pipeline/ratelimit/algos/fixedwin.py +++ b/src/langbot/pkg/pipeline/ratelimit/algos/fixedwin.py @@ -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( diff --git a/src/langbot/pkg/pipeline/resprule/rules/regexp.py b/src/langbot/pkg/pipeline/resprule/rules/regexp.py index 41e1df8e7..b1b13066d 100644 --- a/src/langbot/pkg/pipeline/resprule/rules/regexp.py +++ b/src/langbot/pkg/pipeline/resprule/rules/regexp.py @@ -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) diff --git a/src/langbot/pkg/platform/botmgr.py b/src/langbot/pkg/platform/botmgr.py index e8ce880f4..feaeea7a9 100644 --- a/src/langbot/pkg/platform/botmgr.py +++ b/src/langbot/pkg/platform/botmgr.py @@ -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) diff --git a/src/langbot/pkg/platform/logger.py b/src/langbot/pkg/platform/logger.py index 8fcdc1795..af6a7de76 100644 --- a/src/langbot/pkg/platform/logger.py +++ b/src/langbot/pkg/platform/logger.py @@ -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 = [] diff --git a/src/langbot/pkg/platform/sources/aiocqhttp.py b/src/langbot/pkg/platform/sources/aiocqhttp.py index 4b994c2be..6fc0ca918 100644 --- a/src/langbot/pkg/platform/sources/aiocqhttp.py +++ b/src/langbot/pkg/platform/sources/aiocqhttp.py @@ -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 diff --git a/src/langbot/pkg/platform/sources/dingtalk.py b/src/langbot/pkg/platform/sources/dingtalk.py index 996187f08..ec9de538e 100644 --- a/src/langbot/pkg/platform/sources/dingtalk.py +++ b/src/langbot/pkg/platform/sources/dingtalk.py @@ -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, diff --git a/src/langbot/pkg/platform/sources/discord.py b/src/langbot/pkg/platform/sources/discord.py index fae65f6a0..1be977d61 100644 --- a/src/langbot/pkg/platform/sources/discord.py +++ b/src/langbot/pkg/platform/sources/discord.py @@ -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 diff --git a/src/langbot/pkg/platform/sources/http_bot.py b/src/langbot/pkg/platform/sources/http_bot.py index 6a95bb460..2b8abfbd5 100644 --- a/src/langbot/pkg/platform/sources/http_bot.py +++ b/src/langbot/pkg/platform/sources/http_bot.py @@ -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( { diff --git a/src/langbot/pkg/platform/sources/kook.py b/src/langbot/pkg/platform/sources/kook.py index 5a6bade36..bb895d2e3 100644 --- a/src/langbot/pkg/platform/sources/kook.py +++ b/src/langbot/pkg/platform/sources/kook.py @@ -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: diff --git a/src/langbot/pkg/platform/sources/lark.py b/src/langbot/pkg/platform/sources/lark.py index 96e40469c..3c3476159 100644 --- a/src/langbot/pkg/platform/sources/lark.py +++ b/src/langbot/pkg/platform/sources/lark.py @@ -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 diff --git a/src/langbot/pkg/platform/sources/legacy/gewechat.py b/src/langbot/pkg/platform/sources/legacy/gewechat.py index 68e1bdedd..ee1f40281 100644 --- a/src/langbot/pkg/platform/sources/legacy/gewechat.py +++ b/src/langbot/pkg/platform/sources/legacy/gewechat.py @@ -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: diff --git a/src/langbot/pkg/platform/sources/legacy/nakuru.py b/src/langbot/pkg/platform/sources/legacy/nakuru.py index 1e34af0b9..d69aa2d94 100644 --- a/src/langbot/pkg/platform/sources/legacy/nakuru.py +++ b/src/langbot/pkg/platform/sources/legacy/nakuru.py @@ -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() diff --git a/src/langbot/pkg/platform/sources/legacy/qqbotpy.py b/src/langbot/pkg/platform/sources/legacy/qqbotpy.py index 3a55482c9..21ac5703e 100644 --- a/src/langbot/pkg/platform/sources/legacy/qqbotpy.py +++ b/src/langbot/pkg/platform/sources/legacy/qqbotpy.py @@ -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 diff --git a/src/langbot/pkg/platform/sources/line.py b/src/langbot/pkg/platform/sources/line.py index 3d0f75c7d..496a5ed81 100644 --- a/src/langbot/pkg/platform/sources/line.py +++ b/src/langbot/pkg/platform/sources/line.py @@ -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 diff --git a/src/langbot/pkg/platform/sources/matrix.py b/src/langbot/pkg/platform/sources/matrix.py index da159223f..7cf39dd86 100644 --- a/src/langbot/pkg/platform/sources/matrix.py +++ b/src/langbot/pkg/platform/sources/matrix.py @@ -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') diff --git a/src/langbot/pkg/platform/sources/officialaccount.py b/src/langbot/pkg/platform/sources/officialaccount.py index 288991d62..1506ae8ec 100644 --- a/src/langbot/pkg/platform/sources/officialaccount.py +++ b/src/langbot/pkg/platform/sources/officialaccount.py @@ -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( diff --git a/src/langbot/pkg/platform/sources/openclaw_weixin.py b/src/langbot/pkg/platform/sources/openclaw_weixin.py index dbfa595ee..4726888b1 100644 --- a/src/langbot/pkg/platform/sources/openclaw_weixin.py +++ b/src/langbot/pkg/platform/sources/openclaw_weixin.py @@ -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 diff --git a/src/langbot/pkg/platform/sources/qqofficial.py b/src/langbot/pkg/platform/sources/qqofficial.py index 35fd81f76..bd4f97feb 100644 --- a/src/langbot/pkg/platform/sources/qqofficial.py +++ b/src/langbot/pkg/platform/sources/qqofficial.py @@ -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; ' diff --git a/src/langbot/pkg/platform/sources/satori.py b/src/langbot/pkg/platform/sources/satori.py index 6bfa342ef..aff7bbf21 100644 --- a/src/langbot/pkg/platform/sources/satori.py +++ b/src/langbot/pkg/platform/sources/satori.py @@ -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() diff --git a/src/langbot/pkg/platform/sources/telegram.py b/src/langbot/pkg/platform/sources/telegram.py index e7ad626bb..107894094 100644 --- a/src/langbot/pkg/platform/sources/telegram.py +++ b/src/langbot/pkg/platform/sources/telegram.py @@ -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 diff --git a/src/langbot/pkg/platform/sources/websocket_adapter.py b/src/langbot/pkg/platform/sources/websocket_adapter.py index 9f9d9758d..dcda32120 100644 --- a/src/langbot/pkg/platform/sources/websocket_adapter.py +++ b/src/langbot/pkg/platform/sources/websocket_adapter.py @@ -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 = ( diff --git a/src/langbot/pkg/platform/sources/websocket_manager.py b/src/langbot/pkg/platform/sources/websocket_manager.py index 216043e9a..99ef98bba 100644 --- a/src/langbot/pkg/platform/sources/websocket_manager.py +++ b/src/langbot/pkg/platform/sources/websocket_manager.py @@ -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}') diff --git a/src/langbot/pkg/platform/sources/wechatpad.py b/src/langbot/pkg/platform/sources/wechatpad.py index 1c9c5ee27..9b3e40d7d 100644 --- a/src/langbot/pkg/platform/sources/wechatpad.py +++ b/src/langbot/pkg/platform/sources/wechatpad.py @@ -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 diff --git a/src/langbot/pkg/platform/sources/wecom.py b/src/langbot/pkg/platform/sources/wecom.py index 27c77ad37..bb8ba2b17 100644 --- a/src/langbot/pkg/platform/sources/wecom.py +++ b/src/langbot/pkg/platform/sources/wecom.py @@ -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( diff --git a/src/langbot/pkg/platform/sources/wecombot.py b/src/langbot/pkg/platform/sources/wecombot.py index d0febad9e..0990773ad 100644 --- a/src/langbot/pkg/platform/sources/wecombot.py +++ b/src/langbot/pkg/platform/sources/wecombot.py @@ -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( diff --git a/src/langbot/pkg/platform/sources/wecomcs.py b/src/langbot/pkg/platform/sources/wecomcs.py index ea7d8ef5c..830255a42 100644 --- a/src/langbot/pkg/platform/sources/wecomcs.py +++ b/src/langbot/pkg/platform/sources/wecomcs.py @@ -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: diff --git a/src/langbot/pkg/platform/webhook_pusher.py b/src/langbot/pkg/platform/webhook_pusher.py index 79eb1de24..a6af2284a 100644 --- a/src/langbot/pkg/platform/webhook_pusher.py +++ b/src/langbot/pkg/platform/webhook_pusher.py @@ -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 diff --git a/src/langbot/pkg/plugin/archive.py b/src/langbot/pkg/plugin/archive.py new file mode 100644 index 000000000..f3a8511a4 --- /dev/null +++ b/src/langbot/pkg/plugin/archive.py @@ -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 diff --git a/src/langbot/pkg/plugin/connector.py b/src/langbot/pkg/plugin/connector.py index b92bd2382..00081a43a 100644 --- a/src/langbot/pkg/plugin/connector.py +++ b/src/langbot/pkg/plugin/connector.py @@ -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 {} diff --git a/src/langbot/pkg/plugin/handler.py b/src/langbot/pkg/plugin/handler.py index 99ca0af34..bfc1c575f 100644 --- a/src/langbot/pkg/plugin/handler.py +++ b/src/langbot/pkg/plugin/handler.py @@ -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, } diff --git a/src/langbot/pkg/provider/modelmgr/modelmgr.py b/src/langbot/pkg/provider/modelmgr/modelmgr.py index bfbbfbcbc..925b7de95 100644 --- a/src/langbot/pkg/provider/modelmgr/modelmgr.py +++ b/src/langbot/pkg/provider/modelmgr/modelmgr.py @@ -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: diff --git a/src/langbot/pkg/provider/modelmgr/requester.py b/src/langbot/pkg/provider/modelmgr/requester.py index 7bc6b43f2..0c8232dd5 100644 --- a/src/langbot/pkg/provider/modelmgr/requester.py +++ b/src/langbot/pkg/provider/modelmgr/requester.py @@ -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. diff --git a/src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py b/src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py index efad67250..982f5609b 100644 --- a/src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py +++ b/src/langbot/pkg/provider/modelmgr/requesters/litellmchat.py @@ -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', []): diff --git a/src/langbot/pkg/provider/runner.py b/src/langbot/pkg/provider/runner.py index 987b3a0e9..ca6b86017 100644 --- a/src/langbot/pkg/provider/runner.py +++ b/src/langbot/pkg/provider/runner.py @@ -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.""" diff --git a/src/langbot/pkg/provider/runners/cozeapi.py b/src/langbot/pkg/provider/runners/cozeapi.py index 26980f81e..a7ca53408 100644 --- a/src/langbot/pkg/provider/runners/cozeapi.py +++ b/src/langbot/pkg/provider/runners/cozeapi.py @@ -2,7 +2,6 @@ from __future__ import annotations import typing import json -import base64 from langbot.pkg.provider import runner from langbot.pkg.core import app @@ -11,6 +10,16 @@ from langbot.pkg.utils import image import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query from langbot.libs.coze_server_api.client import AsyncCozeAPIClient +_MAX_COZE_GENERATED_CHARS = 1024 * 1024 +_MAX_COZE_MEDIA_BYTES = 10 * 1024 * 1024 + + +def _append_bounded(current: str, addition: typing.Any) -> str: + addition = str(addition or '') + if len(current) + len(addition) > _MAX_COZE_GENERATED_CHARS: + raise ValueError('Coze response exceeds the runtime limit') + return current + addition + @runner.runner_class('coze-api') class CozeAPIRunner(runner.RequestRunner): @@ -77,7 +86,10 @@ class CozeAPIRunner(runner.RequestRunner): content_parts.append({'type': 'text', 'text': ce.text}) elif ce.type == 'image_base64': image_b64, image_format = await image.extract_b64_and_format(ce.image_base64) - file_bytes = base64.b64decode(image_b64) + file_bytes = await image.decode_base64_limited( + image_b64, + max_bytes=_MAX_COZE_MEDIA_BYTES, + ) file_id = await self._get_file_id(file_bytes) content_parts.append({'type': 'image', 'file_id': file_id}) elif ce.type == 'file': @@ -144,7 +156,7 @@ class CozeAPIRunner(runner.RequestRunner): auto_save_history=self.auto_save_history, stream=True, ): - self.ap.logger.debug(f'coze-chat-stream: {chunk}') + self.ap.logger.debug(f'coze-chat-stream: {str(chunk)[:1000]}') event_type = chunk.get('event') data = chunk.get('data', {}) @@ -153,11 +165,17 @@ class CozeAPIRunner(runner.RequestRunner): if event_type == 'conversation.message.delta': # 收集内容 if 'content' in data: - full_content += data.get('content', '') + full_content = _append_bounded( + full_content, + data.get('content', ''), + ) # 收集推理内容（如果有） if 'reasoning_content' in data: - full_reasoning += data.get('reasoning_content', '') + full_reasoning = _append_bounded( + full_reasoning, + data.get('reasoning_content', ''), + ) elif event_type.split('.')[-1] == 'done': # 本地部署coze时，结束event不为done # 保存会话ID @@ -179,6 +197,8 @@ class CozeAPIRunner(runner.RequestRunner): remove_think = self.pipeline_config.get('output', {}).get('misc', {}).get('remove-think', False) if not remove_think: content = f'\n{full_reasoning}\n\n{content}'.strip() + if len(content) > _MAX_COZE_GENERATED_CHARS: + raise ValueError('Coze response exceeds the runtime limit') # 一次性返回完整内容 yield provider_message.Message( @@ -227,7 +247,7 @@ class CozeAPIRunner(runner.RequestRunner): auto_save_history=self.auto_save_history, stream=True, ): - self.ap.logger.debug(f'coze-chat-stream-chunk: {chunk}') + self.ap.logger.debug(f'coze-chat-stream-chunk: {str(chunk)[:1000]}') event_type = chunk.get('event') data = chunk.get('data', {}) @@ -263,7 +283,7 @@ class CozeAPIRunner(runner.RequestRunner): error_msg = f'Coze API错误: {data.get("message", "未知错误")}' yield provider_message.MessageChunk(role='assistant', content=error_msg, finish_reason='error') return - full_content += content + full_content = _append_bounded(full_content, content) if message_idx % 8 == 0 or is_final: if full_content: yield provider_message.MessageChunk(role='assistant', content=full_content, is_final=is_final) @@ -286,3 +306,6 @@ class CozeAPIRunner(runner.RequestRunner): else: async for msg in self._chat_messages(query): yield msg + + async def aclose(self) -> None: + await self.coze.close() diff --git a/src/langbot/pkg/provider/runners/dashscopeapi.py b/src/langbot/pkg/provider/runners/dashscopeapi.py index a2c593ccc..00af120fb 100644 --- a/src/langbot/pkg/provider/runners/dashscopeapi.py +++ b/src/langbot/pkg/provider/runners/dashscopeapi.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import typing import re @@ -10,6 +11,9 @@ from ...core import app import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query import langbot_plugin.api.entities.builtin.provider.message as provider_message +_MAX_DASHSCOPE_RESPONSE_CHARS = 1024 * 1024 +_MAX_DASHSCOPE_REFERENCES = 1024 + class DashscopeAPIError(Exception): """Dashscope API 请求失败""" @@ -19,6 +23,13 @@ class DashscopeAPIError(Exception): super().__init__(self.message) +def _append_bounded(current: str, addition: typing.Any) -> str: + addition = str(addition or '') + if len(current) + len(addition) > _MAX_DASHSCOPE_RESPONSE_CHARS: + raise DashscopeAPIError('Dashscope response exceeds the runtime limit') + return current + addition + + @runner.runner_class('dashscope-app-api') class DashScopeAPIRunner(runner.RequestRunner): "阿里云百炼DashsscopeAPI对话请求器" @@ -111,18 +122,16 @@ class DashScopeAPIRunner(runner.RequestRunner): if remove_think: has_thoughts = False # 发送对话请求 - response = dashscope.Application.call( - api_key=self.api_key, # 智能体应用的API Key - app_id=self.app_id, # 智能体应用的ID - prompt=plain_text, # 用户输入的文本信息 - stream=True, # 流式输出 - incremental_output=True, # 增量输出，使用流式输出需要开启增量输出 - session_id=query.session.using_conversation.uuid, # 会话ID用于，多轮对话 + response = await asyncio.to_thread( + dashscope.Application.call, + api_key=self.api_key, + app_id=self.app_id, + prompt=plain_text, + stream=True, + incremental_output=True, + session_id=query.session.using_conversation.uuid, enable_thinking=has_thoughts, has_thoughts=has_thoughts, - # rag_options={ # 主要用于文件交互，暂不支持 - # "session_file_ids": ["FILE_ID1"], # FILE_ID1 替换为实际的临时文件ID,逗号隔开多个 - # } ) idx_chunk = 0 try: @@ -131,7 +140,7 @@ class DashScopeAPIRunner(runner.RequestRunner): except AttributeError: is_stream = False if is_stream: - for chunk in response: + async for chunk in runner.iterate_sync(response): if chunk.get('status_code') != 200: raise DashscopeAPIError( f'Dashscope API 请求失败: status_code={chunk.get("status_code")} message={chunk.get("message")} request_id={chunk.get("request_id")} ' @@ -145,15 +154,27 @@ class DashScopeAPIRunner(runner.RequestRunner): if stream_think and stream_think[0].get('thought'): if not think_start: think_start = True - pending_content += f'\n{stream_think[0].get("thought")}' + pending_content = _append_bounded( + pending_content, + f'\n{stream_think[0].get("thought")}', + ) else: # 继续输出 reasoning_content - pending_content += stream_think[0].get('thought') + pending_content = _append_bounded( + pending_content, + stream_think[0].get('thought'), + ) elif think_start and (not stream_think or stream_think[0].get('thought') == '') and not think_end: think_end = True - pending_content += '\n\n' + pending_content = _append_bounded( + pending_content, + '\n\n', + ) if stream_output.get('text') is not None: - pending_content += stream_output.get('text') + pending_content = _append_bounded( + pending_content, + stream_output.get('text'), + ) # 是否是流式最后一个chunk is_final = False if stream_output.get('finish_reason', False) == 'null' else True @@ -162,12 +183,14 @@ class DashScopeAPIRunner(runner.RequestRunner): # 从模型传出的参考资料信息中提取用于替换的字典 if references_dict_list is not None: - for doc in references_dict_list: + for doc in references_dict_list[:_MAX_DASHSCOPE_REFERENCES]: if doc.get('index_id') is not None: references_dict[doc.get('index_id')] = doc.get('doc_name') # 将参考资料替换到文本中 pending_content = self._replace_references(pending_content, references_dict) + if len(pending_content) > _MAX_DASHSCOPE_RESPONSE_CHARS: + raise DashscopeAPIError('Dashscope response exceeds the runtime limit') if idx_chunk % 8 == 0 or is_final: yield provider_message.MessageChunk( @@ -178,7 +201,7 @@ class DashScopeAPIRunner(runner.RequestRunner): # 保存当前会话的session_id用于下次对话的语境 query.session.using_conversation.uuid = stream_output.get('session_id') else: - for chunk in response: + async for chunk in runner.iterate_sync(response): if chunk.get('status_code') != 200: raise DashscopeAPIError( f'Dashscope API 请求失败: status_code={chunk.get("status_code")} message={chunk.get("message")} request_id={chunk.get("request_id")} ' @@ -192,15 +215,27 @@ class DashScopeAPIRunner(runner.RequestRunner): if stream_think and stream_think[0].get('thought'): if not think_start: think_start = True - pending_content += f'\n{stream_think[0].get("thought")}' + pending_content = _append_bounded( + pending_content, + f'\n{stream_think[0].get("thought")}', + ) else: # 继续输出 reasoning_content - pending_content += stream_think[0].get('thought') + pending_content = _append_bounded( + pending_content, + stream_think[0].get('thought'), + ) elif think_start and (not stream_think or stream_think[0].get('thought') == '') and not think_end: think_end = True - pending_content += '\n\n' + pending_content = _append_bounded( + pending_content, + '\n\n', + ) if stream_output.get('text') is not None: - pending_content += stream_output.get('text') + pending_content = _append_bounded( + pending_content, + stream_output.get('text'), + ) # 保存当前会话的session_id用于下次对话的语境 query.session.using_conversation.uuid = stream_output.get('session_id') @@ -210,12 +245,14 @@ class DashScopeAPIRunner(runner.RequestRunner): # 从模型传出的参考资料信息中提取用于替换的字典 if references_dict_list is not None: - for doc in references_dict_list: + for doc in references_dict_list[:_MAX_DASHSCOPE_REFERENCES]: if doc.get('index_id') is not None: references_dict[doc.get('index_id')] = doc.get('doc_name') # 将参考资料替换到文本中 - pending_content = self._replace_references(pending_content, references_dict) + pending_content = self._replace_references(pending_content, references_dict) + if len(pending_content) > _MAX_DASHSCOPE_RESPONSE_CHARS: + raise DashscopeAPIError('Dashscope response exceeds the runtime limit') yield provider_message.Message( role='assistant', @@ -240,18 +277,16 @@ class DashScopeAPIRunner(runner.RequestRunner): biz_params.update(query.variables) # 发送对话请求 - response = dashscope.Application.call( - api_key=self.api_key, # 智能体应用的API Key - app_id=self.app_id, # 智能体应用的ID - prompt=plain_text, # 用户输入的文本信息 - stream=True, # 流式输出 - incremental_output=True, # 增量输出，使用流式输出需要开启增量输出 - session_id=query.session.using_conversation.uuid, # 会话ID用于，多轮对话 - biz_params=biz_params, # 工作流应用的自定义输入参数传递 - flow_stream_mode='message_format', # 消息模式，输出/结束节点的流式结果 - # rag_options={ # 主要用于文件交互，暂不支持 - # "session_file_ids": ["FILE_ID1"], # FILE_ID1 替换为实际的临时文件ID,逗号隔开多个 - # } + response = await asyncio.to_thread( + dashscope.Application.call, + api_key=self.api_key, + app_id=self.app_id, + prompt=plain_text, + stream=True, + incremental_output=True, + session_id=query.session.using_conversation.uuid, + biz_params=biz_params, + flow_stream_mode='message_format', ) # 处理API返回的流式输出 @@ -262,7 +297,7 @@ class DashScopeAPIRunner(runner.RequestRunner): is_stream = False idx_chunk = 0 if is_stream: - for chunk in response: + async for chunk in runner.iterate_sync(response): if chunk.get('status_code') != 200: raise DashscopeAPIError( f'Dashscope API 请求失败: status_code={chunk.get("status_code")} message={chunk.get("message")} request_id={chunk.get("request_id")} ' @@ -273,7 +308,10 @@ class DashScopeAPIRunner(runner.RequestRunner): # 获取流式传输的output stream_output = chunk.get('output', {}) if stream_output.get('workflow_message') is not None: - pending_content += stream_output.get('workflow_message').get('message').get('content') + pending_content = _append_bounded( + pending_content, + stream_output.get('workflow_message').get('message').get('content'), + ) # if stream_output.get('text') is not None: # pending_content += stream_output.get('text') @@ -284,12 +322,14 @@ class DashScopeAPIRunner(runner.RequestRunner): # 从模型传出的参考资料信息中提取用于替换的字典 if references_dict_list is not None: - for doc in references_dict_list: + for doc in references_dict_list[:_MAX_DASHSCOPE_REFERENCES]: if doc.get('index_id') is not None: references_dict[doc.get('index_id')] = doc.get('doc_name') # 将参考资料替换到文本中 pending_content = self._replace_references(pending_content, references_dict) + if len(pending_content) > _MAX_DASHSCOPE_RESPONSE_CHARS: + raise DashscopeAPIError('Dashscope response exceeds the runtime limit') if idx_chunk % 8 == 0 or is_final: yield provider_message.MessageChunk( role='assistant', @@ -301,7 +341,7 @@ class DashScopeAPIRunner(runner.RequestRunner): query.session.using_conversation.uuid = stream_output.get('session_id') else: - for chunk in response: + async for chunk in runner.iterate_sync(response): if chunk.get('status_code') != 200: raise DashscopeAPIError( f'Dashscope API 请求失败: status_code={chunk.get("status_code")} message={chunk.get("message")} request_id={chunk.get("request_id")} ' @@ -312,7 +352,10 @@ class DashScopeAPIRunner(runner.RequestRunner): # 获取流式传输的output stream_output = chunk.get('output', {}) if stream_output.get('text') is not None: - pending_content += stream_output.get('text') + pending_content = _append_bounded( + pending_content, + stream_output.get('text'), + ) is_final = False if stream_output.get('finish_reason', False) == 'null' else True @@ -324,12 +367,14 @@ class DashScopeAPIRunner(runner.RequestRunner): # 从模型传出的参考资料信息中提取用于替换的字典 if references_dict_list is not None: - for doc in references_dict_list: + for doc in references_dict_list[:_MAX_DASHSCOPE_REFERENCES]: if doc.get('index_id') is not None: references_dict[doc.get('index_id')] = doc.get('doc_name') # 将参考资料替换到文本中 - pending_content = self._replace_references(pending_content, references_dict) + pending_content = self._replace_references(pending_content, references_dict) + if len(pending_content) > _MAX_DASHSCOPE_RESPONSE_CHARS: + raise DashscopeAPIError('Dashscope response exceeds the runtime limit') yield provider_message.Message( role='assistant', diff --git a/src/langbot/pkg/provider/runners/difysvapi.py b/src/langbot/pkg/provider/runners/difysvapi.py index 6846e587a..d3c175080 100644 --- a/src/langbot/pkg/provider/runners/difysvapi.py +++ b/src/langbot/pkg/provider/runners/difysvapi.py @@ -1,10 +1,11 @@ from __future__ import annotations +import asyncio +import heapq import typing import json import time import uuid -import base64 import mimetypes import os import re @@ -16,11 +17,9 @@ from langbot.pkg.provider import runner from langbot.pkg.core import app import langbot_plugin.api.entities.builtin.provider.message as provider_message import langbot_plugin.api.entities.builtin.platform.message as platform_message -from langbot.pkg.utils import image +from langbot.pkg.utils import httpclient, image import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query from langbot.libs.dify_service_api.v1 import client, errors -import httpx - # Module-level store for paused-workflow form state. The key includes the full # execution scope before the bot, pipeline, adapter, and launcher dimensions; @@ -29,8 +28,29 @@ import httpx # placement generations. PendingFormKey = tuple[str, str, int, str, str, str, str, str] _PENDING_FORMS: dict[PendingFormKey, 'OrderedDict[str, dict[str, typing.Any]]'] = {} +_PENDING_FORM_EXPIRY_HEAP: list[tuple[float, int, PendingFormKey, str]] = [] +_PENDING_FORM_ACTIVE_COUNT = 0 +_PENDING_FORM_REVISION = 0 _PENDING_FORM_DEFAULT_TTL = 30 * 60 # 30 minutes safety cap +_PENDING_FORM_MAX_SESSIONS = 4096 +_PENDING_FORM_MAX_PER_SESSION = 16 +_PENDING_FORM_HEAP_COMPACT_FLOOR = 64 +_PENDING_FORM_HEAP_MAX_MULTIPLIER = 4 +_PENDING_FORM_REVISION_KEY = '_langbot_cache_revision' _STREAM_FORM_PLACEHOLDER = '\u200b' +_MAX_DIFY_UPLOAD_BYTES = 10 * 1024 * 1024 + + +def _read_local_file_limited(path: str) -> bytes: + """Read a local platform attachment without allowing an oversized allocation.""" + + if os.path.getsize(path) > _MAX_DIFY_UPLOAD_BYTES: + raise ValueError('Dify upload file exceeds the size limit') + with open(path, 'rb') as file: + content = file.read(_MAX_DIFY_UPLOAD_BYTES + 1) + if len(content) > _MAX_DIFY_UPLOAD_BYTES: + raise ValueError('Dify upload file exceeds the size limit') + return content def _merge_stream_text(accumulated: str, incoming: typing.Any) -> str: @@ -65,19 +85,100 @@ def _session_key_from_query(query: pipeline_query.Query) -> PendingFormKey: ) +def _synchronize_pending_form_cache_if_externally_cleared() -> None: + """Keep test/debug direct cache clears from retaining stale heap entries.""" + + global _PENDING_FORM_ACTIVE_COUNT + if _PENDING_FORMS: + return + _PENDING_FORM_EXPIRY_HEAP.clear() + _PENDING_FORM_ACTIVE_COUNT = 0 + + +def _pending_form_entry_is_current( + expires_at: float, + revision: int, + session_key: PendingFormKey, + form_token: str, +) -> bool: + forms = _PENDING_FORMS.get(session_key) + if forms is None: + return False + stored = forms.get(form_token) + if stored is None: + return False + return stored.get(_PENDING_FORM_REVISION_KEY) == revision and stored.get('_expires_at') == expires_at + + +def _peek_valid_pending_form_expiry( + *, + pop: bool = False, +) -> tuple[float, int, PendingFormKey, str] | None: + while _PENDING_FORM_EXPIRY_HEAP: + entry = _PENDING_FORM_EXPIRY_HEAP[0] + if _pending_form_entry_is_current(*entry): + if pop: + heapq.heappop(_PENDING_FORM_EXPIRY_HEAP) + return entry + heapq.heappop(_PENDING_FORM_EXPIRY_HEAP) + return None + + +def _drop_pending_form(session_key: PendingFormKey, form_token: str) -> None: + global _PENDING_FORM_ACTIVE_COUNT + forms = _PENDING_FORMS.get(session_key) + if forms is None or forms.pop(form_token, None) is None: + return + _PENDING_FORM_ACTIVE_COUNT = max(_PENDING_FORM_ACTIVE_COUNT - 1, 0) + if not forms: + _PENDING_FORMS.pop(session_key, None) + + +def _drop_pending_form_session(session_key: PendingFormKey) -> None: + global _PENDING_FORM_ACTIVE_COUNT + forms = _PENDING_FORMS.pop(session_key, None) + if forms is not None: + _PENDING_FORM_ACTIVE_COUNT = max( + _PENDING_FORM_ACTIVE_COUNT - len(forms), + 0, + ) + + +def _compact_pending_form_expiry_heap_if_needed() -> None: + max_heap_entries = max( + _PENDING_FORM_HEAP_COMPACT_FLOOR, + _PENDING_FORM_ACTIVE_COUNT * _PENDING_FORM_HEAP_MAX_MULTIPLIER, + ) + if len(_PENDING_FORM_EXPIRY_HEAP) <= max_heap_entries: + return + _PENDING_FORM_EXPIRY_HEAP[:] = [ + ( + float(stored['_expires_at']), + int(stored[_PENDING_FORM_REVISION_KEY]), + session_key, + form_token, + ) + for session_key, forms in _PENDING_FORMS.items() + for form_token, stored in forms.items() + ] + heapq.heapify(_PENDING_FORM_EXPIRY_HEAP) + + def _prune_pending_forms(now: float | None = None) -> None: + _synchronize_pending_form_cache_if_externally_cleared() if now is None: now = time.time() - for session_key in list(_PENDING_FORMS.keys()): - forms = _PENDING_FORMS[session_key] - expired_tokens = [token for token, data in forms.items() if data.get('_expires_at', 0) <= now] - for token in expired_tokens: - forms.pop(token, None) - if not forms: - _PENDING_FORMS.pop(session_key, None) + while True: + entry = _peek_valid_pending_form_expiry() + if entry is None or entry[0] > now: + break + _, _, session_key, form_token = _peek_valid_pending_form_expiry(pop=True) + _drop_pending_form(session_key, form_token) + _compact_pending_form_expiry_heap_if_needed() def _set_pending_form(session_key: PendingFormKey, form_data: dict[str, typing.Any]) -> None: + global _PENDING_FORM_ACTIVE_COUNT, _PENDING_FORM_REVISION _prune_pending_forms() if isinstance(session_key, tuple) and len(session_key) == 8: form_data['pipeline_uuid'] = session_key[4] @@ -88,11 +189,31 @@ def _set_pending_form(session_key: PendingFormKey, form_data: dict[str, typing.A except (TypeError, ValueError): expiration_ts = 0.0 stored['_expires_at'] = expiration_ts or (time.time() + _PENDING_FORM_DEFAULT_TTL) + _PENDING_FORM_REVISION += 1 + stored[_PENDING_FORM_REVISION_KEY] = _PENDING_FORM_REVISION form_token = str(stored.get('form_token') or '') forms = _PENDING_FORMS.setdefault(session_key, OrderedDict()) # Re-insert at the end so this becomes the "latest" entry - forms.pop(form_token, None) + if forms.pop(form_token, None) is None: + _PENDING_FORM_ACTIVE_COUNT += 1 forms[form_token] = stored + heapq.heappush( + _PENDING_FORM_EXPIRY_HEAP, + ( + stored['_expires_at'], + _PENDING_FORM_REVISION, + session_key, + form_token, + ), + ) + while len(forms) > _PENDING_FORM_MAX_PER_SESSION: + oldest_token = next(iter(forms)) + _drop_pending_form(session_key, oldest_token) + if len(_PENDING_FORMS) > _PENDING_FORM_MAX_SESSIONS: + oldest_entry = _peek_valid_pending_form_expiry() + if oldest_entry is not None: + _drop_pending_form_session(oldest_entry[2]) + _compact_pending_form_expiry_heap_if_needed() def _get_pending_form_by_token(session_key: PendingFormKey, form_token: str) -> dict[str, typing.Any] | None: @@ -144,11 +265,11 @@ def _clear_pending_form(session_key: PendingFormKey, form_token: str | None = No if not forms: return if form_token is None: - _PENDING_FORMS.pop(session_key, None) + _drop_pending_form_session(session_key) + _compact_pending_form_expiry_heap_if_needed() return - forms.pop(form_token, None) - if not forms: - _PENDING_FORMS.pop(session_key, None) + _drop_pending_form(session_key, form_token) + _compact_pending_form_expiry_heap_if_needed() def _format_human_input_text( @@ -721,6 +842,9 @@ class DifyServiceAPIRunner(runner.RequestRunner): base_url=self.pipeline_config['ai']['dify-service-api']['base-url'], ) + async def aclose(self) -> None: + await self.dify_client.aclose() + def _process_thinking_content( self, content: str, @@ -796,13 +920,16 @@ class DifyServiceAPIRunner(runner.RequestRunner): async def download_file(file_url: str) -> tuple[bytes, str]: """Download file from url (supports data url).""" - async with httpx.AsyncClient() as client_session: - resp = await client_session.get(file_url) + client_session = httpclient.get_session() + async with client_session.get(file_url, timeout=120) as resp: resp.raise_for_status() content_type = ( resp.headers.get('content-type') or mimetypes.guess_type(file_url)[0] or 'application/octet-stream' ) - return resp.content, content_type + return ( + await httpclient.read_limited(resp, max_bytes=_MAX_DIFY_UPLOAD_BYTES), + content_type, + ) def _detect_file_type(content_type: str) -> str: """Map MIME to dify file type.""" @@ -820,7 +947,10 @@ class DifyServiceAPIRunner(runner.RequestRunner): plain_text += ce.text elif ce.type == 'image_base64': image_b64, image_format = await image.extract_b64_and_format(ce.image_base64) - file_bytes = base64.b64decode(image_b64) + file_bytes = await image.decode_base64_limited( + image_b64, + max_bytes=_MAX_DIFY_UPLOAD_BYTES, + ) image_id = await upload_file_bytes(f'img.{image_format}', file_bytes, f'image/{image_format}') upload_files.append({'type': 'image', 'id': image_id}) elif ce.type == 'file_url': @@ -840,7 +970,10 @@ class DifyServiceAPIRunner(runner.RequestRunner): content_type = 'application/octet-stream' if ';' in header: content_type = header.split(';')[0][5:] or content_type - file_bytes = base64.b64decode(b64_data) + file_bytes = await image.decode_base64_limited( + b64_data, + max_bytes=_MAX_DIFY_UPLOAD_BYTES, + ) file_id = await upload_file_bytes(file_name, file_bytes, content_type) file_type = _detect_file_type(content_type) upload_files.append({'type': file_type, 'id': file_id}) @@ -865,15 +998,19 @@ class DifyServiceAPIRunner(runner.RequestRunner): } async def _download_file_for_form(self, file_url: str) -> tuple[bytes, str, str]: - async with httpx.AsyncClient() as client_session: - resp = await client_session.get(file_url) + client_session = httpclient.get_session() + async with client_session.get(file_url, timeout=120) as resp: resp.raise_for_status() content_type = ( resp.headers.get('content-type') or mimetypes.guess_type(file_url)[0] or 'application/octet-stream' ) parsed = urlparse(file_url) file_name = os.path.basename(parsed.path) or 'file' - return resp.content, content_type, file_name + return ( + await httpclient.read_limited(resp, max_bytes=_MAX_DIFY_UPLOAD_BYTES), + content_type, + file_name, + ) async def _platform_file_to_dify(self, item: typing.Any, user: str) -> dict | None: try: @@ -890,13 +1027,15 @@ class DifyServiceAPIRunner(runner.RequestRunner): content_type = header.split(';', 1)[0][5:] or content_type return await self._upload_file_bytes_for_user( file_name, - base64.b64decode(b64_data), + await image.decode_base64_limited( + b64_data, + max_bytes=_MAX_DIFY_UPLOAD_BYTES, + ), content_type, user, ) if item.path: - with open(item.path, 'rb') as f: - file_bytes = f.read() + file_bytes = await asyncio.to_thread(_read_local_file_limited, str(item.path)) content_type = mimetypes.guess_type(str(item.path))[0] or 'application/octet-stream' file_name = item.name or os.path.basename(str(item.path)) or 'file' return await self._upload_file_bytes_for_user(file_name, file_bytes, content_type, user) diff --git a/src/langbot/pkg/provider/runners/langflowapi.py b/src/langbot/pkg/provider/runners/langflowapi.py index 8995476d3..10c66df59 100644 --- a/src/langbot/pkg/provider/runners/langflowapi.py +++ b/src/langbot/pkg/provider/runners/langflowapi.py @@ -1,5 +1,6 @@ from __future__ import annotations +import codecs import typing import json import httpx @@ -11,6 +12,44 @@ from ...core import app import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query import langbot_plugin.api.entities.builtin.provider.message as provider_message +_MAX_LANGFLOW_LINE_CHARS = 1024 * 1024 +_MAX_LANGFLOW_TOTAL_BYTES = 16 * 1024 * 1024 +_MAX_LANGFLOW_RESPONSE_BYTES = 1024 * 1024 + + +async def _iter_limited_lines( + response: httpx.Response, +) -> typing.AsyncGenerator[str, None]: + decoder = codecs.getincrementaldecoder('utf-8')('replace') + buffer = '' + total_bytes = 0 + async for chunk in response.aiter_bytes(chunk_size=8192): + total_bytes += len(chunk) + if total_bytes > _MAX_LANGFLOW_TOTAL_BYTES: + raise ValueError('Langflow stream exceeds the runtime limit') + buffer += decoder.decode(chunk) + while '\n' in buffer: + line, buffer = buffer.split('\n', 1) + if len(line) > _MAX_LANGFLOW_LINE_CHARS: + raise ValueError('Langflow event exceeds the runtime limit') + yield line.rstrip('\r') + if len(buffer) > _MAX_LANGFLOW_LINE_CHARS: + raise ValueError('Langflow event exceeds the runtime limit') + buffer += decoder.decode(b'', final=True) + if buffer: + if len(buffer) > _MAX_LANGFLOW_LINE_CHARS: + raise ValueError('Langflow event exceeds the runtime limit') + yield buffer.rstrip('\r') + + +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_LANGFLOW_RESPONSE_BYTES: + raise ValueError('Langflow response exceeds the runtime limit') + return bytes(body) + @runner.runner_class('langflow-api') class LangflowAPIRunner(runner.RequestRunner): @@ -99,7 +138,7 @@ class LangflowAPIRunner(runner.RequestRunner): accumulated_content = '' message_count = 0 - async for line in response.aiter_lines(): + async for line in _iter_limited_lines(response): data_str = line if data_str.startswith('data: '): @@ -144,11 +183,15 @@ class LangflowAPIRunner(runner.RequestRunner): yield provider_message.MessageChunk(role='assistant', content=accumulated_content, is_final=True) else: # 非流式请求 - response = await client.post(url, json=payload, headers=headers, timeout=120.0) - response.raise_for_status() - - # 解析响应 - response_data = response.json() + async with client.stream( + 'POST', + url, + json=payload, + headers=headers, + timeout=120.0, + ) as response: + response.raise_for_status() + response_data = json.loads(await _read_limited_response(response)) # 提取消息内容 # 根据Langflow API文档，响应结构可能在outputs[0].outputs[0].outputs.message.message中 diff --git a/src/langbot/pkg/provider/runners/n8nsvapi.py b/src/langbot/pkg/provider/runners/n8nsvapi.py index 543fd7ef9..24ef7f59c 100644 --- a/src/langbot/pkg/provider/runners/n8nsvapi.py +++ b/src/langbot/pkg/provider/runners/n8nsvapi.py @@ -12,6 +12,8 @@ from ...core import app import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query import langbot_plugin.api.entities.builtin.provider.message as provider_message +_MAX_N8N_RESPONSE_CHARS = 1024 * 1024 + class N8nAPIError(Exception): """N8n API 请求失败""" @@ -94,6 +96,8 @@ class N8nServiceAPIRunner(runner.RequestRunner): else: chunk_str = str(raw_chunk) + if len(full_text) + len(chunk_str) > _MAX_N8N_RESPONSE_CHARS: + raise N8nAPIError('n8n response exceeds the runtime limit') full_text += chunk_str buffer += chunk_str @@ -112,7 +116,9 @@ class N8nServiceAPIRunner(runner.RequestRunner): if obj.get('type') == 'item' and 'content' in obj: chunk_idx += 1 - content = obj['content'] + content = str(obj['content']) + if len(full_content) + len(content) > _MAX_N8N_RESPONSE_CHARS: + raise N8nAPIError('n8n response exceeds the runtime limit') full_content += content elif obj.get('type') == 'end': is_final = True @@ -128,6 +134,8 @@ class N8nServiceAPIRunner(runner.RequestRunner): except json.JSONDecodeError: # buffer 末尾可能是一个不完整的 JSON，等待更多数据 break + except N8nAPIError: + raise except Exception as e: # 记录解析失败并继续接收后续 chunk try: @@ -255,7 +263,12 @@ class N8nServiceAPIRunner(runner.RequestRunner): self.webhook_url, json=payload, headers=headers, auth=auth, timeout=self.timeout ) as response: if response.status != 200: - error_text = await response.text() + error_text = ( + await httpclient.read_limited( + response, + max_bytes=_MAX_N8N_RESPONSE_CHARS, + ) + ).decode('utf-8', errors='replace') self.ap.logger.error(f'n8n webhook call failed: {response.status}, {error_text}') raise Exception(f'n8n webhook call failed: {response.status}, {error_text}') diff --git a/src/langbot/pkg/provider/runners/tboxapi.py b/src/langbot/pkg/provider/runners/tboxapi.py index 0fb22a642..9072daa84 100644 --- a/src/langbot/pkg/provider/runners/tboxapi.py +++ b/src/langbot/pkg/provider/runners/tboxapi.py @@ -1,8 +1,9 @@ from __future__ import annotations +import asyncio import typing import json -import base64 +import logging import tempfile import os @@ -11,10 +12,13 @@ from tboxsdk.model.file import File, FileType from .. import runner from ...core import app -from ...utils import image +from ...utils import bounded_executor, image import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query import langbot_plugin.api.entities.builtin.provider.message as provider_message +_MAX_TBOX_RESPONSE_CHARS = 1024 * 1024 +_MAX_TBOX_MEDIA_BYTES = 10 * 1024 * 1024 + class TboxAPIError(Exception): """TBox API 请求失败""" @@ -24,6 +28,19 @@ class TboxAPIError(Exception): super().__init__(self.message) +def _append_bounded(current: str, addition: typing.Any) -> str: + addition = str(addition or '') + if len(current) + len(addition) > _MAX_TBOX_RESPONSE_CHARS: + raise TboxAPIError('Tbox response exceeds the runtime limit') + return current + addition + + +def _write_temp_media(file_bytes: bytes, suffix: str) -> str: + with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp_file: + tmp_file.write(file_bytes) + return tmp_file.name + + @runner.runner_class('tbox-app-api') class TboxAPIRunner(runner.RequestRunner): "蚂蚁百宝箱API对话请求器" @@ -42,6 +59,7 @@ class TboxAPIRunner(runner.RequestRunner): self.api_key = self.pipeline_config['ai']['tbox-app-api']['api-key'] # 初始化Tbox client + logging.getLogger('tbox.client').setLevel(logging.WARNING) self.tbox_client = TboxClient(authorization=self.api_key) async def _preprocess_user_message(self, query: pipeline_query.Query) -> tuple[str, list[str]]: @@ -59,19 +77,29 @@ class TboxAPIRunner(runner.RequestRunner): plain_text += ce.text elif ce.type == 'image_base64': image_b64, image_format = await image.extract_b64_and_format(ce.image_base64) - # 创建临时文件 - file_bytes = base64.b64decode(image_b64) + file_bytes = await image.decode_base64_limited( + image_b64, + max_bytes=_MAX_TBOX_MEDIA_BYTES, + ) + tmp_file_path: str | None = None try: - with tempfile.NamedTemporaryFile(suffix=f'.{image_format}', delete=False) as tmp_file: - tmp_file.write(file_bytes) - tmp_file_path = tmp_file.name - file_upload_resp = self.tbox_client.upload_file(tmp_file_path) + tmp_file_path = await asyncio.to_thread( + _write_temp_media, + file_bytes, + f'.{image_format}', + ) + file_upload_resp = await asyncio.to_thread( + self.tbox_client.upload_file, + tmp_file_path, + ) image_id = file_upload_resp.get('data', '') image_ids.append(image_id) finally: - # 清理临时文件 - if os.path.exists(tmp_file_path): - os.unlink(tmp_file_path) + if tmp_file_path and os.path.exists(tmp_file_path): + await bounded_executor.run_blocking_cleanup( + os.unlink, + tmp_file_path, + ) elif isinstance(query.user_message.content, str): plain_text = query.user_message.content @@ -98,18 +126,23 @@ class TboxAPIRunner(runner.RequestRunner): files = [File(file_id=image_id, type=FileType.IMAGE) for image_id in image_ids] # 发送对话请求 - response = self.tbox_client.chat( - app_id=self.app_id, # Tbox中智能体应用的ID - user_id=query.bot_uuid, # 用户ID - query=plain_text, # 用户输入的文本信息 - stream=is_stream, # 是否流式输出 - conversation_id=conversation_id, # 会话ID，为None时Tbox会自动创建一个新会话 - files=files, # 图片内容 + response = await asyncio.to_thread( + self.tbox_client.chat, + app_id=self.app_id, + user_id=query.bot_uuid, + query=plain_text, + stream=is_stream, + conversation_id=conversation_id, + files=files, ) if is_stream: # 解析Tbox流式输出内容，并发送给上游 - for chunk in self._process_stream_message(response, query, remove_think): + async for chunk in self._process_stream_message( + response, + query, + remove_think, + ): yield chunk else: message = self._process_non_stream_message(response, query, remove_think) @@ -127,13 +160,16 @@ class TboxAPIRunner(runner.RequestRunner): thinking_content = payload.get('reasoningContent', []) result = '' if thinking_content and not remove_think: - result += f'\n{thinking_content[0].get("text", "")}\n\n' + result = _append_bounded( + result, + f'\n{thinking_content[0].get("text", "")}\n\n', + ) content = payload.get('result', []) if content: - result += content[0].get('chunk', '') + result = _append_bounded(result, content[0].get('chunk', '')) return result - def _process_stream_message( + async def _process_stream_message( self, response: typing.Generator[dict], query: pipeline_query.Query, remove_think: bool ): idx_msg = 0 @@ -141,7 +177,7 @@ class TboxAPIRunner(runner.RequestRunner): conversation_id = None think_start = False think_end = False - for chunk in response: + async for chunk in runner.iterate_sync(response): if chunk.get('type', '') == 'chunk': """ Tbox返回的消息内容chunk结构 @@ -149,7 +185,10 @@ class TboxAPIRunner(runner.RequestRunner): """ # 如果包含思考过程，拼接 if think_start and not think_end: - pending_content += '\n\n' + pending_content = _append_bounded( + pending_content, + '\n\n', + ) think_end = True payload = chunk.get('payload', {}) @@ -158,7 +197,10 @@ class TboxAPIRunner(runner.RequestRunner): query.session.using_conversation.uuid = conversation_id if payload.get('text'): idx_msg += 1 - pending_content += payload.get('text') + pending_content = _append_bounded( + pending_content, + payload.get('text'), + ) elif chunk.get('type', '') == 'thinking' and not remove_think: """ Tbox返回的思考过程chunk结构 @@ -170,9 +212,15 @@ class TboxAPIRunner(runner.RequestRunner): content = payload.get('ext_data', {}).get('text') if not think_start: think_start = True - pending_content += f'\n{content}' + pending_content = _append_bounded( + pending_content, + f'\n{content}', + ) else: - pending_content += content + pending_content = _append_bounded( + pending_content, + content, + ) elif chunk.get('type', '') == 'error': raise TboxAPIError( f'Tbox API 请求失败: status_code={chunk.get("status_code")} message={chunk.get("message")} request_id={chunk.get("request_id")} ' diff --git a/src/langbot/pkg/provider/runners/weknoraapi.py b/src/langbot/pkg/provider/runners/weknoraapi.py index 9d46eebb7..6af41548f 100644 --- a/src/langbot/pkg/provider/runners/weknoraapi.py +++ b/src/langbot/pkg/provider/runners/weknoraapi.py @@ -10,6 +10,15 @@ import langbot_plugin.api.entities.builtin.provider.message as provider_message import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query from langbot.libs.weknora_api import client, errors +_MAX_WEKNORA_GENERATED_CHARS = 1024 * 1024 + + +def _append_bounded(current: str, addition: typing.Any) -> str: + addition = str(addition or '') + if len(current) + len(addition) > _MAX_WEKNORA_GENERATED_CHARS: + raise errors.WeKnoraAPIError('WeKnora response exceeds the runtime limit') + return current + addition + @runner.runner_class('weknora-api') class WeKnoraAPIRunner(runner.RequestRunner): @@ -94,7 +103,7 @@ class WeKnoraAPIRunner(runner.RequestRunner): web_search_enabled=web_search_enabled, timeout=timeout, ): - self.ap.logger.debug('weknora-agent-chunk: ' + str(chunk)) + self.ap.logger.debug('weknora-agent-chunk: ' + str(chunk)[:1000]) response_type = chunk.get('response_type', '') content = chunk.get('content', '') @@ -120,7 +129,7 @@ class WeKnoraAPIRunner(runner.RequestRunner): elif response_type == 'answer': if content: - full_answer += content + full_answer = _append_bounded(full_answer, content) elif response_type == 'error': raise errors.WeKnoraAPIError(f'WeKnora 服务错误: {content}') @@ -158,14 +167,14 @@ class WeKnoraAPIRunner(runner.RequestRunner): knowledge_base_ids=knowledge_base_ids, timeout=timeout, ): - self.ap.logger.debug('weknora-chat-chunk: ' + str(chunk)) + self.ap.logger.debug('weknora-chat-chunk: ' + str(chunk)[:1000]) response_type = chunk.get('response_type', '') content = chunk.get('content', '') if response_type == 'answer': if content: - full_answer += content + full_answer = _append_bounded(full_answer, content) elif response_type == 'error': raise errors.WeKnoraAPIError(f'WeKnora 服务错误: {content}') @@ -207,7 +216,7 @@ class WeKnoraAPIRunner(runner.RequestRunner): web_search_enabled=web_search_enabled, timeout=timeout, ): - self.ap.logger.debug('weknora-agent-chunk: ' + str(chunk)) + self.ap.logger.debug('weknora-agent-chunk: ' + str(chunk)[:1000]) response_type = chunk.get('response_type', '') content = chunk.get('content', '') @@ -235,7 +244,7 @@ class WeKnoraAPIRunner(runner.RequestRunner): elif response_type == 'answer': message_idx += 1 if content: - pending_answer += content + pending_answer = _append_bounded(pending_answer, content) if done: is_final = True @@ -288,7 +297,7 @@ class WeKnoraAPIRunner(runner.RequestRunner): knowledge_base_ids=knowledge_base_ids, timeout=timeout, ): - self.ap.logger.debug('weknora-chat-chunk: ' + str(chunk)) + self.ap.logger.debug('weknora-chat-chunk: ' + str(chunk)[:1000]) response_type = chunk.get('response_type', '') content = chunk.get('content', '') @@ -297,7 +306,7 @@ class WeKnoraAPIRunner(runner.RequestRunner): if response_type == 'answer': message_idx += 1 if content: - pending_answer += content + pending_answer = _append_bounded(pending_answer, content) if done: is_final = True diff --git a/src/langbot/pkg/provider/session/sessionmgr.py b/src/langbot/pkg/provider/session/sessionmgr.py index 4beda1661..74b9a8c40 100644 --- a/src/langbot/pkg/provider/session/sessionmgr.py +++ b/src/langbot/pkg/provider/session/sessionmgr.py @@ -2,6 +2,8 @@ from __future__ import annotations import asyncio import dataclasses +import heapq +import time from langbot_plugin.api.entities.builtin.provider import message as provider_message, prompt as provider_prompt import langbot_plugin.api.entities.builtin.provider.session as provider_session @@ -24,6 +26,10 @@ SessionKey = tuple[ str, int | str, ] +SessionExpiryEntry = tuple[float, int, SessionKey] + +_SESSION_EXPIRY_HEAP_MIN_LIMIT = 64 +_SESSION_EXPIRY_HEAP_ACTIVE_MULTIPLIER = 4 def _query_session_key(query: pipeline_query.Query) -> tuple[SessionKey, ExecutionContext]: @@ -49,11 +55,215 @@ class SessionManager: ap: app.Application - session_list: list[provider_session.Session] - def __init__(self, ap: app.Application): self.ap = ap - self.session_list = [] + self._legacy_sessions: list[provider_session.Session] = [] + self._session_index: dict[SessionKey, provider_session.Session] = {} + self._session_keys_by_workspace: dict[str, set[SessionKey]] = {} + self._session_expiry_heap: list[SessionExpiryEntry] = [] + self._next_access_revision = 0 + + @property + def session_list(self) -> list[provider_session.Session]: + """Compatibility view for API services that enumerate sessions.""" + + return [ + *self._legacy_sessions, + *self._session_index.values(), + ] + + @session_list.setter + def session_list(self, sessions: list[provider_session.Session]) -> None: + """Replace the cache while keeping the O(1) index consistent.""" + + session_values = list(sessions) + self._legacy_sessions = [] + self._session_index = {} + self._session_keys_by_workspace = {} + self._session_expiry_heap = [] + self._next_access_revision = 0 + now = time.monotonic() + for session in session_values: + key = getattr(session, '_langbot_session_key', None) + if isinstance(key, tuple) and len(key) == 6: + self._session_index[key] = session + self._session_keys_by_workspace.setdefault(key[1], set()).add(key) + last_accessed = getattr(session, '_langbot_last_accessed', None) + if last_accessed is None: + last_accessed = now + self._touch_session( + session, + key, + float(last_accessed), + compact=False, + ) + else: + self._legacy_sessions.append(session) + self._compact_session_expiry_heap(force=True) + + def _retention_config(self) -> dict: + instance_config = getattr(getattr(self.ap, 'instance_config', None), 'data', {}) + if not isinstance(instance_config, dict): + return {} + config = instance_config.get('system', {}).get('session_retention', {}) + return config if isinstance(config, dict) else {} + + def _positive_config_int(self, name: str, default: int) -> int: + try: + value = int(self._retention_config().get(name, default)) + except (TypeError, ValueError): + value = default + return max(value, 1) + + @staticmethod + def _session_is_idle(session: provider_session.Session) -> bool: + semaphore = getattr(session, '_semaphore', None) + concurrency = getattr(session, '_langbot_session_concurrency', None) + if semaphore is None or not isinstance(concurrency, int): + return True + return getattr(semaphore, '_value', -1) == concurrency + + def _remove_session(self, session: provider_session.Session) -> None: + key = getattr(session, '_langbot_session_key', None) + if isinstance(key, tuple) and len(key) == 6 and self._session_index.get(key) is session: + self._session_index.pop(key, None) + workspace_keys = self._session_keys_by_workspace.get(key[1]) + if workspace_keys is not None: + workspace_keys.discard(key) + if not workspace_keys: + self._session_keys_by_workspace.pop(key[1], None) + else: + try: + self._legacy_sessions.remove(session) + except ValueError: + pass + + def _touch_session( + self, + session: provider_session.Session, + key: SessionKey, + now: float, + *, + compact: bool = True, + ) -> None: + self._next_access_revision += 1 + revision = self._next_access_revision + object.__setattr__(session, '_langbot_last_accessed', now) + object.__setattr__(session, '_langbot_access_revision', revision) + heapq.heappush( + self._session_expiry_heap, + (now, revision, key), + ) + if compact: + self._compact_session_expiry_heap() + + def _compact_session_expiry_heap(self, *, force: bool = False) -> None: + limit = max( + len(self._session_index) * _SESSION_EXPIRY_HEAP_ACTIVE_MULTIPLIER, + _SESSION_EXPIRY_HEAP_MIN_LIMIT, + ) + if not force and len(self._session_expiry_heap) <= limit: + return + self._session_expiry_heap = [ + ( + float(getattr(session, '_langbot_last_accessed', 0.0)), + int(getattr(session, '_langbot_access_revision', 0)), + key, + ) + for key, session in self._session_index.items() + ] + heapq.heapify(self._session_expiry_heap) + + def _pop_current_expiry_entry( + self, + ) -> tuple[float, int, SessionKey, provider_session.Session] | None: + while self._session_expiry_heap: + last_accessed, revision, key = heapq.heappop(self._session_expiry_heap) + session = self._session_index.get(key) + if session is None: + continue + if getattr(session, '_langbot_access_revision', None) != revision: + continue + return last_accessed, revision, key, session + return None + + def _prune_expired_sessions(self, now: float) -> None: + idle_ttl = self._positive_config_int('idle_ttl_seconds', 86400) + cutoff = now - idle_ttl + while self._session_expiry_heap: + last_accessed, _, _ = self._session_expiry_heap[0] + if last_accessed > cutoff: + break + current = self._pop_current_expiry_entry() + if current is None: + break + last_accessed, revision, key, session = current + if last_accessed > cutoff: + heapq.heappush( + self._session_expiry_heap, + (last_accessed, revision, key), + ) + break + if self._session_is_idle(session): + self._remove_session(session) + continue + # The session became active without another cache lookup. Give it + # a fresh TTL instead of repeatedly examining the same expired + # entry or losing its future expiry record. + self._touch_session(session, key, now) + + def _prune_workspace_capacity( + self, + workspace_uuid: str, + max_entries_per_workspace: int, + ) -> None: + workspace_keys = self._session_keys_by_workspace.get(workspace_uuid, set()) + overflow = len(workspace_keys) - max_entries_per_workspace + 1 + if overflow <= 0: + return + idle_workspace_sessions = sorted( + ( + session + for key in tuple(workspace_keys) + if (session := self._session_index.get(key)) is not None and self._session_is_idle(session) + ), + key=lambda session: float(getattr(session, '_langbot_last_accessed', 0.0)), + ) + for session in idle_workspace_sessions[:overflow]: + self._remove_session(session) + + def _evict_oldest_idle_session(self, now: float) -> bool: + # At most one current entry per active session is examined. Stale heap + # revisions do not count and are discarded in O(log N). + current_probes = 0 + max_probes = len(self._session_index) + while current_probes < max_probes: + current = self._pop_current_expiry_entry() + if current is None: + return False + _, _, key, session = current + current_probes += 1 + if self._session_is_idle(session): + self._remove_session(session) + return True + self._touch_session(session, key, now) + return False + + def _prune_sessions(self, now: float, workspace_uuid: str) -> None: + self._prune_expired_sessions(now) + max_entries_per_workspace = self._positive_config_int('max_entries_per_workspace', 200) + self._prune_workspace_capacity( + workspace_uuid, + max_entries_per_workspace, + ) + + max_entries = self._positive_config_int('max_entries', 2000) + overflow = len(self._session_index) - max_entries + 1 + if overflow <= 0: + return + for _ in range(overflow): + if not self._evict_oldest_idle_session(now): + break async def initialize(self): pass @@ -61,9 +271,25 @@ class SessionManager: async def get_session(self, query: pipeline_query.Query) -> provider_session.Session: """获取会话""" session_key, execution_context = _query_session_key(query) - for session in self.session_list: - if getattr(session, '_langbot_session_key', None) == session_key: - return session + now = time.monotonic() + session = self._session_index.get(session_key) + if session is not None: + self._touch_session(session, session_key, now) + return session + + self._prune_sessions(now, execution_context.workspace_uuid) + max_entries_per_workspace = self._positive_config_int('max_entries_per_workspace', 200) + workspace_entries = len( + self._session_keys_by_workspace.get( + execution_context.workspace_uuid, + (), + ) + ) + if workspace_entries >= max_entries_per_workspace: + raise RuntimeError(f'Workspace session cache capacity reached ({max_entries_per_workspace})') + max_entries = self._positive_config_int('max_entries', 2000) + if len(self._session_index) >= max_entries: + raise RuntimeError(f'Session cache capacity reached ({max_entries})') session_concurrency = self.ap.instance_config.data['concurrency']['session'] @@ -93,8 +319,14 @@ class SessionManager: object.__setattr__(session, 'bot_uuid', query.bot_uuid) object.__setattr__(session, '_execution_context', session_context) object.__setattr__(session, '_langbot_session_key', session_key) + object.__setattr__(session, '_langbot_session_concurrency', session_concurrency) session._semaphore = asyncio.Semaphore(session_concurrency) - self.session_list.append(session) + self._session_index[session_key] = session + self._session_keys_by_workspace.setdefault( + execution_context.workspace_uuid, + set(), + ).add(session_key) + self._touch_session(session, session_key, now) return session async def get_conversation( @@ -144,6 +376,47 @@ class SessionManager: bot_uuid=bot_uuid, ) session.conversations.append(conversation) + max_conversations = self._positive_config_int('max_conversations_per_session', 20) + if len(session.conversations) > max_conversations: + del session.conversations[:-max_conversations] session.using_conversation = conversation return session.using_conversation + + def trim_conversation_messages( + self, + conversation: provider_session.Conversation, + *, + max_rounds: int, + ) -> None: + """Bound retained process-local history after a completed turn.""" + + try: + max_rounds = int(max_rounds) + except (TypeError, ValueError): + max_rounds = 10 + max_rounds = max(max_rounds, 1) + max_messages = self._positive_config_int('max_messages_per_conversation', 100) + + kept_reversed = [] + user_rounds = 0 + for message in reversed(conversation.messages): + if user_rounds >= max_rounds: + break + kept_reversed.append(message) + if getattr(message, 'role', None) == 'user': + user_rounds += 1 + retained = list(reversed(kept_reversed))[-max_messages:] + # Binary payloads are needed for the current model call, but retaining + # them in process-local history makes a few image/file turns consume + # hundreds of MB. Historical URL and text references remain intact. + for message in retained: + content = getattr(message, 'content', None) + if not isinstance(content, list): + continue + for element in content: + if getattr(element, 'image_base64', None) is not None: + element.image_base64 = None + if getattr(element, 'file_base64', None) is not None: + element.file_base64 = None + conversation.messages = retained diff --git a/src/langbot/pkg/provider/tools/loaders/mcp.py b/src/langbot/pkg/provider/tools/loaders/mcp.py index 69dafb4b0..08fefb4aa 100644 --- a/src/langbot/pkg/provider/tools/loaders/mcp.py +++ b/src/langbot/pkg/provider/tools/loaders/mcp.py @@ -1,12 +1,13 @@ from __future__ import annotations -import base64 import enum import json import math import re import time import typing +import ipaddress +from urllib.parse import urlparse from contextlib import AsyncExitStack, asynccontextmanager from datetime import timedelta import traceback @@ -50,6 +51,7 @@ MCP_TOOL_READ_RESOURCE = 'langbot_mcp_read_resource' MCP_RESOURCE_DISCOVERY_MAX_PAGES = 20 MCP_RESOURCE_CACHE_TTL_SECONDS = 30 +MCP_RESOURCE_CACHE_MAX_ENTRIES = 32 MCP_RESOURCE_PREVIEW_MAX_BYTES = 64 * 1024 MCP_RESOURCE_AGENT_READ_MAX_BYTES = 64 * 1024 MCP_RESOURCE_AGENT_READ_MAX_TOKENS = 12000 @@ -134,10 +136,13 @@ def _truncate_text(text: str, max_bytes: int, max_tokens: int | None = None) -> def _blob_size(blob: str) -> int: - try: - return len(base64.b64decode(blob, validate=False)) - except Exception: + # MCP BlobResourceContents is schema-validated base64 without whitespace. + # Compute decoded size in O(1) without allocating a second binary copy. + encoded_chars = len(blob) + if encoded_chars % 4: return len(blob.encode('utf-8', errors='ignore')) + padding = 2 if blob.endswith('==') else 1 if blob.endswith('=') else 0 + return max((encoded_chars // 4) * 3 - padding, 0) def _resource_to_dict(resource: mcp_types.Resource | mcp_types.ResourceLink) -> dict: @@ -434,12 +439,24 @@ class RuntimeMCPSession: await self._box_stdio_runtime.initialize() async def _init_sse_server(self): + trust_env = self._remote_http_trust_env() + + def httpx_client_factory(headers=None, timeout=None, auth=None): + return httpx.AsyncClient( + headers=headers, + timeout=timeout, + auth=auth, + follow_redirects=True, + trust_env=trust_env, + ) + sse_transport = await self.exit_stack.enter_async_context( sse_client( self.server_config['url'], headers=self.server_config.get('headers', {}), timeout=self.server_config.get('timeout', 10), sse_read_timeout=self.server_config.get('ssereadtimeout', 30), + httpx_client_factory=httpx_client_factory, ) ) @@ -449,6 +466,18 @@ class RuntimeMCPSession: await self.session.initialize() + def _remote_http_trust_env(self) -> bool: + configured = self.server_config.get('trust_env') + if isinstance(configured, bool): + return configured + hostname = (urlparse(str(self.server_config.get('url', ''))).hostname or '').lower() + if hostname == 'localhost': + return False + try: + return not ipaddress.ip_address(hostname).is_loopback + except ValueError: + return True + @asynccontextmanager async def _streamable_http_session(self) -> typing.AsyncIterator[ClientSession]: """Enter a fully initialized Streamable HTTP session as one context. @@ -465,6 +494,7 @@ class RuntimeMCPSession: headers=self.server_config.get('headers', {}), timeout=self.server_config.get('timeout', 10), follow_redirects=True, + trust_env=self._remote_http_trust_env(), ) as http_client: async with streamable_http_client( self.server_config['url'], @@ -1215,6 +1245,9 @@ class RuntimeMCPSession: cache_key = (uri, max_bytes, max_tokens, include_blob) now = time.time() + for expired_key, entry in tuple(self._resource_cache.items()): + if now - entry.get('cached_at', 0) > MCP_RESOURCE_CACHE_TTL_SECONDS: + self._resource_cache.pop(expired_key, None) cached = self._resource_cache.get(cache_key) if cached and now - cached.get('cached_at', 0) <= MCP_RESOURCE_CACHE_TTL_SECONDS: envelope = { @@ -1314,6 +1347,12 @@ class RuntimeMCPSession: 'warnings': warnings, } await self._assert_execution_active() + if cache_key not in self._resource_cache and len(self._resource_cache) >= MCP_RESOURCE_CACHE_MAX_ENTRIES: + oldest_key = min( + self._resource_cache, + key=lambda key: self._resource_cache[key].get('cached_at', 0), + ) + self._resource_cache.pop(oldest_key, None) self._resource_cache[cache_key] = {'cached_at': now, 'envelope': envelope} self._record_resource_read_trace(query, envelope) return envelope @@ -1498,17 +1537,249 @@ class MCPLoader(loader.ToolLoader): 在此加载器中管理所有与 MCP Server 的连接。 """ - sessions: dict[tuple[str, str, int, str], RuntimeMCPSession] - - _last_listed_functions: list[resource_tool.LLMTool] + _sessions: dict[tuple[str, str, int, str], RuntimeMCPSession] _hosted_mcp_tasks: list[asyncio.Task] def __init__(self, ap: app.Application): super().__init__(ap) self.sessions = {} - self._last_listed_functions = [] self._hosted_mcp_tasks = [] + self._hosted_mcp_tasks_by_scope: dict[ + tuple[str, str, int], + set[asyncio.Task], + ] = {} + self._host_dispatch_tasks: set[asyncio.Task] = set() + config = getattr(getattr(ap, 'instance_config', None), 'data', {}) + mcp_config = config.get('mcp', {}) if isinstance(config, dict) else {} + raw_lifecycle_concurrency = mcp_config.get('lifecycle_concurrency', 16) if isinstance(mcp_config, dict) else 16 + if ( + isinstance(raw_lifecycle_concurrency, bool) + or not isinstance(raw_lifecycle_concurrency, int) + or raw_lifecycle_concurrency < 1 + ): + raw_lifecycle_concurrency = 16 + self._lifecycle_concurrency = min( + raw_lifecycle_concurrency, + 128, + ) + self._lifecycle_semaphore = asyncio.Semaphore(self._lifecycle_concurrency) + + @property + def sessions( + self, + ) -> dict[tuple[str, str, int, str], RuntimeMCPSession]: + return self._sessions + + @sessions.setter + def sessions(self, sessions: dict) -> None: + """Compatibility setter that rebuilds the per-scope session index.""" + + self._sessions = sessions + self._session_keys_by_scope: dict[ + tuple[str, str, int], + set[tuple[str, str, int, str]], + ] = {} + self._scope_generations: dict[tuple[str, str], int] = {} + for key in sessions: + if not isinstance(key, tuple) or len(key) != 4: + continue + scope_key = key[:3] + self._session_keys_by_scope.setdefault(scope_key, set()).add(key) + self._scope_generations[scope_key[:2]] = scope_key[2] + + def _register_session( + self, + context: TenantContext, + server_name: str, + session: RuntimeMCPSession, + ) -> None: + scope_key = self._scope_key(context) + workspace_scope = scope_key[:2] + previous_generation = self._scope_generations.get(workspace_scope) + if previous_generation is not None and previous_generation != scope_key[2]: + raise WorkspaceInvariantError('MCP session registration crossed a Workspace generation') + key = (*scope_key, server_name) + self._sessions[key] = session + self._session_keys_by_scope.setdefault(scope_key, set()).add(key) + self._scope_generations[workspace_scope] = scope_key[2] + + def _pop_session( + self, + context: TenantContext, + server_name: str, + ) -> RuntimeMCPSession | None: + scope_key = self._scope_key(context) + key = (*scope_key, server_name) + session = self._sessions.pop(key, None) + keys = self._session_keys_by_scope.get(scope_key) + if keys is not None: + keys.discard(key) + if not keys: + self._session_keys_by_scope.pop(scope_key, None) + self._drop_empty_scope(scope_key) + return session + + def _drop_empty_scope(self, scope_key: tuple[str, str, int]) -> None: + if ( + scope_key not in self._session_keys_by_scope + and scope_key not in self._hosted_mcp_tasks_by_scope + and self._scope_generations.get(scope_key[:2]) == scope_key[2] + ): + self._scope_generations.pop(scope_key[:2], None) + + def track_hosted_task( + self, + task: asyncio.Task, + context: TenantContext, + ) -> asyncio.Task: + """Track a host task without retaining it after completion.""" + + scope_key = self._scope_key(context) + workspace_scope = scope_key[:2] + previous_generation = self._scope_generations.get(workspace_scope) + if previous_generation is not None and previous_generation != scope_key[2]: + task.cancel() + raise WorkspaceInvariantError('MCP host task crossed a Workspace generation') + self._scope_generations[workspace_scope] = scope_key[2] + self._hosted_mcp_tasks.append(task) + self._hosted_mcp_tasks_by_scope.setdefault(scope_key, set()).add(task) + + def discard(completed: asyncio.Task) -> None: + try: + self._hosted_mcp_tasks.remove(completed) + except ValueError: + pass + tasks = self._hosted_mcp_tasks_by_scope.get(scope_key) + if tasks is not None: + tasks.discard(completed) + if not tasks: + self._hosted_mcp_tasks_by_scope.pop(scope_key, None) + self._drop_empty_scope(scope_key) + + task.add_done_callback(discard) + return task + + def _track_host_dispatch_task(self, task: asyncio.Task) -> None: + """Track the bounded startup dispatcher without retaining it.""" + + self._host_dispatch_tasks.add(task) + + def discard(completed: asyncio.Task) -> None: + self._host_dispatch_tasks.discard(completed) + if completed.cancelled(): + return + exception = completed.exception() + if exception is not None: + self.ap.logger.error( + f'MCP startup dispatcher failed: {exception}', + ) + + task.add_done_callback(discard) + + async def _retire_runtime_scope( + self, + scope_key: tuple[str, str, int], + ) -> None: + tasks = tuple(self._hosted_mcp_tasks_by_scope.pop(scope_key, ())) + for task in tasks: + if not task.done(): + task.cancel() + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + + keys = tuple(self._session_keys_by_scope.pop(scope_key, ())) + sessions = [session for key in keys if (session := self._sessions.pop(key, None)) is not None] + await self._shutdown_sessions(sessions) + self._scope_generations.pop(scope_key[:2], None) + + async def _observe_execution_context( + self, + context: ExecutionContext, + ) -> None: + workspace_scope = ( + context.instance_uuid, + context.workspace_uuid, + ) + previous_generation = self._scope_generations.get(workspace_scope) + if previous_generation is None: + return + if context.placement_generation < previous_generation: + raise WorkspaceInvariantError('MCP runtime placement generation rolled back') + if context.placement_generation == previous_generation: + return + await self._retire_runtime_scope((*workspace_scope, previous_generation)) + + async def _reset_runtime_state(self) -> None: + """Cancel host tasks and close sessions before reload or shutdown.""" + + dispatch_tasks = tuple(self._host_dispatch_tasks) + self._host_dispatch_tasks.clear() + for task in dispatch_tasks: + if not task.done(): + task.cancel() + if dispatch_tasks: + await asyncio.gather(*dispatch_tasks, return_exceptions=True) + + tasks = tuple(self._hosted_mcp_tasks) + self._hosted_mcp_tasks.clear() + self._hosted_mcp_tasks_by_scope.clear() + for task in tasks: + if not task.done(): + task.cancel() + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + + sessions = tuple(self._sessions.values()) + self.sessions = {} + await self._shutdown_sessions(sessions) + + async def _shutdown_sessions( + self, + sessions: typing.Iterable[RuntimeMCPSession], + ) -> None: + """Close MCP sessions in bounded batches to avoid shutdown storms.""" + + session_list = list(sessions) + for offset in range(0, len(session_list), self._lifecycle_concurrency): + batch = session_list[offset : offset + self._lifecycle_concurrency] + results = await asyncio.gather( + *(session.shutdown() for session in batch), + return_exceptions=True, + ) + for session, result in zip(batch, results, strict=True): + if isinstance(result, BaseException): + self.ap.logger.error(f'Error shutting down MCP session {session.server_name}: {result}') + + async def _host_server_configs_bounded( + self, + server_configs: typing.Sequence[tuple[ExecutionContext, dict],], + ) -> None: + """Create at most one lifecycle batch of MCP host tasks at a time.""" + + for offset in range(0, len(server_configs), self._lifecycle_concurrency): + batch = server_configs[offset : offset + self._lifecycle_concurrency] + tasks: list[asyncio.Task] = [] + for execution_context, config in batch: + task = create_detached_task( + self.host_mcp_server(execution_context, config), + after_commit_manager=getattr( + self.ap, + 'persistence_mgr', + None, + ), + workspace_uuid=execution_context.workspace_uuid, + ) + tasks.append(task) + try: + self.track_hosted_task(task, execution_context) + except WorkspaceInvariantError as exc: + self.ap.logger.warning( + f'Skipping stale MCP startup task for {execution_context.workspace_uuid}: {exc}' + ) + continue + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) async def _assert_execution_active( self, @@ -1523,6 +1794,7 @@ class MCPLoader(loader.ToolLoader): ) if binding.instance_uuid != execution_context.instance_uuid: raise WorkspaceInvariantError('MCP caller instance does not match the active Workspace binding') + await self._observe_execution_context(execution_context) return execution_context async def initialize(self): @@ -1531,9 +1803,36 @@ class MCPLoader(loader.ToolLoader): async def load_mcp_servers_from_db(self): self.ap.logger.info('Loading MCP servers from db...') - self.sessions = {} + await self._reset_runtime_state() + + pending_hosts: list[tuple[ExecutionContext, dict]] = [] + + async def queue_server(binding, server) -> None: + config = self.ap.persistence_mgr.serialize_model( + persistence_mcp.MCPServer, + server, + ) + if config.get('mode') == 'stdio' and not stdio_mcp_enabled(self.ap): + self.ap.logger.info( + f'Skipping disabled stdio MCP server {server.uuid}; ' + 'the persisted configuration is retained but no process is launched' + ) + return + try: + if binding is None: + binding = await self.ap.workspace_service.get_execution_binding(server.workspace_uuid) + execution_context = ExecutionContext( + instance_uuid=binding.instance_uuid, + workspace_uuid=binding.workspace_uuid, + placement_generation=binding.placement_generation, + ) + except Exception as exc: + self.ap.logger.warning( + f'Skipping MCP server {server.uuid}: Workspace execution binding is unavailable: {exc}' + ) + return + pending_hosts.append((execution_context, config)) - server_configs: list[tuple[typing.Any, typing.Any, dict]] = [] list_bindings = getattr(self.ap.workspace_service, 'list_active_execution_bindings', None) tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None) cloud_runtime = getattr(getattr(self.ap.persistence_mgr, 'mode', None), 'value', None) == 'cloud_runtime' @@ -1548,51 +1847,19 @@ class MCPLoader(loader.ToolLoader): .order_by(persistence_mcp.MCPServer.uuid) ) for server in result.all(): - server_configs.append( - ( - binding, - server, - self.ap.persistence_mgr.serialize_model(persistence_mcp.MCPServer, server), - ) - ) + await queue_server(binding, server) else: # Compatibility path for isolated loader tests and older embedders. result = await self.ap.persistence_mgr.execute_async(sqlalchemy.select(persistence_mcp.MCPServer)) for server in result.all(): - server_configs.append( - ( - None, - server, - self.ap.persistence_mgr.serialize_model(persistence_mcp.MCPServer, server), - ) - ) + await queue_server(None, server) - for binding, server, config in server_configs: - if config.get('mode') == 'stdio' and not stdio_mcp_enabled(self.ap): - self.ap.logger.info( - f'Skipping disabled stdio MCP server {server.uuid}; ' - 'the persisted configuration is retained but no process is launched' - ) - continue - try: - if binding is None: - binding = await self.ap.workspace_service.get_execution_binding(server.workspace_uuid) - execution_context = ExecutionContext( - instance_uuid=binding.instance_uuid, - workspace_uuid=binding.workspace_uuid, - placement_generation=binding.placement_generation, - ) - except Exception as exc: - self.ap.logger.warning( - f'Skipping MCP server {server.uuid}: Workspace execution binding is unavailable: {exc}' - ) - continue - - task = create_detached_task( - self.host_mcp_server(execution_context, config), + if pending_hosts: + dispatch_task = create_detached_task( + self._host_server_configs_bounded(pending_hosts), after_commit_manager=getattr(self.ap, 'persistence_mgr', None), ) - self._hosted_mcp_tasks.append(task) + self._track_host_dispatch_task(dispatch_task) @staticmethod def _scope_key(context: TenantContext) -> tuple[str, str, int]: @@ -1609,9 +1876,25 @@ class MCPLoader(loader.ToolLoader): def _sessions_for_context(self, context: TenantContext) -> list[RuntimeMCPSession]: scope_key = self._scope_key(context) - return [session for key, session in self.sessions.items() if key[:3] == scope_key] + return [ + session + for key in self._session_keys_by_scope.get(scope_key, ()) + if (session := self._sessions.get(key)) is not None + ] - async def host_mcp_server(self, context: TenantContext, server_config: dict): + async def host_mcp_server( + self, + context: TenantContext, + server_config: dict, + ) -> None: + async with self._lifecycle_semaphore: + await self._host_mcp_server(context, server_config) + + async def _host_mcp_server( + self, + context: TenantContext, + server_config: dict, + ) -> None: requested_context = _execution_context_from_tenant(context) execution_context = await run_in_workspace_uow( self.ap, @@ -1627,7 +1910,17 @@ class MCPLoader(loader.ToolLoader): try: session = await self.load_mcp_server(execution_context, server_config) await self._assert_execution_active(execution_context) - self.sessions[self._session_key(execution_context, server_config['name'])] = session + old_session = self._pop_session( + execution_context, + server_config['name'], + ) + if old_session is not None: + await old_session.shutdown() + self._register_session( + execution_context, + server_config['name'], + session, + ) except Exception as e: self.ap.logger.error( f'Failed to load MCP server from db: {server_config["name"]}({server_config["uuid"]}): {e}\n{traceback.format_exc()}' @@ -1876,8 +2169,6 @@ class MCPLoader(loader.ToolLoader): if include_resource_tools and self._eligible_resource_sessions_for_bound(context, bound_mcp_servers): all_functions.extend(self._mcp_synthetic_resource_tools()) - self._last_listed_functions = all_functions - return all_functions async def get_tool_catalog( @@ -2140,7 +2431,9 @@ class MCPLoader(loader.ToolLoader): self.ap.logger.warning(f'MCP server {server_name} not found in sessions, skipping removal') return - session = self.sessions.pop(key) + session = self._pop_session(context, server_name) + if session is None: + return await session.shutdown() self.ap.logger.info(f'Removed MCP server: {server_name}') @@ -2180,22 +2473,5 @@ class MCPLoader(loader.ToolLoader): """关闭所有工具""" self.ap.logger.info('Shutting down all MCP sessions...') - hosted_tasks = [task for task in self._hosted_mcp_tasks if not task.done()] - for task in hosted_tasks: - task.cancel() - if hosted_tasks: - await asyncio.gather(*hosted_tasks, return_exceptions=True) - self._hosted_mcp_tasks.clear() - - async def shutdown_session(session: RuntimeMCPSession) -> None: - try: - await session.shutdown() - self.ap.logger.debug(f'Shutdown MCP session: {session.server_name}') - except Exception as e: - self.ap.logger.error( - f'Error shutting down MCP session {session.server_name}: {e}\n{traceback.format_exc()}' - ) - - await asyncio.gather(*(shutdown_session(session) for session in list(self.sessions.values()))) - self.sessions.clear() + await self._reset_runtime_state() self.ap.logger.info('All MCP sessions shutdown complete') diff --git a/src/langbot/pkg/provider/tools/loaders/mcp_stdio.py b/src/langbot/pkg/provider/tools/loaders/mcp_stdio.py index ff72a1c20..110fff431 100644 --- a/src/langbot/pkg/provider/tools/loaders/mcp_stdio.py +++ b/src/langbot/pkg/provider/tools/loaders/mcp_stdio.py @@ -6,10 +6,12 @@ import os import shutil import shlex import threading +import weakref from contextlib import suppress, AsyncExitStack, asynccontextmanager from typing import TYPE_CHECKING, Any import pydantic +from ....utils import bounded_executor from mcp import ClientSession from mcp.client.websocket import websocket_client from ....box.workspace import ( @@ -27,7 +29,7 @@ if TYPE_CHECKING: from .mcp import RuntimeMCPSession -_WORKSPACE_COPY_LOCKS: dict[str, threading.Lock] = {} +_WORKSPACE_COPY_LOCKS: weakref.WeakValueDictionary[str, threading.Lock] = weakref.WeakValueDictionary() _WORKSPACE_COPY_LOCKS_GUARD = threading.Lock() @@ -536,7 +538,11 @@ class BoxStdioSessionRuntime: return try: process_host_root = os.path.join(self._shared_workspace_host_path(), '.mcp', self.process_id) - await asyncio.to_thread(shutil.rmtree, process_host_root, True) + await bounded_executor.run_blocking_cleanup( + shutil.rmtree, + process_host_root, + True, + ) except Exception as exc: self.ap.logger.warning( f'MCP server {self.server_name}: failed to clean staged workspace ' diff --git a/src/langbot/pkg/provider/tools/loaders/native.py b/src/langbot/pkg/provider/tools/loaders/native.py index 67df0a7dd..3b4e68cea 100644 --- a/src/langbot/pkg/provider/tools/loaders/native.py +++ b/src/langbot/pkg/provider/tools/loaders/native.py @@ -1,24 +1,29 @@ from __future__ import annotations +import asyncio import base64 import contextlib import errno +import heapq import json import os import posixpath import stat +import time from collections.abc import Iterator from dataclasses import dataclass from pathlib import PurePosixPath import langbot_plugin.api.entities.builtin.resource.tool as resource_tool from langbot_plugin.api.entities.events import pipeline_query +import regex from .. import loader from ..errors import ToolNotFoundError from .availability import is_box_backend_available from . import skill as skill_loader from ....api.http.context import ExecutionContext +from ....utils.bounded_executor import run_blocking_atomic EXEC_TOOL_NAME = 'exec' READ_TOOL_NAME = 'read' @@ -36,10 +41,18 @@ _DEFAULT_READ_MAX_LINES = 2000 _MAX_READ_MAX_LINES = 10000 _DEFAULT_TOOL_RESULT_MAX_BYTES = 50 * 1024 _BOX_FILE_SCRIPT_MAX_BYTES = 2048 +_MAX_HOST_EDIT_FILE_BYTES = 1024 * 1024 _GLOB_MAX_MATCHES = 100 +_FILE_WALK_MAX_ENTRIES = 100_000 +_DIRECTORY_MAX_ENTRIES = 10_000 _GREP_MAX_MATCHES = 200 _GREP_MAX_FILES = 5000 _GREP_MAX_LINE_CHARS = 500 +_GREP_MAX_SCAN_LINE_CHARS = 1024 * 1024 +_GREP_MAX_FILE_SCAN_CHARS = 10 * 1024 * 1024 +_GREP_MAX_TOTAL_SCAN_CHARS = 50 * 1024 * 1024 +_GREP_MAX_PATTERN_CHARS = 1024 +_GREP_REGEX_TIMEOUT_SECONDS = 0.25 _DIRECTORY_OPEN_FLAGS = ( os.O_RDONLY | getattr(os, 'O_DIRECTORY', 0) | getattr(os, 'O_NOFOLLOW', 0) | getattr(os, 'O_CLOEXEC', 0) @@ -433,7 +446,19 @@ class NativeToolLoader(loader.ToolLoader): with _open_location_fd(root_fd, location.relative_parts, os.O_RDONLY) as target_fd: metadata = os.fstat(target_fd) if stat.S_ISDIR(metadata.st_mode): - return self._build_directory_result(os.listdir(target_fd)) + entries: list[str] = [] + truncated = False + with os.scandir(target_fd) as iterator: + for entry in iterator: + if len(entries) >= _DIRECTORY_MAX_ENTRIES: + truncated = True + break + entries.append(entry.name) + return self._build_directory_result( + entries, + total=len(entries) + int(truncated), + force_truncated_by='entries' if truncated else None, + ) if not stat.S_ISREG(metadata.st_mode): raise ValueError('Path must reference a regular file or directory.') return self._read_text_file_preview(target_fd, parameters, metadata=metadata) @@ -479,10 +504,16 @@ class NativeToolLoader(loader.ToolLoader): with _open_host_root(location, create=False) as root_fd: with _open_location_fd(root_fd, location.relative_parts, os.O_RDWR) as target_fd: - if not stat.S_ISREG(os.fstat(target_fd).st_mode): + metadata = os.fstat(target_fd) + if not stat.S_ISREG(metadata.st_mode): return False, 'File not found.' - with os.fdopen(os.dup(target_fd), 'r', encoding='utf-8', errors='replace') as file_obj: - content = file_obj.read() + if metadata.st_size > _MAX_HOST_EDIT_FILE_BYTES: + return False, f'File exceeds the {_MAX_HOST_EDIT_FILE_BYTES}-byte edit limit.' + with os.fdopen(os.dup(target_fd), 'rb') as file_obj: + raw_content = file_obj.read(_MAX_HOST_EDIT_FILE_BYTES + 1) + if len(raw_content) > _MAX_HOST_EDIT_FILE_BYTES: + return False, f'File exceeds the {_MAX_HOST_EDIT_FILE_BYTES}-byte edit limit.' + content = raw_content.decode('utf-8', errors='replace') count = content.count(old_string) if count == 0: return False, 'old_string not found in file.' @@ -490,6 +521,8 @@ class NativeToolLoader(loader.ToolLoader): return False, f'old_string matches {count} locations; provide a more unique string.' payload = content.replace(old_string, new_string, 1).encode('utf-8') + if len(payload) > _MAX_HOST_EDIT_FILE_BYTES: + return False, f'Edited file exceeds the {_MAX_HOST_EDIT_FILE_BYTES}-byte limit.' os.ftruncate(target_fd, 0) os.lseek(target_fd, 0, os.SEEK_SET) self._write_all(target_fd, payload) @@ -520,11 +553,19 @@ class NativeToolLoader(loader.ToolLoader): return any(candidate and PurePosixPath(relative_path).match(candidate) for candidate in candidates) def _glob_host_location(self, location: _HostLocation, pattern: str, sandbox_base: str) -> dict: - hits: list[tuple[str, float]] = [] + newest_hits: list[tuple[float, str]] = [] + total = 0 + entries_seen = 0 + scan_truncated = False - def walk(directory_fd: int, prefix: str) -> None: + def walk(directory_fd: int, prefix: str) -> bool: + nonlocal entries_seen, scan_truncated, total with os.scandir(directory_fd) as entries: for entry in entries: + entries_seen += 1 + if entries_seen > _FILE_WALK_MAX_ENTRIES: + scan_truncated = True + return True name = entry.name if name in _SKIP_DIRS: continue @@ -536,11 +577,17 @@ class NativeToolLoader(loader.ToolLoader): metadata = os.fstat(child_fd) relative = f'{prefix}/{name}' if prefix else name if self._rglob_matches(relative, pattern): - hits.append((relative, metadata.st_mtime)) - if stat.S_ISDIR(metadata.st_mode): - walk(child_fd, relative) + total += 1 + candidate = (metadata.st_mtime, relative) + if len(newest_hits) < _GLOB_MAX_MATCHES: + heapq.heappush(newest_hits, candidate) + elif candidate > newest_hits[0]: + heapq.heapreplace(newest_hits, candidate) + if stat.S_ISDIR(metadata.st_mode) and walk(child_fd, relative): + return True finally: os.close(child_fd) + return False with _open_host_root(location, create=False) as root_fd: with _open_location_fd(root_fd, location.relative_parts, os.O_RDONLY) as target_fd: @@ -548,12 +595,11 @@ class NativeToolLoader(loader.ToolLoader): return {'ok': False, 'error': f'Path is not a directory: {sandbox_base}'} walk(target_fd, '') - hits.sort(key=lambda item: item[1], reverse=True) - total = len(hits) + hits = sorted(newest_hits, reverse=True) sandbox_paths: list[str] = [] output_bytes = 0 truncated_by_bytes = False - for relative, _mtime in hits[:_GLOB_MAX_MATCHES]: + for _mtime, relative in hits: sandbox_path = self._sandbox_child_path(sandbox_base, relative) entry_bytes = len(sandbox_path.encode('utf-8')) + (1 if sandbox_paths else 0) if output_bytes + entry_bytes > _DEFAULT_TOOL_RESULT_MAX_BYTES: @@ -567,27 +613,57 @@ class NativeToolLoader(loader.ToolLoader): 'matches': sandbox_paths, 'preview': '\n'.join(sandbox_paths), 'total': total, - 'truncated': total > len(sandbox_paths) or truncated_by_bytes, - 'truncated_by': 'bytes' if truncated_by_bytes else ('matches' if total > len(sandbox_paths) else None), + 'truncated': scan_truncated or total > len(sandbox_paths) or truncated_by_bytes, + 'truncated_by': ( + 'scan' + if scan_truncated + else ('bytes' if truncated_by_bytes else ('matches' if total > len(sandbox_paths) else None)) + ), } def _grep_host_location( self, location: _HostLocation, - regex, + pattern: str, include: str | None, sandbox_base: str, ) -> dict: + try: + compiled = regex.compile(pattern) + except regex.error as exc: + return {'ok': False, 'error': f'Invalid regex: {exc}'} + matches: list[dict] = [] output_bytes = 0 truncated_by: str | None = None files_seen = 0 + entries_seen = 0 + total_chars_seen = 0 + deadline = time.monotonic() + _GREP_REGEX_TIMEOUT_SECONDS def grep_file(file_fd: int, sandbox_path: str) -> bool: - nonlocal output_bytes, truncated_by + nonlocal output_bytes, total_chars_seen, truncated_by + file_chars_seen = 0 with os.fdopen(os.dup(file_fd), 'r', encoding='utf-8', errors='ignore') as handle: - for lineno, line in enumerate(handle, 1): - if not regex.search(line): + lineno = 0 + while True: + line = handle.readline(_GREP_MAX_SCAN_LINE_CHARS + 1) + if not line: + break + lineno += 1 + line_chars = len(line) + file_chars_seen += line_chars + total_chars_seen += line_chars + if file_chars_seen > _GREP_MAX_FILE_SCAN_CHARS or total_chars_seen > _GREP_MAX_TOTAL_SCAN_CHARS: + truncated_by = 'scan' + return True + if line_chars > _GREP_MAX_SCAN_LINE_CHARS and not line.endswith('\n'): + truncated_by = truncated_by or 'line' + return False + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError + if not compiled.search(line, timeout=remaining, concurrent=True): continue content, line_truncated = self._truncate_grep_line(line.rstrip()) entry = {'file': sandbox_path, 'line': lineno, 'content': content} @@ -605,9 +681,13 @@ class NativeToolLoader(loader.ToolLoader): return False def walk(directory_fd: int, prefix: str) -> bool: - nonlocal files_seen + nonlocal entries_seen, files_seen, truncated_by with os.scandir(directory_fd) as entries: for entry in entries: + entries_seen += 1 + if entries_seen > _FILE_WALK_MAX_ENTRIES: + truncated_by = 'scan' + return True name = entry.name if name in _SKIP_DIRS: continue @@ -637,13 +717,16 @@ class NativeToolLoader(loader.ToolLoader): with _open_host_root(location, create=False) as root_fd: with _open_location_fd(root_fd, location.relative_parts, os.O_RDONLY) as target_fd: - metadata = os.fstat(target_fd) - if stat.S_ISREG(metadata.st_mode): - grep_file(target_fd, sandbox_base) - elif stat.S_ISDIR(metadata.st_mode): - walk(target_fd, '') - else: - return {'ok': False, 'error': f'Path not found: {sandbox_base}'} + try: + metadata = os.fstat(target_fd) + if stat.S_ISREG(metadata.st_mode): + grep_file(target_fd, sandbox_base) + elif stat.S_ISDIR(metadata.st_mode): + walk(target_fd, '') + else: + return {'ok': False, 'error': f'Path not found: {sandbox_base}'} + except TimeoutError: + return {'ok': False, 'error': 'Regex search timed out'} return { 'ok': True, @@ -701,9 +784,24 @@ if not path.startswith('/workspace'): elif not os.path.exists(path): print(json.dumps({{'ok': False, 'error': f'File not found: {{path}}'}})) elif os.path.isdir(path): - entries = sorted(os.listdir(path)) + entries = [] + directory_truncated = False + with os.scandir(path) as iterator: + for entry in iterator: + if len(entries) >= {_DIRECTORY_MAX_ENTRIES}: + directory_truncated = True + break + entries.append(entry.name) + entries.sort() content = '\\n'.join(entries) - print(json.dumps({{'ok': True, 'content': content, 'is_directory': True, 'total': len(entries), 'truncated': False}})) + print(json.dumps({{ + 'ok': True, + 'content': content, + 'is_directory': True, + 'total': len(entries) + int(directory_truncated), + 'truncated': directory_truncated, + 'truncated_by': 'entries' if directory_truncated else None, + }})) elif encoding == 'base64': size_bytes = os.path.getsize(path) with open(path, 'rb') as f: @@ -824,7 +922,7 @@ else: async def _glob_workspace_via_box(self, path: str, pattern: str, query: pipeline_query.Query) -> dict: script = f""" -import json, os +import heapq, json, os from pathlib import Path path = {json.dumps(path)} pattern = {json.dumps(pattern)} @@ -835,12 +933,28 @@ elif not os.path.isdir(path): print(json.dumps({{'ok': False, 'error': f'Path is not a directory: {{path}}'}})) else: base = Path(path) - hits = [ - item for item in base.rglob(pattern) - if not any(part in skip_dirs for part in item.parts) - ] - hits.sort(key=lambda item: item.stat().st_mtime if item.exists() else 0, reverse=True) - shown = hits[:{_GLOB_MAX_MATCHES}] + newest_hits = [] + total = 0 + entries_seen = 0 + scan_truncated = False + for item in base.rglob(pattern): + entries_seen += 1 + if entries_seen > {_FILE_WALK_MAX_ENTRIES}: + scan_truncated = True + break + if any(part in skip_dirs for part in item.parts): + continue + total += 1 + try: + mtime = item.stat().st_mtime + except OSError: + mtime = 0 + candidate = (mtime, str(item)) + if len(newest_hits) < {_GLOB_MAX_MATCHES}: + heapq.heappush(newest_hits, candidate) + elif candidate > newest_hits[0]: + heapq.heapreplace(newest_hits, candidate) + shown = [Path(item_path) for _mtime, item_path in sorted(newest_hits, reverse=True)] matches = [] output_bytes = 0 truncated_by_bytes = False @@ -857,9 +971,12 @@ else: 'ok': True, 'matches': matches, 'preview': '\\n'.join(matches), - 'total': len(hits), - 'truncated': len(hits) > len(matches) or truncated_by_bytes, - 'truncated_by': 'bytes' if truncated_by_bytes else ('matches' if len(hits) > len(matches) else None), + 'total': total, + 'truncated': scan_truncated or total > len(matches) or truncated_by_bytes, + 'truncated_by': ( + 'scan' if scan_truncated + else ('bytes' if truncated_by_bytes else ('matches' if total > len(matches) else None)) + ), }})) """.strip() return await self._run_workspace_file_script(script, query) @@ -872,12 +989,15 @@ else: query: pipeline_query.Query, ) -> dict: script = f""" -import json, os, re +import json, os, re, signal, time from pathlib import Path path = {json.dumps(path)} pattern = {json.dumps(pattern)} include = {json.dumps(include)} skip_dirs = {json.dumps(sorted(_SKIP_DIRS))} +def regex_timeout(_signum, _frame): + raise TimeoutError +signal.signal(signal.SIGALRM, regex_timeout) try: regex = re.compile(pattern) except re.error as exc: @@ -888,6 +1008,17 @@ else: elif not os.path.exists(path): print(json.dumps({{'ok': False, 'error': f'Path not found: {{path}}'}})) else: + regex_deadline = time.monotonic() + {_GREP_REGEX_TIMEOUT_SECONDS} + def bounded_search(value): + remaining = regex_deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError + signal.setitimer(signal.ITIMER_REAL, remaining) + try: + return regex.search(value) + finally: + signal.setitimer(signal.ITIMER_REAL, 0) + base = Path(path) if base.is_file(): files = [base] @@ -904,14 +1035,37 @@ else: matches = [] output_bytes = 0 truncated_by = None + total_chars_seen = 0 for fp in files: try: handle = fp.open('r', encoding='utf-8', errors='ignore') except OSError: continue + file_chars_seen = 0 with handle: - for lineno, line in enumerate(handle, 1): - if regex.search(line): + lineno = 0 + while True: + line = handle.readline({_GREP_MAX_SCAN_LINE_CHARS + 1}) + if not line: + break + lineno += 1 + file_chars_seen += len(line) + total_chars_seen += len(line) + if ( + file_chars_seen > {_GREP_MAX_FILE_SCAN_CHARS} + or total_chars_seen > {_GREP_MAX_TOTAL_SCAN_CHARS} + ): + truncated_by = 'scan' + break + if len(line) > {_GREP_MAX_SCAN_LINE_CHARS} and not line.endswith('\\n'): + truncated_by = truncated_by or 'line' + break + try: + matched = bounded_search(line) + except TimeoutError: + print(json.dumps({{'ok': False, 'error': 'Regex search timed out'}})) + raise SystemExit(0) + if matched: if base.is_file(): file_path = path else: @@ -934,9 +1088,9 @@ else: if len(matches) >= {_GREP_MAX_MATCHES}: truncated_by = truncated_by or 'matches' break - if truncated_by == 'bytes' or len(matches) >= {_GREP_MAX_MATCHES}: + if truncated_by in ('bytes', 'scan') or len(matches) >= {_GREP_MAX_MATCHES}: break - if truncated_by == 'bytes' or len(matches) >= {_GREP_MAX_MATCHES}: + if truncated_by in ('bytes', 'scan') or len(matches) >= {_GREP_MAX_MATCHES}: break print(json.dumps({{ @@ -966,7 +1120,7 @@ else: host_location = None if host_location is not None: try: - return self._read_host_location(host_location, parameters) + return await asyncio.to_thread(self._read_host_location, host_location, parameters) except FileNotFoundError: pass @@ -998,7 +1152,7 @@ else: if self._should_use_box_workspace_files(host_location.selected_skill): return await self._read_workspace_via_box(path, parameters, query) try: - return self._read_host_location(host_location, parameters) + return await asyncio.to_thread(self._read_host_location, host_location, parameters) except (FileNotFoundError, NotADirectoryError): return {'ok': False, 'error': f'File not found: {path}'} @@ -1031,7 +1185,7 @@ else: if self._should_use_box_workspace_files(host_location.selected_skill): return await self._write_workspace_via_box(path, content, parameters, query) try: - self._write_host_location(host_location, content, parameters) + await run_blocking_atomic(self._write_host_location, host_location, content, parameters) except ValueError as exc: return {'ok': False, 'error': str(exc)} self._refresh_skill_from_disk(query, host_location.selected_skill) @@ -1091,7 +1245,12 @@ else: if self._should_use_box_workspace_files(host_location.selected_skill): return await self._edit_workspace_via_box(path, old_string, new_string, query) try: - changed, error = self._edit_host_location(host_location, old_string, new_string) + changed, error = await run_blocking_atomic( + self._edit_host_location, + host_location, + old_string, + new_string, + ) except (FileNotFoundError, NotADirectoryError): return {'ok': False, 'error': f'File not found: {path}'} if not changed: @@ -1364,7 +1523,7 @@ else: if self._should_use_box_workspace_files(host_location.selected_skill): return await self._glob_workspace_via_box(path, pattern, query) try: - return self._glob_host_location(host_location, pattern, path) + return await asyncio.to_thread(self._glob_host_location, host_location, pattern, path) except (FileNotFoundError, NotADirectoryError): return {'ok': False, 'error': f'Path is not a directory: {path}'} @@ -1374,12 +1533,8 @@ else: include = parameters.get('include') self.ap.logger.info(f'grep tool invoked: query_id={query.query_id} pattern={pattern} path={path}') - import re - - try: - regex = re.compile(pattern) - except re.error as e: - return {'ok': False, 'error': f'Invalid regex: {e}'} + if not isinstance(pattern, str) or len(pattern) > _GREP_MAX_PATTERN_CHARS: + return {'ok': False, 'error': f'Regex patterns may contain at most {_GREP_MAX_PATTERN_CHARS} characters'} host_location = self._resolve_host_location( query, @@ -1390,7 +1545,13 @@ else: if self._should_use_box_workspace_files(host_location.selected_skill): return await self._grep_workspace_via_box(path, pattern, include, query) try: - return self._grep_host_location(host_location, regex, include, path) + return await asyncio.to_thread( + self._grep_host_location, + host_location, + pattern, + include, + path, + ) except (FileNotFoundError, NotADirectoryError): return {'ok': False, 'error': f'Path not found: {path}'} @@ -1430,18 +1591,24 @@ else: normalized['truncated_by'] = 'bytes' return normalized - def _build_directory_result(self, entries: list[str]) -> dict: + def _build_directory_result( + self, + entries: list[str], + *, + total: int | None = None, + force_truncated_by: str | None = None, + ) -> dict: sorted_entries = sorted(str(entry) for entry in entries) content = '\n'.join(sorted_entries) preview = self._truncate_text_to_bytes(content, _DEFAULT_TOOL_RESULT_MAX_BYTES) - truncated = preview != content + truncated_by = force_truncated_by or ('bytes' if preview != content else None) return { 'ok': True, 'content': preview, 'is_directory': True, - 'total': len(sorted_entries), - 'truncated': truncated, - 'truncated_by': 'bytes' if truncated else None, + 'total': len(sorted_entries) if total is None else total, + 'truncated': truncated_by is not None, + 'truncated_by': truncated_by, } def _read_text_file_preview(self, file_fd: int, parameters: dict, *, metadata: os.stat_result) -> dict: diff --git a/src/langbot/pkg/rag/knowledge/kbmgr.py b/src/langbot/pkg/rag/knowledge/kbmgr.py index 9d87a6e9c..8042b3271 100644 --- a/src/langbot/pkg/rag/knowledge/kbmgr.py +++ b/src/langbot/pkg/rag/knowledge/kbmgr.py @@ -17,11 +17,18 @@ from langbot.pkg.api.http.service.tenant import TenantContext, require_workspace from langbot.pkg.core import app, taskmgr from langbot.pkg.core.task_boundary import run_in_workspace_uow from langbot.pkg.entity.persistence import rag as persistence_rag -from langbot.pkg.workspace.errors import WorkspaceNotFoundError +from langbot.pkg.workspace.errors import WorkspaceInvariantError, WorkspaceNotFoundError from .base import KnowledgeBaseInterface +_MAX_ZIP_ARCHIVE_ENTRIES = 1024 +_MAX_ZIP_DOCUMENTS = 8 +_MAX_ZIP_FILE_BYTES = 10 * 1024 * 1024 +_MAX_ZIP_UNCOMPRESSED_BYTES = 40 * 1024 * 1024 +_MAX_ZIP_COMPRESSION_RATIO = 100 + + class RuntimeKnowledgeBase(KnowledgeBaseInterface): ap: app.Application @@ -261,21 +268,29 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface): # run background task asynchronously ctx = taskmgr.TaskContext.new() - wrapper = self.ap.task_mgr.create_user_task( - self._store_file_task( - execution_context, - file_obj, - task_context=ctx, - parser_plugin_id=parser_plugin_id, - ), - kind='knowledge-operation', - name=f'knowledge-store-file-{file_id}', - label=f'Store file {file_id}', - 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( + self._store_file_task( + execution_context, + file_obj, + task_context=ctx, + parser_plugin_id=parser_plugin_id, + ), + kind='knowledge-operation', + name=f'knowledge-store-file-{file_id}', + label=f'Store file {file_id}', + context=ctx, + instance_uuid=execution_context.instance_uuid, + workspace_uuid=execution_context.workspace_uuid, + placement_generation=execution_context.placement_generation, + ) + except taskmgr.TaskCapacityError: + await self.ap.persistence_mgr.execute_async( + sqlalchemy.delete(persistence_rag.File) + .where(persistence_rag.File.workspace_uuid == execution_context.workspace_uuid) + .where(persistence_rag.File.uuid == file_uuid) + ) + raise return wrapper.id async def _store_zip_file( @@ -301,9 +316,21 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface): try: # use utf-8 encoding with zipfile.ZipFile(io.BytesIO(zip_bytes), 'r', metadata_encoding='utf-8') as zip_ref: + if len(zip_ref.filelist) > _MAX_ZIP_ARCHIVE_ENTRIES: + raise ValueError('ZIP archive contains too many entries') + + supported_files: list[zipfile.ZipInfo] = [] + total_uncompressed_bytes = 0 for file_info in zip_ref.filelist: # skip directories and hidden files - if file_info.is_dir() or file_info.filename.startswith('.'): + normalized_name = file_info.filename.replace('\\', '/').strip('/') + path_parts = normalized_name.split('/') + if ( + file_info.is_dir() + or not normalized_name + or any(part.startswith('.') for part in path_parts) + or '__MACOSX' in path_parts + ): continue _, file_ext = os.path.splitext(file_info.filename) @@ -311,17 +338,30 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface): if file_extension not in supported_extensions: self.ap.logger.debug(f'Skipping unsupported file in ZIP: {file_info.filename}') continue + if file_info.flag_bits & 0x1: + raise ValueError('Encrypted ZIP entries are not supported') + if file_info.file_size > _MAX_ZIP_FILE_BYTES: + raise ValueError(f'ZIP document exceeds the file size limit: {file_info.filename}') + if ( + file_info.file_size + and file_info.file_size > max(file_info.compress_size, 1) * _MAX_ZIP_COMPRESSION_RATIO + ): + raise ValueError(f'ZIP document exceeds the compression-ratio limit: {file_info.filename}') + total_uncompressed_bytes += file_info.file_size + if total_uncompressed_bytes > _MAX_ZIP_UNCOMPRESSED_BYTES: + raise ValueError('ZIP documents exceed the uncompressed size limit') + supported_files.append(file_info) + if len(supported_files) > _MAX_ZIP_DOCUMENTS: + raise ValueError('ZIP archive contains too many supported documents') + for file_info in supported_files: try: - file_content = zip_ref.read(file_info.filename) + file_content = await asyncio.to_thread(zip_ref.read, file_info) base_name = file_info.filename.replace('/', '_').replace('\\', '_') file_stem, file_ext = os.path.splitext(base_name) extension = file_ext.lstrip('.') - if file_stem.startswith('__MACOSX'): - continue - extracted_file_id = file_stem + '_' + str(uuid.uuid4())[:8] + '.' + extension extracted_object_key = await self.ap.storage_mgr.save_scoped( execution_context, @@ -350,6 +390,8 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface): f'Extracted and stored file from ZIP: {file_info.filename} -> {extracted_object_key}' ) + except taskmgr.TaskCapacityError: + raise except Exception as e: self.ap.logger.warning(f'Failed to extract file {file_info.filename} from ZIP: {e}') continue @@ -580,6 +622,53 @@ class RAGManager: def __init__(self, ap: app.Application): self.ap = ap self.knowledge_bases = {} + self._scope_generations: dict[tuple[str, str], int] = {} + self._knowledge_keys_by_scope: dict[ + tuple[str, str], + set[tuple[str, str]], + ] = {} + + def _cache_runtime( + self, + runtime: RuntimeKnowledgeBase, + ) -> None: + context = runtime.execution_context + self._observe_execution_context(context) + key = ( + context.workspace_uuid, + runtime.get_uuid(), + ) + self.knowledge_bases[key] = runtime + scope = (context.instance_uuid, context.workspace_uuid) + self._knowledge_keys_by_scope.setdefault(scope, set()).add(key) + + def _pop_runtime( + self, + context: ExecutionContext, + kb_uuid: str, + ) -> RuntimeKnowledgeBase | None: + key = (context.workspace_uuid, kb_uuid) + runtime = self.knowledge_bases.pop(key, None) + scope = (context.instance_uuid, context.workspace_uuid) + keys = self._knowledge_keys_by_scope.get(scope) + if keys is not None: + keys.discard(key) + if not keys: + self._knowledge_keys_by_scope.pop(scope, None) + self._scope_generations.pop(scope, None) + return runtime + + 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('RAG runtime placement generation rolled back') + if previous_generation == context.placement_generation: + return + if previous_generation is not None: + for key in self._knowledge_keys_by_scope.pop(scope, ()): + self.knowledge_bases.pop(key, None) + self._scope_generations[scope] = context.placement_generation async def initialize(self): await self.load_knowledge_bases_from_db() @@ -587,6 +676,8 @@ class RAGManager: async def _to_execution_context( self, context: RequestContext | ExecutionContext, + *, + _binding_validated: bool = False, ) -> ExecutionContext: if isinstance(context, RequestContext): execution_context = ExecutionContext.from_request(context) @@ -595,12 +686,19 @@ class RAGManager: else: raise WorkspaceRequiredError('RequestContext or ExecutionContext is required') - binding = await self.ap.workspace_service.get_execution_binding( + if not _binding_validated: + binding = await self.ap.workspace_service.get_execution_binding( + execution_context.workspace_uuid, + expected_generation=execution_context.placement_generation, + ) + if binding.instance_uuid != execution_context.instance_uuid: + raise WorkspaceNotFoundError('Workspace not found') + scope = ( + execution_context.instance_uuid, execution_context.workspace_uuid, - expected_generation=execution_context.placement_generation, ) - if binding.instance_uuid != execution_context.instance_uuid: - raise WorkspaceNotFoundError('Workspace not found') + if scope in self._scope_generations: + self._observe_execution_context(execution_context) return execution_context async def _get_engine_map(self, context: TenantContext) -> dict[str, dict]: @@ -748,7 +846,7 @@ class RAGManager: try: await runtime_kb._on_kb_create(execution_context) except Exception: - self.knowledge_bases.pop((execution_context.workspace_uuid, kb_uuid), None) + self._pop_runtime(execution_context, kb_uuid) await self.ap.persistence_mgr.execute_async( sqlalchemy.delete(persistence_rag.KnowledgeBase) .where(persistence_rag.KnowledgeBase.workspace_uuid == execution_context.workspace_uuid) @@ -763,6 +861,8 @@ class RAGManager: self.ap.logger.info('Loading knowledge bases from db...') self.knowledge_bases = {} + self._scope_generations = {} + self._knowledge_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) @@ -786,6 +886,7 @@ class RAGManager: placement_generation=binding.placement_generation, ), knowledge_base, + _binding_validated=True, ) except Exception as e: self.ap.logger.error( @@ -815,6 +916,8 @@ class RAGManager: self, context: RequestContext | ExecutionContext, knowledge_base_entity: persistence_rag.KnowledgeBase | sqlalchemy.Row | dict, + *, + _binding_validated: bool = False, ) -> RuntimeKnowledgeBase: if isinstance(knowledge_base_entity, sqlalchemy.Row): # Safe access to _mapping for SQLAlchemy 1.4+ @@ -826,7 +929,10 @@ class RAGManager: } knowledge_base_entity = persistence_rag.KnowledgeBase(**filtered_dict) - execution_context = await self._to_execution_context(context) + execution_context = await self._to_execution_context( + context, + _binding_validated=_binding_validated, + ) if knowledge_base_entity.workspace_uuid != execution_context.workspace_uuid: raise WorkspaceNotFoundError('Knowledge base not found') runtime_knowledge_base = RuntimeKnowledgeBase( @@ -837,9 +943,7 @@ class RAGManager: await runtime_knowledge_base.initialize() - self.knowledge_bases[(execution_context.workspace_uuid, runtime_knowledge_base.get_uuid())] = ( - runtime_knowledge_base - ) + self._cache_runtime(runtime_knowledge_base) return runtime_knowledge_base @@ -857,7 +961,7 @@ class RAGManager: kb_uuid: str, ) -> None: execution_context = await self._to_execution_context(context) - self.knowledge_bases.pop((execution_context.workspace_uuid, kb_uuid), None) + self._pop_runtime(execution_context, kb_uuid) async def delete_knowledge_base( self, @@ -865,7 +969,7 @@ class RAGManager: kb_uuid: str, ) -> None: execution_context = await self._to_execution_context(context) - kb = self.knowledge_bases.pop((execution_context.workspace_uuid, kb_uuid), None) + kb = self._pop_runtime(execution_context, kb_uuid) if kb is not None: await kb.dispose(execution_context) else: diff --git a/src/langbot/pkg/skill/manager.py b/src/langbot/pkg/skill/manager.py index 42ce7acda..aeb688b7f 100644 --- a/src/langbot/pkg/skill/manager.py +++ b/src/langbot/pkg/skill/manager.py @@ -58,6 +58,9 @@ class SkillManager: async def reload_skills(self, context: TenantContext) -> None: execution_context = self._execution_context(context) key = self._scope_key(execution_context) + for existing_key in tuple(self._skills_by_scope): + if existing_key[:2] == key[:2] and existing_key != key: + self._skills_by_scope.pop(existing_key, None) self._skills_by_scope[key] = {} box_service = getattr(self.ap, 'box_service', None) diff --git a/src/langbot/pkg/storage/mgr.py b/src/langbot/pkg/storage/mgr.py index 354bd2f29..1ec20375c 100644 --- a/src/langbot/pkg/storage/mgr.py +++ b/src/langbot/pkg/storage/mgr.py @@ -6,6 +6,7 @@ import re from pathlib import PurePath from ..core import app +from ..utils import bounded_executor from ..api.http.authz import WorkspaceRequiredError from ..api.http.context import ExecutionContext, RequestContext from . import provider @@ -212,16 +213,17 @@ class StorageMgr: return None workspace_uuid = match.group('workspace') generation = int(match.group('generation')) - try: - await self.ap.workspace_service.get_execution_binding( - workspace_uuid, - expected_generation=generation, - ) - except Exception: - return None - if not await self.storage_provider.exists(object_key): - return None - return await self.storage_provider.load(object_key) + with bounded_executor.blocking_work_scope(workspace_uuid): + try: + await self.ap.workspace_service.get_execution_binding( + workspace_uuid, + expected_generation=generation, + ) + except Exception: + return None + if not await self.storage_provider.exists(object_key): + return None + return await self.storage_provider.load(object_key) @classmethod def require_scoped_object_key( @@ -331,3 +333,8 @@ class StorageMgr: self.ap.logger.info('Initialized local storage backend.') await self.storage_provider.initialize() + + async def shutdown(self) -> None: + storage_provider = getattr(self, 'storage_provider', None) + if storage_provider is not None: + await storage_provider.shutdown() diff --git a/src/langbot/pkg/storage/provider.py b/src/langbot/pkg/storage/provider.py index e24dcbf97..4e46f09c7 100644 --- a/src/langbot/pkg/storage/provider.py +++ b/src/langbot/pkg/storage/provider.py @@ -14,6 +14,11 @@ class StorageProvider(abc.ABC): async def initialize(self): pass + async def shutdown(self) -> None: + """Release provider-owned clients or pools.""" + + return None + @abc.abstractmethod async def save( self, diff --git a/src/langbot/pkg/storage/providers/localstorage.py b/src/langbot/pkg/storage/providers/localstorage.py index 43566fc7b..c64698232 100644 --- a/src/langbot/pkg/storage/providers/localstorage.py +++ b/src/langbot/pkg/storage/providers/localstorage.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import os import aiofiles import shutil @@ -40,10 +41,9 @@ class LocalStorageProvider(provider.StorageProvider): key: str, value: bytes, ): - resolved = _safe_resolve(LOCAL_STORAGE_PATH, key) + resolved = await asyncio.to_thread(_safe_resolve, LOCAL_STORAGE_PATH, key) parent = os.path.dirname(resolved) - if not os.path.exists(parent): - os.makedirs(parent) + await asyncio.to_thread(os.makedirs, parent, exist_ok=True) async with aiofiles.open(resolved, 'wb') as f: await f.write(value) @@ -51,7 +51,7 @@ class LocalStorageProvider(provider.StorageProvider): self, key: str, ) -> bytes: - resolved = _safe_resolve(LOCAL_STORAGE_PATH, key) + resolved = await asyncio.to_thread(_safe_resolve, LOCAL_STORAGE_PATH, key) async with aiofiles.open(resolved, 'rb') as f: return await f.read() @@ -59,28 +59,31 @@ class LocalStorageProvider(provider.StorageProvider): self, key: str, ) -> bool: - resolved = _safe_resolve(LOCAL_STORAGE_PATH, key) - return os.path.exists(resolved) + resolved = await asyncio.to_thread(_safe_resolve, LOCAL_STORAGE_PATH, key) + return await asyncio.to_thread(os.path.exists, resolved) async def delete( self, key: str, ): - resolved = _safe_resolve(LOCAL_STORAGE_PATH, key) - os.remove(resolved) + resolved = await asyncio.to_thread(_safe_resolve, LOCAL_STORAGE_PATH, key) + await asyncio.to_thread(os.remove, resolved) async def size( self, key: str, ) -> int: - resolved = _safe_resolve(LOCAL_STORAGE_PATH, key) - return os.path.getsize(resolved) + resolved = await asyncio.to_thread(_safe_resolve, LOCAL_STORAGE_PATH, key) + return await asyncio.to_thread(os.path.getsize, resolved) async def delete_dir_recursive( self, dir_path: str, ): - resolved = _safe_resolve(LOCAL_STORAGE_PATH, dir_path) - # 直接删除整个目录 - if os.path.exists(resolved): - shutil.rmtree(resolved) + resolved = await asyncio.to_thread( + _safe_resolve, + LOCAL_STORAGE_PATH, + dir_path, + ) + if await asyncio.to_thread(os.path.exists, resolved): + await asyncio.to_thread(shutil.rmtree, resolved) diff --git a/src/langbot/pkg/storage/providers/s3storage.py b/src/langbot/pkg/storage/providers/s3storage.py index 43cc2e966..8729da494 100644 --- a/src/langbot/pkg/storage/providers/s3storage.py +++ b/src/langbot/pkg/storage/providers/s3storage.py @@ -1,9 +1,12 @@ from __future__ import annotations +import asyncio + import boto3 from botocore.exceptions import ClientError from ...core import app +from ...utils import bounded_executor from .. import provider @@ -14,6 +17,7 @@ class S3StorageProvider(provider.StorageProvider): super().__init__(ap) self.s3_client = None self.bucket_name = None + self._io_semaphore = asyncio.Semaphore(16) async def initialize(self): """Initialize S3 client with configuration from config.yaml""" @@ -26,6 +30,11 @@ class S3StorageProvider(provider.StorageProvider): secret_access_key = s3_config.get('secret_access_key', '') region_name = s3_config.get('region', 'us-east-1') self.bucket_name = s3_config.get('bucket', 'langbot-storage') + try: + max_concurrency = int(s3_config.get('max_concurrency', 16)) + except (TypeError, ValueError): + max_concurrency = 16 + self._io_semaphore = asyncio.Semaphore(max(1, min(max_concurrency, 128))) # Initialize S3 client session = boto3.session.Session() @@ -37,7 +46,25 @@ class S3StorageProvider(provider.StorageProvider): aws_secret_access_key=secret_access_key, ) - # Ensure bucket exists + await self._run_io(self._ensure_bucket) + + async def shutdown(self) -> None: + """Close the botocore HTTP connection pool without blocking the loop.""" + + client = self.s3_client + self.s3_client = None + if client is not None: + await bounded_executor.run_blocking_cleanup(client.close) + + async def _run_io(self, operation, /, *args, **kwargs): + """Run one blocking boto3 operation behind a bounded concurrency gate.""" + + async with self._io_semaphore: + return await asyncio.to_thread(operation, *args, **kwargs) + + def _ensure_bucket(self) -> None: + """Probe/create the bucket without blocking the application event loop.""" + try: self.s3_client.head_bucket(Bucket=self.bucket_name) except ClientError as e: @@ -61,7 +88,8 @@ class S3StorageProvider(provider.StorageProvider): ): """Save bytes to S3""" try: - self.s3_client.put_object( + await self._run_io( + self.s3_client.put_object, Bucket=self.bucket_name, Key=key, Body=value, @@ -76,22 +104,30 @@ class S3StorageProvider(provider.StorageProvider): ) -> bytes: """Load bytes from S3""" try: - response = self.s3_client.get_object( - Bucket=self.bucket_name, - Key=key, - ) - return response['Body'].read() + return await self._run_io(self._load_sync, key) except Exception as e: self.ap.logger.error(f'Failed to load from S3: {e}') raise + def _load_sync(self, key: str) -> bytes: + response = self.s3_client.get_object( + Bucket=self.bucket_name, + Key=key, + ) + body = response['Body'] + try: + return body.read() + finally: + body.close() + async def exists( self, key: str, ) -> bool: """Check if object exists in S3""" try: - self.s3_client.head_object( + await self._run_io( + self.s3_client.head_object, Bucket=self.bucket_name, Key=key, ) @@ -109,7 +145,8 @@ class S3StorageProvider(provider.StorageProvider): ): """Delete object from S3""" try: - self.s3_client.delete_object( + await self._run_io( + self.s3_client.delete_object, Bucket=self.bucket_name, Key=key, ) @@ -123,7 +160,8 @@ class S3StorageProvider(provider.StorageProvider): ) -> int: """Get object size from S3 without downloading it""" try: - response = self.s3_client.head_object( + response = await self._run_io( + self.s3_client.head_object, Bucket=self.bucket_name, Key=key, ) @@ -138,23 +176,23 @@ class S3StorageProvider(provider.StorageProvider): ): """Delete all objects with the given prefix (directory)""" try: - # Ensure dir_path ends with / - if not dir_path.endswith('/'): - dir_path = dir_path + '/' - - # List all objects with the prefix - paginator = self.s3_client.get_paginator('list_objects_v2') - pages = paginator.paginate(Bucket=self.bucket_name, Prefix=dir_path) - - # Delete all objects - for page in pages: - if 'Contents' in page: - objects_to_delete = [{'Key': obj['Key']} for obj in page['Contents']] - if objects_to_delete: - self.s3_client.delete_objects( - Bucket=self.bucket_name, - Delete={'Objects': objects_to_delete}, - ) + await self._run_io(self._delete_dir_recursive_sync, dir_path) except Exception as e: self.ap.logger.error(f'Failed to delete directory from S3: {e}') raise + + def _delete_dir_recursive_sync(self, dir_path: str) -> None: + if not dir_path.endswith('/'): + dir_path = dir_path + '/' + + paginator = self.s3_client.get_paginator('list_objects_v2') + pages = paginator.paginate(Bucket=self.bucket_name, Prefix=dir_path) + for page in pages: + if 'Contents' not in page: + continue + objects_to_delete = [{'Key': obj['Key']} for obj in page['Contents']] + if objects_to_delete: + self.s3_client.delete_objects( + Bucket=self.bucket_name, + Delete={'Objects': objects_to_delete}, + ) diff --git a/src/langbot/pkg/survey/manager.py b/src/langbot/pkg/survey/manager.py index da227c285..51bba58a7 100644 --- a/src/langbot/pkg/survey/manager.py +++ b/src/langbot/pkg/survey/manager.py @@ -2,7 +2,6 @@ from __future__ import annotations -import asyncio import contextlib import json import typing @@ -10,9 +9,10 @@ import httpx import sqlalchemy from ..core import app as core_app +from ..core import entities as core_entities from ..entity.persistence.metadata import Metadata from ..persistence.tenant_uow import CrossScopeTransactionError -from ..utils import constants +from ..utils import constants, httpclient SURVEY_TRIGGERED_KEY = 'survey_triggered_events' BOT_RESPONSE_COUNT_KEY = 'survey_bot_response_count' @@ -162,7 +162,13 @@ class SurveyManager: await self._save_triggered_events() # Check for pending survey asynchronously - asyncio.create_task(self._fetch_pending_survey(event)) + self.ap.task_mgr.create_task( + self._fetch_pending_survey(event), + kind='survey-fetch', + name=f'survey-fetch-{event}', + scopes=[core_entities.LifecycleControlScope.APPLICATION], + instance_uuid=self.ap.workspace_service.instance_uuid, + ) async def _fetch_pending_survey(self, event: str): """Fetch pending survey from Space for this event.""" @@ -172,10 +178,13 @@ class SurveyManager: 'instance_id': constants.instance_id, 'event': event, } - async with httpx.AsyncClient(timeout=httpx.Timeout(10)) as client: + async with httpx.AsyncClient( + timeout=httpx.Timeout(10), + event_hooks=httpclient.httpx_response_limit_hooks(), + ) as client: resp = await client.post(url, json=payload) if resp.status_code == 200: - data = resp.json() + data = await httpclient.parse_json_response(resp) if data.get('code') == 0 and data.get('data', {}).get('survey'): self._pending_survey = data['data']['survey'] self.ap.logger.info(f'Survey pending: {self._pending_survey.get("survey_id")}') @@ -218,7 +227,10 @@ class SurveyManager: 'metadata': await self._build_base_metadata(), 'completed': completed, } - async with httpx.AsyncClient(timeout=httpx.Timeout(10)) as client: + async with httpx.AsyncClient( + timeout=httpx.Timeout(10), + event_hooks=httpclient.httpx_response_limit_hooks(), + ) as client: resp = await client.post(url, json=payload) if resp.status_code == 200: self.clear_pending_survey() @@ -245,11 +257,15 @@ class SurveyManager: 'attachments': attachments, 'metadata': metadata, } - async with httpx.AsyncClient(timeout=httpx.Timeout(30)) as client: + async with httpx.AsyncClient( + timeout=httpx.Timeout(30), + event_hooks=httpclient.httpx_response_limit_hooks(), + ) as client: resp = await client.post(url, json=payload) if resp.status_code == 200: return True - self.ap.logger.warning(f'Failed to submit feedback: {resp.status_code} {resp.text[:200]}') + body = await httpclient.response_text(resp, max_chars=200) + self.ap.logger.warning(f'Failed to submit feedback: {resp.status_code} {body}') except Exception as e: self.ap.logger.warning(f'Failed to submit feedback: {e}') return False @@ -264,7 +280,10 @@ class SurveyManager: 'survey_id': survey_id, 'instance_id': constants.instance_id, } - async with httpx.AsyncClient(timeout=httpx.Timeout(10)) as client: + async with httpx.AsyncClient( + timeout=httpx.Timeout(10), + event_hooks=httpclient.httpx_response_limit_hooks(), + ) as client: resp = await client.post(url, json=payload) if resp.status_code == 200: self.clear_pending_survey() diff --git a/src/langbot/pkg/telemetry/heartbeat.py b/src/langbot/pkg/telemetry/heartbeat.py index 094818b62..3c3efb507 100644 --- a/src/langbot/pkg/telemetry/heartbeat.py +++ b/src/langbot/pkg/telemetry/heartbeat.py @@ -27,23 +27,25 @@ if typing.TYPE_CHECKING: HEARTBEAT_INTERVAL_SECONDS = 24 * 3600 -async def _count(ap: core_app.Application, table) -> int: +async def _count( + ap: core_app.Application, + table, + *, + cloud_counter: typing.Callable[[], int] | None = None, +) -> int: """Count rows in a persistence table; -1 when unavailable.""" try: persistence_mgr = ap.persistence_mgr cloud_runtime = getattr(getattr(persistence_mgr, 'mode', None), 'value', None) == 'cloud_runtime' if cloud_runtime: - tenant_uow = getattr(persistence_mgr, 'tenant_uow', None) - if not callable(tenant_uow): + # The Cloud runtime role deliberately cannot bypass RLS. Counting + # every tenant by opening one UoW per Workspace turns a best-effort + # daily heartbeat into thousands of serial SQL statements. The + # already-loaded runtime registries are authoritative for this + # process and provide an O(1), connection-free operational count. + if cloud_counter is None: return -1 - total = 0 - for binding in await ap.workspace_service.list_active_execution_bindings(): - async with tenant_uow(binding.workspace_uuid): - result = await persistence_mgr.execute_async( - sqlalchemy.select(sqlalchemy.func.count()).select_from(table) - ) - total += int(result.scalar() or 0) - return total + return max(int(cloud_counter()), 0) result = await ap.persistence_mgr.execute_async(sqlalchemy.select(sqlalchemy.func.count()).select_from(table)) return int(result.scalar() or 0) except Exception: @@ -95,11 +97,27 @@ async def build_heartbeat_payload(ap: core_app.Application) -> dict: pass # Resource counts - features['pipeline_count'] = await _count(ap, persistence_pipeline.LegacyPipeline) - features['mcp_server_count'] = await _count(ap, persistence_mcp.MCPServer) - features['knowledge_base_count'] = await _count(ap, persistence_rag.KnowledgeBase) + features['pipeline_count'] = await _count( + ap, + persistence_pipeline.LegacyPipeline, + cloud_counter=lambda: len(ap.pipeline_mgr._pipelines_by_key), + ) + features['mcp_server_count'] = await _count( + ap, + persistence_mcp.MCPServer, + cloud_counter=lambda: len(ap.tool_mgr.mcp_tool_loader._sessions), + ) + features['knowledge_base_count'] = await _count( + ap, + persistence_rag.KnowledgeBase, + cloud_counter=lambda: len(ap.rag_mgr.knowledge_bases), + ) if 'bot_count' not in features: - features['bot_count'] = await _count(ap, persistence_bot.Bot) + features['bot_count'] = await _count( + ap, + persistence_bot.Bot, + cloud_counter=lambda: len(ap.platform_mgr._bots_by_key), + ) # Plugin count (from plugin runtime) try: diff --git a/src/langbot/pkg/telemetry/telemetry.py b/src/langbot/pkg/telemetry/telemetry.py index 2948e268a..6c3b8ba29 100644 --- a/src/langbot/pkg/telemetry/telemetry.py +++ b/src/langbot/pkg/telemetry/telemetry.py @@ -1,8 +1,13 @@ from __future__ import annotations import asyncio +import contextlib import httpx from ..core import app as core_app +from ..utils import httpclient + + +_MAX_INFLIGHT_TELEMETRY_TASKS = 8 class TelemetryManager: @@ -18,13 +23,45 @@ class TelemetryManager: self.telemetry_config = {} self.send_tasks: list[asyncio.Task] = [] + self._client: httpx.AsyncClient | None = None async def initialize(self): self.telemetry_config = self.ap.instance_config.data.get('space', {}) async def start_send_task(self, payload: dict): + self.send_tasks = [task for task in self.send_tasks if not task.done()] + if len(self.send_tasks) >= _MAX_INFLIGHT_TELEMETRY_TASKS: + self.ap.logger.debug('Telemetry queue is full; dropping best-effort event') + return task = asyncio.create_task(self.send(payload)) self.send_tasks.append(task) + task.add_done_callback(self._send_task_done) + + def _send_task_done(self, task: asyncio.Task) -> None: + try: + self.send_tasks.remove(task) + except ValueError: + pass + + async def shutdown(self) -> None: + tasks = list(self.send_tasks) + for task in tasks: + task.cancel() + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + self.send_tasks.clear() + if self._client is not None: + await self._client.aclose() + self._client = None + + @contextlib.asynccontextmanager + async def _client_context(self): + if self._client is None or self._client.is_closed: + self._client = httpx.AsyncClient( + timeout=httpx.Timeout(10), + event_hooks=httpclient.httpx_response_limit_hooks(), + ) + yield self._client async def send(self, payload: dict): """Send telemetry payload to configured telemetry server (non-blocking). @@ -91,20 +128,21 @@ class TelemetryManager: except Exception: sanitized['duration_ms'] = 0 - async with httpx.AsyncClient(timeout=httpx.Timeout(10)) as client: + async with self._client_context() as client: try: # Use asyncio.wait_for to ensure we always bound the total time resp = await asyncio.wait_for(client.post(url, json=sanitized), timeout=10 + 1) if resp.status_code >= 400: + body = await httpclient.response_text(resp, max_chars=200) self.ap.logger.warning( - f'Telemetry post to {url} returned status {resp.status_code} - {resp.text}' + f'Telemetry post to {url} returned status {resp.status_code} - {body}' ) else: # Detect application-level errors inside HTTP 200 responses app_err = False try: - j = resp.json() + j = await httpclient.parse_json_response(resp) if isinstance(j, dict) and j.get('code') is not None and int(j.get('code')) >= 400: app_err = True self.ap.logger.warning( @@ -114,12 +152,14 @@ class TelemetryManager: pass if app_err: + body = await httpclient.response_text(resp, max_chars=200) self.ap.logger.warning( - f'Telemetry post to {url} returned app-level error - response: {resp.text[:200]}' + f'Telemetry post to {url} returned app-level error - response: {body}' ) else: + body = await httpclient.response_text(resp, max_chars=200) self.ap.logger.debug( - f'Telemetry posted to {url}, status {resp.status_code} - response: {resp.text[:200]}' + f'Telemetry posted to {url}, status {resp.status_code} - response: {body}' ) except asyncio.TimeoutError: self.ap.logger.warning(f'Telemetry post to {url} timed out') diff --git a/src/langbot/pkg/utils/bounded_executor.py b/src/langbot/pkg/utils/bounded_executor.py new file mode 100644 index 000000000..006714181 --- /dev/null +++ b/src/langbot/pkg/utils/bounded_executor.py @@ -0,0 +1,312 @@ +from __future__ import annotations + +import asyncio +import concurrent.futures +import contextlib +import contextvars +import threading +from collections.abc import Callable +from typing import Any + + +DEFAULT_MAX_WORKERS = 8 +DEFAULT_MAX_PENDING = 128 +DEFAULT_MAX_INFLIGHT_PER_SCOPE = 4 +HARD_MAX_WORKERS = 64 +HARD_MAX_PENDING = 4096 +BLOCKING_CLEANUP_SCOPE = 'system:cleanup' +_CLEANUP_RETRY_INITIAL_SECONDS = 0.01 +_CLEANUP_RETRY_MAX_SECONDS = 0.25 + +_blocking_work_scope: contextvars.ContextVar[str | None] = contextvars.ContextVar( + 'langbot_blocking_work_scope', + default=None, +) + + +class BlockingWorkCapacityError(RuntimeError): + """Raised before unbounded blocking work can enter the executor queue.""" + + def __init__(self, message: str, *, scope: str | None = None) -> None: + super().__init__(message) + self.scope = scope + + +@contextlib.contextmanager +def blocking_work_scope(scope: str | None): + """Attribute blocking submissions to one trusted tenant scope.""" + + normalized = str(scope).strip() if scope is not None else None + if not normalized: + yield + return + token = _blocking_work_scope.set(normalized) + try: + yield + finally: + _blocking_work_scope.reset(token) + + +def current_blocking_work_scope() -> str | None: + """Return the active trusted blocking-work scope, if any.""" + + return _blocking_work_scope.get() + + +async def run_blocking_atomic( + fn: Callable[..., Any], + /, + *args: Any, + **kwargs: Any, +) -> Any: + """Let an admitted filesystem operation finish before propagating cancel.""" + + task = asyncio.create_task(asyncio.to_thread(fn, *args, **kwargs)) + try: + return await asyncio.shield(task) + except asyncio.CancelledError: + await asyncio.gather(task, return_exceptions=True) + raise + + +async def run_blocking_cleanup( + fn: Callable[..., Any], + /, + *args: Any, + **kwargs: Any, +) -> Any: + """Wait for bounded executor capacity and complete cleanup atomically.""" + + retry_delay = _CLEANUP_RETRY_INITIAL_SECONDS + while True: + try: + with blocking_work_scope(BLOCKING_CLEANUP_SCOPE): + return await run_blocking_atomic(fn, *args, **kwargs) + except BlockingWorkCapacityError as exc: + if exc.scope != BLOCKING_CLEANUP_SCOPE: + raise + await asyncio.sleep(retry_delay) + retry_delay = min( + retry_delay * 2, + _CLEANUP_RETRY_MAX_SECONDS, + ) + + +async def run_in_blocking_work_scope( + coro, + scope: str | None, +): + """Run a coroutine with blocking-work fairness attribution.""" + + with blocking_work_scope(scope): + return await coro + + +def _bounded_integer( + value: Any, + *, + name: str, + minimum: int, + maximum: int, +) -> int: + if isinstance(value, bool): + raise ValueError(f'{name} must be an integer') + try: + parsed = int(value) + except (TypeError, ValueError) as exc: + raise ValueError(f'{name} must be an integer') from exc + if parsed < minimum or parsed > maximum: + raise ValueError(f'{name} must be between {minimum} and {maximum}') + return parsed + + +def _validated_limits( + max_workers: Any, + max_pending: Any, + max_inflight_per_scope: Any | None, +) -> tuple[int, int, int]: + workers = _bounded_integer( + max_workers, + name='blocking_executor.max_workers', + minimum=1, + maximum=HARD_MAX_WORKERS, + ) + pending = _bounded_integer( + max_pending, + name='blocking_executor.max_pending', + minimum=0, + maximum=HARD_MAX_PENDING, + ) + fair_share = max(1, workers // 2) + scope_limit = ( + min(DEFAULT_MAX_INFLIGHT_PER_SCOPE, fair_share) + if max_inflight_per_scope is None + else _bounded_integer( + max_inflight_per_scope, + name='blocking_executor.max_inflight_per_scope', + minimum=1, + maximum=HARD_MAX_PENDING, + ) + ) + if scope_limit > fair_share: + raise ValueError(f'blocking_executor.max_inflight_per_scope must not exceed half of max_workers ({fair_share})') + return workers, pending, scope_limit + + +class BoundedThreadPoolExecutor(concurrent.futures.ThreadPoolExecutor): + """Thread pool with a hard cap on running plus queued submissions.""" + + def __init__( + self, + *, + max_workers: int = DEFAULT_MAX_WORKERS, + max_pending: int = DEFAULT_MAX_PENDING, + max_inflight_per_scope: int | None = None, + thread_name_prefix: str = 'langbot-blocking', + ) -> None: + max_workers, max_pending, max_inflight_per_scope = _validated_limits( + max_workers, + max_pending, + max_inflight_per_scope, + ) + super().__init__( + max_workers=max_workers, + thread_name_prefix=thread_name_prefix, + ) + self.max_workers = max_workers + self.max_pending = max_pending + self.max_inflight_per_scope = max_inflight_per_scope + self._capacity = threading.BoundedSemaphore(max_workers + max_pending) + self._stats_lock = threading.Lock() + self._inflight_by_scope: dict[str, int] = {} + self._inflight = 0 + self._running = 0 + self._submitted_total = 0 + self._completed_total = 0 + self._rejected_total = 0 + self._global_rejected_total = 0 + self._scope_rejected_total = 0 + + def submit( + self, + fn: Callable[..., Any], + /, + *args: Any, + **kwargs: Any, + ) -> concurrent.futures.Future: + scope = current_blocking_work_scope() + if not self._capacity.acquire(blocking=False): + with self._stats_lock: + self._rejected_total += 1 + self._global_rejected_total += 1 + raise BlockingWorkCapacityError( + 'Blocking executor capacity reached', + scope=scope, + ) + + with self._stats_lock: + if scope is not None and self._inflight_by_scope.get(scope, 0) >= self.max_inflight_per_scope: + self._rejected_total += 1 + self._scope_rejected_total += 1 + self._capacity.release() + raise BlockingWorkCapacityError( + 'Workspace blocking executor capacity reached', + scope=scope, + ) + self._inflight += 1 + self._submitted_total += 1 + if scope is not None: + self._inflight_by_scope[scope] = self._inflight_by_scope.get(scope, 0) + 1 + + def run() -> Any: + with self._stats_lock: + self._running += 1 + try: + return fn(*args, **kwargs) + finally: + with self._stats_lock: + self._running -= 1 + + try: + future = super().submit(run) + except BaseException: + with self._stats_lock: + self._inflight -= 1 + self._release_scope_locked(scope) + self._capacity.release() + raise + + def complete(_future: concurrent.futures.Future) -> None: + with self._stats_lock: + self._inflight -= 1 + self._completed_total += 1 + self._release_scope_locked(scope) + self._capacity.release() + + future.add_done_callback(complete) + return future + + def _release_scope_locked(self, scope: str | None) -> None: + if scope is None: + return + remaining = self._inflight_by_scope.get(scope, 0) - 1 + if remaining > 0: + self._inflight_by_scope[scope] = remaining + else: + self._inflight_by_scope.pop(scope, None) + + def snapshot(self) -> dict[str, int]: + with self._stats_lock: + inflight = self._inflight + running = self._running + return { + 'max_workers': self.max_workers, + 'max_pending': self.max_pending, + 'max_inflight_per_scope': self.max_inflight_per_scope, + 'inflight': inflight, + 'running': running, + 'pending': max(inflight - running, 0), + 'active_scopes': len(self._inflight_by_scope), + 'submitted_total': self._submitted_total, + 'completed_total': self._completed_total, + 'rejected_total': self._rejected_total, + 'global_rejected_total': self._global_rejected_total, + 'scope_rejected_total': self._scope_rejected_total, + } + + +def configure_bounded_default_executor( + loop: asyncio.AbstractEventLoop, + *, + max_workers: int = DEFAULT_MAX_WORKERS, + max_pending: int = DEFAULT_MAX_PENDING, + max_inflight_per_scope: int | None = None, + thread_name_prefix: str = 'langbot-blocking', +) -> BoundedThreadPoolExecutor: + """Install one bounded owner for every ``asyncio.to_thread`` call.""" + + max_workers, max_pending, max_inflight_per_scope = _validated_limits( + max_workers, + max_pending, + max_inflight_per_scope, + ) + existing = getattr(loop, '_default_executor', None) + if isinstance(existing, BoundedThreadPoolExecutor): + if ( + existing.max_workers != max_workers + or existing.max_pending != max_pending + or existing.max_inflight_per_scope != max_inflight_per_scope + ): + raise RuntimeError('The blocking executor is already configured with different limits') + return existing + if existing is not None: + raise RuntimeError('The event loop default executor was initialized before LangBot resource limits') + + executor = BoundedThreadPoolExecutor( + max_workers=max_workers, + max_pending=max_pending, + max_inflight_per_scope=max_inflight_per_scope, + thread_name_prefix=thread_name_prefix, + ) + loop.set_default_executor(executor) + return executor diff --git a/src/langbot/pkg/utils/event_loop_monitor.py b/src/langbot/pkg/utils/event_loop_monitor.py new file mode 100644 index 000000000..f38f10379 --- /dev/null +++ b/src/langbot/pkg/utils/event_loop_monitor.py @@ -0,0 +1,97 @@ +from __future__ import annotations + +import asyncio +import contextlib +import math +from collections import deque + + +DEFAULT_SAMPLE_INTERVAL_SECONDS = 1.0 +DEFAULT_RECENT_SAMPLE_COUNT = 120 + + +class EventLoopLagMonitor: + """Measure event-loop scheduling delay with fixed, bounded state.""" + + def __init__( + self, + *, + sample_interval_seconds: float = DEFAULT_SAMPLE_INTERVAL_SECONDS, + recent_sample_count: int = DEFAULT_RECENT_SAMPLE_COUNT, + ) -> None: + interval = float(sample_interval_seconds) + if not math.isfinite(interval) or interval <= 0: + raise ValueError('sample_interval_seconds must be greater than zero') + sample_count = int(recent_sample_count) + if sample_count < 2 or sample_count > 3600: + raise ValueError('recent_sample_count must be between 2 and 3600') + self.sample_interval_seconds = interval + self.recent_sample_count = sample_count + self._recent_lag_ms: deque[float] = deque(maxlen=sample_count) + self._samples_total = 0 + self._max_lag_ms = 0.0 + self._last_lag_ms = 0.0 + self._task: asyncio.Task[None] | None = None + + @property + def running(self) -> bool: + return self._task is not None and not self._task.done() + + def start(self) -> None: + """Start sampling on the current event loop; repeated calls are safe.""" + + if self.running: + return + self._task = asyncio.create_task( + self._run(), + name='event-loop-lag-monitor', + ) + + async def stop(self) -> None: + """Cancel and await the owned sampler task.""" + + task = self._task + self._task = None + if task is None: + return + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + + async def _run(self) -> None: + loop = asyncio.get_running_loop() + expected_at = loop.time() + self.sample_interval_seconds + while True: + await asyncio.sleep(max(expected_at - loop.time(), 0.0)) + observed_at = loop.time() + self._record_lag_seconds(max(observed_at - expected_at, 0.0)) + # One observation captures a long stall; do not replay every + # missed interval in a tight loop after the scheduler recovers. + expected_at = observed_at + self.sample_interval_seconds + + def _record_lag_seconds(self, lag_seconds: float) -> None: + lag_ms = max(float(lag_seconds), 0.0) * 1000 + self._last_lag_ms = lag_ms + self._max_lag_ms = max(self._max_lag_ms, lag_ms) + self._recent_lag_ms.append(lag_ms) + self._samples_total += 1 + + def snapshot(self) -> dict[str, int | float | bool]: + """Return aggregate metrics without exposing task or tenant state.""" + + recent = sorted(self._recent_lag_ms) + if recent: + p95_index = max(math.ceil(len(recent) * 0.95) - 1, 0) + recent_p95_ms = recent[p95_index] + recent_max_ms = recent[-1] + else: + recent_p95_ms = 0.0 + recent_max_ms = 0.0 + return { + 'running': self.running, + 'samples_total': self._samples_total, + 'last_lag_ms': self._last_lag_ms, + 'recent_p95_lag_ms': recent_p95_ms, + 'recent_max_lag_ms': recent_max_ms, + 'max_lag_ms': self._max_lag_ms, + } diff --git a/src/langbot/pkg/utils/httpclient.py b/src/langbot/pkg/utils/httpclient.py index 3c7df9afe..26c75d16f 100644 --- a/src/langbot/pkg/utils/httpclient.py +++ b/src/langbot/pkg/utils/httpclient.py @@ -11,9 +11,64 @@ reuses the same underlying SSL context and connection pool. from __future__ import annotations +import asyncio +import inspect +import json +import typing + import aiohttp +import httpx _sessions: dict[str, aiohttp.ClientSession] = {} +DEFAULT_REMOTE_BODY_LIMIT = 10 * 1024 * 1024 + + +class RemoteResponseTooLargeError(ValueError): + """Raised before an untrusted remote response can exhaust process memory.""" + + +class _LimitedHTTPXAsyncByteStream(httpx.AsyncByteStream): + def __init__(self, inner: httpx.AsyncByteStream, max_bytes: int) -> None: + self._inner = inner + self._max_bytes = max_bytes + self._read_bytes = 0 + + async def __aiter__(self): + async for chunk in self._inner: + self._read_bytes += len(chunk) + if self._read_bytes > self._max_bytes: + raise RemoteResponseTooLargeError(f'Remote response exceeds the {self._max_bytes}-byte limit') + yield chunk + + async def aclose(self) -> None: + await self._inner.aclose() + + +def httpx_response_limit_hooks( + max_bytes: int = DEFAULT_REMOTE_BODY_LIMIT, +) -> dict[str, list]: + """Return hooks that cap HTTPX bodies before its automatic buffering.""" + + max_bytes = max(int(max_bytes), 1) + + async def limit_response(response: httpx.Response) -> None: + content_length = response.headers.get('Content-Length') + if content_length is not None: + try: + declared_size = int(content_length) + except (TypeError, ValueError): + declared_size = None + if declared_size is not None and declared_size > max_bytes: + await response.aclose() + raise RemoteResponseTooLargeError(f'Remote response exceeds the {max_bytes}-byte limit') + + if response.is_stream_consumed: + if len(response.content) > max_bytes: + raise RemoteResponseTooLargeError(f'Remote response exceeds the {max_bytes}-byte limit') + return + response.stream = _LimitedHTTPXAsyncByteStream(response.stream, max_bytes) + + return {'response': [limit_response]} def get_session(*, trust_env: bool = False) -> aiohttp.ClientSession: @@ -47,3 +102,69 @@ async def close_all(): if not session.closed: await session.close() _sessions.clear() + + +async def read_limited( + response: aiohttp.ClientResponse, + *, + max_bytes: int = DEFAULT_REMOTE_BODY_LIMIT, +) -> bytes: + """Read an HTTP response incrementally with a strict byte limit.""" + + max_bytes = max(int(max_bytes), 1) + content_length = response.headers.get('Content-Length') + if content_length is not None: + try: + declared_size = int(content_length) + except (TypeError, ValueError): + declared_size = None + if declared_size is not None and declared_size > max_bytes: + raise RemoteResponseTooLargeError(f'Remote response exceeds the {max_bytes}-byte limit') + + body = bytearray() + async for chunk in response.content.iter_chunked(64 * 1024): + body.extend(chunk) + if len(body) > max_bytes: + raise RemoteResponseTooLargeError(f'Remote response exceeds the {max_bytes}-byte limit') + return bytes(body) + + +async def read_text_limited( + response: aiohttp.ClientResponse, + *, + max_bytes: int = DEFAULT_REMOTE_BODY_LIMIT, +) -> str: + body = await read_limited(response, max_bytes=max_bytes) + return body.decode(response.charset or 'utf-8', errors='replace') + + +async def read_json_limited( + response: aiohttp.ClientResponse, + *, + max_bytes: int = DEFAULT_REMOTE_BODY_LIMIT, +) -> typing.Any: + body = await read_limited(response, max_bytes=max_bytes) + return await asyncio.to_thread(json.loads, body) + + +async def parse_json_response(response: typing.Any) -> typing.Any: + """Parse an already bounded HTTP response without blocking the event loop.""" + + parsed = await asyncio.to_thread(response.json) + if inspect.isawaitable(parsed): + parsed = await parsed + return parsed + + +async def response_text( + response: typing.Any, + *, + max_chars: int = 4096, +) -> str: + """Decode an already bounded response body off-loop and cap diagnostics.""" + + text = await asyncio.to_thread(lambda: str(response.text)) + max_chars = max(int(max_chars), 1) + if len(text) <= max_chars: + return text + return f'{text[:max_chars]}... [truncated]' diff --git a/src/langbot/pkg/utils/image.py b/src/langbot/pkg/utils/image.py index 0296ba05f..435671ae8 100644 --- a/src/langbot/pkg/utils/image.py +++ b/src/langbot/pkg/utils/image.py @@ -8,10 +8,46 @@ import aiohttp from langbot.pkg.utils import httpclient import PIL.Image -import httpx import asyncio +_INSECURE_SSL_CONTEXT = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) +_INSECURE_SSL_CONTEXT.check_hostname = False +_INSECURE_SSL_CONTEXT.verify_mode = ssl.CERT_NONE +DEFAULT_BASE64_MEDIA_LIMIT = 10 * 1024 * 1024 + + +def _detect_image_format(file_bytes: bytes) -> str: + with PIL.Image.open(io.BytesIO(file_bytes)) as image: + return str(image.format or 'jpeg').lower() + + +def _decode_base64_limited(value: str, max_bytes: int) -> bytes: + max_bytes = max(int(max_bytes), 1) + max_encoded_chars = 4 * ((max_bytes + 2) // 3) + 4 + if len(value) > max_encoded_chars: + raise ValueError(f'Base64 media exceeds the {max_bytes}-byte limit') + decoded = base64.b64decode(value) + if len(decoded) > max_bytes: + raise ValueError(f'Base64 media exceeds the {max_bytes}-byte limit') + return decoded + + +async def decode_base64_limited( + value: str, + *, + max_bytes: int = DEFAULT_BASE64_MEDIA_LIMIT, +) -> bytes: + """Decode bounded media outside the event loop.""" + + return await asyncio.to_thread(_decode_base64_limited, value, max_bytes) + + +async def encode_base64(data: bytes) -> str: + """Encode a bounded byte payload outside the event loop.""" + + return (await asyncio.to_thread(base64.b64encode, data)).decode('utf-8') + async def get_gewechat_image_base64( gewechat_url: str, @@ -59,10 +95,10 @@ async def get_gewechat_image_base64( timeout=timeout, ) as response: if response.status != 200: - # print(response) - raise Exception(f'获取gewechat图片下载失败: {await response.text()}') + error = await httpclient.read_text_limited(response) + raise Exception(f'获取gewechat图片下载失败: {error}') - resp_data = await response.json() + resp_data = await httpclient.read_json_limited(response) if resp_data.get('ret') != 200: raise Exception(f'获取gewechat图片下载链接失败: {resp_data}') @@ -80,9 +116,10 @@ async def get_gewechat_image_base64( try: async with session.get(download_url) as img_response: if img_response.status != 200: - raise Exception(f'下载图片失败: {await img_response.text()}, URL: {download_url}') + error = await httpclient.read_text_limited(img_response) + raise Exception(f'下载图片失败: {error}, URL: {download_url}') - image_data = await img_response.read() + image_data = await httpclient.read_limited(img_response) content_type = img_response.headers.get('Content-Type', '') if content_type: @@ -90,7 +127,7 @@ async def get_gewechat_image_base64( else: image_format = file_url.split('.')[-1] - base64_str = base64.b64encode(image_data).decode('utf-8') + base64_str = await encode_base64(image_data) return base64_str, image_format except asyncio.TimeoutError: @@ -113,16 +150,13 @@ async def get_wecom_image_base64(pic_url: str) -> tuple[str, str]: raise Exception(f'Failed to download image: {response.status}') # 读取图片数据 - image_data = await response.read() + image_data = await httpclient.read_limited(response) # 获取图片格式 content_type = response.headers.get('Content-Type', '') image_format = content_type.split('/')[-1] # 例如 'image/jpeg' -> 'jpeg' - # 转换为 base64 - import base64 - - image_base64 = base64.b64encode(image_data).decode('utf-8') + image_base64 = await encode_base64(image_data) return image_base64, image_format @@ -132,11 +166,11 @@ async def get_qq_official_image_base64(pic_url: str, content_type: str) -> tuple 下载QQ官方图片， 并且转换为base64格式 """ - async with httpx.AsyncClient() as client: - response = await client.get(pic_url) - response.raise_for_status() # 确保请求成功 - image_data = response.content - base64_data = base64.b64encode(image_data).decode('utf-8') + session = httpclient.get_session() + async with session.get(pic_url) as response: + response.raise_for_status() + image_data = await httpclient.read_limited(response) + base64_data = await encode_base64(image_data) return f'data:{content_type};base64,{base64_data}' @@ -153,19 +187,20 @@ async def get_qq_image_bytes(image_url: str, query: dict = {}) -> tuple[bytes, s """[弃用]获取QQ图片的bytes""" image_url, query_in_url = get_qq_image_downloadable_url(image_url) query = {**query, **query_in_url} - ssl_context = ssl.create_default_context() - ssl_context.check_hostname = False - ssl_context.verify_mode = ssl.CERT_NONE session = httpclient.get_session() - async with session.get(image_url, params=query, ssl=ssl_context, timeout=aiohttp.ClientTimeout(total=30.0)) as resp: + async with session.get( + image_url, + params=query, + ssl=_INSECURE_SSL_CONTEXT, + timeout=aiohttp.ClientTimeout(total=30.0), + ) as resp: resp.raise_for_status() - file_bytes = await resp.read() + file_bytes = await httpclient.read_limited(resp) content_type = resp.headers.get('Content-Type') if not content_type: image_format = 'jpeg' elif not content_type.startswith('image/'): - pil_img = PIL.Image.open(io.BytesIO(file_bytes)) - image_format = pil_img.format.lower() + image_format = await asyncio.to_thread(_detect_image_format, file_bytes) else: image_format = content_type.split('/')[-1] return file_bytes, image_format @@ -187,7 +222,7 @@ async def qq_image_url_to_base64(image_url: str) -> typing.Tuple[str, str]: file_bytes, image_format = await get_qq_image_bytes(image_url, query) - base64_str = base64.b64encode(file_bytes).decode() + base64_str = await encode_base64(file_bytes) return base64_str, image_format @@ -209,8 +244,8 @@ async def get_slack_image_to_base64(pic_url: str, bot_token: str): session = httpclient.get_session() async with session.get(pic_url, headers=headers) as resp: mime_type = resp.headers.get('Content-Type', 'application/octet-stream') - file_bytes = await resp.read() - base64_str = base64.b64encode(file_bytes).decode('utf-8') + file_bytes = await httpclient.read_limited(resp) + base64_str = await encode_base64(file_bytes) return f'data:{mime_type};base64,{base64_str}' except Exception as e: raise (e) diff --git a/src/langbot/pkg/utils/logcache.py b/src/langbot/pkg/utils/logcache.py index 84c58f557..99c6c76a8 100644 --- a/src/langbot/pkg/utils/logcache.py +++ b/src/langbot/pkg/utils/logcache.py @@ -3,6 +3,7 @@ from __future__ import annotations LOG_PAGE_SIZE = 20 MAX_CACHED_PAGES = 10 +MAX_LOG_LINE_CHARS = 20000 class LogPage: @@ -40,6 +41,10 @@ class LogCache: def add_log(self, log: str): """添加日志""" + log = str(log) + if len(log) > MAX_LOG_LINE_CHARS: + marker = '\n[log truncated]' + log = log[: MAX_LOG_LINE_CHARS - len(marker)] + marker if self.log_pages[-1].add_log(log): self.log_pages.append(LogPage(number=self.log_pages[-1].number + 1)) diff --git a/src/langbot/pkg/utils/safe_regex.py b/src/langbot/pkg/utils/safe_regex.py new file mode 100644 index 000000000..32ae27843 --- /dev/null +++ b/src/langbot/pkg/utils/safe_regex.py @@ -0,0 +1,177 @@ +from __future__ import annotations + +import asyncio +import time +from collections.abc import Sequence + +import regex + + +MAX_PATTERN_COUNT = 64 +MAX_PATTERN_CHARS = 1024 +MAX_INPUT_CHARS = 1024 * 1024 +MAX_REPLACEMENT_CHARS = 64 +MAX_MASKED_OUTPUT_CHARS = 2 * 1024 * 1024 +DEFAULT_OPERATION_TIMEOUT_SECONDS = 0.05 + + +class SafeRegexError(ValueError): + """Base class for rejected, invalid, or timed-out tenant regex work.""" + + +class SafeRegexLimitError(SafeRegexError): + """Raised when a regex operation exceeds a deterministic resource limit.""" + + +class SafeRegexTimeoutError(SafeRegexError): + """Raised when the regex engine exhausts the operation CPU budget.""" + + +def _validate_patterns(patterns: Sequence[str]) -> tuple[str, ...]: + normalized = tuple(patterns) + if len(normalized) > MAX_PATTERN_COUNT: + raise SafeRegexLimitError(f'At most {MAX_PATTERN_COUNT} regex patterns are allowed') + for pattern in normalized: + if not isinstance(pattern, str): + raise SafeRegexError('Regex patterns must be strings') + if len(pattern) > MAX_PATTERN_CHARS: + raise SafeRegexLimitError(f'Regex patterns may contain at most {MAX_PATTERN_CHARS} characters') + return normalized + + +def _validate_input(value: str) -> None: + if not isinstance(value, str): + raise SafeRegexError('Regex input must be a string') + if len(value) > MAX_INPUT_CHARS: + raise SafeRegexLimitError(f'Regex input may contain at most {MAX_INPUT_CHARS} characters') + + +def _remaining_seconds(deadline: float) -> float: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise SafeRegexTimeoutError('Regex operation timed out') + return remaining + + +def _compile(pattern: str): + try: + return regex.compile(pattern) + except regex.error as exc: + raise SafeRegexError(f'Invalid regex: {exc}') from exc + + +def _matches_any_sync( + patterns: Sequence[str], + value: str, + *, + mode: str, + timeout_seconds: float, +) -> bool: + normalized_patterns = _validate_patterns(patterns) + _validate_input(value) + if mode not in {'match', 'search'}: + raise ValueError(f'Unsupported safe regex mode: {mode}') + + deadline = time.monotonic() + timeout_seconds + try: + for pattern in normalized_patterns: + compiled = _compile(pattern) + matcher = compiled.match if mode == 'match' else compiled.search + if matcher( + value, + timeout=_remaining_seconds(deadline), + concurrent=True, + ): + return True + except TimeoutError as exc: + raise SafeRegexTimeoutError('Regex operation timed out') from exc + return False + + +async def matches_any( + patterns: Sequence[str], + value: str, + *, + mode: str = 'search', + timeout_seconds: float = DEFAULT_OPERATION_TIMEOUT_SECONDS, +) -> bool: + """Match untrusted patterns without blocking the shared event loop.""" + + if timeout_seconds <= 0: + raise ValueError('timeout_seconds must be positive') + return await asyncio.to_thread( + _matches_any_sync, + patterns, + value, + mode=mode, + timeout_seconds=timeout_seconds, + ) + + +def _mask_patterns_sync( + patterns: Sequence[str], + value: str, + *, + mask: str, + mask_word: str, + timeout_seconds: float, +) -> tuple[bool, str]: + normalized_patterns = _validate_patterns(patterns) + _validate_input(value) + if len(mask) > MAX_REPLACEMENT_CHARS or len(mask_word) > MAX_REPLACEMENT_CHARS: + raise SafeRegexLimitError(f'Regex replacements may contain at most {MAX_REPLACEMENT_CHARS} characters') + + # Reject amplification before invoking a replacement callback. This is + # deliberately conservative: a hostile replacement must not allocate tens + # of megabytes before the post-operation output check can run. + replacement_width = len(mask_word) if mask_word else len(mask) + if replacement_width * max(1, len(value)) > MAX_MASKED_OUTPUT_CHARS: + raise SafeRegexLimitError('Regex replacement could exceed the masked output limit') + + deadline = time.monotonic() + timeout_seconds + found = False + current = value + + def replace(match) -> str: + nonlocal found + found = True + if mask_word: + return mask_word + return mask * len(match.group(0)) + + try: + for pattern in normalized_patterns: + compiled = _compile(pattern) + current = compiled.sub( + replace, + current, + timeout=_remaining_seconds(deadline), + concurrent=True, + ) + if len(current) > MAX_MASKED_OUTPUT_CHARS: + raise SafeRegexLimitError('Regex replacement exceeded the masked output limit') + except TimeoutError as exc: + raise SafeRegexTimeoutError('Regex operation timed out') from exc + return found, current + + +async def mask_patterns( + patterns: Sequence[str], + value: str, + *, + mask: str, + mask_word: str, + timeout_seconds: float = DEFAULT_OPERATION_TIMEOUT_SECONDS, +) -> tuple[bool, str]: + """Apply untrusted masking patterns with bounded CPU and output growth.""" + + if timeout_seconds <= 0: + raise ValueError('timeout_seconds must be positive') + return await asyncio.to_thread( + _mask_patterns_sync, + patterns, + value, + mask=mask, + mask_word=mask_word, + timeout_seconds=timeout_seconds, + ) diff --git a/src/langbot/pkg/utils/version.py b/src/langbot/pkg/utils/version.py index 1e19420db..0ee49de14 100644 --- a/src/langbot/pkg/utils/version.py +++ b/src/langbot/pkg/utils/version.py @@ -1,12 +1,13 @@ from __future__ import annotations +import asyncio import typing import logging import requests from ..core import app -from . import constants +from . import constants, httpclient class VersionManager: @@ -26,13 +27,14 @@ class VersionManager: async def get_release_list(self) -> list: """Fetch release list from Space API (cached GitHub releases).""" try: - rls_list_resp = requests.get( - url='https://space.langbot.app/api/v1/dist/info/releases', + rls_list_resp = await asyncio.to_thread( + requests.get, + 'https://space.langbot.app/api/v1/dist/info/releases', proxies=self.ap.proxy_mgr.get_forward_proxies(), timeout=10, ) rls_list_resp.raise_for_status() - resp_json = rls_list_resp.json() + resp_json = await httpclient.parse_json_response(rls_list_resp) if resp_json.get('code') == 0 and isinstance(resp_json.get('data'), list): return resp_json['data'] self.ap.logger.warning(f'Failed to fetch release list: unexpected response: {resp_json.get("msg", "")}') diff --git a/src/langbot/pkg/vector/mgr.py b/src/langbot/pkg/vector/mgr.py index 78107a32f..7a748d325 100644 --- a/src/langbot/pkg/vector/mgr.py +++ b/src/langbot/pkg/vector/mgr.py @@ -113,6 +113,14 @@ class VectorDBManager: self.vector_db = ChromaVectorDatabase(self.ap) self.ap.logger.warning('No vector database backend configured, defaulting to Chroma.') + async def shutdown(self) -> None: + """Release the active vector backend deterministically.""" + + vector_db = self.vector_db + self.vector_db = None + if vector_db is not None: + await vector_db.close() + def get_supported_search_types(self) -> list[str]: """Return the search types supported by the current VDB backend.""" if self.vector_db is None: diff --git a/src/langbot/pkg/vector/vdb.py b/src/langbot/pkg/vector/vdb.py index df9bfacbf..67bc011f6 100644 --- a/src/langbot/pkg/vector/vdb.py +++ b/src/langbot/pkg/vector/vdb.py @@ -1,9 +1,37 @@ from __future__ import annotations import abc import enum -from typing import Any, Dict +from typing import Any, Dict, MutableMapping, MutableSet, TypeVar import numpy as np +_CacheValue = TypeVar('_CacheValue') + + +def runtime_cache_limit(ap: Any, default: int = 1024) -> int: + try: + value = int(ap.instance_config.data.get('vdb', {}).get('runtime_cache_limit', default)) + except (AttributeError, TypeError, ValueError): + value = default + return max(value, 1) + + +def remember_bounded_mapping( + cache: MutableMapping[str, _CacheValue], + key: str, + value: _CacheValue, + limit: int, +) -> None: + cache.pop(key, None) + cache[key] = value + while len(cache) > limit: + cache.pop(next(iter(cache))) + + +def remember_bounded_set(cache: MutableSet[str], key: str, limit: int) -> None: + cache.add(key) + while len(cache) > limit: + cache.pop() + class SearchType(str, enum.Enum): """Supported search types for vector databases.""" @@ -14,6 +42,9 @@ class SearchType(str, enum.Enum): class VectorDatabase(abc.ABC): + async def close(self) -> None: + """Release backend clients and in-process caches.""" + @classmethod def supported_search_types(cls) -> list[SearchType]: """Return the search types supported by this VDB backend. diff --git a/src/langbot/pkg/vector/vdbs/chroma.py b/src/langbot/pkg/vector/vdbs/chroma.py index b4d8af19c..6965443a5 100644 --- a/src/langbot/pkg/vector/vdbs/chroma.py +++ b/src/langbot/pkg/vector/vdbs/chroma.py @@ -2,7 +2,12 @@ from __future__ import annotations import asyncio from typing import Any from chromadb import PersistentClient -from langbot.pkg.vector.vdb import VectorDatabase, SearchType +from langbot.pkg.vector.vdb import ( + VectorDatabase, + SearchType, + remember_bounded_mapping, + runtime_cache_limit, +) from langbot.pkg.core import app import chromadb import chromadb.errors @@ -16,6 +21,13 @@ class ChromaVectorDatabase(VectorDatabase): self.ap = ap self.client = PersistentClient(path=base_path) self._collections = {} + self._runtime_cache_limit = runtime_cache_limit(ap) + + async def close(self) -> None: + # Chroma's PersistentClient has no public close API. Collection + # wrappers are safe to discard and otherwise retain every collection + # touched during the lifetime of this application object. + self._collections.clear() @classmethod def supported_search_types(cls) -> list[SearchType]: @@ -23,8 +35,12 @@ class ChromaVectorDatabase(VectorDatabase): async def get_or_create_collection(self, collection: str) -> chromadb.Collection: if collection not in self._collections: - self._collections[collection] = await asyncio.to_thread( - self.client.get_or_create_collection, name=collection + runtime_collection = await asyncio.to_thread(self.client.get_or_create_collection, name=collection) + remember_bounded_mapping( + self._collections, + collection, + runtime_collection, + self._runtime_cache_limit, ) self.ap.logger.info(f"Chroma collection '{collection}' accessed/created.") return self._collections[collection] diff --git a/src/langbot/pkg/vector/vdbs/milvus.py b/src/langbot/pkg/vector/vdbs/milvus.py index 1b01b1101..ac6012451 100644 --- a/src/langbot/pkg/vector/vdbs/milvus.py +++ b/src/langbot/pkg/vector/vdbs/milvus.py @@ -3,9 +3,10 @@ import asyncio from typing import Any, Dict from pymilvus import MilvusClient, DataType, CollectionSchema, FieldSchema from pymilvus.milvus_client.index import IndexParams -from langbot.pkg.vector.vdb import VectorDatabase +from langbot.pkg.vector.vdb import VectorDatabase, remember_bounded_set, runtime_cache_limit from langbot.pkg.vector.filter_utils import normalize_filter, strip_unsupported_fields from langbot.pkg.core import app +from langbot.pkg.utils import bounded_executor # Milvus schema only stores these metadata fields; filter on other fields is # silently dropped with a warning. @@ -71,8 +72,15 @@ class MilvusVectorDatabase(VectorDatabase): self.db_name = db_name self.client = None self._collections: set[str] = set() + self._runtime_cache_limit = runtime_cache_limit(ap) self._initialize_client() + async def close(self) -> None: + self._collections.clear() + if self.client is not None: + await bounded_executor.run_blocking_cleanup(self.client.close) + self.client = None + def _initialize_client(self): """Initialize Milvus client connection""" try: @@ -169,7 +177,7 @@ class MilvusVectorDatabase(VectorDatabase): await self._ensure_index_if_missing(collection) self.ap.logger.info(f"Milvus collection '{collection}' already exists") - self._collections.add(collection) + remember_bounded_set(self._collections, collection, self._runtime_cache_limit) return collection async def _ensure_index_if_missing(self, collection: str) -> None: diff --git a/src/langbot/pkg/vector/vdbs/qdrant.py b/src/langbot/pkg/vector/vdbs/qdrant.py index 6da4c79ef..dea6a8676 100644 --- a/src/langbot/pkg/vector/vdbs/qdrant.py +++ b/src/langbot/pkg/vector/vdbs/qdrant.py @@ -4,7 +4,7 @@ from typing import Any, Dict, List from qdrant_client import AsyncQdrantClient, models from langbot.pkg.core import app -from langbot.pkg.vector.vdb import VectorDatabase +from langbot.pkg.vector.vdb import VectorDatabase, remember_bounded_set, runtime_cache_limit from langbot.pkg.vector.filter_utils import normalize_filter @@ -52,6 +52,11 @@ class QdrantVectorDatabase(VectorDatabase): self.client = AsyncQdrantClient(host=host, port=int(port), api_key=api_key) self._collections: set[str] = set() + self._runtime_cache_limit = runtime_cache_limit(ap) + + async def close(self) -> None: + self._collections.clear() + await self.client.close() async def _ensure_collection(self, collection: str, vector_size: int) -> None: if collection in self._collections: @@ -59,14 +64,14 @@ class QdrantVectorDatabase(VectorDatabase): exists = await self.client.collection_exists(collection) if exists: - self._collections.add(collection) + remember_bounded_set(self._collections, collection, self._runtime_cache_limit) return await self.client.create_collection( collection_name=collection, vectors_config=models.VectorParams(size=vector_size, distance=models.Distance.COSINE), ) - self._collections.add(collection) + remember_bounded_set(self._collections, collection, self._runtime_cache_limit) self.ap.logger.info(f"Qdrant collection '{collection}' created with dim={vector_size}.") async def get_or_create_collection(self, collection: str): diff --git a/src/langbot/pkg/vector/vdbs/seekdb.py b/src/langbot/pkg/vector/vdbs/seekdb.py index c96a9ca77..fc82298e0 100644 --- a/src/langbot/pkg/vector/vdbs/seekdb.py +++ b/src/langbot/pkg/vector/vdbs/seekdb.py @@ -7,7 +7,13 @@ from typing import Any, Dict, List from langbot.pkg.core import app -from langbot.pkg.vector.vdb import VectorDatabase, SearchType +from langbot.pkg.utils import bounded_executor +from langbot.pkg.vector.vdb import ( + VectorDatabase, + SearchType, + remember_bounded_mapping, + runtime_cache_limit, +) try: import pyseekdb @@ -90,6 +96,7 @@ class SeekDBVectorDatabase(VectorDatabase): self._collections: Dict[str, Any] = {} self._collection_configs: Dict[str, HNSWConfiguration] = {} + self._runtime_cache_limit = runtime_cache_limit(ap) self._escape_table = str.maketrans( { @@ -103,6 +110,13 @@ class SeekDBVectorDatabase(VectorDatabase): } ) + async def close(self) -> None: + self._collections.clear() + self._collection_configs.clear() + close = getattr(self.client, 'close', None) + if callable(close): + await bounded_executor.run_blocking_cleanup(close) + def _normalize_collection_name(self, collection: str) -> str: """SeekDB only accepts [a-zA-Z0-9_], while LangBot uses UUID-like KB IDs.""" normalized = re.sub(r'[^A-Za-z0-9_]', '_', collection) @@ -132,7 +146,12 @@ class SeekDBVectorDatabase(VectorDatabase): if await asyncio.to_thread(self.client.has_collection, collection): # Collection exists, get it coll = await asyncio.to_thread(self.client.get_collection, collection, embedding_function=None) - self._collections[collection] = coll + remember_bounded_mapping( + self._collections, + collection, + coll, + self._runtime_cache_limit, + ) self.ap.logger.info(f"SeekDB collection '{collection}' retrieved.") return coll @@ -145,7 +164,12 @@ class SeekDBVectorDatabase(VectorDatabase): # Create HNSW configuration config = HNSWConfiguration(dimension=vector_size, distance='cosine') - self._collection_configs[collection] = config + remember_bounded_mapping( + self._collection_configs, + collection, + config, + self._runtime_cache_limit, + ) # Create collection without embedding function (we manage embeddings externally) coll = await asyncio.to_thread( @@ -155,7 +179,12 @@ class SeekDBVectorDatabase(VectorDatabase): embedding_function=None, # Disable automatic embedding ) - self._collections[collection] = coll + remember_bounded_mapping( + self._collections, + collection, + coll, + self._runtime_cache_limit, + ) self.ap.logger.info(f"SeekDB collection '{collection}' created with dimension={vector_size}, distance='cosine'") return coll @@ -243,7 +272,12 @@ class SeekDBVectorDatabase(VectorDatabase): # Get collection if collection not in self._collections: coll = await asyncio.to_thread(self.client.get_collection, collection, embedding_function=None) - self._collections[collection] = coll + remember_bounded_mapping( + self._collections, + collection, + coll, + self._runtime_cache_limit, + ) else: coll = self._collections[collection] @@ -349,7 +383,12 @@ class SeekDBVectorDatabase(VectorDatabase): # Get collection if collection not in self._collections: coll = await asyncio.to_thread(self.client.get_collection, collection, embedding_function=None) - self._collections[collection] = coll + remember_bounded_mapping( + self._collections, + collection, + coll, + self._runtime_cache_limit, + ) else: coll = self._collections[collection] @@ -374,7 +413,12 @@ class SeekDBVectorDatabase(VectorDatabase): if collection not in self._collections: coll = await asyncio.to_thread(self.client.get_collection, collection, embedding_function=None) - self._collections[collection] = coll + remember_bounded_mapping( + self._collections, + collection, + coll, + self._runtime_cache_limit, + ) else: coll = self._collections[collection] @@ -396,7 +440,12 @@ class SeekDBVectorDatabase(VectorDatabase): if collection not in self._collections: coll = await asyncio.to_thread(self.client.get_collection, collection, embedding_function=None) - self._collections[collection] = coll + remember_bounded_mapping( + self._collections, + collection, + coll, + self._runtime_cache_limit, + ) else: coll = self._collections[collection] diff --git a/src/langbot/pkg/vector/vdbs/valkey_search.py b/src/langbot/pkg/vector/vdbs/valkey_search.py index 66be43664..6f96fccb9 100644 --- a/src/langbot/pkg/vector/vdbs/valkey_search.py +++ b/src/langbot/pkg/vector/vdbs/valkey_search.py @@ -6,7 +6,12 @@ import struct from typing import Any from langbot.pkg.core import app -from langbot.pkg.vector.vdb import VectorDatabase, SearchType +from langbot.pkg.vector.vdb import ( + VectorDatabase, + SearchType, + remember_bounded_set, + runtime_cache_limit, +) from langbot.pkg.vector.filter_utils import normalize_filter, strip_unsupported_fields try: @@ -153,6 +158,7 @@ class ValkeySearchVectorDatabase(VectorDatabase): self._client_lock = asyncio.Lock() # Index names we have already ensured this process lifetime. self._ensured_indexes: set[str] = set() + self._runtime_cache_limit = runtime_cache_limit(ap) # Whether we have already warned about the non-honored vector_weight. self._vector_weight_warned = False @@ -364,7 +370,7 @@ class ValkeySearchVectorDatabase(VectorDatabase): # check-then-create TOCTOU window. try: await ft.info(client, index) - self._ensured_indexes.add(index) + remember_bounded_set(self._ensured_indexes, index, self._runtime_cache_limit) return except RequestError: pass @@ -389,7 +395,7 @@ class ValkeySearchVectorDatabase(VectorDatabase): ] options = FtCreateOptions(data_type=DataType.HASH, prefixes=[self._key_prefix(collection)]) await ft.create(client, index, schema, options) - self._ensured_indexes.add(index) + remember_bounded_set(self._ensured_indexes, index, self._runtime_cache_limit) self.ap.logger.info( f"Valkey Search index '{index}' created (dim={dim}, algo={self._algorithm.value}, " f'metric={self._distance_metric.value})' @@ -772,7 +778,7 @@ class ValkeySearchVectorDatabase(VectorDatabase): # was being paid on the first query to each collection. try: await ft.info(client, index) - self._ensured_indexes.add(index) + remember_bounded_set(self._ensured_indexes, index, self._runtime_cache_limit) return True except RequestError: return False diff --git a/src/langbot/pkg/workspace/invitation_delivery.py b/src/langbot/pkg/workspace/invitation_delivery.py index 42decbf6f..a3e2e8c86 100644 --- a/src/langbot/pkg/workspace/invitation_delivery.py +++ b/src/langbot/pkg/workspace/invitation_delivery.py @@ -11,6 +11,8 @@ from urllib.parse import quote import httpx +from ..utils import httpclient + if typing.TYPE_CHECKING: from ..core.app import Application @@ -99,7 +101,11 @@ class InvitationDeliveryService: 'text': self._plain_text(workspace_name, invitation_link), 'html': self._html(workspace_name, invitation_link), } - async with httpx.AsyncClient(timeout=httpx.Timeout(config.timeout), trust_env=True) as client: + async with httpx.AsyncClient( + timeout=httpx.Timeout(config.timeout), + trust_env=True, + event_hooks=httpclient.httpx_response_limit_hooks(), + ) as client: response = await client.post( config.resend_api_url, headers={'Authorization': f'Bearer {config.resend_api_key}'}, diff --git a/src/langbot/pkg/workspace/service.py b/src/langbot/pkg/workspace/service.py index 963d26f30..bab30d4f2 100644 --- a/src/langbot/pkg/workspace/service.py +++ b/src/langbot/pkg/workspace/service.py @@ -53,6 +53,13 @@ class WorkspaceService: self.ap = ap self.policy = policy or SingleWorkspacePolicy() self._instance_uuid = instance_uuid + self._startup_execution_bindings: ( + tuple[ + WorkspaceExecutionBinding, + ..., + ] + | None + ) = None @property def instance_uuid(self) -> str: @@ -129,6 +136,38 @@ class WorkspaceService: of work. """ + self._require_deployment_admission() + self._require_directory_projection() + if self._startup_execution_bindings is not None: + return list(self._startup_execution_bindings) + return await self._discover_active_execution_bindings() + + async def prime_startup_execution_bindings( + self, + ) -> list[WorkspaceExecutionBinding]: + """Freeze one validated binding snapshot during the serial boot graph. + + Directory synchronization starts only after ``BuildAppStage`` has + finished, so all runtime managers would otherwise rediscover and + revalidate the same active Workspace set independently. A boot-scoped + immutable snapshot removes those repeated tenant transactions without + weakening request-time generation checks. + """ + + self._require_deployment_admission() + self._require_directory_projection() + if self._startup_execution_bindings is None: + self._startup_execution_bindings = tuple(await self._discover_active_execution_bindings()) + return list(self._startup_execution_bindings) + + def release_startup_execution_bindings(self) -> None: + """Release the boot-only projection snapshot before background sync.""" + + self._startup_execution_bindings = None + + async def _discover_active_execution_bindings( + self, + ) -> list[WorkspaceExecutionBinding]: instance_discovery_uow = getattr(self.ap.persistence_mgr, 'instance_discovery_uow', None) statement = ( sqlalchemy.select(WorkspaceExecutionState.workspace_uuid) diff --git a/src/langbot/templates/config.yaml b/src/langbot/templates/config.yaml index 23f5aa679..8b00fb87b 100644 --- a/src/langbot/templates/config.yaml +++ b/src/langbot/templates/config.yaml @@ -47,6 +47,9 @@ command: concurrency: pipeline: 20 session: 1 + # Hard admission limits for queued + running pipeline queries. + pending_queries: 1000 + pending_queries_per_workspace: 100 proxy: http: '' https: '' @@ -56,6 +59,15 @@ system: recovery_key: '' allow_modify_login_info: true disabled_adapters: [] + blocking_executor: + # All asyncio.to_thread work shares this process-wide bounded pool. + # Both running threads and queued calls are capped to prevent tenant + # bursts from creating an unbounded queue of retained request objects. + max_workers: 8 + max_pending: 128 + # One trusted Workspace can occupy at most this many running + queued + # slots. This must not exceed half of max_workers. + max_inflight_per_scope: 4 # Public outbound IP addresses of this LangBot deployment. Some platforms # (e.g. WeCom, WeChat Official Account, QQ Official API) require the # caller's IPs to be added to their trusted-IP / IP-whitelist settings. @@ -76,6 +88,32 @@ system: task_retention: # Keep at most this many completed async task records in memory completed_limit: 200 + # Bound progress output retained by one task, including running tasks. + max_log_chars: 200000 + # Protect the shared process from user-triggered operation storms. + max_active_user_tasks: 256 + max_active_user_tasks_per_workspace: 8 + session_retention: + # Process-local conversation sessions are a cache, not durable history. + max_entries: 2000 + max_entries_per_workspace: 200 + idle_ttl_seconds: 86400 + max_conversations_per_session: 20 + max_messages_per_conversation: 100 + websocket_retention: + # Bound live browser sockets and per-Workspace fan-out in the shared process. + max_connections: 1024 + max_connections_per_workspace: 32 + # Idle proxy runtimes are evicted when this process-local cache fills. + max_workspace_proxies: 1024 + max_conversations_per_workspace: 200 + max_messages_per_conversation: 100 + conversation_idle_ttl_seconds: 86400 + send_queue_size: 100 + response_limits: + # Defense in depth for tenant-configured upstream providers. + max_generated_chars: 1048576 + max_stream_chunks: 100000 jwt: expire: 604800 secret: '' @@ -92,6 +130,11 @@ database: user: 'postgres' password: 'postgres' database: 'postgres' + # One bounded pool is shared by business data and Cloud pgvector. + pool_size: 10 + max_overflow: 10 + pool_timeout_seconds: 30 + pool_recycle_seconds: 1800 cloud_migration: # `langbot migrate --cloud` reads an operator-only PostgreSQL DSN from # this environment variable. The operator role must differ from the @@ -99,6 +142,8 @@ database: operator_dsn_env: 'LANGBOT_CLOUD_MIGRATION_DSN' vdb: use: chroma + # Bound process-local collection/index handles across all Workspaces. + runtime_cache_limit: 1024 qdrant: url: '' host: localhost @@ -157,6 +202,9 @@ storage: secret_access_key: '' region: 'us-east-1' bucket: 'langbot-storage' + # boto3 is synchronous; bound the number of operations delegated to + # worker threads so an S3 slowdown cannot saturate the process. + max_concurrency: 16 plugin: enable: true runtime_ws_url: 'ws://langbot_plugin_runtime:5400/control/ws' @@ -170,6 +218,15 @@ plugin: max_pids: 128 max_open_files: 256 max_file_size_mb: 512 + # Instance-wide admission budgets. The effective worker count is the + # lowest of max_workers, max_total_cpus/max_cpus and + # max_total_memory_mb/max_memory_mb. + max_workers: 16 + max_total_cpus: 8.0 + max_total_memory_mb: 8192 + # Includes disabled and historical installation fences retained to + # reject stale desired-state replay. + max_installations: 10000 # Cloud shared Runtime sets this to true and fails closed unless # delegated cgroup v2 controllers are available. require_hard_limits: false @@ -177,6 +234,9 @@ plugin: # Max bytes for a single plugin binary storage value max_value_bytes: 10485760 mcp: + # Bound instance-wide MCP startup and shutdown bursts. Supports + # MCP__LIFECYCLE_CONCURRENCY and is clamped to a maximum of 128. + lifecycle_concurrency: 16 stdio: # Independent gate for local stdio MCP transports. Cloud v2 sets # MCP__STDIO__ENABLED=false even when Box Runtime is available. @@ -204,6 +264,15 @@ box: # External WebSocket runtimes also require LANGBOT_BOX_CONTROL_TOKEN in # both LangBot and Box. Keep the shared secret out of this config file. endpoint: '' # External Box Runtime base URL, e.g. 'ws://127.0.0.1:5410'. Leave empty for local auto-managed runtime. + limits: + max_sessions: 64 + max_managed_processes: 64 + max_completed_processes: 256 + # Retained admission fences prevent replay after entitlement expiry or + # revocation. Fail closed before that monotonic state can grow without + # bound; Cloud may override this with BOX__LIMITS__MAX_ADMISSION_RECORDS. + max_admission_records: 100000 + max_rpc_file_bytes: 20971520 # Cloud v2 overrides these values through the instance config/environment. # OSS keeps admission disabled and preserves the existing multi-session # local behavior. These limits are Runtime-owned and cannot be relaxed by diff --git a/tests/e2e/test_startup.py b/tests/e2e/test_startup.py index e63150b4c..955e33af6 100644 --- a/tests/e2e/test_startup.py +++ b/tests/e2e/test_startup.py @@ -27,6 +27,21 @@ class TestStartupFlow: """Verify LangBot API is responding.""" assert langbot_process.health_check() + def test_health_check_exposes_bounded_blocking_executor(self, e2e_client): + """The production startup path installs blocking-work admission.""" + response = e2e_client.get('/healthz') + + assert response.status_code == 200 + executor = response.json()['resources']['blocking_executor'] + assert executor['max_workers'] == 8 + assert executor['max_pending'] == 128 + assert executor['max_inflight_per_scope'] == 4 + assert executor['inflight'] >= 0 + assert executor['rejected_total'] >= 0 + event_loop = response.json()['resources']['event_loop'] + assert event_loop['running'] is True + assert event_loop['recent_max_lag_ms'] >= 0 + def test_system_info_endpoint(self, e2e_client): """Test /api/v1/system/info endpoint.""" response = e2e_client.get('/api/v1/system/info') diff --git a/tests/integration/api/test_smoke.py b/tests/integration/api/test_smoke.py index 92087fca4..9c2927503 100644 --- a/tests/integration/api/test_smoke.py +++ b/tests/integration/api/test_smoke.py @@ -160,7 +160,7 @@ class TestHealthEndpoint: assert response.status_code == 200 data = await response.get_json() - assert data == {'code': 0, 'msg': 'ok'} + assert data == {'code': 0, 'msg': 'ok', 'resources': {}} @pytest.mark.asyncio async def test_healthz_no_auth_required(self, quart_test_client): @@ -327,9 +327,7 @@ class TestUserInitEndpoint: ) assert response.status_code == 200 - fake_api_app.user_service.reset_password.assert_awaited_once_with( - 'member@example.com', 'new-member-password' - ) + fake_api_app.user_service.reset_password.assert_awaited_once_with('member@example.com', 'new-member-password') @pytest.mark.usefixtures('mock_circular_import_chain') diff --git a/tests/integration/persistence/test_migrations_postgres.py b/tests/integration/persistence/test_migrations_postgres.py index 023a0de11..4eec1668c 100644 --- a/tests/integration/persistence/test_migrations_postgres.py +++ b/tests/integration/persistence/test_migrations_postgres.py @@ -20,6 +20,7 @@ import asyncio import contextlib import datetime import hashlib +import time import typing from types import SimpleNamespace from unittest.mock import AsyncMock @@ -72,6 +73,12 @@ from langbot.pkg.api.http.service.user import UserService from langbot.pkg.api.mcp.context import get_request_context as get_mcp_request_context from langbot.pkg.api.mcp.mount import MCPMount from langbot.pkg.entity.persistence.apikey import ApiKey +from langbot.pkg.entity.persistence import bot as persistence_bot +from langbot.pkg.entity.persistence import mcp as persistence_mcp +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 plugin as persistence_plugin +from langbot.pkg.entity.persistence import rag as persistence_rag from langbot.pkg.entity.persistence.metadata import WorkspaceMetadata from langbot.pkg.entity.persistence.monitoring import MonitoringFeedback from langbot.pkg.entity.persistence.workspace import ( @@ -81,6 +88,9 @@ from langbot.pkg.entity.persistence.workspace import ( ) from langbot.pkg.platform.botmgr import PlatformManager from langbot.pkg.pipeline.pipelinemgr import PipelineManager +from langbot.pkg.plugin.connector import PluginRuntimeConnector +from langbot.pkg.provider.modelmgr import requester as model_requester +from langbot.pkg.provider.modelmgr import token as model_token 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 @@ -99,6 +109,39 @@ class _NoopDirectoryProjectionProvider: raise AssertionError(f'Unexpected delta fetch for {instance_uuid}: {workspace_uuids!r}') +class _CapacityPluginRuntimeHandler: + """Minimal shared Runtime control-plane surface for the startup probe.""" + + def __init__(self) -> None: + self.bindings: dict[str, typing.Any] = {} + self.reconciled: tuple[typing.Any, ...] = () + + def register_installation_binding( + self, + binding, + *, + plugin_author: str, + plugin_name: str, + ) -> None: + self.bindings[binding.installation_uuid] = ( + binding, + plugin_author, + plugin_name, + ) + + def unregister_installation_binding(self, binding) -> None: + self.bindings.pop(binding.installation_uuid, None) + + async def reconcile_plugin_installations(self, desired_states) -> dict: + self.reconciled = tuple(desired_states) + return { + 'applied': [], + 'removed': [], + 'missing_artifacts': [], + 'failed_installations': [], + } + + def _get_script_head() -> str: """Resolve the current Alembic head revision from the script directory. @@ -653,6 +696,385 @@ class TestPostgreSQLTenantRuntime: finally: await _dispose_manager(manager) + @pytest.mark.asyncio + async def test_populated_cloud_startup_is_linear_and_task_bounded( + self, + postgres_url, + postgres_engine, + clean_tables, + clean_alembic_version, + monkeypatch, + ): + """Run the real Cloud startup query graph against populated RLS tenants. + + The default is intentionally small enough for CI. Audit runs can raise + ``LANGBOT_PG_CAPACITY_WORKSPACES`` without changing the test contract. + Every tenant owns one representative startup resource of each kind. + """ + + workspace_count = int(os.environ.get('LANGBOT_PG_CAPACITY_WORKSPACES', '25')) + if not 1 <= workspace_count <= 2_000: + raise ValueError('LANGBOT_PG_CAPACITY_WORKSPACES must be between 1 and 2000') + max_elapsed_raw = os.environ.get( + 'LANGBOT_PG_CAPACITY_MAX_SECONDS', + ) + max_elapsed = float(max_elapsed_raw) if max_elapsed_raw is not None else None + instance_uuid = 'cloud-populated-startup-capacity-test' + role_name = f'lb_capacity_{uuid.uuid4().hex[:12]}' + role_password = f'Lb{uuid.uuid4().hex}' + quote = postgres_engine.dialect.identifier_preparer.quote + managers: list[PersistenceManager] = [] + role_created = False + statement_counts = { + table_name: 0 + for table_name in ( + 'model_providers', + 'llm_models', + 'embedding_models', + 'rerank_models', + 'bots', + 'legacy_pipelines', + 'knowledge_bases', + 'mcp_servers', + 'plugin_settings', + ) + } + measured_engine = None + model_manager = None + platform_manager = None + mcp_loader = None + + _restore_postgres_manager_registry(monkeypatch) + monkeypatch.setattr(constants, 'instance_id', instance_uuid) + release_manager = PersistenceManager( + _application_for_postgres_url( + postgres_url, + 'postgres-capacity-release-test', + ), + mode=PersistenceMode.RELEASE_MIGRATION, + ) + managers.append(release_manager) + + def role_url() -> str: + return ( + sa.engine.make_url(postgres_url) + .set(username=role_name, password=role_password) + .render_as_string(hide_password=False) + ) + + def count_resource_statements( + _conn, + _cursor, + statement, + _parameters, + _context, + _executemany, + ) -> None: + normalized = ' '.join(str(statement).lower().split()) + for table_name in statement_counts: + if f' from {table_name}' in normalized or f' from "{table_name}"' in normalized: + statement_counts[table_name] += 1 + + try: + await release_manager.initialize() + for index in range(workspace_count): + workspace_uuid = f'ca{index:06x}-0000-4000-8000-{index:012x}' + suffix = f'{index:06d}' + provider_uuid = f'capacity-provider-{suffix}' + async with release_manager.tenant_uow(workspace_uuid) as uow: + uow.session.add( + Workspace( + uuid=workspace_uuid, + instance_uuid=instance_uuid, + name=f'Capacity {suffix}', + slug=f'capacity-{suffix}', + type='team', + status='active', + source='cloud_projection', + projection_revision=1, + ) + ) + await uow.session.flush() + uow.session.add_all( + [ + WorkspaceExecutionState( + workspace_uuid=workspace_uuid, + instance_uuid=instance_uuid, + active_generation=1, + state='active', + write_fenced=False, + source='cloud', + desired_state_revision=1, + ), + persistence_model.ModelProvider( + uuid=provider_uuid, + workspace_uuid=workspace_uuid, + name='Capacity Provider', + requester='capacity-probe', + base_url='https://capacity.invalid', + api_keys=[], + ), + ] + ) + await uow.session.flush() + uow.session.add_all( + [ + persistence_model.LLMModel( + uuid=f'capacity-llm-{suffix}', + workspace_uuid=workspace_uuid, + name='Capacity LLM', + provider_uuid=provider_uuid, + abilities=[], + extra_args={}, + ), + persistence_model.EmbeddingModel( + uuid=f'capacity-embedding-{suffix}', + workspace_uuid=workspace_uuid, + name='Capacity Embedding', + provider_uuid=provider_uuid, + extra_args={}, + ), + persistence_model.RerankModel( + uuid=f'capacity-rerank-{suffix}', + workspace_uuid=workspace_uuid, + name='Capacity Rerank', + provider_uuid=provider_uuid, + extra_args={}, + ), + persistence_bot.Bot( + uuid=f'capacity-bot-{suffix}', + workspace_uuid=workspace_uuid, + name='Capacity Bot', + description='', + adapter='capacity-probe', + adapter_config={}, + enable=False, + pipeline_routing_rules=[], + ), + persistence_pipeline.LegacyPipeline( + uuid=f'capacity-pipeline-{suffix}', + workspace_uuid=workspace_uuid, + name='Capacity Pipeline', + description='', + for_version='capacity-probe', + is_default=True, + stages=[], + config={}, + extensions_preferences={}, + ), + persistence_rag.KnowledgeBase( + uuid=f'capacity-kb-{suffix}', + workspace_uuid=workspace_uuid, + name='Capacity Knowledge Base', + description='', + collection_id=f'capacity-collection-{suffix}', + legacy_vector_collection=False, + ), + persistence_mcp.MCPServer( + uuid=f'capacity-mcp-{suffix}', + workspace_uuid=workspace_uuid, + name=f'capacity-mcp-{suffix}', + enable=False, + mode='remote', + extra_args={}, + ), + persistence_plugin.PluginSetting( + workspace_uuid=workspace_uuid, + plugin_author='capacity', + plugin_name=f'plugin-{suffix}', + installation_uuid=str( + uuid.uuid5( + uuid.NAMESPACE_URL, + f'langbot-capacity:{workspace_uuid}', + ) + ), + artifact_digest=hashlib.sha256(workspace_uuid.encode()).hexdigest(), + runtime_revision=1, + enabled=False, + config={}, + install_source='github', + install_info={}, + ), + ] + ) + + async with postgres_engine.connect() as conn: + await conn.execute(text(f"CREATE ROLE {quote(role_name)} LOGIN PASSWORD '{role_password}'")) + role_created = True + await conn.execute( + text( + f'GRANT CONNECT ON DATABASE ' + f'{quote(sa.engine.make_url(postgres_url).database)} ' + f'TO {quote(role_name)}' + ) + ) + await conn.execute(text(f'GRANT USAGE ON SCHEMA public TO {quote(role_name)}')) + await _grant_runtime_role_business_objects( + conn, + role_name, + quote, + ) + + runtime_application = _application_for_postgres_url( + role_url(), + 'postgres-capacity-runtime-test', + ) + runtime_application.instance_config.data.update( + { + 'plugin': {'enable': True}, + 'mcp': {'lifecycle_concurrency': 8}, + } + ) + runtime_application.deployment = SimpleNamespace( + mode='cloud', + multi_workspace_enabled=False, + ) + runtime_application.task_mgr = SimpleNamespace( + cancel_by_scope=lambda *_args, **_kwargs: None, + ) + cloud_manager = PersistenceManager( + runtime_application, + mode=PersistenceMode.CLOUD_RUNTIME, + ) + managers.append(cloud_manager) + await cloud_manager.initialize() + runtime_application.persistence_mgr = cloud_manager + cloud_manager.ap = runtime_application + runtime_application.workspace_service = WorkspaceService( + runtime_application, + policy=CloudWorkspacePolicy(), + instance_uuid=instance_uuid, + ) + + measured_engine = cloud_manager.get_db_engine().sync_engine + sa.event.listen( + measured_engine, + 'before_cursor_execute', + count_resource_statements, + ) + wall_started = time.monotonic() + cpu_started = time.process_time() + + bindings = await runtime_application.workspace_service.prime_startup_execution_bindings() + assert len(bindings) == workspace_count + + model_manager = ModelManager(runtime_application) + + async def build_capacity_provider( + context, + provider_entity, + ): + return model_requester.RuntimeProvider( + context, + provider_entity, + model_token.TokenManager( + provider_entity.uuid, + provider_entity.api_keys or [], + ), + SimpleNamespace(aclose=AsyncMock()), + ) + + model_manager._build_provider = build_capacity_provider + await model_manager.load_models_from_db() + + platform_manager = PlatformManager(runtime_application) + platform_manager.load_bot = AsyncMock() + await platform_manager.load_bots_from_db() + + pipeline_manager = PipelineManager(runtime_application) + pipeline_manager.stage_dict = {} + await pipeline_manager.load_pipelines_from_db() + + rag_manager = RAGManager(runtime_application) + await rag_manager.load_knowledge_bases_from_db() + + mcp_loader = MCPLoader(runtime_application) + mcp_loader.host_mcp_server = AsyncMock() + await mcp_loader.load_mcp_servers_from_db() + dispatch_tasks = tuple(mcp_loader._host_dispatch_tasks) + if dispatch_tasks: + await asyncio.gather(*dispatch_tasks) + await asyncio.sleep(0) + + plugin_connector = PluginRuntimeConnector( + runtime_application, + AsyncMock(), + ) + plugin_handler = _CapacityPluginRuntimeHandler() + plugin_connector.handler = plugin_handler + contexts = [ + ExecutionContext( + instance_uuid=binding.instance_uuid, + workspace_uuid=binding.workspace_uuid, + placement_generation=binding.placement_generation, + ) + for binding in bindings + ] + await plugin_connector.reconcile_projected_workspaces(contexts) + + elapsed = time.monotonic() - wall_started + cpu_seconds = time.process_time() - cpu_started + logging.getLogger('postgres-capacity-runtime-test').info( + 'Populated Cloud startup capacity: workspaces=%d elapsed=%.3fs cpu=%.3fs statements=%s', + workspace_count, + elapsed, + cpu_seconds, + statement_counts, + ) + + assert len(model_manager.provider_dict) == workspace_count + assert len(model_manager.llm_model_dict) == workspace_count + assert len(model_manager.embedding_model_dict) == workspace_count + assert len(model_manager.rerank_model_dict) == workspace_count + assert platform_manager.load_bot.await_count == workspace_count + assert len(pipeline_manager.pipelines) == workspace_count + assert len(rag_manager.knowledge_bases) == workspace_count + assert mcp_loader.host_mcp_server.await_count == workspace_count + assert not mcp_loader._host_dispatch_tasks + assert not mcp_loader._hosted_mcp_tasks + assert len(plugin_handler.reconciled) == workspace_count + assert len(plugin_handler.bindings) == workspace_count + assert all(count == workspace_count for count in statement_counts.values()), statement_counts + if max_elapsed is not None: + assert elapsed <= max_elapsed + finally: + cleanup_errors: list[BaseException] = [] + if measured_engine is not None: + sa.event.remove( + measured_engine, + 'before_cursor_execute', + count_resource_statements, + ) + if mcp_loader is not None: + try: + await mcp_loader.shutdown() + except BaseException as exc: + cleanup_errors.append(exc) + if platform_manager is not None: + try: + await platform_manager.shutdown() + except BaseException as exc: + cleanup_errors.append(exc) + if model_manager is not None: + try: + await model_manager.shutdown() + except BaseException as exc: + cleanup_errors.append(exc) + for manager in reversed(managers): + try: + await _dispose_manager(manager) + except BaseException as exc: + cleanup_errors.append(exc) + if role_created: + try: + async with postgres_engine.connect() as conn: + await conn.execute(text(f'DROP OWNED BY {quote(role_name)}')) + await conn.execute(text(f'DROP ROLE {quote(role_name)}')) + except BaseException as exc: + cleanup_errors.append(exc) + if cleanup_errors: + raise cleanup_errors[0] + @pytest.mark.asyncio async def test_release_bootstrap_and_runtime_isolation( self, diff --git a/tests/integration/persistence/test_release_migration_postgres.py b/tests/integration/persistence/test_release_migration_postgres.py index e1cec1cf7..b6637afda 100644 --- a/tests/integration/persistence/test_release_migration_postgres.py +++ b/tests/integration/persistence/test_release_migration_postgres.py @@ -367,6 +367,12 @@ async def test_release_entrypoint_holds_lock_migrates_validates_and_disposes( with pytest.raises(RuntimeError, match="table 'metadata' grants are incomplete"): await manager._validate_configured_runtime_postgres_role(require_grants=True) finally: + manager = getattr(ap, 'persistence_mgr', None) + if isinstance(manager, PersistenceManager): + # The one-shot entrypoint disposes its pool before returning. This + # test deliberately reuses the manager for catalog mutation checks, + # which can open a fresh pool and therefore owns a second shutdown. + await manager.shutdown() async with postgres_engine.connect() as conn: await conn.execute(text(f'DROP OWNED BY {quote(runtime_role)}')) await conn.execute(text(f'DROP ROLE IF EXISTS {quote(runtime_role)}')) @@ -631,6 +637,11 @@ async def test_runtime_role_catalog_validator_rejects_delegation_and_escape_hatc async with postgres_engine.connect() as conn: await conn.execute(text(f'ALTER EXTENSION vector DROP FUNCTION public.{quote(security_definer)}()')) finally: + manager = getattr(ap, 'persistence_mgr', None) + if isinstance(manager, PersistenceManager): + # The entrypoint disposed the release pool before returning; all + # validator calls above happened afterward and can reopen it. + await manager.shutdown() async with postgres_engine.connect() as conn: await conn.execute(text(f'ALTER DATABASE {quote(database_name)} RESET search_path')) await conn.execute(text(f'ALTER ROLE {quote(runtime_role)} RESET search_path')) diff --git a/tests/unit_tests/api/http/test_bounded_json_request.py b/tests/unit_tests/api/http/test_bounded_json_request.py new file mode 100644 index 000000000..6fbdb7f97 --- /dev/null +++ b/tests/unit_tests/api/http/test_bounded_json_request.py @@ -0,0 +1,39 @@ +from __future__ import annotations + +import quart + +from langbot.pkg.api.http.controller import main as controller_main +from langbot.pkg.utils import bounded_executor + + +async def test_bounded_json_request_decodes_off_loop_in_workspace_scope( + monkeypatch, +): + app = quart.Quart(__name__) + app.request_class = controller_main.BoundedJSONRequest + observed_scopes: list[str | None] = [] + + async def fake_to_thread(fn, *args, **kwargs): + observed_scopes.append(bounded_executor.current_blocking_work_scope()) + return fn(*args, **kwargs) + + monkeypatch.setattr( + controller_main.asyncio, + 'to_thread', + fake_to_thread, + ) + + @app.post('/json') + async def parse_json(): + with bounded_executor.blocking_work_scope('workspace-a'): + payload = await quart.request.get_json() + return quart.jsonify(payload) + + response = await app.test_client().post( + '/json', + json={'nested': {'value': 1}}, + ) + + assert response.status_code == 200 + assert await response.get_json() == {'nested': {'value': 1}} + assert observed_scopes == ['workspace-a'] diff --git a/tests/unit_tests/api/http/test_internal_error_responses.py b/tests/unit_tests/api/http/test_internal_error_responses.py index f5e7eb272..25c0e83ef 100644 --- a/tests/unit_tests/api/http/test_internal_error_responses.py +++ b/tests/unit_tests/api/http/test_internal_error_responses.py @@ -9,6 +9,10 @@ import quart from langbot.pkg.api.http.controller import group from langbot.pkg.api.http.controller.groups.webhooks import WebhookRouterGroup +from langbot.pkg.utils.bounded_executor import ( + BlockingWorkCapacityError, + current_blocking_work_scope, +) pytestmark = pytest.mark.asyncio @@ -34,6 +38,16 @@ class _AuthenticatedRouterGroup(group.RouterGroup): return self.success() +class _BlockingCapacityRouterGroup(group.RouterGroup): + name = 'blocking-capacity-test' + path = '/blocking-capacity-test' + + async def initialize(self) -> None: + @self.route('', methods=['GET'], auth_type=group.AuthType.NONE) + async def _(): + raise BlockingWorkCapacityError('Workspace blocking executor capacity reached') + + class _InvalidAccountRouterGroup(group.RouterGroup): name = 'invalid-account-test' path = '/invalid-account-test' @@ -102,6 +116,20 @@ async def test_public_webhook_error_uses_same_generic_error_contract(): assert 'do-not-return' not in (await response.get_data(as_text=True)) +async def test_blocking_work_capacity_maps_to_retryable_http_response(): + application = SimpleNamespace(logger=Mock()) + quart_app = quart.Quart(__name__) + await _BlockingCapacityRouterGroup(application, quart_app).initialize() + + response = await quart_app.test_client().get('/blocking-capacity-test') + + assert response.status_code == 429 + assert await response.get_json() == { + 'code': 'blocking_work_capacity_exceeded', + 'msg': 'Workspace blocking executor capacity reached', + } + + async def test_public_webhook_carries_scope_without_holding_database_session(): class ScopeOnlyPersistenceManager: mode = SimpleNamespace(value='cloud_runtime') @@ -128,6 +156,7 @@ async def test_public_webhook_carries_scope_without_holding_database_session(): async def handle_unified_webhook(self, **_kwargs): assert persistence_mgr.active_workspace == workspace_uuid assert persistence_mgr.current_session() is None + assert current_blocking_work_scope() == workspace_uuid return {'ok': True} async def get_execution_binding(resolved_workspace_uuid, expected_generation=None): @@ -156,6 +185,41 @@ async def test_public_webhook_carries_scope_without_holding_database_session(): assert persistence_mgr.active_workspace is None +async def test_public_webhook_blocking_capacity_is_retryable(): + workspace_uuid = '00000000-0000-0000-0000-00000000000a' + bot_uuid = '11111111-1111-4111-8111-111111111111' + + class Adapter: + async def handle_unified_webhook(self, **_kwargs): + raise BlockingWorkCapacityError( + 'Workspace blocking executor capacity reached', + scope=workspace_uuid, + ) + + runtime_bot = SimpleNamespace( + workspace_uuid=workspace_uuid, + placement_generation=4, + enable=True, + adapter=Adapter(), + ) + application = SimpleNamespace( + logger=Mock(), + persistence_mgr=SimpleNamespace(mode=SimpleNamespace(value='oss')), + platform_mgr=SimpleNamespace(resolve_public_bot=AsyncMock(return_value=runtime_bot)), + workspace_service=SimpleNamespace(get_execution_binding=AsyncMock(return_value=None)), + ) + quart_app = quart.Quart(__name__) + await WebhookRouterGroup(application, quart_app).initialize() + + response = await quart_app.test_client().post(f'/bots/{bot_uuid}') + + assert response.status_code == 429 + assert await response.get_json() == { + 'code': 'blocking_work_capacity_exceeded', + 'msg': 'Workspace blocking executor capacity reached', + } + + async def test_authentication_failure_does_not_return_internal_exception_text(): logger = Mock() application = SimpleNamespace( diff --git a/tests/unit_tests/api/service/test_mcp_service.py b/tests/unit_tests/api/service/test_mcp_service.py index 018bf77fe..31c531e98 100644 --- a/tests/unit_tests/api/service/test_mcp_service.py +++ b/tests/unit_tests/api/service/test_mcp_service.py @@ -29,6 +29,7 @@ from langbot.pkg.api.http.context import ( WorkspaceContext, ) from langbot.pkg.api.http.service.mcp import MCPService, redact_mcp_secrets, restore_mcp_secret_placeholders +from langbot.pkg.core.taskmgr import TaskCapacityError from langbot.pkg.entity.persistence.mcp import MCPServer from langbot.pkg.provider.tools.loaders.mcp_policy import MCPStdioDisabledError from langbot.pkg.workspace.errors import WorkspaceNotFoundError @@ -1050,3 +1051,23 @@ class TestMCPServiceTestMCPServer: # Verify - load_mcp_server called ap.tool_mgr.mcp_tool_loader.load_mcp_server.assert_called_once() assert task_id == 456 + + async def test_rejected_transient_test_session_is_shut_down(self): + ap = SimpleNamespace() + mock_session = MagicMock() + mock_session.shutdown = AsyncMock() + ap.tool_mgr = SimpleNamespace( + mcp_tool_loader=SimpleNamespace(load_mcp_server=AsyncMock(return_value=mock_session)) + ) + + def reject(coroutine, **_kwargs): + coroutine.close() + raise TaskCapacityError('capacity') + + ap.task_mgr = SimpleNamespace(create_user_task=Mock(side_effect=reject)) + service = _service(ap) + + with pytest.raises(TaskCapacityError, match='capacity'): + await service.test_mcp_server(_CONTEXT, '_', {'name': 'New Server'}) + + mock_session.shutdown.assert_awaited_once_with() diff --git a/tests/unit_tests/api/service/test_space_service.py b/tests/unit_tests/api/service/test_space_service.py index b9a833bd3..6fb74d20b 100644 --- a/tests/unit_tests/api/service/test_space_service.py +++ b/tests/unit_tests/api/service/test_space_service.py @@ -13,6 +13,8 @@ Source: src/langbot/pkg/api/http/service/space.py from __future__ import annotations +from collections import OrderedDict +import json from urllib.parse import parse_qs, urlsplit import pytest @@ -28,6 +30,23 @@ from langbot.pkg.entity.persistence.user import User pytestmark = pytest.mark.asyncio +def _set_response_body(response: MagicMock, body: dict | str) -> None: + """Configure an aiohttp-like streaming body on an HTTP response mock.""" + + raw_body = body.encode() if isinstance(body, str) else json.dumps(body).encode() + + class Content: + async def iter_chunked(self, _chunk_size: int): + midpoint = max(len(raw_body) // 2, 1) + yield raw_body[:midpoint] + if midpoint < len(raw_body): + yield raw_body[midpoint:] + + response.headers = {} + response.content = Content() + response.charset = 'utf-8' + + def _create_mock_user( email: str = 'test@example.com', account_type: str = 'space', @@ -292,6 +311,40 @@ class TestSpaceServiceGetCredits: # Verify - returns cached value without API call assert result == 100 + async def test_cached_credit_lookup_does_not_scan_all_users(self): + ap = SimpleNamespace() + ap.instance_config = SimpleNamespace(data={}) + ap.persistence_mgr = SimpleNamespace() + service = SpaceService(ap) + + class AtMostOneStepOrderedDict(OrderedDict): + def __iter__(self): + iterator = super().__iter__() + yielded = False + + def next_entry(): + nonlocal yielded + if yielded: + raise AssertionError('credits cache scanned all users') + yielded = True + return next(iterator) + + class AtMostOneStepIterator: + def __iter__(self): + return self + + def __next__(self): + return next_entry() + + return AtMostOneStepIterator() + + now = time.time() + service._credits_cache = AtMostOneStepOrderedDict( + (f'user-{index}@example.com', (index, now)) for index in range(512) + ) + + assert await service.get_credits('user-511@example.com') == 511 + async def test_get_credits_cache_expired_refreshes(self): """Refreshes expired cache.""" # Setup @@ -406,6 +459,7 @@ class TestSpaceServiceRefreshToken: }, } ) + _set_response_body(mock_response, mock_response.json.return_value) with patch('langbot.pkg.api.http.service.space.httpclient.get_session') as mock_session: mock_session_obj = MagicMock() @@ -441,6 +495,7 @@ class TestSpaceServiceRefreshToken: } ) mock_response.text = AsyncMock(return_value='{"code":1,"msg":"Invalid refresh token"}') + _set_response_body(mock_response, mock_response.json.return_value) with patch('langbot.pkg.api.http.service.space.httpclient.get_session') as mock_session: mock_session_obj = MagicMock() @@ -467,6 +522,7 @@ class TestSpaceServiceRefreshToken: mock_response = MagicMock() mock_response.status = 500 mock_response.text = AsyncMock(return_value='Internal Server Error') + _set_response_body(mock_response, mock_response.text.return_value) with patch('langbot.pkg.api.http.service.space.httpclient.get_session') as mock_session: mock_session_obj = MagicMock() @@ -506,6 +562,7 @@ class TestSpaceServiceExchangeOAuthCode: }, } ) + _set_response_body(mock_response, mock_response.json.return_value) with patch('langbot.pkg.api.http.service.space.httpclient.get_session') as mock_session: mock_session_obj = MagicMock() @@ -535,6 +592,7 @@ class TestSpaceServiceExchangeOAuthCode: mock_response.status = 200 mock_response.json = AsyncMock(return_value={'code': 1, 'msg': 'Invalid code'}) mock_response.text = AsyncMock(return_value='{"code":1,"msg":"Invalid code"}') + _set_response_body(mock_response, mock_response.json.return_value) with patch('langbot.pkg.api.http.service.space.httpclient.get_session') as mock_session: mock_session_obj = MagicMock() @@ -573,6 +631,7 @@ class TestSpaceServiceGetUserInfoRaw: }, } ) + _set_response_body(mock_response, mock_response.json.return_value) with patch('langbot.pkg.api.http.service.space.httpclient.get_session') as mock_session: mock_session_obj = MagicMock() @@ -603,6 +662,7 @@ class TestSpaceServiceGetUserInfoRaw: mock_response.status = 200 mock_response.json = AsyncMock(return_value={'code': 1, 'msg': 'Unauthorized'}) mock_response.text = AsyncMock(return_value='{"code":1,"msg":"Unauthorized"}') + _set_response_body(mock_response, mock_response.json.return_value) with patch('langbot.pkg.api.http.service.space.httpclient.get_session') as mock_session: mock_session_obj = MagicMock() @@ -703,6 +763,7 @@ class TestSpaceServiceGetModels: }, } ) + _set_response_body(mock_response, mock_response.json.return_value) with patch('langbot.pkg.api.http.service.space.httpclient.get_session') as mock_session: mock_session_obj = MagicMock() @@ -732,6 +793,7 @@ class TestSpaceServiceGetModels: mock_response.status = 200 mock_response.json = AsyncMock(return_value={'code': 1, 'msg': 'Unauthorized'}) mock_response.text = AsyncMock(return_value='{"code":1,"msg":"Unauthorized"}') + _set_response_body(mock_response, mock_response.json.return_value) with patch('langbot.pkg.api.http.service.space.httpclient.get_session') as mock_session: mock_session_obj = MagicMock() diff --git a/tests/unit_tests/api/service/test_user_service.py b/tests/unit_tests/api/service/test_user_service.py index 8a00fad94..1e9dfa270 100644 --- a/tests/unit_tests/api/service/test_user_service.py +++ b/tests/unit_tests/api/service/test_user_service.py @@ -29,11 +29,25 @@ from langbot.pkg.entity.errors.account import ( SpaceAccountBindingRequiredError, SpaceAccountNotRegisteredError, ) +from langbot.pkg.utils.bounded_executor import BlockingWorkCapacityError pytestmark = pytest.mark.asyncio +async def test_password_hashing_rejects_concurrent_waiters() -> None: + service = UserService(SimpleNamespace()) + await service._password_hash_lock.acquire() + try: + with pytest.raises( + BlockingWorkCapacityError, + match='Password hashing capacity reached', + ): + await service._hash_password('secret') + finally: + service._password_hash_lock.release() + + class TestSpaceOAuthState: async def test_login_state_is_opaque_single_use(self): service = UserService(SimpleNamespace()) @@ -91,6 +105,32 @@ class TestSpaceOAuthState: assert consumed.account is None assert consumed.launch_workspace_uuid == 'workspace-a' + async def test_issue_state_does_not_scan_all_live_states(self, monkeypatch): + service = UserService(SimpleNamespace()) + for _ in range(512): + await service.issue_space_oauth_state('login') + + class NoGlobalIterationDict(dict): + def __iter__(self): + raise AssertionError('OAuth state issuance scanned all live states') + + def keys(self): + raise AssertionError('OAuth state issuance scanned all live states') + + def items(self): + raise AssertionError('OAuth state issuance scanned all live states') + + def values(self): + raise AssertionError('OAuth state issuance scanned all live states') + + guarded_states = NoGlobalIterationDict(service._space_oauth_states) + monkeypatch.setattr(service, '_space_oauth_states', guarded_states) + + state = await service.issue_space_oauth_state('login') + + assert await service.consume_space_oauth_state(state, 'login') is None + assert len(guarded_states) == 512 + def _create_mock_user( email: str = 'test@example.com', diff --git a/tests/unit_tests/api/test_adapter_session_scoping.py b/tests/unit_tests/api/test_adapter_session_scoping.py index cdbcf8cca..3e8642aad 100644 --- a/tests/unit_tests/api/test_adapter_session_scoping.py +++ b/tests/unit_tests/api/test_adapter_session_scoping.py @@ -2,7 +2,7 @@ from __future__ import annotations import asyncio from types import SimpleNamespace -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, Mock import lark_oapi import pytest @@ -19,7 +19,9 @@ from langbot.pkg.api.http.controller.groups.platform.adapters import ( _AdapterSessionScope, _bind_session_scope, _get_owned_session, + _make_room_for_session, _pop_owned_session, + _start_adapter_session_task, ) @@ -94,6 +96,10 @@ async def _create_client(*, role: str = 'developer'): ), ) + class TestTaskManager: + def create_user_task(self, coro, **_kwargs): + return SimpleNamespace(task=asyncio.create_task(coro)) + application = SimpleNamespace( user_service=SimpleNamespace( get_authenticated_account=AsyncMock(side_effect=get_authenticated_account), @@ -102,6 +108,7 @@ async def _create_client(*, role: str = 'developer'): resolve_account_workspace=AsyncMock(side_effect=resolve_account_workspace), ), platform_mgr=SimpleNamespace(), + task_mgr=TestTaskManager(), ) router = AdaptersRouterGroup(application, quart_app) await router.initialize() @@ -142,6 +149,58 @@ async def test_session_scope_matches_exact_tenant_placement_and_principal(): assert sessions == {} +async def test_session_capacity_evicts_oldest_session_in_same_workspace(): + owner_context = _request_context() + sessions: dict[str, dict] = {} + tasks = [] + for index in range(10): + task = SimpleNamespace(done=Mock(return_value=False), cancel=Mock()) + tasks.append(task) + session = {'created_at': float(index), 'task': task} + _bind_session_scope(session, owner_context) + sessions[f'session-{index}'] = session + + _make_room_for_session(sessions, owner_context) + + assert 'session-0' not in sessions + assert len(sessions) == 9 + tasks[0].cancel.assert_called_once_with() + + +async def test_adapter_session_task_uses_tenant_task_admission(): + blocker = asyncio.Event() + + async def credential_exchange(): + await blocker.wait() + + task_manager = SimpleNamespace(create_user_task=Mock()) + + def create_user_task(coro, **_kwargs): + return SimpleNamespace(task=asyncio.create_task(coro)) + + task_manager.create_user_task.side_effect = create_user_task + application = SimpleNamespace(task_mgr=task_manager) + request_context = _request_context() + + returned = _start_adapter_session_task( + application, + credential_exchange(), + adapter='lark', + session_id='session-test', + request_context=request_context, + ) + + assert returned is not None + task_manager.create_user_task.assert_called_once() + kwargs = task_manager.create_user_task.call_args.kwargs + assert kwargs['kind'] == 'platform-adapter-credential-exchange' + assert kwargs['instance_uuid'] == request_context.instance_uuid + assert kwargs['workspace_uuid'] == request_context.workspace_uuid + assert kwargs['placement_generation'] == request_context.placement_generation + blocker.set() + await returned + + async def test_lark_session_status_and_delete_hide_cross_scope_sessions(monkeypatch): registration_blocker = asyncio.Event() diff --git a/tests/unit_tests/api/test_websocket_chat_uow.py b/tests/unit_tests/api/test_websocket_chat_uow.py index 9322eb67c..5f08563d2 100644 --- a/tests/unit_tests/api/test_websocket_chat_uow.py +++ b/tests/unit_tests/api/test_websocket_chat_uow.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio from contextlib import asynccontextmanager from types import SimpleNamespace from unittest.mock import AsyncMock, Mock @@ -13,6 +14,11 @@ from langbot.pkg.api.http.context import ( WorkspaceContext, ) from langbot.pkg.api.http.controller.groups.pipelines.websocket_chat import WebSocketChatRouterGroup +from langbot.pkg.api.http.controller.groups.pipelines.websocket_chat import ( + create_scoped_duplex_tasks, +) +from langbot.pkg.api.http.controller.groups.pipelines.websocket_chat import wait_for_duplex_tasks +from langbot.pkg.utils.bounded_executor import current_blocking_work_scope @pytest.mark.asyncio @@ -67,3 +73,72 @@ async def test_websocket_pipeline_lookup_opens_workspace_uow_after_auth_scope_cl assert result is adapter assert scopes == [workspace_uuid] + + +@pytest.mark.asyncio +async def test_duplex_websocket_tasks_cancel_blocked_peer_when_one_direction_ends() -> None: + blocked = asyncio.Event() + + async def receive_forever() -> None: + blocked.set() + await asyncio.Future() + + async def send_finishes() -> None: + await blocked.wait() + + receive_task = asyncio.create_task(receive_forever()) + send_task = asyncio.create_task(send_finishes()) + + await asyncio.wait_for( + wait_for_duplex_tasks(receive_task, send_task), + timeout=1, + ) + + assert receive_task.cancelled() + assert send_task.done() + + +@pytest.mark.asyncio +async def test_duplex_websocket_tasks_allow_terminal_send_to_drain() -> None: + receive_finished = asyncio.Event() + send_drained = asyncio.Event() + + async def receive_finishes() -> None: + receive_finished.set() + + async def send_terminal_frame() -> None: + await receive_finished.wait() + await asyncio.sleep(0) + send_drained.set() + + receive_task = asyncio.create_task(receive_finishes()) + send_task = asyncio.create_task(send_terminal_frame()) + + await wait_for_duplex_tasks(receive_task, send_task) + + assert send_drained.is_set() + assert send_task.done() + assert not send_task.cancelled() + + +@pytest.mark.asyncio +async def test_duplex_websocket_tasks_share_trusted_workspace_budget() -> None: + observed: list[tuple[str, str | None]] = [] + + async def observe(direction: str) -> None: + await asyncio.sleep(0) + observed.append((direction, current_blocking_work_scope())) + + receive_task, send_task = create_scoped_duplex_tasks( + observe('receive'), + observe('send'), + 'workspace-a', + ) + + await asyncio.gather(receive_task, send_task) + + assert sorted(observed) == [ + ('receive', 'workspace-a'), + ('send', 'workspace-a'), + ] + assert current_blocking_work_scope() is None diff --git a/tests/unit_tests/box/test_box_service.py b/tests/unit_tests/box/test_box_service.py index 5f31734ef..bce4a03e5 100644 --- a/tests/unit_tests/box/test_box_service.py +++ b/tests/unit_tests/box/test_box_service.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio import datetime as dt import os +import pathlib import tempfile from types import SimpleNamespace from unittest.mock import AsyncMock, Mock @@ -2050,6 +2051,17 @@ class TestAttachmentHelpers: component = SimpleNamespace(base64=None, url=None, path=None) assert await BoxService._component_to_bytes(component) is None + @pytest.mark.asyncio + async def test_component_to_bytes_rejects_oversized_base64(self, monkeypatch): + monkeypatch.setattr(BoxService, '_ATTACHMENT_MAX_BYTES', 4) + component = SimpleNamespace( + base64='data:application/octet-stream;base64,' + ('A' * 12), + url=None, + path=None, + ) + + assert await BoxService._component_to_bytes(component) is None + class TestInboundOutboundRoundTrip: def _service(self) -> BoxService: @@ -2206,7 +2218,7 @@ class TestAttachmentHostPath: # File actually landed on the host workspace. host_file = os.path.join(ws, 'inbox', str(query.query_uuid), d['name']) assert os.path.isfile(host_file) - assert open(host_file, 'rb').read() == big + assert pathlib.Path(host_file).read_bytes() == big @pytest.mark.asyncio async def test_inbound_host_clears_stale_query_dir(self, tmp_path): @@ -2218,7 +2230,7 @@ class TestAttachmentHostPath: # Seed a stale file under the same query_id (simulates webchat id reuse). stale_dir = os.path.join(ws, 'inbox', 'query-42') os.makedirs(stale_dir, exist_ok=True) - open(os.path.join(stale_dir, 'image_1.png'), 'wb').write(b'STALE-OLD-IMAGE') + pathlib.Path(stale_dir, 'image_1.png').write_bytes(b'STALE-OLD-IMAGE') new = b'\x89PNG\r\n\x1a\n NEW' b64 = 'data:image/png;base64,' + base64.b64encode(new).decode() @@ -2228,9 +2240,10 @@ class TestAttachmentHostPath: descriptors = await service.materialize_inbound_attachments(query) # The new write recreated the dir; the stale file is gone, new bytes present. host_file = os.path.join(stale_dir, descriptors[0]['name']) - assert open(host_file, 'rb').read() == new + host_bytes = pathlib.Path(host_file).read_bytes() + assert host_bytes == new # No leftover content from the stale image. - assert b'STALE-OLD-IMAGE' not in open(host_file, 'rb').read() + assert b'STALE-OLD-IMAGE' not in host_bytes @pytest.mark.asyncio async def test_inbound_host_replaces_query_symlink_without_touching_other_workspace(self, tmp_path): @@ -2259,7 +2272,7 @@ class TestAttachmentHostPath: assert descriptors[0]['path'] == '/workspace/inbox/query-42/input.bin' assert protected.read_bytes() == b'workspace-b-secret' assert not os.path.islink(os.path.join(inbox, 'query-42')) - assert open(os.path.join(inbox, 'query-42', 'input.bin'), 'rb').read() == payload + assert pathlib.Path(inbox, 'query-42', 'input.bin').read_bytes() == payload @pytest.mark.asyncio async def test_outbound_reads_host_and_clears(self, tmp_path): @@ -2269,8 +2282,8 @@ class TestAttachmentHostPath: os.makedirs(outbox, exist_ok=True) # A large file that would be truncated on the exec/stdout path: big_png = b'\x89PNG\r\n\x1a\n' + b'y' * (400 * 1024) - open(os.path.join(outbox, 'result.png'), 'wb').write(big_png) - open(os.path.join(outbox, 'notes.txt'), 'wb').write(b'hello') + pathlib.Path(outbox, 'result.png').write_bytes(big_png) + pathlib.Path(outbox, 'notes.txt').write_bytes(b'hello') service.execute_tool = AsyncMock(side_effect=AssertionError('exec must not be used on host path')) attachments = await service.collect_outbound_attachments(query) @@ -2354,7 +2367,7 @@ class TestAttachmentHostPath: # it and then wipe the dir. Simulate "nothing produced this turn" by # treating any present file as stale and asserting it is not re-sent # across a second, genuinely-empty collection. - open(os.path.join(outbox, 'stale.png'), 'wb').write(b'\x89PNG\r\n\x1a\n old') + pathlib.Path(outbox, 'stale.png').write_bytes(b'\x89PNG\r\n\x1a\n old') service.execute_tool = AsyncMock(side_effect=AssertionError('exec must not be used on host path')) # First collection drains + clears the dir. @@ -2376,7 +2389,7 @@ class TestAttachmentHostPath: for sub in ('inbox', 'outbox'): d = os.path.join(ws, sub, '0') os.makedirs(d, exist_ok=True) - open(os.path.join(d, 'leftover.bin'), 'wb').write(b'from a previous process') + pathlib.Path(d, 'leftover.bin').write_bytes(b'from a previous process') service.execute_tool = AsyncMock(side_effect=AssertionError('exec must not be used for host-owned files')) await service._purge_attachment_dirs() diff --git a/tests/unit_tests/cloud/test_entitlements.py b/tests/unit_tests/cloud/test_entitlements.py index 5e3d385b7..9268f8a13 100644 --- a/tests/unit_tests/cloud/test_entitlements.py +++ b/tests/unit_tests/cloud/test_entitlements.py @@ -1,5 +1,7 @@ from __future__ import annotations +import asyncio + import pytest from unittest.mock import AsyncMock @@ -107,3 +109,53 @@ async def test_resolver_checks_deployment_admission_before_and_after_provider_ca with pytest.raises(RuntimeError, match='expired during provider call'): await resolver.resolve('workspace-a', now=150) assert checks == 2 + + +@pytest.mark.asyncio +async def test_directory_activity_reconciliation_drops_historical_snapshots(): + provider = AsyncMock() + provider.get_workspace_entitlement = AsyncMock(return_value=_snapshot()) + resolver = EntitlementResolver('instance-a', provider) + await resolver.reconcile_active_workspaces({'workspace-a', 'workspace-b'}) + await resolver.resolve('workspace-a', now=150) + + await resolver.reconcile_active_workspaces({'workspace-b'}) + + assert resolver.snapshot_counts() == { + 'active_workspaces': 1, + 'cached_snapshots': 0, + } + with pytest.raises(EntitlementUnavailableError, match='directory projection'): + await resolver.resolve('workspace-a', now=150) + provider.get_workspace_entitlement.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_directory_fence_wins_race_with_inflight_entitlement_fetch(): + provider_started = asyncio.Event() + release_provider = asyncio.Event() + + async def fetch(_workspace_uuid: str) -> EntitlementSnapshot: + provider_started.set() + await release_provider.wait() + return _snapshot() + + provider = AsyncMock() + provider.get_workspace_entitlement = AsyncMock(side_effect=fetch) + resolver = EntitlementResolver('instance-a', provider) + await resolver.reconcile_active_workspaces({'workspace-a'}) + resolve_task = asyncio.create_task(resolver.resolve('workspace-a', now=150)) + await provider_started.wait() + + await resolver.update_workspace_activity( + active_workspace_uuids=set(), + inactive_workspace_uuids={'workspace-a'}, + ) + release_provider.set() + + with pytest.raises(EntitlementUnavailableError, match='directory projection'): + await resolve_task + assert resolver.snapshot_counts() == { + 'active_workspaces': 0, + 'cached_snapshots': 0, + } diff --git a/tests/unit_tests/cloud/test_space_launch.py b/tests/unit_tests/cloud/test_space_launch.py index feca2ac65..14b6b7f17 100644 --- a/tests/unit_tests/cloud/test_space_launch.py +++ b/tests/unit_tests/cloud/test_space_launch.py @@ -86,6 +86,50 @@ async def test_consumes_valid_workspace_launch_assertion_once(): await service.consume_assertion(token, expected_workspace_uuid=WORKSPACE_UUID) +async def test_replay_cache_does_not_scan_all_live_assertions(monkeypatch): + private_key = Ed25519PrivateKey.generate() + now = int(time.time()) + service = _service(private_key, now=now) + for index in range(512): + await service._consume_jti(f'jti-{index}', now + 90) + + class NoGlobalIterationDict(dict): + def __iter__(self): + raise AssertionError('replay admission scanned all live assertions') + + def keys(self): + raise AssertionError('replay admission scanned all live assertions') + + def items(self): + raise AssertionError('replay admission scanned all live assertions') + + def values(self): + raise AssertionError('replay admission scanned all live assertions') + + guarded_jtis = NoGlobalIterationDict(service._consumed_jtis) + monkeypatch.setattr(service, '_consumed_jtis', guarded_jtis) + + await service._consume_jti('jti-new', now + 90) + + assert len(guarded_jtis) == 513 + + +async def test_replay_cache_fails_closed_at_capacity(monkeypatch): + from langbot.pkg.cloud import launch + + private_key = Ed25519PrivateKey.generate() + now = int(time.time()) + service = _service(private_key, now=now) + monkeypatch.setattr(launch, '_CONSUMED_JTI_MAX_ENTRIES', 2) + await service._consume_jti('jti-1', now + 90) + await service._consume_jti('jti-2', now + 90) + + with pytest.raises(SpaceLaunchError, match='replay cache capacity'): + await service._consume_jti('jti-3', now + 90) + with pytest.raises(SpaceLaunchError, match='already been consumed'): + await service._consume_jti('jti-1', now + 90) + + async def test_rejects_expired_wrong_workspace_and_wrong_instance_assertions(): private_key = Ed25519PrivateKey.generate() now = int(time.time()) diff --git a/tests/unit_tests/core/test_app_shutdown.py b/tests/unit_tests/core/test_app_shutdown.py new file mode 100644 index 000000000..7f17ab825 --- /dev/null +++ b/tests/unit_tests/core/test_app_shutdown.py @@ -0,0 +1,122 @@ +from __future__ import annotations + +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from langbot.pkg.core.app import Application + + +@pytest.mark.asyncio +async def test_shutdown_closes_mcp_session_manager_once() -> None: + app = Application() + stop_session_manager = AsyncMock() + app.platform_mgr = SimpleNamespace(shutdown=AsyncMock()) + app.tool_mgr = SimpleNamespace(shutdown=AsyncMock()) + app.model_mgr = SimpleNamespace(shutdown=AsyncMock()) + app.box_service = SimpleNamespace(shutdown=AsyncMock()) + app.plugin_connector = SimpleNamespace(aclose=AsyncMock()) + app.telemetry = SimpleNamespace(shutdown=AsyncMock()) + app.vector_db_mgr = SimpleNamespace(shutdown=AsyncMock()) + app.storage_mgr = SimpleNamespace(shutdown=AsyncMock()) + manifest_provider = SimpleNamespace(aclose=AsyncMock()) + app.deployment = SimpleNamespace(manifest_provider=manifest_provider) + persistence_engine = SimpleNamespace(dispose=AsyncMock()) + app.persistence_mgr = SimpleNamespace(db=SimpleNamespace(engine=persistence_engine)) + app.http_ctrl = SimpleNamespace(mcp_mount=SimpleNamespace(stop_session_manager=stop_session_manager)) + + await app.shutdown() + await app.shutdown() + + stop_session_manager.assert_awaited_once() + app.platform_mgr.shutdown.assert_awaited_once() + app.tool_mgr.shutdown.assert_awaited_once() + app.model_mgr.shutdown.assert_awaited_once() + app.box_service.shutdown.assert_awaited_once() + app.plugin_connector.aclose.assert_awaited_once() + app.telemetry.shutdown.assert_awaited_once() + app.vector_db_mgr.shutdown.assert_awaited_once() + app.storage_mgr.shutdown.assert_awaited_once() + manifest_provider.aclose.assert_awaited_once() + persistence_engine.dispose.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_dispose_tracks_only_one_shutdown_task() -> None: + app = Application() + app.event_loop = asyncio.get_running_loop() + + app.dispose() + shutdown_task = app._shutdown_task + app.dispose() + + assert shutdown_task is not None + assert app._shutdown_task is shutdown_task + await shutdown_task + + app.dispose() + assert app._shutdown_task is shutdown_task + + +@pytest.mark.asyncio +async def test_runtime_resource_stats_are_aggregate_and_constant_time() -> None: + app = Application() + app.event_loop = asyncio.get_running_loop() + app.blocking_executor = SimpleNamespace( + snapshot=lambda: { + 'inflight': 3, + 'running': 2, + 'pending': 1, + 'rejected_total': 4, + } + ) + app.task_mgr = SimpleNamespace(get_stats=lambda: {'total': 5, 'completed': 2}) + app.query_pool = SimpleNamespace( + queries=[object()], + cached_queries={}, + active_query_count_by_workspace={'workspace-a': 1}, + ) + app.model_mgr = SimpleNamespace( + provider_dict={'provider': object()}, + llm_model_dict={}, + embedding_model_dict={}, + rerank_model_dict={}, + ) + app.platform_mgr = SimpleNamespace(_bots_by_key={}) + app.pipeline_mgr = SimpleNamespace(_pipelines_by_key={}) + app.rag_mgr = SimpleNamespace(knowledge_bases={}) + app.plugin_connector = SimpleNamespace(_known_desired_states={'installation': object()}) + app.tool_mgr = SimpleNamespace( + mcp_tool_loader=SimpleNamespace( + _sessions={}, + _hosted_mcp_tasks=[], + _host_dispatch_tasks=set(), + ) + ) + app.telemetry = SimpleNamespace(send_tasks=[]) + + stats = app.get_runtime_resource_stats() + + assert stats['asyncio_tasks'] >= 1 + assert stats['event_loop'] == { + 'running': False, + 'samples_total': 0, + 'last_lag_ms': 0, + 'recent_p95_lag_ms': 0, + 'recent_max_lag_ms': 0, + 'max_lag_ms': 0, + } + assert stats['blocking_executor']['rejected_total'] == 4 + assert stats['application_tasks'] == { + 'total': 5, + 'completed': 2, + } + assert stats['query_pool'] == { + 'queued': 1, + 'cached': 0, + 'active_workspaces': 1, + } + assert stats['models']['providers'] == 1 + assert stats['runtimes']['plugin_installations'] == 1 diff --git a/tests/unit_tests/core/test_boot.py b/tests/unit_tests/core/test_boot.py index b2b717006..03c82a732 100644 --- a/tests/unit_tests/core/test_boot.py +++ b/tests/unit_tests/core/test_boot.py @@ -2,13 +2,37 @@ from __future__ import annotations import signal from types import SimpleNamespace -from unittest.mock import Mock +from unittest.mock import AsyncMock, Mock import pytest from langbot.pkg.core import boot +@pytest.mark.asyncio +async def test_make_app_shuts_down_partially_built_application(monkeypatch): + app_inst = SimpleNamespace( + event_loop=None, + shutdown=AsyncMock(), + initialize=AsyncMock(), + ) + + class FailingStage: + async def run(self, ap): + assert ap is app_inst + raise RuntimeError('startup failed') + + monkeypatch.setattr(boot.app, 'Application', lambda: app_inst) + monkeypatch.setattr(boot, 'stage_order', ['FailingStage']) + monkeypatch.setitem(boot.stage.preregistered_stages, 'FailingStage', FailingStage) + + with pytest.raises(RuntimeError, match='startup failed'): + await boot.make_app(SimpleNamespace()) + + app_inst.shutdown.assert_awaited_once() + app_inst.initialize.assert_not_awaited() + + @pytest.mark.asyncio async def test_main_signal_handler_handles_sigint_before_app_created(monkeypatch): captured_handler = {} diff --git a/tests/unit_tests/core/test_load_config.py b/tests/unit_tests/core/test_load_config.py index a32f30dc3..83c5697ff 100644 --- a/tests/unit_tests/core/test_load_config.py +++ b/tests/unit_tests/core/test_load_config.py @@ -333,11 +333,23 @@ class TestApplyEnvOverridesToConfig: { 'PLUGIN__WORKER__MAX_MEMORY_MB': '768', 'MCP__STDIO__ENABLED': 'false', + 'SYSTEM__BLOCKING_EXECUTOR__MAX_WORKERS': '12', + 'SYSTEM__BLOCKING_EXECUTOR__MAX_PENDING': '256', + 'SYSTEM__BLOCKING_EXECUTOR__MAX_INFLIGHT_PER_SCOPE': '3', }, clear=True, ): result = load_config._apply_env_overrides_to_config(completed) + assert result['system']['blocking_executor'] == { + 'max_workers': 12, + 'max_pending': 256, + 'max_inflight_per_scope': 3, + } + assert isinstance( + result['system']['blocking_executor']['max_workers'], + int, + ) assert result['plugin']['worker']['max_memory_mb'] == 768 assert isinstance(result['plugin']['worker']['max_memory_mb'], int) assert result['mcp']['stdio']['enabled'] is False diff --git a/tests/unit_tests/core/test_taskmgr.py b/tests/unit_tests/core/test_taskmgr.py index ada445709..44503de9a 100644 --- a/tests/unit_tests/core/test_taskmgr.py +++ b/tests/unit_tests/core/test_taskmgr.py @@ -13,6 +13,7 @@ from __future__ import annotations import pytest import asyncio import contextvars +import inspect import sys from unittest.mock import Mock, MagicMock from contextlib import contextmanager @@ -265,6 +266,28 @@ class TestTaskWrapper: wrapper.cancel() + @pytest.mark.asyncio + async def test_workspace_task_sets_blocking_work_scope(self): + """Detached tasks recover tenant fairness from durable ownership.""" + _, TaskWrapper, _ = get_taskmgr_classes() + from langbot.pkg.utils.bounded_executor import ( + current_blocking_work_scope, + ) + + mock_app = create_mock_app() + + async def read_scope(): + return current_blocking_work_scope() + + wrapper = TaskWrapper( + mock_app, + read_scope(), + workspace_uuid='workspace-a', + ) + + assert await wrapper.task == 'workspace-a' + assert current_blocking_work_scope() is None + @pytest.mark.asyncio async def test_to_dict_serialization(self): """Test TaskWrapper.to_dict serialization.""" @@ -530,6 +553,56 @@ class TestAsyncTaskManager: wrapper.cancel() + @pytest.mark.asyncio + async def test_create_user_task_enforces_workspace_active_limit_and_closes_rejected_coroutine(self): + """A noisy Workspace cannot accumulate unbounded background work.""" + _, _, AsyncTaskManager = get_taskmgr_classes() + mock_app = create_mock_app() + mock_app.instance_config.data['system']['task_retention'].update( + { + 'max_active_user_tasks': 10, + 'max_active_user_tasks_per_workspace': 1, + } + ) + manager = AsyncTaskManager(mock_app) + + async def long_coro(): + await asyncio.sleep(10) + + first = manager.create_user_task(long_coro(), workspace_uuid='workspace-a') + rejected = long_coro() + with pytest.raises(RuntimeError, match='Workspace has too many active user operations'): + manager.create_user_task(rejected, workspace_uuid='workspace-a') + + assert inspect.getcoroutinestate(rejected) == inspect.CORO_CLOSED + other_workspace = manager.create_user_task(long_coro(), workspace_uuid='workspace-b') + first.cancel() + other_workspace.cancel() + + @pytest.mark.asyncio + async def test_create_user_task_enforces_instance_active_limit(self): + """The shared process retains a hard cap even across Workspaces.""" + _, _, AsyncTaskManager = get_taskmgr_classes() + mock_app = create_mock_app() + mock_app.instance_config.data['system']['task_retention'].update( + { + 'max_active_user_tasks': 1, + 'max_active_user_tasks_per_workspace': 10, + } + ) + manager = AsyncTaskManager(mock_app) + + async def long_coro(): + await asyncio.sleep(10) + + first = manager.create_user_task(long_coro(), workspace_uuid='workspace-a') + rejected = long_coro() + with pytest.raises(RuntimeError, match='instance has too many active user operations'): + manager.create_user_task(rejected, workspace_uuid='workspace-b') + + assert inspect.getcoroutinestate(rejected) == inspect.CORO_CLOSED + first.cancel() + @pytest.mark.asyncio async def test_get_task_by_id(self): """Test get_task_by_id returns correct task.""" diff --git a/tests/unit_tests/persistence/test_postgresql_database.py b/tests/unit_tests/persistence/test_postgresql_database.py index 29175a202..63a3b751b 100644 --- a/tests/unit_tests/persistence/test_postgresql_database.py +++ b/tests/unit_tests/persistence/test_postgresql_database.py @@ -13,11 +13,13 @@ from langbot.pkg.persistence.databases import postgresql @pytest.mark.asyncio async def test_postgresql_manager_parses_explicit_url_without_string_reassembly(monkeypatch) -> None: captured = None + captured_options = None sentinel_engine = object() - def create_engine(url): - nonlocal captured + def create_engine(url, **options): + nonlocal captured, captured_options captured = url + captured_options = options return sentinel_engine monkeypatch.setattr(postgresql.sqlalchemy_asyncio, 'create_async_engine', create_engine) @@ -40,6 +42,13 @@ async def test_postgresql_manager_parses_explicit_url_without_string_reassembly( assert captured.password == 'p@ss' assert captured.query['ssl'] == 'require' assert 'sslmode' not in captured.query + assert captured_options == { + 'pool_size': 10, + 'max_overflow': 10, + 'pool_timeout': 30, + 'pool_recycle': 1800, + 'pool_pre_ping': True, + } assert manager.engine is sentinel_engine @@ -47,7 +56,7 @@ async def test_postgresql_manager_parses_explicit_url_without_string_reassembly( async def test_postgresql_manager_builds_structured_url_with_special_password(monkeypatch) -> None: captured = None - def create_engine(url): + def create_engine(url, **_options): nonlocal captured captured = url return object() @@ -76,6 +85,60 @@ async def test_postgresql_manager_builds_structured_url_with_special_password(mo assert captured.database == 'langbot' +@pytest.mark.asyncio +async def test_postgresql_manager_applies_explicit_bounded_pool_options(monkeypatch) -> None: + captured_options = None + + def create_engine(_url, **options): + nonlocal captured_options + captured_options = options + return object() + + monkeypatch.setattr(postgresql.sqlalchemy_asyncio, 'create_async_engine', create_engine) + ap = SimpleNamespace( + instance_config=SimpleNamespace( + data={ + 'database': { + 'postgresql': { + 'pool_size': 24, + 'max_overflow': 0, + 'pool_timeout_seconds': 7, + 'pool_recycle_seconds': 600, + } + } + } + ) + ) + + await postgresql.PostgreSQLDatabaseManager(ap).initialize() + + assert captured_options == { + 'pool_size': 24, + 'max_overflow': 0, + 'pool_timeout': 7, + 'pool_recycle': 600, + 'pool_pre_ping': True, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ('name', 'value'), + [ + ('pool_size', 0), + ('pool_size', True), + ('max_overflow', -1), + ('pool_timeout_seconds', 0), + ('pool_recycle_seconds', '1800'), + ], +) +async def test_postgresql_manager_rejects_invalid_pool_options(name, value) -> None: + ap = SimpleNamespace(instance_config=SimpleNamespace(data={'database': {'postgresql': {name: value}}})) + + with pytest.raises(ValueError, match=rf'database\.postgresql\.{name}'): + await postgresql.PostgreSQLDatabaseManager(ap).initialize() + + @pytest.mark.asyncio async def test_postgresql_manager_rejects_non_postgresql_url_without_echoing_secret() -> None: ap = SimpleNamespace( diff --git a/tests/unit_tests/persistence/test_release_migration.py b/tests/unit_tests/persistence/test_release_migration.py index d622270a0..bcbaf4ffa 100644 --- a/tests/unit_tests/persistence/test_release_migration.py +++ b/tests/unit_tests/persistence/test_release_migration.py @@ -104,6 +104,7 @@ async def test_release_migration_disposes_operator_engine_on_failure(monkeypatch manager = SimpleNamespace( db=SimpleNamespace(engine=engine), initialize=AsyncMock(side_effect=RuntimeError('migration failed')), + shutdown=AsyncMock(side_effect=engine.dispose), ) def manager_factory(*args, **kwargs): @@ -121,6 +122,7 @@ async def test_release_migration_disposes_operator_engine_on_failure(monkeypatch await release_migration.run_cloud_release_migration(ap, environ=_operator_environ()) assert ap.persistence_mgr is manager + manager.shutdown.assert_awaited_once() engine.dispose.assert_awaited_once() diff --git a/tests/unit_tests/pipeline/test_chat_handler.py b/tests/unit_tests/pipeline/test_chat_handler.py index c8a923d78..2301ec23c 100644 --- a/tests/unit_tests/pipeline/test_chat_handler.py +++ b/tests/unit_tests/pipeline/test_chat_handler.py @@ -464,3 +464,13 @@ class TestChatHandlerHelper: handler = chat.ChatMessageHandler(fake_app) result = handler.cut_str('first line\nsecond line') assert '...' in result + + def test_response_size_limit_uses_instance_config(self, fake_app): + from langbot_plugin.api.entities.builtin.provider.message import Message + + fake_app.instance_config.data['system'] = {'response_limits': {'max_generated_chars': 4}} + chat = get_chat_handler() + handler = chat.ChatMessageHandler(fake_app) + + with pytest.raises(RuntimeError, match='configured limit'): + handler._check_response_size(Message(role='assistant', content='12345')) diff --git a/tests/unit_tests/pipeline/test_controller_tenancy.py b/tests/unit_tests/pipeline/test_controller_tenancy.py index f41b9668a..b99e54bb2 100644 --- a/tests/unit_tests/pipeline/test_controller_tenancy.py +++ b/tests/unit_tests/pipeline/test_controller_tenancy.py @@ -36,6 +36,7 @@ async def test_controller_drops_stale_query_before_pipeline_lookup( query_pool, session = _prepare_scheduler(mock_app) mock_app.workspace_service.get_execution_binding.side_effect = WorkspaceGenerationMismatchError('stale generation') controller = Controller(mock_app) + initial_slots = controller.semaphore._value await controller._process_query(sample_query) @@ -47,6 +48,7 @@ async def test_controller_drops_stale_query_before_pipeline_lookup( query_pool.remove_query.assert_awaited_once_with(sample_query) session._semaphore.release.assert_called_once_with() query_pool.condition.notify_all.assert_called_once_with() + assert controller.semaphore._value == initial_slots @pytest.mark.asyncio diff --git a/tests/unit_tests/pipeline/test_longtext_image.py b/tests/unit_tests/pipeline/test_longtext_image.py new file mode 100644 index 000000000..c02c8ecef --- /dev/null +++ b/tests/unit_tests/pipeline/test_longtext_image.py @@ -0,0 +1,17 @@ +from unittest.mock import Mock + +from langbot.pkg.pipeline.longtext.strategies.image import Text2ImageStrategy + + +class _WideFont: + def getlength(self, text: str) -> int: + return len(text) * 100 + + +def test_image_strategy_line_split_always_consumes_input(): + strategy = Text2ImageStrategy(Mock()) + + lines = strategy._split_text_lines('abc', 1, _WideFont()) + + assert lines == ['a', 'b', 'c'] + assert ''.join(lines) == 'abc' diff --git a/tests/unit_tests/pipeline/test_n8nsvapi.py b/tests/unit_tests/pipeline/test_n8nsvapi.py index 787472375..54266aec4 100644 --- a/tests/unit_tests/pipeline/test_n8nsvapi.py +++ b/tests/unit_tests/pipeline/test_n8nsvapi.py @@ -37,7 +37,11 @@ _saved_modules = {name: sys.modules.get(name) for name in _import_stubs} for _name, _stub in _import_stubs.items(): sys.modules[_name] = _stub try: - from langbot.pkg.provider.runners.n8nsvapi import N8nServiceAPIRunner + from langbot.pkg.provider.runners.n8nsvapi import ( + N8nAPIError, + N8nServiceAPIRunner, + _MAX_N8N_RESPONSE_CHARS, + ) finally: for _name, _original in _saved_modules.items(): if _original is None: @@ -230,6 +234,17 @@ async def test_plain_json_non_dict_response(): assert chunks[0].content == '["a", "b"]' +@pytest.mark.asyncio +async def test_response_size_is_bounded(): + runner = make_runner() + + with pytest.raises(N8nAPIError, match='exceeds the runtime limit'): + await collect_chunks( + runner, + [b'x' * (_MAX_N8N_RESPONSE_CHARS + 1)], + ) + + @pytest.mark.asyncio async def test_invalid_json_returns_raw_text(): """Non-JSON response returns raw text as-is.""" diff --git a/tests/unit_tests/pipeline/test_pipelinemgr.py b/tests/unit_tests/pipeline/test_pipelinemgr.py index 2033df402..5a290a8e1 100644 --- a/tests/unit_tests/pipeline/test_pipelinemgr.py +++ b/tests/unit_tests/pipeline/test_pipelinemgr.py @@ -3,11 +3,13 @@ PipelineManager unit tests """ import pytest +from types import SimpleNamespace from unittest.mock import AsyncMock, Mock from importlib import import_module from langbot.pkg.api.http.context import ExecutionContext -from langbot.pkg.workspace.errors import WorkspaceGenerationMismatchError +from langbot.pkg.workspace.entities import WorkspaceExecutionBinding +from langbot.pkg.workspace.errors import WorkspaceGenerationMismatchError, WorkspaceInvariantError def _context(pipeline_uuid: str = 'test-uuid') -> ExecutionContext: @@ -49,6 +51,95 @@ async def test_pipeline_manager_initialize(mock_app): assert len(manager.pipelines) == 0 +@pytest.mark.asyncio +async def test_cloud_startup_reuses_validated_pipeline_binding(mock_app): + class TenantUow: + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return False + + binding = WorkspaceExecutionBinding( + instance_uuid='test-instance', + workspace_uuid='test-workspace', + placement_generation=1, + write_fenced=False, + state='active', + ) + pipeline_entity = Mock( + uuid='test-uuid', + workspace_uuid='test-workspace', + stages=[], + config={}, + extensions_preferences={}, + ) + mock_app.persistence_mgr.mode = SimpleNamespace(value='cloud_runtime') + mock_app.persistence_mgr.tenant_uow = lambda _workspace_uuid: TenantUow() + mock_app.persistence_mgr.execute_async = AsyncMock(return_value=Mock(all=Mock(return_value=[pipeline_entity]))) + mock_app.workspace_service.list_active_execution_bindings = AsyncMock(return_value=[binding]) + mock_app.workspace_service.get_execution_binding = AsyncMock( + side_effect=AssertionError('startup pipeline loader repeated a validated binding lookup') + ) + manager = get_pipelinemgr_module().PipelineManager(mock_app) + manager.stage_dict = {} + + await manager.load_pipelines_from_db() + + assert len(manager.pipelines) == 1 + mock_app.workspace_service.get_execution_binding.assert_not_awaited() + + +def test_generation_advance_prunes_superseded_workspace_pipelines(mock_app): + class NoGlobalIterationDict(dict): + def __iter__(self): + raise AssertionError('generation advance scanned every pipeline') + + def items(self): + raise AssertionError('generation advance scanned every pipeline') + + def values(self): + raise AssertionError('generation advance scanned every pipeline') + + pipelinemgr = get_pipelinemgr_module() + manager = pipelinemgr.PipelineManager(mock_app) + old_context = _context() + next_context = ExecutionContext( + instance_uuid=old_context.instance_uuid, + workspace_uuid=old_context.workspace_uuid, + placement_generation=2, + pipeline_uuid=old_context.pipeline_uuid, + ) + old_pipeline = SimpleNamespace( + execution_context=old_context, + workspace_uuid=old_context.workspace_uuid, + placement_generation=old_context.placement_generation, + ) + other_pipelines = [ + SimpleNamespace( + execution_context=ExecutionContext( + instance_uuid='test-instance', + workspace_uuid=f'workspace-{index}', + placement_generation=1, + pipeline_uuid=f'pipeline-{index}', + ), + workspace_uuid=f'workspace-{index}', + placement_generation=1, + ) + for index in range(1_000) + ] + manager.pipelines = [old_pipeline, *other_pipelines] + + manager._observe_execution_context(old_context) + manager._pipelines_by_key = NoGlobalIterationDict(manager._pipelines_by_key) + manager._observe_execution_context(next_context) + manager._pipelines_by_key = dict(manager._pipelines_by_key) + + assert manager.pipelines == other_pipelines + with pytest.raises(WorkspaceInvariantError, match='rolled back'): + manager._observe_execution_context(old_context) + + @pytest.mark.asyncio async def test_load_pipeline(mock_app): """Test loading a single pipeline""" diff --git a/tests/unit_tests/pipeline/test_pool.py b/tests/unit_tests/pipeline/test_pool.py index 15d4bffb9..853738512 100644 --- a/tests/unit_tests/pipeline/test_pool.py +++ b/tests/unit_tests/pipeline/test_pool.py @@ -18,6 +18,7 @@ from langbot.pkg.pipeline.pool import ( ExecutionContextRequiredError, QueryNotFoundError, QueryPool, + QueryPoolCapacityError, get_query_execution_context, ) @@ -433,3 +434,55 @@ class TestQueryPoolWorkspaceIsolation: assert pool.get_query_count(workspace_b) == 1 assert pool.get_query_count(next_generation) == 0 assert pool.query_id_counter == 3 + + async def test_workspace_capacity_discards_oldest_queued_query(self): + pool = QueryPool(max_queries=3, max_queries_per_workspace=2) + first = await add_scoped_mock_query(pool, TEST_CONTEXT) + second = await add_scoped_mock_query(pool, TEST_CONTEXT) + third = await add_scoped_mock_query(pool, TEST_CONTEXT) + + assert await pool.get_query(TEST_CONTEXT.workspace_uuid, first.query_uuid) is None + assert await pool.get_query(TEST_CONTEXT.workspace_uuid, second.query_uuid) is second + assert await pool.get_query(TEST_CONTEXT.workspace_uuid, third.query_uuid) is third + assert pool.active_query_count_by_workspace == {TEST_CONTEXT.workspace_uuid: 2} + assert pool.get_dropped_query_count(TEST_CONTEXT) == 1 + + async def test_capacity_rejects_when_every_query_is_already_running(self): + pool = QueryPool(max_queries=1, max_queries_per_workspace=1) + running = await add_scoped_mock_query(pool, TEST_CONTEXT) + async with pool: + pool.mark_query_running_locked(running) + + with pytest.raises(QueryPoolCapacityError): + await add_scoped_mock_query(pool, TEST_CONTEXT) + + assert pool.active_query_count_by_workspace == {TEST_CONTEXT.workspace_uuid: 1} + + async def test_mark_query_running_keeps_active_indexes_but_removes_queue_entry(self): + pool = QueryPool(max_queries=1, max_queries_per_workspace=1) + running = await add_scoped_mock_query(pool, TEST_CONTEXT) + + async with pool: + pool.mark_query_running_locked(running) + + assert running not in pool.queries + assert await pool.get_query(TEST_CONTEXT.workspace_uuid, running.query_uuid) is running + assert pool.active_query_count_by_workspace == {TEST_CONTEXT.workspace_uuid: 1} + + async def test_historical_workspace_counters_are_bounded(self): + pool = QueryPool(max_queries=2, max_queries_per_workspace=1) + contexts = [ + ExecutionContext( + instance_uuid='instance-test', + workspace_uuid=f'workspace-{index}', + placement_generation=1, + ) + for index in range(3) + ] + + for context in contexts: + query = await add_scoped_mock_query(pool, context) + await pool.remove_query(query) + + assert len(pool.query_count_by_scope) == 2 + assert (contexts[0].instance_uuid, contexts[0].workspace_uuid, 1) not in pool.query_count_by_scope diff --git a/tests/unit_tests/pipeline/test_ratelimit.py b/tests/unit_tests/pipeline/test_ratelimit.py index be767ad56..e23636ab7 100644 --- a/tests/unit_tests/pipeline/test_ratelimit.py +++ b/tests/unit_tests/pipeline/test_ratelimit.py @@ -152,8 +152,18 @@ class TestFixedWindowAlgo: # First request creates container await algo.require_access(sample_query_with_rate_limit, provider_session.LauncherTypes.PERSON, '12345') - # Key format: 'LauncherTypes.PERSON_12345' (enum string representation) - expected_key = 'LauncherTypes.PERSON_12345' + context = sample_query_with_rate_limit._execution_context + expected_key = ':'.join( + ( + context.instance_uuid, + context.workspace_uuid, + str(context.placement_generation), + str(sample_query_with_rate_limit.bot_uuid), + str(sample_query_with_rate_limit.pipeline_uuid), + str(provider_session.LauncherTypes.PERSON), + '12345', + ) + ) assert expected_key in algo.containers container = algo.containers[expected_key] @@ -191,8 +201,18 @@ class TestFixedWindowAlgo: for i in range(5): await algo.require_access(sample_query, provider_session.LauncherTypes.PERSON, 'test') - # Key format: 'LauncherTypes.PERSON_test' - expected_key = 'LauncherTypes.PERSON_test' + context = sample_query._execution_context + expected_key = ':'.join( + ( + context.instance_uuid, + context.workspace_uuid, + str(context.placement_generation), + str(sample_query.bot_uuid), + str(sample_query.pipeline_uuid), + str(provider_session.LauncherTypes.PERSON), + 'test', + ) + ) container = algo.containers[expected_key] assert window_start in container.records assert container.records[window_start] == 5 diff --git a/tests/unit_tests/platform/test_aiocqhttp_message_converter.py b/tests/unit_tests/platform/test_aiocqhttp_message_converter.py index b41e82631..1ee1730d0 100644 --- a/tests/unit_tests/platform/test_aiocqhttp_message_converter.py +++ b/tests/unit_tests/platform/test_aiocqhttp_message_converter.py @@ -68,6 +68,25 @@ async def test_connection_listener_only_suppresses_exact_duplicates(): ] +@pytest.mark.asyncio +async def test_connection_event_cache_is_bounded(): + adapter, _ = _make_adapter() + + for index in range(150): + await adapter._on_websocket_connection(aiocqhttp.Event({'self_id': index, 'time': index})) + + assert len(adapter.on_websocket_connection_event_cache) == 100 + + +def test_group_lookup_caches_are_bounded(): + converter = AiocqhttpEventConverter() + converter._group_name_cache = {index: (str(index), 10_000.0) for index in range(5000)} + + converter._prune_caches(1.0) + + assert len(converter._group_name_cache) == 4096 + + def test_unregister_listener_removes_registered_wrapper(): adapter, _ = _make_adapter() diff --git a/tests/unit_tests/platform/test_botmgr_tenancy.py b/tests/unit_tests/platform/test_botmgr_tenancy.py index 1824ecbff..e3c055eb0 100644 --- a/tests/unit_tests/platform/test_botmgr_tenancy.py +++ b/tests/unit_tests/platform/test_botmgr_tenancy.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import contextlib from types import SimpleNamespace from unittest.mock import AsyncMock @@ -8,7 +9,10 @@ import pytest from langbot.pkg.api.http.authz import WorkspaceRequiredError from langbot.pkg.api.http.context import ExecutionContext +from langbot.pkg.entity.persistence.bot import Bot from langbot.pkg.platform.botmgr import PlatformManager, RuntimeBot +from langbot.pkg.workspace.entities import WorkspaceExecutionBinding +from langbot.pkg.workspace.errors import WorkspaceInvariantError import langbot_plugin.api.entities.builtin.platform.events as platform_events @@ -83,7 +87,220 @@ async def test_public_route_key_resolves_bound_runtime_and_rejects_non_opaque_in @pytest.mark.asyncio async def test_stale_runtime_generation_is_not_returned(manager): - assert await manager.get_bot_by_uuid(_context(WORKSPACE_A, BOT_A, generation=5), BOT_A) is None + with pytest.raises(ValueError, match='stale'): + await manager.get_bot_by_uuid(_context(WORKSPACE_A, BOT_A, generation=5), BOT_A) + + +@pytest.mark.asyncio +async def test_generation_advance_shuts_down_and_prunes_old_workspace_bots(): + class NoGlobalIterationDict(dict): + def __iter__(self): + raise AssertionError('generation advance scanned every bot runtime') + + def items(self): + raise AssertionError('generation advance scanned every bot runtime') + + def values(self): + raise AssertionError('generation advance scanned every bot runtime') + + manager = PlatformManager(SimpleNamespace()) + old_bot = SimpleNamespace( + workspace_uuid=WORKSPACE_A, + placement_generation=4, + enable=True, + shutdown=AsyncMock(), + ) + other_bot = SimpleNamespace( + workspace_uuid=WORKSPACE_B, + placement_generation=4, + enable=True, + shutdown=AsyncMock(), + ) + unrelated_bots = [ + SimpleNamespace( + workspace_uuid=f'workspace-{index}', + placement_generation=4, + enable=False, + shutdown=AsyncMock(), + ) + for index in range(1_000) + ] + manager.bots = [old_bot, other_bot, *unrelated_bots] + old_context = _context(WORKSPACE_A, BOT_A, generation=4) + next_context = _context(WORKSPACE_A, BOT_A, generation=5) + + await manager._observe_execution_context(old_context) + manager._bots_by_key = NoGlobalIterationDict(manager._bots_by_key) + await manager._observe_execution_context(next_context) + manager._bots_by_key = dict(manager._bots_by_key) + + old_bot.shutdown.assert_awaited_once_with() + assert manager.bots == [other_bot, *unrelated_bots] + with pytest.raises(WorkspaceInvariantError, match='rolled back'): + await manager._observe_execution_context(old_context) + + +@pytest.mark.asyncio +async def test_concurrent_websocket_proxy_creation_reuses_one_runtime(): + created_adapters = [] + + class WebsocketAdapter: + def __init__(self, *_args, **_kwargs): + created_adapters.append(self) + + def register_listener(self, *_args): + pass + + application = SimpleNamespace(workspace_service=_WorkspaceService()) + manager = PlatformManager(application) + manager.adapter_dict = {'websocket': WebsocketAdapter} + context = ExecutionContext( + instance_uuid='instance', + workspace_uuid=WORKSPACE_A, + placement_generation=4, + ) + + runtimes = await asyncio.gather(*(manager.get_websocket_proxy_bot(context) for _ in range(20))) + + assert len(created_adapters) == 1 + assert len({id(runtime) for runtime in runtimes}) == 1 + assert manager.websocket_proxy_bots == {WORKSPACE_A: runtimes[0]} + + +@pytest.mark.asyncio +async def test_websocket_proxy_cache_evicts_oldest_idle_workspace(): + created_adapters = [] + + class WebsocketAdapter: + def __init__(self, *_args, **_kwargs): + self.kill = AsyncMock() + self.inbound_listener_tasks = set() + created_adapters.append(self) + + def register_listener(self, *_args): + pass + + application = SimpleNamespace( + workspace_service=_WorkspaceService(), + instance_config=SimpleNamespace( + data={ + 'system': { + 'websocket_retention': {'max_workspace_proxies': 1}, + } + } + ), + ) + manager = PlatformManager(application) + manager.adapter_dict = {'websocket': WebsocketAdapter} + + await manager.get_websocket_proxy_bot( + ExecutionContext( + instance_uuid='instance', + workspace_uuid=WORKSPACE_A, + placement_generation=4, + ) + ) + second = await manager.get_websocket_proxy_bot( + ExecutionContext( + instance_uuid='instance', + workspace_uuid=WORKSPACE_B, + placement_generation=4, + ) + ) + + created_adapters[0].kill.assert_awaited_once_with() + assert manager.websocket_proxy_bots == {WORKSPACE_B: second} + assert WORKSPACE_A not in manager._proxy_last_accessed + + +@pytest.mark.asyncio +async def test_reload_stops_and_drops_existing_platform_runtimes(): + old_bot = SimpleNamespace(enable=True, shutdown=AsyncMock()) + old_proxy = SimpleNamespace(enable=True, shutdown=AsyncMock()) + persistence_mgr = SimpleNamespace( + execute_async=AsyncMock(return_value=SimpleNamespace(all=lambda: [])), + ) + application = SimpleNamespace( + logger=SimpleNamespace(info=lambda *_args: None, warning=lambda *_args: None), + persistence_mgr=persistence_mgr, + workspace_service=SimpleNamespace(), + ) + manager = PlatformManager(application) + manager.bots = [old_bot] + manager.websocket_proxy_bots = {WORKSPACE_A: old_proxy} + manager._scope_generations = {('instance', WORKSPACE_A): 4} + + await manager.load_bots_from_db() + + old_bot.shutdown.assert_awaited_once_with() + old_proxy.shutdown.assert_awaited_once_with() + assert manager.bots == [] + assert manager.websocket_proxy_bots == {} + assert manager._scope_generations == {} + + +@pytest.mark.asyncio +async def test_cloud_startup_reuses_validated_platform_binding(): + class TenantUow: + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return False + + class ProbeAdapter: + def __init__(self, _config, _logger): + self.listeners = [] + + def register_listener(self, event_type, listener): + self.listeners.append((event_type, listener)) + + async def kill(self): + return None + + binding = WorkspaceExecutionBinding( + instance_uuid='instance', + workspace_uuid=WORKSPACE_A, + placement_generation=4, + write_fenced=False, + state='active', + ) + bot = Bot( + uuid=BOT_A, + workspace_uuid=WORKSPACE_A, + name='Probe', + description='', + adapter='probe', + adapter_config={}, + enable=False, + pipeline_routing_rules=[], + ) + workspace_service = SimpleNamespace( + list_active_execution_bindings=AsyncMock(return_value=[binding]), + get_execution_binding=AsyncMock( + side_effect=AssertionError('startup platform loader repeated a validated binding lookup') + ), + ) + application = SimpleNamespace( + logger=SimpleNamespace( + info=lambda *_args, **_kwargs: None, + warning=lambda *_args, **_kwargs: None, + error=lambda *_args, **_kwargs: None, + ), + persistence_mgr=SimpleNamespace( + mode=SimpleNamespace(value='cloud_runtime'), + tenant_uow=lambda _workspace_uuid: TenantUow(), + execute_async=AsyncMock(return_value=SimpleNamespace(all=lambda: [bot])), + ), + workspace_service=workspace_service, + ) + manager = PlatformManager(application) + manager.adapter_dict = {'probe': ProbeAdapter} + + await manager.load_bots_from_db() + + assert len(manager.bots) == 1 + workspace_service.get_execution_binding.assert_not_awaited() @pytest.mark.asyncio diff --git a/tests/unit_tests/platform/test_dingtalk_adapter.py b/tests/unit_tests/platform/test_dingtalk_adapter.py index df2087866..b2974837a 100644 --- a/tests/unit_tests/platform/test_dingtalk_adapter.py +++ b/tests/unit_tests/platform/test_dingtalk_adapter.py @@ -5,6 +5,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from langbot.pkg.platform import botmgr as _botmgr # noqa: F401 from langbot.pkg.platform.sources.dingtalk import ( DingTalkAdapter, _dingtalk_card_markdown, @@ -17,6 +18,17 @@ from langbot.pkg.platform.sources.dingtalk import ( ) +def test_dingtalk_auxiliary_tasks_are_bounded(): + adapter = DingTalkAdapter.model_construct() + adapter._background_tasks = {MagicMock(done=MagicMock(return_value=False)) for _ in range(100)} + + async def callback(): + raise AssertionError('rejected callback must not run') + + assert adapter._start_background_task(callback()) is False + assert len(adapter._background_tasks) == 100 + + def test_dingtalk_select_component_params_expose_options(): params = _dingtalk_form_component_params( { diff --git a/tests/unit_tests/platform/test_discord_limits.py b/tests/unit_tests/platform/test_discord_limits.py new file mode 100644 index 000000000..57656b348 --- /dev/null +++ b/tests/unit_tests/platform/test_discord_limits.py @@ -0,0 +1,12 @@ +from __future__ import annotations + +import pytest + +from langbot.pkg.platform.sources import discord + + +def test_discord_base64_decode_is_bounded(monkeypatch): + monkeypatch.setattr(discord, '_MAX_DISCORD_MEDIA_BYTES', 4) + + with pytest.raises(ValueError, match='exceeds'): + discord._decode_discord_base64_limited('A' * 12) diff --git a/tests/unit_tests/platform/test_http_bot_tenancy.py b/tests/unit_tests/platform/test_http_bot_tenancy.py index b00b2364d..42d947e79 100644 --- a/tests/unit_tests/platform/test_http_bot_tenancy.py +++ b/tests/unit_tests/platform/test_http_bot_tenancy.py @@ -1,11 +1,13 @@ from __future__ import annotations +import asyncio from types import SimpleNamespace import pytest from langbot.pkg.api.http.context import ExecutionContext from langbot.pkg.platform.sources.http_bot import HttpBotAdapter +from langbot.pkg.platform.sources import http_bot as http_bot_module def _session(key): @@ -22,6 +24,7 @@ def _adapter(app, execution_context) -> HttpBotAdapter: outbound_states={}, idempotency_cache={}, sync_waiters={}, + inbound_tasks=set(), ) object.__setattr__(adapter, 'ap', app) return adapter @@ -64,3 +67,28 @@ async def test_http_bot_reset_fails_closed_without_trusted_scope(): with pytest.raises(RuntimeError, match='trusted execution scope'): await adapter._reset_session('person', 'shared-session') + + +@pytest.mark.asyncio +async def test_http_bot_bounds_inbound_listener_tasks(monkeypatch): + monkeypatch.setattr(http_bot_module, '_INBOUND_TASK_MAX', 1) + adapter = _adapter(SimpleNamespace(), None) + started = asyncio.Event() + release = asyncio.Event() + + async def blocking_listener(): + started.set() + await release.wait() + + first = adapter._start_inbound_task(blocking_listener()) + await started.wait() + rejected = adapter._start_inbound_task(blocking_listener()) + + assert first is not None + assert rejected is None + assert len(adapter.inbound_tasks) == 1 + + release.set() + await first + await asyncio.sleep(0) + assert adapter.inbound_tasks == set() diff --git a/tests/unit_tests/platform/test_kook_limits.py b/tests/unit_tests/platform/test_kook_limits.py new file mode 100644 index 000000000..669ab4517 --- /dev/null +++ b/tests/unit_tests/platform/test_kook_limits.py @@ -0,0 +1,24 @@ +from __future__ import annotations + +import json +import zlib + +import pytest + +from langbot.pkg.platform.sources import kook + + +def test_kook_gateway_decoder_accepts_raw_and_compressed_json(): + payload = {'s': 1, 'd': {'session_id': 'session-a'}} + encoded = json.dumps(payload).encode() + + assert kook._decode_gateway_message(encoded) == payload + assert kook._decode_gateway_message(zlib.compress(encoded)) == payload + + +def test_kook_gateway_decoder_rejects_decompression_bomb(monkeypatch): + monkeypatch.setattr(kook, '_KOOK_MAX_GATEWAY_MESSAGE_BYTES', 1024) + compressed = zlib.compress(b'x' * 1025) + + with pytest.raises(ValueError, match='decompressed size limit'): + kook._decode_gateway_message(compressed) diff --git a/tests/unit_tests/platform/test_lark_adapter.py b/tests/unit_tests/platform/test_lark_adapter.py index 59b91ee60..a8e1fe3d9 100644 --- a/tests/unit_tests/platform/test_lark_adapter.py +++ b/tests/unit_tests/platform/test_lark_adapter.py @@ -1,7 +1,13 @@ """Tests for Lark adapter helper behavior.""" +import threading +from unittest.mock import MagicMock + +import pytest + from langbot.pkg.platform.sources.lark import ( LarkAdapter, + _decode_lark_base64_limited, _lark_clean_form_content, _lark_completed_input_lines, _lark_current_input_defs, @@ -11,6 +17,27 @@ from langbot.pkg.platform.sources.lark import ( ) +def test_lark_base64_decode_is_bounded(monkeypatch): + import langbot.pkg.platform.sources.lark as lark_module + + monkeypatch.setattr(lark_module, '_MAX_LARK_MEDIA_BYTES', 4) + + with pytest.raises(ValueError, match='exceeds'): + _decode_lark_base64_limited('A' * 12) + + +def test_lark_threadsafe_callbacks_are_bounded(): + adapter = LarkAdapter.model_construct() + adapter.threadsafe_event_lock = threading.Lock() + adapter.threadsafe_event_futures = {MagicMock(done=MagicMock(return_value=False)) for _ in range(100)} + + async def callback(): + raise AssertionError('rejected callback must not run') + + assert adapter._schedule_threadsafe_event(callback()) is None + assert len(adapter.threadsafe_event_futures) == 100 + + def test_lark_current_input_defs_only_returns_active_stage(): input_defs = [ {'output_variable_name': 'us_input', 'type': 'paragraph'}, diff --git a/tests/unit_tests/platform/test_line_limits.py b/tests/unit_tests/platform/test_line_limits.py new file mode 100644 index 000000000..53ef12e44 --- /dev/null +++ b/tests/unit_tests/platform/test_line_limits.py @@ -0,0 +1,30 @@ +from __future__ import annotations + +import pytest +from unittest.mock import MagicMock + +from langbot.pkg.platform import botmgr as _botmgr # noqa: F401 +from langbot.pkg.platform.sources import line + + +def test_line_media_content_accepts_limit_boundary(monkeypatch) -> None: + monkeypatch.setattr(line, 'MAX_LINE_MEDIA_BYTES', 4) + content = b'1234' + + assert line._validate_line_media_content(content) is content + + +def test_line_media_content_rejects_oversized_payload(monkeypatch) -> None: + monkeypatch.setattr(line, 'MAX_LINE_MEDIA_BYTES', 4) + + with pytest.raises(ValueError, match='LINE media exceeds'): + line._validate_line_media_content(b'12345') + + +@pytest.mark.asyncio +async def test_line_kill_closes_api_client() -> None: + api_client = MagicMock() + adapter = line.LINEAdapter.model_construct(api_client=api_client) + + assert await adapter.kill() is True + api_client.close.assert_called_once_with() diff --git a/tests/unit_tests/platform/test_matrix_limits.py b/tests/unit_tests/platform/test_matrix_limits.py new file mode 100644 index 000000000..1f8356b3f --- /dev/null +++ b/tests/unit_tests/platform/test_matrix_limits.py @@ -0,0 +1,21 @@ +from __future__ import annotations + +import pytest + +from langbot.pkg.platform.sources import matrix + + +def test_matrix_base64_decode_is_bounded(monkeypatch): + monkeypatch.setattr(matrix, '_MAX_MATRIX_MEDIA_BYTES', 4) + + with pytest.raises(ValueError, match='exceeds'): + matrix._decode_matrix_base64_limited('A' * 12) + + +def test_matrix_local_file_read_is_bounded(tmp_path, monkeypatch): + monkeypatch.setattr(matrix, '_MAX_MATRIX_MEDIA_BYTES', 4) + path = tmp_path / 'large.bin' + path.write_bytes(b'12345') + + with pytest.raises(ValueError, match='exceeds'): + matrix._read_matrix_file_limited(str(path)) diff --git a/tests/unit_tests/platform/test_openclaw_weixin_client.py b/tests/unit_tests/platform/test_openclaw_weixin_client.py new file mode 100644 index 000000000..6c8dcadc5 --- /dev/null +++ b/tests/unit_tests/platform/test_openclaw_weixin_client.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +import pytest + +from langbot.libs.openclaw_weixin_api.client import ( + MAX_CDN_MEDIA_BYTES, + OpenClawWeixinClient, + _decrypt_cdn_payload, + _encrypt_cdn_payload, +) +from langbot.libs.openclaw_weixin_api.types import ApiError + + +def test_cdn_crypto_helpers_round_trip(): + original = b'tenant-media' * 128 + + aes_key_hex, _encoded_key, encrypted, _raw_md5 = _encrypt_cdn_payload(original) + + assert _decrypt_cdn_payload(encrypted, bytes.fromhex(aes_key_hex)) == original + + +@pytest.mark.asyncio +async def test_upload_media_rejects_oversized_input_before_network_access(): + client = OpenClawWeixinClient('https://example.invalid', 'token') + + with pytest.raises(ApiError, match='exceeds the size limit'): + await client.upload_media( + b'x' * (MAX_CDN_MEDIA_BYTES + 1), + 'recipient', + 3, + ) diff --git a/tests/unit_tests/platform/test_openclaw_weixin_tenancy.py b/tests/unit_tests/platform/test_openclaw_weixin_tenancy.py index 3bb1298a6..8a2f53e9b 100644 --- a/tests/unit_tests/platform/test_openclaw_weixin_tenancy.py +++ b/tests/unit_tests/platform/test_openclaw_weixin_tenancy.py @@ -2,8 +2,10 @@ from types import SimpleNamespace from unittest.mock import AsyncMock, Mock import pytest +import langbot_plugin.api.entities.builtin.platform.message as platform_message from langbot.pkg.api.http.context import ExecutionContext +from langbot.pkg.platform.sources import openclaw_weixin from langbot.pkg.platform.sources.openclaw_weixin import OpenClawWeixinAdapter @@ -82,3 +84,12 @@ async def test_persist_config_fails_closed_without_matching_execution_context(ex app.persistence_mgr.execute_async.assert_not_awaited() logger.warning.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_component_base64_decode_is_bounded(monkeypatch): + monkeypatch.setattr(openclaw_weixin, '_MAX_OPENCLAW_COMPONENT_BYTES', 4) + component = platform_message.File(base64='MTIzNDU=') + + with pytest.raises(ValueError, match='exceeds'): + await OpenClawWeixinAdapter._get_component_bytes(component) diff --git a/tests/unit_tests/platform/test_qqofficial_api.py b/tests/unit_tests/platform/test_qqofficial_api.py index 0f791f612..704dd142c 100644 --- a/tests/unit_tests/platform/test_qqofficial_api.py +++ b/tests/unit_tests/platform/test_qqofficial_api.py @@ -10,6 +10,7 @@ import langbot_plugin.api.entities.builtin.platform.message as platform_message from langbot.libs.qq_official_api.api import ( QQ_SELECT_ACTION_PREFIX, + QQOfficialClient, build_keyboard_from_select_field, get_select_field_options, resolve_select_button_action, @@ -49,6 +50,28 @@ def test_qq_select_button_resolves_field_and_value(): assert resolve_select_button_action(form_data, f'{QQ_SELECT_ACTION_PREFIX}99') is None +@pytest.mark.asyncio +async def test_qq_seed_rejects_empty_secret_without_spinning(): + client = QQOfficialClient('', 'token', 'app-id', AsyncMock()) + + with pytest.raises(ValueError, match='must not be empty'): + await asyncio.wait_for(client.repeat_seed(''), timeout=0.1) + + +def test_qq_auxiliary_tasks_are_bounded(): + import langbot.pkg.core.app # noqa: F401 + from langbot.pkg.platform.sources.qqofficial import QQOfficialAdapter + + adapter = QQOfficialAdapter.model_construct() + adapter._background_tasks = {MagicMock(done=MagicMock(return_value=False)) for _ in range(100)} + + async def callback(): + raise AssertionError('rejected callback must not run') + + assert adapter._start_background_task(callback()) is False + assert len(adapter._background_tasks) == 100 + + def test_qq_select_keyboard_fits_twenty_five_options(): form_data = _select_form_data() form_data['input_defs'][0]['option_source']['value'] = [f'Option {idx}' for idx in range(25)] diff --git a/tests/unit_tests/platform/test_telegram_adapter.py b/tests/unit_tests/platform/test_telegram_adapter.py index 6e2262283..885310914 100644 --- a/tests/unit_tests/platform/test_telegram_adapter.py +++ b/tests/unit_tests/platform/test_telegram_adapter.py @@ -11,11 +11,21 @@ import langbot_plugin.api.entities.builtin.platform.events as platform_events import langbot_plugin.api.entities.builtin.platform.message as platform_message from langbot.pkg.platform.sources.telegram import ( TelegramAdapter, + _decode_telegram_base64_limited, _telegram_form_action_from_callback, _telegram_select_field_options, ) +def test_telegram_base64_decode_is_bounded(monkeypatch): + import langbot.pkg.platform.sources.telegram as telegram_module + + monkeypatch.setattr(telegram_module, '_MAX_TELEGRAM_MEDIA_BYTES', 4) + + with pytest.raises(ValueError, match='exceeds'): + _decode_telegram_base64_limited('A' * 12) + + def _select_form_data() -> dict: return { '_current_input_field': 'choice', @@ -88,6 +98,18 @@ def test_telegram_form_callback_cache_preserves_pipeline_uuid(): ) +def test_telegram_form_callback_cache_is_bounded(): + adapter = TelegramAdapter.model_construct() + adapter._form_action_titles = {} + + adapter._cache_form_action_titles( + {f'callback-{index}': str(index) for index in range(5000)}, + now=100.0, + ) + + assert len(adapter._form_action_titles) == adapter._MAX_FORM_ACTION_TITLES + + @pytest.mark.asyncio async def test_telegram_select_field_sends_two_column_inline_keyboard(): bot = MagicMock() diff --git a/tests/unit_tests/platform/test_websocket_session_isolation.py b/tests/unit_tests/platform/test_websocket_session_isolation.py index 110e6e0c3..240e7acfc 100644 --- a/tests/unit_tests/platform/test_websocket_session_isolation.py +++ b/tests/unit_tests/platform/test_websocket_session_isolation.py @@ -108,6 +108,71 @@ async def test_pipeline_indexes_and_broadcasts_are_workspace_scoped(): assert manager.get_stats(scope=SCOPE_A)['total_connections'] == 1 +@pytest.mark.asyncio +async def test_connection_admission_is_bounded_globally_and_per_workspace(): + manager = WebSocketConnectionManager() + await manager.add_connection( + websocket=Mock(), + scope=SCOPE_A, + pipeline_uuid='pipeline-1', + session_type='person', + max_connections=2, + max_connections_per_workspace=1, + ) + + with pytest.raises(RuntimeError, match='Workspace WebSocket'): + await manager.add_connection( + websocket=Mock(), + scope=SCOPE_A, + pipeline_uuid='pipeline-2', + session_type='person', + max_connections=2, + max_connections_per_workspace=1, + ) + + await manager.add_connection( + websocket=Mock(), + scope=SCOPE_B, + pipeline_uuid='pipeline-1', + session_type='person', + max_connections=2, + max_connections_per_workspace=1, + ) + with pytest.raises(RuntimeError, match='WebSocket connection capacity'): + await manager.add_connection( + websocket=Mock(), + scope=WebSocketScope('instance-a', 'workspace-c', 1), + pipeline_uuid='pipeline-1', + session_type='person', + max_connections=2, + max_connections_per_workspace=1, + ) + + +@pytest.mark.asyncio +async def test_close_scope_closes_and_removes_only_matching_connections(): + manager = WebSocketConnectionManager() + websocket_a = Mock(close=AsyncMock()) + connection_a = await manager.add_connection( + websocket=websocket_a, + scope=SCOPE_A, + pipeline_uuid='pipeline-1', + session_type='person', + ) + connection_b = await manager.add_connection( + websocket=Mock(close=AsyncMock()), + scope=SCOPE_B, + pipeline_uuid='pipeline-1', + session_type='person', + ) + + await manager.close_scope(SCOPE_A) + + websocket_a.close.assert_awaited_once() + assert await manager.get_connection(connection_a.connection_id, scope=SCOPE_A) is None + assert await manager.get_connection(connection_b.connection_id, scope=SCOPE_B) is connection_b + + @pytest.mark.asyncio async def test_embed_event_uses_stable_session_launcher(monkeypatch): manager = WebSocketConnectionManager() diff --git a/tests/unit_tests/platform/test_wechatpad_limits.py b/tests/unit_tests/platform/test_wechatpad_limits.py new file mode 100644 index 000000000..621cf7b0e --- /dev/null +++ b/tests/unit_tests/platform/test_wechatpad_limits.py @@ -0,0 +1,39 @@ +from __future__ import annotations + +import httpx +import pytest + +from langbot.libs.wechatpad_api.api import downloadpai +from langbot.libs.wechatpad_api.util import http_util + + +class _Response: + headers = {} + + def __init__(self, chunks: list[bytes]): + self._chunks = chunks + + def iter_content(self, chunk_size=None): + del chunk_size + yield from self._chunks + + +def test_wechatpad_response_reader_is_bounded(monkeypatch): + monkeypatch.setattr(http_util, '_MAX_WECHATPAD_RESPONSE_BYTES', 4) + + with pytest.raises(RuntimeError, match='exceeds the runtime limit'): + http_util._read_requests_response_limited(_Response([b'1234', b'5'])) + + +def test_wechatpad_response_reader_requires_json_object(): + with pytest.raises(RuntimeError, match='non-object'): + http_util._read_requests_response_limited(_Response([b'[]'])) + + +@pytest.mark.asyncio +async def test_wechatpad_media_reader_is_bounded(monkeypatch): + monkeypatch.setattr(downloadpai, '_MAX_WECHATPAD_MEDIA_BYTES', 4) + response = httpx.Response(200, content=b'oversized') + + with pytest.raises(RuntimeError, match='exceeds'): + await downloadpai._read_media_limited(response) diff --git a/tests/unit_tests/platform/test_wecom_api_limits.py b/tests/unit_tests/platform/test_wecom_api_limits.py new file mode 100644 index 000000000..9a8b17f1f --- /dev/null +++ b/tests/unit_tests/platform/test_wecom_api_limits.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +import pytest + +from langbot.libs.wecom_api.api import ( + _EXTENDED_HTTP_TIMEOUT_SECONDS, + _decode_media_base64_limited, + WecomClient, +) + + +@pytest.mark.asyncio +async def test_wecom_extended_client_timeout_is_still_bounded() -> None: + client = object.__new__(WecomClient) + client._http_clients = {} + + try: + async with client._http_client_context(unbounded_timeout=True) as http_client: + assert http_client.timeout.read == _EXTENDED_HTTP_TIMEOUT_SECONDS + finally: + await client.close() + + +@pytest.mark.asyncio +async def test_wecom_base64_decode_is_bounded(monkeypatch) -> None: + import langbot.libs.wecom_api.api as wecom_api + + monkeypatch.setattr(wecom_api, '_MAX_MEDIA_BYTES', 4) + + with pytest.raises(ValueError, match='exceeds'): + await _decode_media_base64_limited('MTIzNDU=') diff --git a/tests/unit_tests/platform/test_wecombot_template_card.py b/tests/unit_tests/platform/test_wecombot_template_card.py index c7ebfbee3..03d8791b1 100644 --- a/tests/unit_tests/platform/test_wecombot_template_card.py +++ b/tests/unit_tests/platform/test_wecombot_template_card.py @@ -1,5 +1,6 @@ import sys import types +from unittest.mock import Mock import pytest @@ -24,6 +25,25 @@ from langbot.libs.wecom_ai_bot_api.api import ( # noqa: E402 from langbot.libs.wecom_ai_bot_api.ws_client import WecomBotWsClient # noqa: E402 +def test_ws_callback_tasks_are_bounded(): + client = WecomBotWsClient('bot-id', 'secret', object()) + client._callback_tasks = {Mock(done=Mock(return_value=False)) for _ in range(100)} + + async def callback(): + raise AssertionError('rejected callback must not run') + + assert client._start_callback_task(callback()) is False + assert len(client._callback_tasks) == 100 + + +def test_webhook_dispatch_tasks_are_bounded(): + client = WecomBotClient('', '', '', object(), unified_mode=True) + client._dispatch_tasks = {Mock(done=Mock(return_value=False)) for _ in range(100)} + + assert client._start_dispatch_task(Mock()) is False + assert len(client._dispatch_tasks) == 100 + + def test_extract_template_card_action_supports_nested_button_key(): task_id, event_key, card_type = extract_template_card_action( { @@ -287,16 +307,9 @@ async def test_webhook_stream_queues_cumulative_snapshots_for_followups(): assert await client.push_stream_chunk('msg-1', '你好', is_final=False) assert await client.push_stream_chunk('msg-1', '你好', is_final=True) - chunks = [ - await client.stream_sessions.consume(session.stream_id), - await client.stream_sessions.consume(session.stream_id), - await client.stream_sessions.consume(session.stream_id), - ] - assert [(chunk.content, chunk.is_final) for chunk in chunks] == [ - ('你', False), - ('你好', False), - ('你好', True), - ] + assert session.queue.qsize() == 1 + chunk = await client.stream_sessions.consume(session.stream_id) + assert (chunk.content, chunk.is_final) == ('你好', True) def test_human_input_payload_keeps_action_select_stage_as_buttons(): diff --git a/tests/unit_tests/plugin/test_connector_pure.py b/tests/unit_tests/plugin/test_connector_pure.py index 13ba29b50..5cc863940 100644 --- a/tests/unit_tests/plugin/test_connector_pure.py +++ b/tests/unit_tests/plugin/test_connector_pure.py @@ -12,6 +12,7 @@ import zipfile from types import SimpleNamespace from unittest.mock import MagicMock +import httpx import pytest @@ -123,6 +124,17 @@ class TestExtractDepsMetadata: # Should find requirements.txt in subdirectory assert task_context.metadata['deps_total'] == 2 + def test_archive_preview_rejects_extreme_compression_ratio(self): + from langbot.pkg.plugin.connector import inspect_plugin_archive_metadata + + zip_buffer = io.BytesIO() + with zipfile.ZipFile(zip_buffer, 'w', compression=zipfile.ZIP_DEFLATED) as zf: + zf.writestr('manifest.yaml', 'kind: Plugin\nmetadata: {}\n') + zf.writestr('bomb.py', b'A' * (1024 * 1024)) + + with pytest.raises(ValueError, match='compression-ratio limit'): + inspect_plugin_archive_metadata(zip_buffer.getvalue()) + class TestParsePluginId: """Tests for _parse_plugin_id static method.""" @@ -141,3 +153,13 @@ class TestParsePluginId: with pytest.raises(ValueError): PluginRuntimeConnector._parse_plugin_id('') + + +@pytest.mark.asyncio +async def test_marketplace_response_reader_is_bounded(): + from langbot.pkg.plugin.connector import _read_httpx_response_limited + + response = httpx.Response(200, content=b'oversized') + + with pytest.raises(ValueError, match='exceeds'): + await _read_httpx_response_limited(response, max_bytes=4) diff --git a/tests/unit_tests/plugin/test_connector_reconcile.py b/tests/unit_tests/plugin/test_connector_reconcile.py index afd2bd206..b22e8e13e 100644 --- a/tests/unit_tests/plugin/test_connector_reconcile.py +++ b/tests/unit_tests/plugin/test_connector_reconcile.py @@ -136,6 +136,23 @@ async def test_shared_reconnect_replays_two_workspaces_and_removes_missing_proje assert set(connector._known_desired_states) == {setting_a.installation_uuid} +@pytest.mark.asyncio +async def test_empty_projected_workspaces_do_not_retain_installation_sets(): + binding_a = execution_binding('workspace-a') + binding_b = execution_binding('workspace-b') + connector = shared_connector( + [[binding_a, binding_b]], + {'workspace-a': [], 'workspace-b': []}, + ) + connector.handler = runtime_handler() + + await connector._prepare_connected_runtime() + + assert connector._workspace_installations == {} + assert connector._known_desired_states == {} + connector.handler.reconcile_plugin_installations.assert_awaited_once_with(()) + + @pytest.mark.asyncio async def test_fresh_shared_runtime_cache_replays_persisted_local_package(): package = b'local-lbpkg-bytes' diff --git a/tests/unit_tests/plugin/test_handler_actions.py b/tests/unit_tests/plugin/test_handler_actions.py index 63e9ea48c..c215670c1 100644 --- a/tests/unit_tests/plugin/test_handler_actions.py +++ b/tests/unit_tests/plugin/test_handler_actions.py @@ -257,7 +257,7 @@ class TestSetBinaryStorage: ) assert response.code != 0 - assert '2048 > 1024 bytes' in response.message + assert '1024-byte limit' in response.message app.persistence_mgr.execute_async.assert_not_awaited() @pytest.mark.asyncio @@ -312,21 +312,25 @@ class TestSetBinaryStorage: ) assert response.code != 0 - assert '10485761 > 10485760 bytes' in response.message + assert '10485760' in response.message app.persistence_mgr.execute_async.assert_not_awaited() @pytest.mark.asyncio - async def test_negative_limit_disables_size_check(self, app): - """Negative max_value_bytes allows values larger than the normal default.""" + async def test_negative_limit_falls_back_to_bounded_default(self, app, monkeypatch): + """Negative max_value_bytes cannot disable the process memory boundary.""" + import langbot.pkg.plugin.handler as handler_module + runtime_handler = make_handler(app) app.instance_config.data['plugin']['binary_storage']['max_value_bytes'] = -1 + monkeypatch.setattr(handler_module, '_DEFAULT_BINARY_STORAGE_VALUE_BYTES', 1024) response = await runtime_handler.actions[RuntimeToLangBotAction.SET_BINARY_STORAGE.value]( self.payload(b'x' * 2048) ) - assert response.code == 0 - assert app.persistence_mgr.execute_async.await_count == 2 + assert response.code != 0 + assert '1024-byte limit' in response.message + app.persistence_mgr.execute_async.assert_not_awaited() @pytest.mark.asyncio async def test_zero_limit_rejects_non_empty_values(self, app): diff --git a/tests/unit_tests/provider/runners/test_difysvapi_runner.py b/tests/unit_tests/provider/runners/test_difysvapi_runner.py index f6d414f34..028d278be 100644 --- a/tests/unit_tests/provider/runners/test_difysvapi_runner.py +++ b/tests/unit_tests/provider/runners/test_difysvapi_runner.py @@ -6,6 +6,7 @@ Tests the helper methods that don't require real Dify API calls. from __future__ import annotations import pytest +import time from unittest.mock import AsyncMock, MagicMock import langbot_plugin.api.entities.builtin.platform.message as platform_message @@ -18,13 +19,12 @@ class TestDifyWorkflowSubmitClient: class FakeResponse: status_code = 503 + headers = {} - async def aread(self): - return b'' - - async def aiter_lines(self): - raise AssertionError('error responses must not enter the SSE loop') - yield + async def aiter_bytes(self, chunk_size=None): + del chunk_size + if False: + yield b'' class FakeStreamContext: async def __aenter__(self): @@ -66,6 +66,33 @@ class TestDifyWorkflowSubmitClient: ) ) + @pytest.mark.asyncio + async def test_sse_parser_rejects_an_unbounded_line(self): + from langbot.libs.dify_service_api.v1 import client, errors + + class FakeResponse: + async def aiter_bytes(self, chunk_size=None): + del chunk_size + for _ in range(129): + yield b'x' * 8192 + + with pytest.raises(errors.DifyAPIError, match='SSE event exceeds'): + await anext(client._iter_sse_json(FakeResponse())) + + @pytest.mark.asyncio + async def test_upload_rejects_oversized_local_file(self, tmp_path): + from langbot.libs.dify_service_api.v1 import client + + file_path = tmp_path / 'large.bin' + file_path.write_bytes(b'x' * (client._MAX_DIFY_UPLOAD_BYTES + 1)) + dify_client = client.AsyncDifyServiceClient( + 'test-key', + 'https://dify.example/v1', + ) + + with pytest.raises(ValueError, match='exceeds the size limit'): + await dify_client.upload_file(file_path, 'person_user-1') + class TestDifyExtractTextOutput: """Tests for _extract_dify_text_output method.""" @@ -320,6 +347,130 @@ class TestDifyHumanInputForms: assert difysvapi._dify_user_from_query(query_a) == difysvapi._dify_user_from_query(query_c) difysvapi._PENDING_FORMS.clear() + def test_pending_form_lookup_does_not_scan_unrelated_sessions(self, monkeypatch): + from langbot.pkg.provider.runners import difysvapi + + def session_key(index: int): + return ( + 'instance', + f'workspace-{index}', + 1, + 'bot', + 'pipeline', + 'adapter', + 'person', + f'user-{index}', + ) + + difysvapi._PENDING_FORMS.clear() + for index in range(512): + difysvapi._set_pending_form( + session_key(index), + { + 'form_token': f'token-{index}', + 'workflow_run_id': f'run-{index}', + }, + ) + + class NoGlobalIterationDict(dict): + def __iter__(self): + raise AssertionError('pending form lookup scanned all sessions') + + def keys(self): + raise AssertionError('pending form lookup scanned all sessions') + + def items(self): + raise AssertionError('pending form lookup scanned all sessions') + + def values(self): + raise AssertionError('pending form lookup scanned all sessions') + + guarded_forms = NoGlobalIterationDict(difysvapi._PENDING_FORMS) + monkeypatch.setattr(difysvapi, '_PENDING_FORMS', guarded_forms) + + assert difysvapi._get_pending_form_by_token(session_key(511), 'token-511')['workflow_run_id'] == 'run-511' + difysvapi._set_pending_form( + session_key(512), + {'form_token': 'token-512', 'workflow_run_id': 'run-512'}, + ) + assert len(guarded_forms) == 513 + + def test_pending_form_expiry_heap_ignores_stale_overwrite_and_stays_bounded(self): + from langbot.pkg.provider.runners import difysvapi + + session_key = ( + 'instance', + 'workspace', + 1, + 'bot', + 'pipeline', + 'adapter', + 'person', + 'user', + ) + difysvapi._PENDING_FORMS.clear() + now = time.time() + difysvapi._set_pending_form( + session_key, + { + 'form_token': 'token', + 'workflow_run_id': 'stale', + 'expiration_time': now + 1, + }, + ) + for revision in range(500): + difysvapi._set_pending_form( + session_key, + { + 'form_token': 'token', + 'workflow_run_id': f'current-{revision}', + 'expiration_time': now + 3600 + revision, + }, + ) + + difysvapi._prune_pending_forms(now + 2) + + assert difysvapi._get_pending_form_by_token(session_key, 'token')['workflow_run_id'] == 'current-499' + assert difysvapi._PENDING_FORM_ACTIVE_COUNT == 1 + assert len(difysvapi._PENDING_FORM_EXPIRY_HEAP) <= max( + difysvapi._PENDING_FORM_HEAP_COMPACT_FLOOR, + difysvapi._PENDING_FORM_ACTIVE_COUNT * difysvapi._PENDING_FORM_HEAP_MAX_MULTIPLIER, + ) + + def test_pending_form_capacity_evicts_earliest_session_without_full_scan( + self, + monkeypatch, + ): + from langbot.pkg.provider.runners import difysvapi + + def session_key(index: int): + return ( + 'instance', + f'workspace-{index}', + 1, + 'bot', + 'pipeline', + 'adapter', + 'person', + f'user-{index}', + ) + + difysvapi._PENDING_FORMS.clear() + monkeypatch.setattr(difysvapi, '_PENDING_FORM_MAX_SESSIONS', 2) + now = time.time() + for index, expires_in in ((1, 300), (2, 100), (3, 200)): + difysvapi._set_pending_form( + session_key(index), + { + 'form_token': f'token-{index}', + 'expiration_time': now + expires_in, + }, + ) + + assert session_key(1) in difysvapi._PENDING_FORMS + assert session_key(2) not in difysvapi._PENDING_FORMS + assert session_key(3) in difysvapi._PENDING_FORMS + def test_interactive_form_data_preserves_pipeline_uuid(self): from langbot.pkg.provider.runners import difysvapi diff --git a/tests/unit_tests/provider/runners/test_remote_stream_limits.py b/tests/unit_tests/provider/runners/test_remote_stream_limits.py new file mode 100644 index 000000000..6791999d6 --- /dev/null +++ b/tests/unit_tests/provider/runners/test_remote_stream_limits.py @@ -0,0 +1,49 @@ +from __future__ import annotations + +import pytest + +from langbot.libs.deerflow_api.client import ( + ERROR_BODY_MAX_BYTES, + _read_error_body, +) +from langbot.libs.deerflow_api.errors import DeerFlowAPIError +from langbot.pkg.provider.runners.langflowapi import ( + _MAX_LANGFLOW_LINE_CHARS, + _MAX_LANGFLOW_RESPONSE_BYTES, + _iter_limited_lines, + _read_limited_response, +) + + +class _ChunkedResponse: + def __init__(self, chunks: list[bytes]): + self._chunks = chunks + + async def aiter_bytes(self, chunk_size=None): + del chunk_size + for chunk in self._chunks: + yield chunk + + +@pytest.mark.asyncio +async def test_langflow_rejects_oversized_stream_event(): + response = _ChunkedResponse([b'x' * (_MAX_LANGFLOW_LINE_CHARS + 1)]) + + with pytest.raises(ValueError, match='event exceeds'): + await anext(_iter_limited_lines(response)) + + +@pytest.mark.asyncio +async def test_langflow_rejects_oversized_blocking_response(): + response = _ChunkedResponse([b'x' * (_MAX_LANGFLOW_RESPONSE_BYTES + 1)]) + + with pytest.raises(ValueError, match='response exceeds'): + await _read_limited_response(response) + + +@pytest.mark.asyncio +async def test_deerflow_rejects_oversized_error_body(): + response = _ChunkedResponse([b'x' * (ERROR_BODY_MAX_BYTES + 1)]) + + with pytest.raises(DeerFlowAPIError, match='response exceeds'): + await _read_error_body(response) diff --git a/tests/unit_tests/provider/runners/test_runner_resource_limits.py b/tests/unit_tests/provider/runners/test_runner_resource_limits.py new file mode 100644 index 000000000..2aeebdee8 --- /dev/null +++ b/tests/unit_tests/provider/runners/test_runner_resource_limits.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +import asyncio +import threading +from unittest.mock import AsyncMock + +import pytest + +from langbot.pkg.provider import runner +from langbot.pkg.provider.runners import ( + cozeapi, + dashscopeapi, + tboxapi, + weknoraapi, +) + + +@pytest.mark.asyncio +async def test_blocking_provider_iterator_runs_outside_event_loop(): + release = threading.Event() + + def values(): + release.wait(timeout=2) + yield 'ready' + + task = asyncio.create_task(anext(runner.iterate_sync(values()))) + await asyncio.sleep(0) + assert not task.done() + + release.set() + assert await asyncio.wait_for(task, timeout=1) == 'ready' + + +@pytest.mark.asyncio +async def test_sync_provider_iterator_has_event_limit(): + with pytest.raises(RuntimeError, match='event limit'): + async for _ in runner.iterate_sync(iter([1, 2]), max_items=1): + pass + + +@pytest.mark.asyncio +async def test_coze_runner_closes_request_scoped_client(): + request_runner = object.__new__(cozeapi.CozeAPIRunner) + request_runner.coze = AsyncMock() + + await request_runner.aclose() + + request_runner.coze.close.assert_awaited_once() + + +@pytest.mark.parametrize( + ('append', 'exception_type'), + [ + (cozeapi._append_bounded, ValueError), + (dashscopeapi._append_bounded, dashscopeapi.DashscopeAPIError), + (tboxapi._append_bounded, tboxapi.TboxAPIError), + (weknoraapi._append_bounded, weknoraapi.errors.WeKnoraAPIError), + ], +) +def test_provider_accumulators_reject_oversized_output(append, exception_type): + with pytest.raises(exception_type, match='exceeds the runtime limit'): + append('x' * (1024 * 1024), 'y') diff --git a/tests/unit_tests/provider/test_mcp_resources.py b/tests/unit_tests/provider/test_mcp_resources.py index afaabd7d6..efc039cce 100644 --- a/tests/unit_tests/provider/test_mcp_resources.py +++ b/tests/unit_tests/provider/test_mcp_resources.py @@ -97,7 +97,11 @@ def _query(variables: dict | None = None, context: ExecutionContext = TEST_EXECU def _register_session(loader: MCPLoader, session: RuntimeMCPSession) -> None: - loader.sessions[loader._session_key(session.execution_context, session.server_name)] = session + loader._register_session( + session.execution_context, + session.server_name, + session, + ) def _http_status_error(status_code: int) -> httpx.HTTPStatusError: @@ -615,3 +619,165 @@ async def test_mcp_loader_shutdown_cancels_startup_tasks_and_closes_sessions_con assert started == {'one', 'two'} assert loader._hosted_mcp_tasks == [] assert loader.sessions == {} + + +@pytest.mark.asyncio +async def test_completed_mcp_host_tasks_do_not_accumulate(): + loader = MCPLoader(_app()) + task = asyncio.create_task(asyncio.sleep(0)) + + loader.track_hosted_task(task, TEST_EXECUTION_CONTEXT) + await task + await asyncio.sleep(0) + + assert loader._hosted_mcp_tasks == [] + assert loader._hosted_mcp_tasks_by_scope == {} + assert loader._scope_generations == {} + + +@pytest.mark.asyncio +async def test_generation_advance_cancels_host_tasks_and_closes_old_sessions(): + loader = MCPLoader(_app()) + old_session = SimpleNamespace( + server_name='old', + shutdown=AsyncMock(), + ) + loader._register_session( + TEST_EXECUTION_CONTEXT, + old_session.server_name, + old_session, + ) + + async def pending_host(): + await asyncio.Event().wait() + + hosted_task = asyncio.create_task(pending_host()) + loader.track_hosted_task(hosted_task, TEST_EXECUTION_CONTEXT) + await asyncio.sleep(0) + next_context = ExecutionContext( + instance_uuid=TEST_EXECUTION_CONTEXT.instance_uuid, + workspace_uuid=TEST_EXECUTION_CONTEXT.workspace_uuid, + placement_generation=2, + ) + loader.ap.workspace_service.get_execution_binding = AsyncMock( + return_value=SimpleNamespace( + instance_uuid=next_context.instance_uuid, + workspace_uuid=next_context.workspace_uuid, + placement_generation=next_context.placement_generation, + ) + ) + + await loader._assert_execution_active(next_context) + + assert hosted_task.cancelled() + old_session.shutdown.assert_awaited_once_with() + assert loader.sessions == {} + assert loader._session_keys_by_scope == {} + assert loader._hosted_mcp_tasks_by_scope == {} + assert loader._scope_generations == {} + + +def test_session_lookup_uses_scope_index_without_global_iteration(): + class NoGlobalIterationDict(dict): + def __iter__(self): + raise AssertionError('MCP lookup scanned every tenant session') + + def items(self): + raise AssertionError('MCP lookup scanned every tenant session') + + def values(self): + raise AssertionError('MCP lookup scanned every tenant session') + + loader = MCPLoader(_app()) + target_context = None + target_session = None + for index in range(1_000): + context = ExecutionContext( + instance_uuid='instance-a', + workspace_uuid=f'workspace-{index}', + placement_generation=1, + ) + session = SimpleNamespace(server_name=f'server-{index}') + loader._register_session(context, session.server_name, session) + if index == 777: + target_context = context + target_session = session + loader._sessions = NoGlobalIterationDict(loader._sessions) + + assert loader._sessions_for_context(target_context) == [target_session] + + +@pytest.mark.asyncio +async def test_mcp_startup_concurrency_is_instance_bounded(): + app = _app() + app.instance_config = SimpleNamespace(data={'mcp': {'lifecycle_concurrency': 2}}) + loader = MCPLoader(app) + active = 0 + maximum_active = 0 + release = asyncio.Event() + + async def fake_host(_context, _config): + nonlocal active, maximum_active + active += 1 + maximum_active = max(maximum_active, active) + if maximum_active == 2: + release.set() + await release.wait() + await asyncio.sleep(0) + active -= 1 + + loader._host_mcp_server = fake_host + + await asyncio.gather( + *( + loader.host_mcp_server( + TEST_EXECUTION_CONTEXT, + {'name': f'server-{index}'}, + ) + for index in range(20) + ) + ) + + assert loader._lifecycle_concurrency == 2 + assert maximum_active == 2 + + +@pytest.mark.asyncio +async def test_mcp_startup_dispatcher_does_not_create_every_server_task_at_once(): + app = _app() + app.instance_config = SimpleNamespace(data={'mcp': {'lifecycle_concurrency': 2}}) + loader = MCPLoader(app) + started = 0 + first_batch_started = asyncio.Event() + release = asyncio.Event() + + async def fake_host(_context, _config): + nonlocal started + started += 1 + if started == 2: + first_batch_started.set() + await release.wait() + + loader.host_mcp_server = fake_host + configs = [(TEST_EXECUTION_CONTEXT, {'name': f'server-{index}'}) for index in range(20)] + + dispatch_task = asyncio.create_task(loader._host_server_configs_bounded(configs)) + await asyncio.wait_for(first_batch_started.wait(), timeout=1) + + assert started == 2 + assert len(loader._hosted_mcp_tasks) == 2 + + release.set() + await dispatch_task + await asyncio.sleep(0) + + assert started == 20 + assert loader._hosted_mcp_tasks == [] + assert loader._hosted_mcp_tasks_by_scope == {} + + +def test_invalid_mcp_lifecycle_concurrency_uses_safe_default(): + app = _app() + app.instance_config = SimpleNamespace(data={'mcp': {'lifecycle_concurrency': True}}) + + assert MCPLoader(app)._lifecycle_concurrency == 16 diff --git a/tests/unit_tests/provider/test_model_manager.py b/tests/unit_tests/provider/test_model_manager.py index 5c2005ccd..de343140b 100644 --- a/tests/unit_tests/provider/test_model_manager.py +++ b/tests/unit_tests/provider/test_model_manager.py @@ -7,6 +7,7 @@ and error handling without calling real LLM APIs. from __future__ import annotations +import dataclasses import pytest from types import SimpleNamespace from unittest.mock import AsyncMock, Mock @@ -18,7 +19,7 @@ from langbot.pkg.entity.errors import provider as provider_errors from langbot.pkg.provider.modelmgr import token from langbot.pkg.api.http.context import ExecutionContext from langbot.pkg.workspace.entities import WorkspaceExecutionBinding -from langbot.pkg.workspace.errors import WorkspaceGenerationMismatchError +from langbot.pkg.workspace.errors import WorkspaceGenerationMismatchError, WorkspaceInvariantError from tests.unit_tests.provider.conftest import ( TEST_EXECUTION_CONTEXT, TEST_WORKSPACE_UUID, @@ -125,6 +126,24 @@ async def test_model_manager_load_models_from_db(fake_requester_registry, fake_p assert len(model_mgr.rerank_model_dict) == 1 +@pytest.mark.asyncio +async def test_empty_cloud_workspace_does_not_retain_generation( + mock_app_for_modelmgr, +): + model_mgr = ModelManager(mock_app_for_modelmgr) + + await model_mgr._load_workspace_models(TEST_EXECUTION_CONTEXT) + + assert model_mgr.provider_dict == {} + assert model_mgr.llm_model_dict == {} + assert model_mgr.embedding_model_dict == {} + assert model_mgr.rerank_model_dict == {} + assert model_mgr._scope_generations == {} + + await model_mgr.resolve_execution_context(TEST_EXECUTION_CONTEXT) + assert model_mgr._scope_generations == {} + + @pytest.mark.asyncio async def test_model_manager_load_provider_unknown_requester(mock_app_for_modelmgr): """Test ModelManager raises RequesterNotFoundError for unknown requester.""" @@ -939,6 +958,110 @@ async def test_runtime_cache_rejects_stale_placement_generation(fake_requester_r await model_mgr.get_model_by_uuid(stale_context, 'any-model') +def test_generation_advance_prunes_superseded_model_runtime_objects(): + class NoGlobalIterationDict(dict): + def __iter__(self): + raise AssertionError('generation advance scanned every model runtime') + + def items(self): + raise AssertionError('generation advance scanned every model runtime') + + def keys(self): + raise AssertionError('generation advance scanned every model runtime') + + model_mgr = ModelManager(Mock()) + old_context = ExecutionContext( + instance_uuid='instance-a', + workspace_uuid='workspace-a', + placement_generation=1, + ) + new_context = dataclasses.replace(old_context, placement_generation=2) + model_mgr._observe_execution_context(old_context) + for cache in ( + model_mgr.provider_dict, + model_mgr.llm_model_dict, + model_mgr.embedding_model_dict, + model_mgr.rerank_model_dict, + ): + model_mgr._cache_set( + cache, + ('instance-a', 'workspace-a', 1, 'resource-a'), + object(), + ) + model_mgr._cache_set( + cache, + ('instance-a', 'workspace-b', 1, 'resource-b'), + object(), + ) + + model_mgr.provider_dict = NoGlobalIterationDict(model_mgr.provider_dict) + model_mgr.llm_model_dict = NoGlobalIterationDict(model_mgr.llm_model_dict) + model_mgr.embedding_model_dict = NoGlobalIterationDict(model_mgr.embedding_model_dict) + model_mgr.rerank_model_dict = NoGlobalIterationDict(model_mgr.rerank_model_dict) + + model_mgr._observe_execution_context(new_context) + + for cache in ( + model_mgr.provider_dict, + model_mgr.llm_model_dict, + model_mgr.embedding_model_dict, + model_mgr.rerank_model_dict, + ): + assert ('instance-a', 'workspace-a', 1, 'resource-a') not in cache + assert ('instance-a', 'workspace-b', 1, 'resource-b') in cache + with pytest.raises(WorkspaceInvariantError, match='rolled back'): + model_mgr._observe_execution_context(old_context) + + +@pytest.mark.asyncio +async def test_generation_advance_closes_retired_provider_requester( + fake_requester_registry, + runtime_provider, +): + model_mgr = fake_requester_registry + runtime_provider.requester.aclose = AsyncMock() + await model_mgr.cache_provider(TEST_EXECUTION_CONTEXT, runtime_provider) + + next_context = dataclasses.replace( + TEST_EXECUTION_CONTEXT, + placement_generation=2, + ) + model_mgr.ap.workspace_service.get_execution_binding = AsyncMock( + return_value=WorkspaceExecutionBinding( + instance_uuid=next_context.instance_uuid, + workspace_uuid=next_context.workspace_uuid, + placement_generation=next_context.placement_generation, + write_fenced=False, + state='active', + ) + ) + + await model_mgr.resolve_execution_context(next_context) + + runtime_provider.requester.aclose.assert_awaited_once_with() + assert model_mgr.provider_dict == {} + assert model_mgr._scope_generations == {} + + +@pytest.mark.asyncio +async def test_model_manager_shutdown_closes_all_requesters_once( + fake_requester_registry, + runtime_provider, +): + model_mgr = fake_requester_registry + runtime_provider.requester.aclose = AsyncMock() + await model_mgr.cache_provider(TEST_EXECUTION_CONTEXT, runtime_provider) + + await model_mgr.shutdown() + await model_mgr.shutdown() + + runtime_provider.requester.aclose.assert_awaited_once_with() + assert model_mgr.provider_dict == {} + assert model_mgr.llm_model_dict == {} + assert model_mgr.embedding_model_dict == {} + assert model_mgr.rerank_model_dict == {} + + def test_provider_not_found_error_str(): """Test ProviderNotFoundError string representation.""" error = provider_errors.ProviderNotFoundError('test-provider') diff --git a/tests/unit_tests/provider/test_session_manager.py b/tests/unit_tests/provider/test_session_manager.py index 018917430..469743c5b 100644 --- a/tests/unit_tests/provider/test_session_manager.py +++ b/tests/unit_tests/provider/test_session_manager.py @@ -16,6 +16,8 @@ from importlib import import_module import langbot_plugin.api.entities.builtin.provider.session as provider_session import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query +import langbot_plugin.api.entities.builtin.provider.message as provider_message +import langbot_plugin.api.entities.builtin.provider.prompt as provider_prompt from langbot.pkg.api.http.context import ExecutionContext from langbot.pkg.pipeline.pool import ( @@ -426,3 +428,131 @@ class TestSessionManagerWorkspaceIsolation: with pytest.raises(ExecutionContextMismatchError): await manager.get_conversation(query, session, [], 'pipeline-1', 'bot-b') + + @pytest.mark.asyncio + async def test_per_workspace_capacity_evicts_oldest_idle_session(self): + manager = self.manager() + manager.ap.instance_config.data['system'] = { + 'session_retention': { + 'max_entries': 10, + 'max_entries_per_workspace': 2, + } + } + queries = [] + sessions = [] + for index in range(3): + query = scoped_query() + query.launcher_id = f'launcher-{index}' + queries.append(query) + sessions.append(await manager.get_session(query)) + + assert len(manager.session_list) == 2 + assert sessions[0] not in manager.session_list + assert sessions[1:] == manager.session_list + assert await manager.get_session(queries[-1]) is manager.session_list[-1] + + @pytest.mark.asyncio + async def test_new_session_does_not_scan_other_workspace_sessions(self): + manager = self.manager() + manager.ap.instance_config.data['system'] = { + 'session_retention': { + 'max_entries': 600, + 'max_entries_per_workspace': 2, + } + } + for index in range(512): + query = scoped_query(workspace_uuid=f'workspace-{index}') + query.launcher_id = f'launcher-{index}' + await manager.get_session(query) + + class NoGlobalIterationDict(dict): + def __iter__(self): + raise AssertionError('global session index iteration is forbidden') + + def items(self): + raise AssertionError('global session index iteration is forbidden') + + def values(self): + raise AssertionError('global session index iteration is forbidden') + + manager._session_index = NoGlobalIterationDict(manager._session_index) + query = scoped_query(workspace_uuid='workspace-new') + query.launcher_id = 'launcher-new' + + session = await manager.get_session(query) + + assert session.workspace_uuid == 'workspace-new' + assert len(manager._session_index) == 513 + + @pytest.mark.asyncio + async def test_stale_expiry_revision_does_not_evict_recent_session( + self, + monkeypatch, + ): + sessionmgr = get_session_module() + manager = self.manager() + manager.ap.instance_config.data['system'] = { + 'session_retention': { + 'max_entries': 10, + 'max_entries_per_workspace': 10, + 'idle_ttl_seconds': 1, + } + } + clock = [0.0] + monkeypatch.setattr(sessionmgr.time, 'monotonic', lambda: clock[0]) + first_query = scoped_query() + first_query.launcher_id = 'first' + first = await manager.get_session(first_query) + + clock[0] = 0.5 + assert await manager.get_session(first_query) is first + + clock[0] = 1.25 + second_query = scoped_query() + second_query.launcher_id = 'second' + await manager.get_session(second_query) + assert first in manager.session_list + + clock[0] = 2.0 + third_query = scoped_query() + third_query.launcher_id = 'third' + await manager.get_session(third_query) + assert first not in manager.session_list + + @pytest.mark.asyncio + async def test_access_revision_heap_stays_bounded(self): + manager = self.manager() + query = scoped_query() + await manager.get_session(query) + + for _ in range(1000): + await manager.get_session(query) + + assert len(manager._session_expiry_heap) <= 64 + + def test_trim_conversation_drops_retained_binary_payloads(self): + manager = self.manager() + conversation = provider_session.Conversation( + prompt=provider_prompt.Prompt(name='test', messages=[]), + messages=[ + provider_message.Message( + role='user', + content=[ + provider_message.ContentElement.from_text('hello'), + provider_message.ContentElement.from_image_base64('x' * 1000000), + provider_message.ContentElement.from_file_base64( + 'y' * 1000000, + 'large.bin', + ), + ], + ) + ], + pipeline_uuid='pipeline-1', + bot_uuid=TEST_BOT_UUID, + ) + + manager.trim_conversation_messages(conversation, max_rounds=10) + + content = conversation.messages[0].content + assert content[1].image_base64 is None + assert content[2].file_base64 is None diff --git a/tests/unit_tests/provider/test_tool_manager_native.py b/tests/unit_tests/provider/test_tool_manager_native.py index 8ef217a1e..d18a830dc 100644 --- a/tests/unit_tests/provider/test_tool_manager_native.py +++ b/tests/unit_tests/provider/test_tool_manager_native.py @@ -4,6 +4,7 @@ import base64 import contextlib import os import tempfile +import threading from types import SimpleNamespace from unittest.mock import AsyncMock, Mock @@ -459,6 +460,24 @@ async def test_edit_rejects_missing_string(): assert 'not found' in result['error'].lower() +@pytest.mark.asyncio +async def test_edit_rejects_oversized_host_file(monkeypatch): + with tempfile.TemporaryDirectory() as tmpdir: + loader, _ = _make_loader_with_workspace(tmpdir) + with open(os.path.join(tmpdir, 'large.txt'), 'wb') as f: + f.write(b'12345') + monkeypatch.setattr(native_loader, '_MAX_HOST_EDIT_FILE_BYTES', 4) + + result = await loader.invoke_tool( + 'edit', + {'path': '/workspace/large.txt', 'old_string': '1', 'new_string': 'x'}, + _make_query(), + ) + + assert result['ok'] is False + assert 'edit limit' in result['error'] + + @pytest.mark.asyncio async def test_path_escape_blocked(): with tempfile.TemporaryDirectory() as tmpdir: @@ -678,6 +697,34 @@ async def test_glob_caps_match_count_and_returns_preview(): assert result['truncated_by'] == 'matches' +@pytest.mark.asyncio +async def test_glob_runs_off_event_loop_and_caps_directory_walk(monkeypatch): + monkeypatch.setattr(native_loader, '_FILE_WALK_MAX_ENTRIES', 10) + + with tempfile.TemporaryDirectory() as tmpdir: + loader, _ = _make_loader_with_workspace(tmpdir) + event_loop_thread = threading.get_ident() + observed_threads: list[int] = [] + original = loader._glob_host_location + + def observe(*args, **kwargs): + observed_threads.append(threading.get_ident()) + return original(*args, **kwargs) + + monkeypatch.setattr(loader, '_glob_host_location', observe) + for index in range(12): + with open(os.path.join(tmpdir, f'file-{index:03d}.txt'), 'w', encoding='utf-8') as f: + f.write(str(index)) + + result = await loader.invoke_tool('glob', {'path': '/workspace', 'pattern': '*.txt'}, _make_query()) + + assert result['ok'] is True + assert result['total'] == 10 + assert result['truncated'] is True + assert result['truncated_by'] == 'scan' + assert observed_threads and observed_threads[0] != event_loop_thread + + @pytest.mark.asyncio async def test_grep_reports_invalid_regex_and_truncates_long_matching_lines(): with tempfile.TemporaryDirectory() as tmpdir: @@ -695,3 +742,21 @@ async def test_grep_reports_invalid_regex_and_truncates_long_matching_lines(): assert result['truncated_by'] == 'line' assert result['matches'][0]['file'] == '/workspace/data.txt' assert result['matches'][0]['content'].endswith('... [truncated]') + + +@pytest.mark.asyncio +async def test_grep_interrupts_catastrophic_regex(monkeypatch): + monkeypatch.setattr(native_loader, '_GREP_REGEX_TIMEOUT_SECONDS', 0.001) + + with tempfile.TemporaryDirectory() as tmpdir: + loader, _ = _make_loader_with_workspace(tmpdir) + with open(os.path.join(tmpdir, 'data.txt'), 'w', encoding='utf-8') as f: + f.write(('a' * 100_000) + '!') + + result = await loader.invoke_tool( + 'grep', + {'path': '/workspace', 'pattern': r'(a+)+$'}, + _make_query(), + ) + + assert result == {'ok': False, 'error': 'Regex search timed out'} diff --git a/tests/unit_tests/rag/test_file_storage.py b/tests/unit_tests/rag/test_file_storage.py index 4d82dda1f..9adec28f0 100644 --- a/tests/unit_tests/rag/test_file_storage.py +++ b/tests/unit_tests/rag/test_file_storage.py @@ -12,6 +12,7 @@ from unittest.mock import AsyncMock, Mock import pytest from langbot.pkg.api.http.context import ExecutionContext +from langbot.pkg.core.taskmgr import TaskCapacityError from langbot.pkg.rag.knowledge.kbmgr import RuntimeKnowledgeBase from langbot.pkg.storage.mgr import StorageMgr from langbot.pkg.workspace.errors import WorkspaceNotFoundError @@ -34,9 +35,9 @@ def _upload_key(logical_key: str, *, context: ExecutionContext = CONTEXT) -> str ) -def _make_zip_bytes(entries: dict[str, bytes]) -> bytes: +def _make_zip_bytes(entries: dict[str, bytes], *, compression: int = zipfile.ZIP_STORED) -> bytes: buffer = io.BytesIO() - with zipfile.ZipFile(buffer, 'w') as zf: + with zipfile.ZipFile(buffer, 'w', compression=compression) as zf: for name, content in entries.items(): zf.writestr(name, content) zf.mkdir('emptydir') @@ -151,6 +152,24 @@ class TestStoreFile: kb.ap.storage_mgr.storage_provider.exists.assert_not_awaited() + @pytest.mark.asyncio + async def test_store_file_rolls_back_pending_record_when_task_capacity_is_exhausted(self): + kb = _make_kb() + object_key = _upload_key('queued.pdf') + + def reject(coro, **_kwargs): + coro.close() + raise TaskCapacityError('capacity') + + kb.ap.task_mgr.create_user_task.side_effect = reject + + with pytest.raises(TaskCapacityError, match='capacity'): + await kb.store_file(CONTEXT, object_key) + + statements = [str(call.args[0]) for call in kb.ap.persistence_mgr.execute_async.await_args_list] + assert any(statement.startswith('INSERT') for statement in statements) + assert any(statement.startswith('DELETE') for statement in statements) + class TestStoreZipFile: @pytest.mark.asyncio @@ -202,6 +221,39 @@ class TestStoreZipFile: kb.store_file.assert_not_awaited() kb.ap.storage_mgr.storage_provider.delete.assert_awaited_once_with(_upload_key('archive.zip')) + @pytest.mark.asyncio + async def test_store_zip_file_rejects_too_many_documents_before_extracting(self): + kb = _make_kb() + kb.ap.storage_mgr.storage_provider.load = AsyncMock( + return_value=_make_zip_bytes({f'doc-{index}.txt': b'text' for index in range(9)}) + ) + kb.store_file = AsyncMock() + + with pytest.raises(ValueError, match='too many supported documents'): + await kb._store_zip_file(CONTEXT, _upload_key('archive.zip')) + + kb.store_file.assert_not_awaited() + kb.ap.storage_mgr.storage_provider.save.assert_not_awaited() + kb.ap.storage_mgr.storage_provider.delete.assert_awaited_once_with(_upload_key('archive.zip')) + + @pytest.mark.asyncio + async def test_store_zip_file_rejects_extreme_compression_ratio_before_extracting(self): + kb = _make_kb() + kb.ap.storage_mgr.storage_provider.load = AsyncMock( + return_value=_make_zip_bytes( + {'bomb.txt': b'A' * (1024 * 1024)}, + compression=zipfile.ZIP_DEFLATED, + ) + ) + kb.store_file = AsyncMock() + + with pytest.raises(ValueError, match='compression-ratio limit'): + await kb._store_zip_file(CONTEXT, _upload_key('archive.zip')) + + kb.store_file.assert_not_awaited() + kb.ap.storage_mgr.storage_provider.save.assert_not_awaited() + kb.ap.storage_mgr.storage_provider.delete.assert_awaited_once_with(_upload_key('archive.zip')) + class TestStoreFileTask: @pytest.mark.asyncio diff --git a/tests/unit_tests/rag/test_kbmgr.py b/tests/unit_tests/rag/test_kbmgr.py index ba8fa0faa..9699269d9 100644 --- a/tests/unit_tests/rag/test_kbmgr.py +++ b/tests/unit_tests/rag/test_kbmgr.py @@ -2,6 +2,7 @@ from __future__ import annotations +import dataclasses from types import SimpleNamespace from unittest.mock import AsyncMock, Mock @@ -10,7 +11,8 @@ import pytest from langbot.pkg.api.http.context import ExecutionContext from langbot.pkg.entity.persistence.rag import KnowledgeBase from langbot.pkg.rag.knowledge.kbmgr import RAGManager, RuntimeKnowledgeBase -from langbot.pkg.workspace.errors import WorkspaceNotFoundError +from langbot.pkg.workspace.entities import WorkspaceExecutionBinding +from langbot.pkg.workspace.errors import WorkspaceInvariantError, WorkspaceNotFoundError CONTEXT_A = ExecutionContext( @@ -511,6 +513,37 @@ class TestRAGManagerLoadKnowledgeBasesFromDB: ('workspace-a', 'kb-2'), } + @pytest.mark.asyncio + async def test_cloud_startup_reuses_validated_binding(self): + class TenantUow: + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return False + + app = _app() + binding = WorkspaceExecutionBinding( + instance_uuid='instance-a', + workspace_uuid='workspace-a', + placement_generation=5, + write_fenced=False, + state='active', + ) + app.persistence_mgr.mode = SimpleNamespace(value='cloud_runtime') + app.persistence_mgr.tenant_uow = lambda _workspace_uuid: TenantUow() + app.persistence_mgr.execute_async.return_value = _Result([_entity()]) + app.workspace_service.list_active_execution_bindings = AsyncMock(return_value=[binding]) + app.workspace_service.get_execution_binding = AsyncMock( + side_effect=AssertionError('startup RAG loader repeated a validated binding lookup') + ) + manager = RAGManager(app) + + await manager.load_knowledge_bases_from_db() + + assert set(manager.knowledge_bases) == {('workspace-a', 'kb-a')} + app.workspace_service.get_execution_binding.assert_not_awaited() + @pytest.mark.asyncio async def test_handles_load_error_gracefully(self): app = _app() @@ -643,6 +676,29 @@ class TestRAGManagerInit: def test_init_creates_empty_knowledge_bases_dict(self): assert RAGManager(_app()).knowledge_bases == {} + def test_generation_advance_prunes_superseded_runtime_knowledge_bases(self): + class NoGlobalItemsScan(dict): + def items(self): + raise AssertionError('generation advance scanned every knowledge runtime') + + app = _app() + manager = RAGManager(app) + manager._cache_runtime( + RuntimeKnowledgeBase( + app, + _entity(), + CONTEXT_A, + ) + ) + manager.knowledge_bases = NoGlobalItemsScan(manager.knowledge_bases) + + next_context = dataclasses.replace(CONTEXT_A, placement_generation=6) + manager._observe_execution_context(next_context) + + assert manager.knowledge_bases == {} + with pytest.raises(WorkspaceInvariantError, match='rolled back'): + manager._observe_execution_context(CONTEXT_A) + class TestRAGManagerGetKnowledgeBase: @pytest.mark.asyncio diff --git a/tests/unit_tests/storage/test_s3storage.py b/tests/unit_tests/storage/test_s3storage.py index eb3d0f7e6..a325e5c61 100644 --- a/tests/unit_tests/storage/test_s3storage.py +++ b/tests/unit_tests/storage/test_s3storage.py @@ -74,6 +74,19 @@ class TestS3StorageProviderInit: assert provider.s3_client is None assert provider.bucket_name is None + @pytest.mark.asyncio + async def test_shutdown_closes_client_once(self): + s3storage = get_s3storage_module() + provider = s3storage.S3StorageProvider(Mock()) + client = Mock() + provider.s3_client = client + + await provider.shutdown() + await provider.shutdown() + + client.close.assert_called_once_with() + assert provider.s3_client is None + class TestS3StorageProviderWithMoto: """Tests using moto to mock AWS S3.""" diff --git a/tests/unit_tests/storage/test_storage_manager.py b/tests/unit_tests/storage/test_storage_manager.py index d96f1cb04..e17c317bd 100644 --- a/tests/unit_tests/storage/test_storage_manager.py +++ b/tests/unit_tests/storage/test_storage_manager.py @@ -92,6 +92,15 @@ class TestStorageMgr: await storage_mgr.initialize() mock_init.assert_called_once() + @pytest.mark.asyncio + async def test_shutdown_delegates_to_active_provider(self): + storage_mgr = StorageMgr(Mock()) + storage_mgr.storage_provider = Mock(shutdown=AsyncMock()) + + await storage_mgr.shutdown() + + storage_mgr.storage_provider.shutdown.assert_awaited_once() + class TestStorageProviderBase: """Test StorageProvider base class methods.""" diff --git a/tests/unit_tests/storage/test_workspace_scoping.py b/tests/unit_tests/storage/test_workspace_scoping.py index 7d8407248..79869a85c 100644 --- a/tests/unit_tests/storage/test_workspace_scoping.py +++ b/tests/unit_tests/storage/test_workspace_scoping.py @@ -8,6 +8,7 @@ import pytest from langbot.pkg.api.http.authz import WorkspaceRequiredError from langbot.pkg.api.http.context import ExecutionContext from langbot.pkg.storage.mgr import StorageMgr +from langbot.pkg.utils.bounded_executor import current_blocking_work_scope WORKSPACE_A = '00000000-0000-0000-0000-00000000000a' @@ -33,14 +34,17 @@ def _context_for_instance(instance_uuid: str) -> ExecutionContext: class _Provider: def __init__(self): self.values: dict[str, bytes] = {} + self.observed_public_scopes: list[str | None] = [] async def save(self, key: str, value: bytes): self.values[key] = value async def load(self, key: str) -> bytes: + self.observed_public_scopes.append(current_blocking_work_scope()) return self.values[key] async def exists(self, key: str) -> bool: + self.observed_public_scopes.append(current_blocking_work_scope()) return key in self.values async def size(self, key: str) -> int: @@ -120,6 +124,10 @@ async def test_public_object_route_derives_trusted_workspace(manager): value=b'image-a', ) assert await manager.resolve_public_object(object_key, expected_owner_type='upload') == b'image-a' + assert manager.storage_provider.observed_public_scopes[-2:] == [ + WORKSPACE_A, + WORKSPACE_A, + ] assert await manager.resolve_public_object(object_key, expected_owner_type='plugin') is None guessed_workspace_key = object_key.replace(WORKSPACE_A, WORKSPACE_B) diff --git a/tests/unit_tests/survey/test_survey_manager.py b/tests/unit_tests/survey/test_survey_manager.py index f040b19cb..da67e8219 100644 --- a/tests/unit_tests/survey/test_survey_manager.py +++ b/tests/unit_tests/survey/test_survey_manager.py @@ -29,6 +29,13 @@ def create_mock_app(): mock_app.instance_config.data = {'space': {'url': 'https://space.example.com'}} mock_app.persistence_mgr = AsyncMock() mock_app.persistence_mgr.execute_async = AsyncMock() + mock_app.workspace_service.instance_uuid = 'instance-test' + + def close_scheduled_coroutine(coro, **kwargs): + coro.close() + return Mock() + + mock_app.task_mgr.create_task = Mock(side_effect=close_scheduled_coroutine) return mock_app diff --git a/tests/unit_tests/telemetry/test_heartbeat.py b/tests/unit_tests/telemetry/test_heartbeat.py index f13aa9a39..a8dfae936 100644 --- a/tests/unit_tests/telemetry/test_heartbeat.py +++ b/tests/unit_tests/telemetry/test_heartbeat.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +from types import SimpleNamespace import pytest from unittest.mock import AsyncMock, Mock @@ -96,6 +97,35 @@ class TestBuildHeartbeatPayload: payload = await heartbeat.build_heartbeat_payload(ap) assert payload['features']['pipeline_count'] == -1 + @pytest.mark.asyncio + async def test_cloud_counts_loaded_registries_without_tenant_sql(self): + heartbeat = get_heartbeat_module() + ap = make_app() + ap.persistence_mgr.mode = SimpleNamespace(value='cloud_runtime') + ap.persistence_mgr.execute_async = AsyncMock( + side_effect=AssertionError('Cloud heartbeat must not issue per-tenant COUNTs') + ) + ap.pipeline_mgr = SimpleNamespace( + _pipelines_by_key={'pipeline-a': object(), 'pipeline-b': object()}, + ) + ap.tool_mgr = SimpleNamespace( + mcp_tool_loader=SimpleNamespace( + _sessions={'mcp-a': object(), 'mcp-b': object(), 'mcp-c': object()}, + ), + ) + ap.rag_mgr = SimpleNamespace( + knowledge_bases={'kb-a': object()}, + ) + + payload = await heartbeat.build_heartbeat_payload(ap) + + features = payload['features'] + assert features['pipeline_count'] == 2 + assert features['mcp_server_count'] == 3 + assert features['knowledge_base_count'] == 1 + assert features['bot_count'] == 1 + ap.persistence_mgr.execute_async.assert_not_awaited() + @pytest.mark.asyncio async def test_no_user_content_fields(self): """The heartbeat must never carry message content / credentials keys.""" diff --git a/tests/unit_tests/test_telemetry.py b/tests/unit_tests/test_telemetry.py index e9159337c..ed3994923 100644 --- a/tests/unit_tests/test_telemetry.py +++ b/tests/unit_tests/test_telemetry.py @@ -1,6 +1,7 @@ from __future__ import annotations from types import SimpleNamespace +import asyncio import pytest @@ -20,7 +21,9 @@ async def test_send_tasks_are_scoped_to_manager_instance(monkeypatch): assert first.send_tasks is not second.send_tasks await first.start_send_task({'event': 'first'}) - await first.send_tasks[0] + task = first.send_tasks[0] + await task + await asyncio.sleep(0) - assert len(first.send_tasks) == 1 + assert first.send_tasks == [] assert second.send_tasks == [] diff --git a/tests/unit_tests/utils/test_bounded_executor.py b/tests/unit_tests/utils/test_bounded_executor.py new file mode 100644 index 000000000..34581734d --- /dev/null +++ b/tests/unit_tests/utils/test_bounded_executor.py @@ -0,0 +1,278 @@ +from __future__ import annotations + +import asyncio +import threading + +import pytest + +from langbot.pkg.utils.bounded_executor import ( + BlockingWorkCapacityError, + BoundedThreadPoolExecutor, + blocking_work_scope, + configure_bounded_default_executor, + run_blocking_atomic, + run_blocking_cleanup, +) + + +def test_bounded_executor_rejects_instead_of_queueing_without_limit(): + executor = BoundedThreadPoolExecutor( + max_workers=1, + max_pending=1, + max_inflight_per_scope=1, + ) + started = threading.Event() + release = threading.Event() + + def block() -> str: + started.set() + release.wait(timeout=5) + return 'done' + + first = executor.submit(block) + assert started.wait(timeout=1) + second = executor.submit(lambda: 'queued') + + with pytest.raises( + BlockingWorkCapacityError, + match='capacity reached', + ): + executor.submit(lambda: 'rejected') + + assert executor.snapshot() == { + 'max_workers': 1, + 'max_pending': 1, + 'max_inflight_per_scope': 1, + 'inflight': 2, + 'running': 1, + 'pending': 1, + 'active_scopes': 0, + 'submitted_total': 2, + 'completed_total': 0, + 'rejected_total': 1, + 'global_rejected_total': 1, + 'scope_rejected_total': 0, + } + + release.set() + assert first.result(timeout=1) == 'done' + assert second.result(timeout=1) == 'queued' + assert executor.snapshot()['inflight'] == 0 + executor.shutdown() + + +def test_workspace_scope_cannot_monopolize_global_workers(): + executor = BoundedThreadPoolExecutor( + max_workers=2, + max_pending=2, + max_inflight_per_scope=1, + ) + release = threading.Event() + workspace_a_started = threading.Event() + workspace_b_started = threading.Event() + + def block(started: threading.Event) -> str: + started.set() + release.wait(timeout=5) + return 'done' + + try: + with blocking_work_scope('workspace-a'): + workspace_a = executor.submit(block, workspace_a_started) + assert workspace_a_started.wait(timeout=1) + with pytest.raises( + BlockingWorkCapacityError, + match='Workspace blocking executor capacity reached', + ): + executor.submit(lambda: 'rejected') + + with blocking_work_scope('workspace-b'): + workspace_b = executor.submit(block, workspace_b_started) + assert workspace_b_started.wait(timeout=1) + + snapshot = executor.snapshot() + assert snapshot['inflight'] == 2 + assert snapshot['active_scopes'] == 2 + assert snapshot['scope_rejected_total'] == 1 + assert snapshot['global_rejected_total'] == 0 + finally: + release.set() + assert workspace_a.result(timeout=1) == 'done' + assert workspace_b.result(timeout=1) == 'done' + executor.shutdown() + + +def test_default_executor_bounds_asyncio_to_thread(): + loop = asyncio.new_event_loop() + executor = configure_bounded_default_executor( + loop, + max_workers=2, + max_pending=3, + ) + try: + assert loop.run_until_complete(asyncio.to_thread(lambda: 'bounded')) == 'bounded' + assert executor.snapshot()['completed_total'] == 1 + assert ( + configure_bounded_default_executor( + loop, + max_workers=2, + max_pending=3, + ) + is executor + ) + finally: + executor.shutdown() + loop.close() + + +def test_workspace_scope_is_enforced_for_asyncio_to_thread(): + loop = asyncio.new_event_loop() + executor = configure_bounded_default_executor( + loop, + max_workers=2, + max_pending=2, + max_inflight_per_scope=1, + ) + started = threading.Event() + release = threading.Event() + + def block() -> str: + started.set() + release.wait(timeout=5) + return 'workspace-a' + + async def exercise() -> None: + with blocking_work_scope('workspace-a'): + workspace_a = asyncio.create_task(asyncio.to_thread(block)) + try: + while not started.is_set(): + await asyncio.sleep(0) + + with blocking_work_scope('workspace-a'): + with pytest.raises( + BlockingWorkCapacityError, + match='Workspace blocking executor capacity reached', + ): + await asyncio.to_thread(lambda: 'rejected') + + with blocking_work_scope('workspace-b'): + assert await asyncio.to_thread(lambda: 'workspace-b') == 'workspace-b' + finally: + release.set() + assert await workspace_a == 'workspace-a' + + try: + loop.run_until_complete(exercise()) + finally: + executor.shutdown() + loop.close() + + +def test_blocking_cleanup_waits_for_capacity_instead_of_leaking_work(): + loop = asyncio.new_event_loop() + executor = configure_bounded_default_executor( + loop, + max_workers=1, + max_pending=0, + max_inflight_per_scope=1, + ) + started = threading.Event() + release = threading.Event() + cleaned = threading.Event() + + def block() -> None: + started.set() + release.wait(timeout=5) + + async def exercise() -> None: + blocker = asyncio.create_task(asyncio.to_thread(block)) + while not started.is_set(): + await asyncio.sleep(0) + cleanup = asyncio.create_task(run_blocking_cleanup(cleaned.set)) + await asyncio.sleep(0.03) + assert not cleanup.done() + release.set() + await blocker + await cleanup + + try: + loop.run_until_complete(exercise()) + assert cleaned.is_set() + assert executor.snapshot()['global_rejected_total'] >= 1 + finally: + release.set() + executor.shutdown() + loop.close() + + +def test_blocking_atomic_waits_for_thread_before_propagating_cancellation(): + loop = asyncio.new_event_loop() + executor = configure_bounded_default_executor( + loop, + max_workers=1, + max_pending=1, + max_inflight_per_scope=1, + ) + started = threading.Event() + release = threading.Event() + completed = threading.Event() + + def block() -> None: + started.set() + release.wait(timeout=5) + completed.set() + + async def exercise() -> None: + operation = asyncio.create_task(run_blocking_atomic(block)) + while not started.is_set(): + await asyncio.sleep(0) + operation.cancel() + await asyncio.sleep(0) + assert not operation.done() + release.set() + with pytest.raises(asyncio.CancelledError): + await operation + + try: + loop.run_until_complete(exercise()) + assert completed.is_set() + finally: + release.set() + executor.shutdown() + loop.close() + + +@pytest.mark.parametrize( + ('max_workers', 'max_pending'), + [ + (0, 1), + (65, 1), + (1, -1), + (1, 4097), + (True, 1), + ], +) +def test_bounded_executor_rejects_unsafe_limits( + max_workers, + max_pending, +): + with pytest.raises(ValueError): + BoundedThreadPoolExecutor( + max_workers=max_workers, + max_pending=max_pending, + ) + + +@pytest.mark.parametrize( + ('max_workers', 'max_inflight_per_scope'), + [(8, 0), (8, 4097), (8, True), (8, 5), (2, 2)], +) +def test_bounded_executor_rejects_unsafe_scope_limits( + max_workers, + max_inflight_per_scope, +): + with pytest.raises(ValueError): + BoundedThreadPoolExecutor( + max_workers=max_workers, + max_inflight_per_scope=max_inflight_per_scope, + ) diff --git a/tests/unit_tests/utils/test_cloud_runtime_soak.py b/tests/unit_tests/utils/test_cloud_runtime_soak.py new file mode 100644 index 000000000..de0bf7531 --- /dev/null +++ b/tests/unit_tests/utils/test_cloud_runtime_soak.py @@ -0,0 +1,446 @@ +from __future__ import annotations + +import io +import json +from collections import deque +from pathlib import Path + +import pytest + +from scripts import cloud_runtime_soak as soak + + +def _write(path: Path, value: str) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(value, encoding='utf-8') + + +def _process_stat(pid: int, *, user_ticks: int, system_ticks: int) -> str: + fields = [ + 'S', + '0', + '0', + '0', + '0', + '0', + '0', + '0', + '0', + '0', + '0', + str(user_ticks), + str(system_ticks), + '0', + '0', + '0', + '0', + '0', + '0', + '0', + ] + return f'{pid} (worker with spaces) ' + ' '.join(fields) + + +def _sample(timestamp: float, **metrics: float) -> soak.MetricSample: + return soak.MetricSample( + monotonic_seconds=timestamp, + wall_time=f'sample-{timestamp}', + metrics=metrics, + ) + + +def _state( + kind: str, + samples: list[soak.MetricSample], + *, + baseline: dict[str, float] | None = None, + latest: dict[str, float] | None = None, +) -> soak.TargetState: + return soak.TargetState( + target=soak.Target(name='target', kind=kind, location='/target'), + samples=deque(samples), + baseline_metrics=baseline or dict(samples[0].metrics), + last_metrics=latest or dict(samples[-1].metrics), + attempted_samples=len(samples), + successful_samples=len(samples), + ) + + +def _thresholds(**overrides) -> soak.Thresholds: + values = { + 'max_memory_growth_bytes': 64 * soak.BYTES_PER_MIB, + 'max_memory_slope_bytes_per_hour': 32 * soak.BYTES_PER_MIB, + 'max_tail_cpu_cores': 0.5, + 'max_throttled_period_ratio': 0.25, + 'allow_rejections': False, + 'max_transient_gauge_growth': 0, + 'require_hard_limits': False, + 'max_event_loop_lag_ms': 1000, + 'max_event_loop_p95_lag_ms': 250, + 'require_event_loop_metrics': True, + } + values.update(overrides) + return soak.Thresholds(**values) + + +@pytest.mark.parametrize( + ('value', 'expected'), + [ + ('1', 1), + ('1.5s', 1.5), + ('2m', 120), + ('3H', 10_800), + ('1d', 86_400), + ], +) +def test_parse_duration(value: str, expected: float) -> None: + assert soak.parse_duration(value) == expected + + +@pytest.mark.parametrize('value', ['', '0s', '-1s', 'wat']) +def test_parse_duration_rejects_invalid_values(value: str) -> None: + with pytest.raises(Exception): + soak.parse_duration(value) + + +def test_build_targets_rejects_secret_bearing_health_url() -> None: + with pytest.raises(ValueError, match='must not contain credentials'): + soak.build_targets( + endpoints=['core=https://user:secret@example.test/healthz'], + cgroups=[], + pids=[], + ) + + with pytest.raises(ValueError, match='query parameters'): + soak.build_targets( + endpoints=['core=https://example.test/healthz?token=secret'], + cgroups=[], + pids=[], + ) + + +def test_read_cgroup_snapshot_reads_v2_pressure_and_limits(tmp_path: Path) -> None: + _write(tmp_path / 'memory.current', '1048576\n') + _write(tmp_path / 'memory.peak', '2097152\n') + _write(tmp_path / 'memory.swap.current', '4096\n') + _write(tmp_path / 'memory.max', '1073741824\n') + _write(tmp_path / 'memory.swap.max', 'max\n') + _write(tmp_path / 'pids.current', '7\n') + _write(tmp_path / 'pids.max', '128\n') + _write( + tmp_path / 'cpu.stat', + 'usage_usec 1234\nnr_periods 20\nnr_throttled 2\nthrottled_usec 99\n', + ) + _write(tmp_path / 'memory.events', 'high 1\nmax 2\noom 0\noom_kill 0\n') + _write(tmp_path / 'pids.events', 'max 3\n') + _write(tmp_path / 'cpu.max', '100000 100000\n') + + metrics = soak.read_cgroup_snapshot(tmp_path) + + assert metrics['memory.current_bytes'] == 1_048_576 + assert metrics['memory.max_bytes'] == 1_073_741_824 + assert 'memory.swap.max_bytes' not in metrics + assert metrics['cpu.usage_usec'] == 1234 + assert metrics['cpu.nr_throttled'] == 2 + assert metrics['memory.events.max'] == 2 + assert metrics['pids.events.max'] == 3 + assert metrics['cpu.quota_usec'] == 100_000 + + +def test_read_process_snapshot_aggregates_descendants(tmp_path: Path) -> None: + for pid, rss_kib, threads, user_ticks, system_ticks in ( + (100, 1000, 2, 100, 50), + (200, 500, 1, 20, 10), + ): + process_root = tmp_path / str(pid) + _write( + process_root / 'status', + f'Name:\tworker\nVmRSS:\t{rss_kib} kB\nThreads:\t{threads}\n', + ) + _write( + process_root / 'stat', + _process_stat( + pid, + user_ticks=user_ticks, + system_ticks=system_ticks, + ), + ) + (process_root / 'fd').mkdir() + (process_root / 'fd' / '0').touch() + (process_root / 'fd' / '1').touch() + (process_root / 'task' / str(pid)).mkdir(parents=True) + _write(tmp_path / '100' / 'task' / '100' / 'children', '200\n') + _write(tmp_path / '200' / 'task' / '200' / 'children', '\n') + + metrics = soak.read_process_snapshot(100, proc_root=tmp_path, clock_ticks=100) + + assert metrics == { + 'rss_bytes': 1500 * 1024, + 'cpu_seconds': 1.8, + 'threads': 3, + 'open_fds': 4, + 'processes': 2, + } + + +class _FakeResponse: + def __init__(self, payload: dict) -> None: + self.status = 200 + self.headers = {'Content-Type': 'application/json'} + self._body = json.dumps(payload).encode() + + def __enter__(self): + return self + + def __exit__(self, *_args): + return None + + def getcode(self) -> int: + return self.status + + def read(self, limit: int) -> bytes: + return self._body[:limit] + + +def test_read_endpoint_snapshot_flattens_resource_metrics() -> None: + def opener(_request, *, timeout: float): + assert timeout == 2 + return _FakeResponse( + { + 'code': 0, + 'resources': { + 'blocking_executor': { + 'pending': 0, + 'global_rejected_total': 2, + } + }, + } + ) + + metrics = soak.read_endpoint_snapshot( + 'http://langbot.test/healthz', + timeout_seconds=2, + opener=opener, + ) + + assert metrics['http.ok'] == 1 + assert metrics['body.resources.blocking_executor.pending'] == 0 + assert metrics['body.resources.blocking_executor.global_rejected_total'] == 2 + + +def test_read_endpoint_snapshot_fails_closed_on_not_ready() -> None: + def opener(_request, *, timeout: float): + return _FakeResponse({'ready': False}) + + with pytest.raises(RuntimeError, match='not ready'): + soak.read_endpoint_snapshot( + 'http://box.test/readyz', + timeout_seconds=2, + opener=opener, + ) + + +def test_evaluate_gate_accepts_stable_process_tail() -> None: + state = _state( + 'process', + [ + _sample(0, rss_bytes=100 * soak.BYTES_PER_MIB, cpu_seconds=0), + _sample(1800, rss_bytes=101 * soak.BYTES_PER_MIB, cpu_seconds=10), + _sample(3600, rss_bytes=100 * soak.BYTES_PER_MIB, cpu_seconds=20), + ], + ) + + result = soak.evaluate_gate( + [state], + analysis_start_seconds=0, + thresholds=_thresholds(), + ) + + assert result.passed + assert result.targets['process:target']['cpu.average_cores'] < 0.01 + + +def test_evaluate_gate_detects_material_memory_leak_and_idle_cpu() -> None: + state = _state( + 'process', + [ + _sample(0, rss_bytes=100 * soak.BYTES_PER_MIB, cpu_seconds=0), + _sample(1800, rss_bytes=150 * soak.BYTES_PER_MIB, cpu_seconds=1800), + _sample(3600, rss_bytes=200 * soak.BYTES_PER_MIB, cpu_seconds=3600), + ], + ) + + result = soak.evaluate_gate( + [state], + analysis_start_seconds=0, + thresholds=_thresholds(), + ) + + assert not result.passed + assert any('grew 100.00 MiB' in failure for failure in result.failures) + assert any('tail CPU averaged 1.000 cores' in failure for failure in result.failures) + + +def test_evaluate_gate_counts_oom_and_throttling_across_workload() -> None: + baseline = { + 'memory.current_bytes': 100, + 'memory.events.high': 0, + 'memory.events.max': 0, + 'memory.events.oom': 0, + 'memory.events.oom_kill': 0, + 'pids.events.max': 0, + 'cpu.usage_usec': 0, + 'cpu.nr_periods': 0, + 'cpu.nr_throttled': 0, + } + latest = { + **baseline, + 'memory.events.oom_kill': 1, + 'pids.events.max': 2, + 'cpu.usage_usec': 1_000_000, + 'cpu.nr_periods': 100, + 'cpu.nr_throttled': 30, + } + state = _state( + 'cgroup', + [ + _sample(100, **{**baseline, 'cpu.usage_usec': 500_000}), + _sample(200, **latest), + ], + baseline=baseline, + latest=latest, + ) + + result = soak.evaluate_gate( + [state], + analysis_start_seconds=100, + thresholds=_thresholds(max_tail_cpu_cores=100), + ) + + assert any('memory.events.oom_kill by 1' in failure for failure in result.failures) + assert any('pids.events.max by 2' in failure for failure in result.failures) + assert any('throttled-period ratio 0.300' in failure for failure in result.failures) + + +def test_evaluate_gate_can_require_all_hard_cgroup_limits() -> None: + metrics = { + 'memory.current_bytes': 100, + 'memory.max_bytes': 1000, + 'memory.events.high': 0, + 'memory.events.max': 0, + 'memory.events.oom': 0, + 'memory.events.oom_kill': 0, + 'pids.current': 1, + 'pids.events.max': 0, + 'cpu.usage_usec': 0, + 'cpu.nr_periods': 0, + 'cpu.nr_throttled': 0, + } + state = _state( + 'cgroup', + [_sample(0, **metrics), _sample(60, **metrics)], + ) + + result = soak.evaluate_gate( + [state], + analysis_start_seconds=0, + thresholds=_thresholds(require_hard_limits=True), + ) + + assert any('missing hard cgroup limits: cpu, pids, swap' in failure for failure in result.failures) + + +def test_evaluate_gate_detects_executor_rejection_and_stuck_pending() -> None: + prefix = 'body.resources.blocking_executor' + state = _state( + 'endpoint', + [ + _sample( + 0, + **{ + f'{prefix}.pending': 1, + f'{prefix}.global_rejected_total': 0, + }, + ), + _sample( + 60, + **{ + f'{prefix}.pending': 2, + f'{prefix}.global_rejected_total': 1, + }, + ), + ], + ) + + result = soak.evaluate_gate( + [state], + analysis_start_seconds=0, + thresholds=_thresholds(), + ) + + assert any('global_rejected_total by 1' in failure for failure in result.failures) + assert any('pending above zero' in failure for failure in result.failures) + + +def test_evaluate_gate_detects_event_loop_stall_and_sustained_lag() -> None: + prefix = 'body.resources.event_loop' + samples = [ + _sample( + 0, + **{ + f'{prefix}.running': 1, + f'{prefix}.samples_total': 10, + f'{prefix}.recent_max_lag_ms': 20, + f'{prefix}.recent_p95_lag_ms': 10, + }, + ), + _sample( + 60, + **{ + f'{prefix}.running': 1, + f'{prefix}.samples_total': 70, + f'{prefix}.recent_max_lag_ms': 1500, + f'{prefix}.recent_p95_lag_ms': 300, + }, + ), + ] + state = _state('endpoint', samples) + state.observed_max_metrics = { + metric: max(sample.metrics[metric] for sample in samples) for metric in samples[0].metrics + } + + result = soak.evaluate_gate( + [state], + analysis_start_seconds=0, + thresholds=_thresholds(), + ) + + assert any('event-loop lag reached 1500.00 ms' in item for item in result.failures) + assert any('recent p95 reached 300.00 ms' in item for item in result.failures) + + +def test_evaluate_gate_requires_running_event_loop_monitor() -> None: + state = _state( + 'endpoint', + [_sample(0, **{'http.ok': 1}), _sample(60, **{'http.ok': 1})], + ) + + result = soak.evaluate_gate( + [state], + analysis_start_seconds=0, + thresholds=_thresholds(), + ) + + assert any('did not expose event-loop health metrics' in item for item in result.failures) + + +def test_write_json_line_streams_one_record() -> None: + stream = io.StringIO() + soak._write_json_line(stream, {'z': 1, 'a': 2}) + assert stream.getvalue() == '{"a":2,"z":1}\n' + + +def test_main_requires_a_target() -> None: + with pytest.raises(SystemExit) as exc_info: + soak.main(['--duration', '2s', '--sample-interval', '1s', '--startup-grace', '1s']) + assert exc_info.value.code == 2 diff --git a/tests/unit_tests/utils/test_event_loop_monitor.py b/tests/unit_tests/utils/test_event_loop_monitor.py new file mode 100644 index 000000000..f80186d94 --- /dev/null +++ b/tests/unit_tests/utils/test_event_loop_monitor.py @@ -0,0 +1,75 @@ +from __future__ import annotations + +import asyncio +import time + +import pytest + +from langbot.pkg.utils.event_loop_monitor import EventLoopLagMonitor + + +@pytest.mark.parametrize( + 'kwargs', + [ + {'sample_interval_seconds': 0}, + {'sample_interval_seconds': float('inf')}, + {'recent_sample_count': 1}, + {'recent_sample_count': 3601}, + ], +) +def test_event_loop_monitor_rejects_unbounded_configuration(kwargs) -> None: + with pytest.raises(ValueError): + EventLoopLagMonitor(**kwargs) + + +def test_event_loop_monitor_snapshot_is_bounded_and_reports_p95() -> None: + monitor = EventLoopLagMonitor(recent_sample_count=4) + for lag_seconds in (0.001, 0.002, 0.003, 0.004, 0.100): + monitor._record_lag_seconds(lag_seconds) + + snapshot = monitor.snapshot() + + assert snapshot == { + 'running': False, + 'samples_total': 5, + 'last_lag_ms': 100, + 'recent_p95_lag_ms': 100, + 'recent_max_lag_ms': 100, + 'max_lag_ms': 100, + } + assert len(monitor._recent_lag_ms) == 4 + + +async def test_event_loop_monitor_start_and_stop_are_idempotent() -> None: + monitor = EventLoopLagMonitor( + sample_interval_seconds=0.001, + recent_sample_count=4, + ) + + monitor.start() + task = monitor._task + monitor.start() + assert monitor._task is task + await asyncio.sleep(0.005) + assert monitor.snapshot()['samples_total'] > 0 + assert monitor.snapshot()['running'] is True + + await monitor.stop() + await monitor.stop() + assert monitor.snapshot()['running'] is False + assert task is not None and task.done() + + +async def test_event_loop_monitor_observes_real_scheduler_stall() -> None: + monitor = EventLoopLagMonitor( + sample_interval_seconds=0.005, + recent_sample_count=8, + ) + monitor.start() + try: + await asyncio.sleep(0.01) + time.sleep(0.05) + await asyncio.sleep(0.01) + assert monitor.snapshot()['recent_max_lag_ms'] >= 35 + finally: + await monitor.stop() diff --git a/tests/unit_tests/utils/test_httpclient.py b/tests/unit_tests/utils/test_httpclient.py index da110afe2..c045f07b3 100644 --- a/tests/unit_tests/utils/test_httpclient.py +++ b/tests/unit_tests/utils/test_httpclient.py @@ -6,8 +6,13 @@ Tests session management, reuse, and cleanup. from __future__ import annotations +import threading +from types import SimpleNamespace +from unittest.mock import AsyncMock + import pytest import aiohttp +import httpx from aiohttp import web from langbot.pkg.utils import httpclient @@ -88,6 +93,89 @@ class TestCloseAll: assert len(httpclient._sessions) == 0 + +class TestReadLimited: + async def test_rejects_oversized_content_length_before_reading(self): + content = SimpleNamespace(iter_chunked=None) + response = SimpleNamespace(headers={'Content-Length': '11'}, content=content) + + with pytest.raises(httpclient.RemoteResponseTooLargeError): + await httpclient.read_limited(response, max_bytes=10) + + async def test_rejects_chunked_body_that_crosses_limit(self): + class Content: + async def iter_chunked(self, _chunk_size): + yield b'12345' + yield b'678901' + + response = SimpleNamespace(headers={}, content=Content()) + + with pytest.raises(httpclient.RemoteResponseTooLargeError): + await httpclient.read_limited(response, max_bytes=10) + + async def test_returns_body_within_limit(self): + class Content: + async def iter_chunked(self, _chunk_size): + yield b'12345' + yield b'67890' + + response = SimpleNamespace(headers={}, content=Content()) + + assert await httpclient.read_limited(response, max_bytes=10) == b'1234567890' + + async def test_json_reader_uses_same_limit(self): + class Content: + async def iter_chunked(self, _chunk_size): + yield b'{"ok":true}' + + response = SimpleNamespace( + headers={}, + content=Content(), + ) + + assert await httpclient.read_json_limited(response, max_bytes=16) == {'ok': True} + + async def test_response_json_parse_runs_off_event_loop(self): + event_loop_thread = threading.get_ident() + response = SimpleNamespace(json=lambda: threading.get_ident()) + + assert await httpclient.parse_json_response(response) != event_loop_thread + + async def test_response_json_parse_supports_async_test_doubles(self): + response = SimpleNamespace(json=AsyncMock(return_value={'ok': True})) + + assert await httpclient.parse_json_response(response) == {'ok': True} + + async def test_response_text_runs_off_loop_and_caps_diagnostics(self): + event_loop_thread = threading.get_ident() + + class Response: + @property + def text(self): + return f'{threading.get_ident()}:abcdef' + + value = await httpclient.response_text(Response(), max_chars=4) + + assert not value.startswith(str(event_loop_thread)) + assert value.endswith('[truncated]') + + async def test_httpx_hook_rejects_before_automatic_buffer_grows(self): + class Source(httpx.AsyncByteStream): + async def __aiter__(self): + yield b'123' + yield b'45' + + async def aclose(self): + return None + + transport = httpx.MockTransport(lambda _request: httpx.Response(200, stream=Source())) + async with httpx.AsyncClient( + transport=transport, + event_hooks=httpclient.httpx_response_limit_hooks(max_bytes=4), + ) as client: + with pytest.raises(httpclient.RemoteResponseTooLargeError, match='4-byte'): + await client.get('https://example.invalid') + async def test_close_all_handles_already_closed(self): """close_all handles already closed sessions gracefully.""" session = httpclient.get_session() diff --git a/tests/unit_tests/utils/test_image.py b/tests/unit_tests/utils/test_image.py index 4a42717ba..5c4f9eb50 100644 --- a/tests/unit_tests/utils/test_image.py +++ b/tests/unit_tests/utils/test_image.py @@ -10,11 +10,28 @@ import pytest import base64 from langbot.pkg.utils.image import ( + decode_base64_limited, + encode_base64, get_qq_image_downloadable_url, extract_b64_and_format, ) +@pytest.mark.asyncio +async def test_base64_media_helpers_round_trip_within_limit(): + encoded = await encode_base64(b'1234') + + assert await decode_base64_limited(encoded, max_bytes=4) == b'1234' + + +@pytest.mark.asyncio +async def test_base64_media_decode_rejects_oversized_payload(): + encoded = base64.b64encode(b'12345').decode() + + with pytest.raises(ValueError, match='exceeds'): + await decode_base64_limited(encoded, max_bytes=4) + + class TestGetQQImageDownloadableUrl: """Tests for get_qq_image_downloadable_url function.""" diff --git a/tests/unit_tests/utils/test_logcache.py b/tests/unit_tests/utils/test_logcache.py index ed05d0ccd..775a3ed27 100644 --- a/tests/unit_tests/utils/test_logcache.py +++ b/tests/unit_tests/utils/test_logcache.py @@ -7,7 +7,13 @@ Tests log page management and pointer-based retrieval. from __future__ import annotations -from langbot.pkg.utils.logcache import LogPage, LogCache, LOG_PAGE_SIZE, MAX_CACHED_PAGES +from langbot.pkg.utils.logcache import ( + LogPage, + LogCache, + LOG_PAGE_SIZE, + MAX_CACHED_PAGES, + MAX_LOG_LINE_CHARS, +) class TestLogPage: @@ -208,3 +214,11 @@ class TestLogCache: """LOG_PAGE_SIZE is defined and reasonable.""" assert LOG_PAGE_SIZE > 0 assert LOG_PAGE_SIZE <= 1000 # Reasonable upper bound + + def test_single_log_line_is_bounded(self): + cache = LogCache() + + cache.add_log('x' * (MAX_LOG_LINE_CHARS * 2)) + + assert len(cache.log_pages[0].logs[0]) == MAX_LOG_LINE_CHARS + assert cache.log_pages[0].logs[0].endswith('[log truncated]') diff --git a/tests/unit_tests/utils/test_safe_regex.py b/tests/unit_tests/utils/test_safe_regex.py new file mode 100644 index 000000000..6f0f985f5 --- /dev/null +++ b/tests/unit_tests/utils/test_safe_regex.py @@ -0,0 +1,73 @@ +from __future__ import annotations + +import threading + +import pytest + +from langbot.pkg.utils import safe_regex +from langbot.pkg.utils.bounded_executor import blocking_work_scope, current_blocking_work_scope + + +@pytest.mark.asyncio +async def test_matches_any_runs_off_event_loop_and_preserves_workspace_scope(monkeypatch): + event_loop_thread = threading.get_ident() + observed: dict[str, object] = {} + original = safe_regex._matches_any_sync + + def observe(*args, **kwargs): + observed['thread'] = threading.get_ident() + observed['scope'] = current_blocking_work_scope() + return original(*args, **kwargs) + + monkeypatch.setattr(safe_regex, '_matches_any_sync', observe) + + with blocking_work_scope('workspace-a'): + assert await safe_regex.matches_any(['^hello'], 'hello world') is True + + assert observed['scope'] == 'workspace-a' + assert observed['thread'] != event_loop_thread + + +@pytest.mark.asyncio +async def test_matches_any_interrupts_catastrophic_backtracking(): + with pytest.raises(safe_regex.SafeRegexTimeoutError): + await safe_regex.matches_any( + [r'(a+)+$'], + ('a' * 100_000) + '!', + timeout_seconds=0.001, + ) + + +@pytest.mark.asyncio +async def test_matches_any_rejects_pattern_and_input_amplification(): + with pytest.raises(safe_regex.SafeRegexLimitError): + await safe_regex.matches_any( + ['a'] * (safe_regex.MAX_PATTERN_COUNT + 1), + 'a', + ) + + with pytest.raises(safe_regex.SafeRegexLimitError): + await safe_regex.matches_any( + ['a'], + 'a' * (safe_regex.MAX_INPUT_CHARS + 1), + ) + + +@pytest.mark.asyncio +async def test_mask_patterns_bounds_replacement_growth_and_masks_matches(): + found, masked = await safe_regex.mask_patterns( + [r'secret-\d+'], + 'a secret-42 value', + mask='*', + mask_word='[hidden]', + ) + assert found is True + assert masked == 'a [hidden] value' + + with pytest.raises(safe_regex.SafeRegexLimitError): + await safe_regex.mask_patterns( + ['a'], + 'a' * safe_regex.MAX_INPUT_CHARS, + mask='0123456789', + mask_word='', + ) diff --git a/tests/unit_tests/vector/test_mgr.py b/tests/unit_tests/vector/test_mgr.py index 17193ac69..608f35121 100644 --- a/tests/unit_tests/vector/test_mgr.py +++ b/tests/unit_tests/vector/test_mgr.py @@ -6,7 +6,9 @@ based on configuration, without actually creating real VDB instances. from __future__ import annotations -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock + +import pytest from tests.utils.import_isolation import isolated_sys_modules @@ -383,3 +385,21 @@ class TestVectorDBManagerProxies: result = mgr.get_supported_search_types() assert result == ['vector', 'full_text'] + + @pytest.mark.asyncio + async def test_shutdown_closes_backend_and_releases_reference(self): + mock_app = MagicMock() + mocks = {'langbot.pkg.core.app': MagicMock()} + + with isolated_sys_modules(mocks): + from langbot.pkg.vector.mgr import VectorDBManager + + mgr = VectorDBManager(mock_app) + backend = MagicMock() + backend.close = AsyncMock() + mgr.vector_db = backend + + await mgr.shutdown() + + backend.close.assert_awaited_once_with() + assert mgr.vector_db is None diff --git a/tests/unit_tests/vector/test_valkey_search_filter.py b/tests/unit_tests/vector/test_valkey_search_filter.py index f2f35f074..19c5bf138 100644 --- a/tests/unit_tests/vector/test_valkey_search_filter.py +++ b/tests/unit_tests/vector/test_valkey_search_filter.py @@ -31,6 +31,7 @@ def make_backend(): # _ensure_client serializes creation through this lock; set it here since # __init__ (which normally creates it) is bypassed. backend._client_lock = asyncio.Lock() + backend._runtime_cache_limit = 1024 return backend diff --git a/tests/unit_tests/vector/test_vdb_base.py b/tests/unit_tests/vector/test_vdb_base.py index 427df9f19..662ed022b 100644 --- a/tests/unit_tests/vector/test_vdb_base.py +++ b/tests/unit_tests/vector/test_vdb_base.py @@ -5,7 +5,12 @@ from __future__ import annotations from unittest.mock import AsyncMock import pytest -from langbot.pkg.vector.vdb import SearchType, VectorDatabase +from langbot.pkg.vector.vdb import ( + SearchType, + VectorDatabase, + remember_bounded_mapping, + remember_bounded_set, +) class TestSearchType: @@ -29,6 +34,18 @@ class TestSearchType: assert SearchType('hybrid') == SearchType.HYBRID +def test_runtime_cache_helpers_bound_mapping_and_set(): + mapping = {} + values = set() + + for index in range(100): + remember_bounded_mapping(mapping, str(index), object(), 8) + remember_bounded_set(values, str(index), 8) + + assert len(mapping) == 8 + assert len(values) == 8 + + class TestVectorDatabaseAbstractMethods: """Tests for VectorDatabase abstract methods.""" diff --git a/tests/unit_tests/workspace/test_workspace_service.py b/tests/unit_tests/workspace/test_workspace_service.py index 32f8875e8..8657998a3 100644 --- a/tests/unit_tests/workspace/test_workspace_service.py +++ b/tests/unit_tests/workspace/test_workspace_service.py @@ -2,6 +2,7 @@ from __future__ import annotations import uuid from types import SimpleNamespace +from unittest.mock import AsyncMock import pytest import sqlalchemy @@ -18,6 +19,7 @@ from langbot.pkg.entity.persistence.workspace import ( ) from langbot.pkg.workspace import ( WorkspaceExecutionUnavailableError, + WorkspaceExecutionBinding, WorkspaceGenerationMismatchError, WorkspaceInvariantError, WorkspaceLimitExceededError, @@ -235,3 +237,28 @@ async def test_cloud_policy_never_creates_or_guesses_a_workspace(workspace_test_ await cloud_service.create_local_workspace(name='Forbidden', slug='forbidden') assert cloud_service.policy.multi_workspace_enabled is True + + +async def test_startup_binding_snapshot_avoids_repeated_discovery(): + service = WorkspaceService( + SimpleNamespace(persistence_mgr=SimpleNamespace()), + policy=CloudWorkspacePolicy(), + instance_uuid='instance-service-test', + ) + binding = WorkspaceExecutionBinding( + instance_uuid='instance-service-test', + workspace_uuid='workspace-a', + placement_generation=1, + write_fenced=False, + state='active', + ) + service._discover_active_execution_bindings = AsyncMock(return_value=[binding]) + + assert await service.prime_startup_execution_bindings() == [binding] + assert await service.list_active_execution_bindings() == [binding] + assert await service.list_active_execution_bindings() == [binding] + service._discover_active_execution_bindings.assert_awaited_once() + + service.release_startup_execution_bindings() + assert await service.list_active_execution_bindings() == [binding] + assert service._discover_active_execution_bindings.await_count == 2 diff --git a/uv.lock b/uv.lock index 9c8c81c12..7dee07888 100644 --- a/uv.lock +++ b/uv.lock @@ -2065,6 +2065,7 @@ dependencies = [ { name = "qrcode" }, { name = "quart" }, { name = "quart-cors" }, + { name = "regex" }, { name = "requests" }, { name = "ruff" }, { name = "slack-sdk" }, @@ -2115,7 +2116,7 @@ requires-dist = [ { name = "ebooklib", specifier = ">=0.18" }, { name = "gewechat-client", specifier = ">=0.1.5" }, { name = "html2text", specifier = ">=2024.2.26" }, - { name = "langbot-plugin", git = "https://github.com/langbot-app/langbot-plugin-sdk.git?rev=0cddf3c2bea5939c67b71e488a719e9903c28d17" }, + { name = "langbot-plugin", git = "https://github.com/langbot-app/langbot-plugin-sdk.git?rev=67be7c332ded4c1afa131b28232590cf9728f817" }, { name = "langchain", specifier = ">=1.3.9" }, { name = "langchain-core", specifier = ">=1.3.3" }, { name = "langchain-text-splitters", specifier = ">=1.1.2" }, @@ -2154,6 +2155,7 @@ requires-dist = [ { name = "qrcode", specifier = ">=7.4" }, { name = "quart", specifier = ">=0.20.0" }, { name = "quart-cors", specifier = ">=0.8.0" }, + { name = "regex", specifier = ">=2026.1.15" }, { name = "requests", specifier = ">=2.33.0" }, { name = "ruff", specifier = ">=0.11.9" }, { name = "slack-sdk", specifier = ">=3.35.0" }, @@ -2181,7 +2183,7 @@ dev = [ [[package]] name = "langbot-plugin" version = "0.4.18" -source = { git = "https://github.com/langbot-app/langbot-plugin-sdk.git?rev=0cddf3c2bea5939c67b71e488a719e9903c28d17#0cddf3c2bea5939c67b71e488a719e9903c28d17" } +source = { git = "https://github.com/langbot-app/langbot-plugin-sdk.git?rev=67be7c332ded4c1afa131b28232590cf9728f817#67be7c332ded4c1afa131b28232590cf9728f817" } dependencies = [ { name = "aiofiles" }, { name = "aiohttp" },