修复生产发布内存持续增长 (#203)
修复 release 内存持续增长问题。 本次范围： - 收口备份扫描与历史维护的内存峰值。 - 限制外部生成 worker 脱管任务的实际并发。 - 为 API 内存态和历史数据增加有界留存。 验收： - 定向测试、cargo 检查和生产运维门禁通过。 - 备份不再触发全局 OOM。 - worker/API/SpacetimeDB 内存曲线在空闲期停止单调增长。 Reviewed-on: http://192.168.35.82/git/GenarrativeAI/Genarrative/pulls/203 Co-authored-by: kdletters <kdletters@qq.com> Co-committed-by: kdletters <kdletters@qq.com>
This commit was merged in pull request #203.
This commit is contained in:
@@ -11,6 +11,17 @@ WorkingDirectory=/opt/genarrative/current
|
||||
EnvironmentFile=/etc/genarrative/api-server.env
|
||||
ExecStart=/usr/bin/node -- /opt/genarrative/current/scripts/database-backup-to-oss.mjs --env-file /etc/genarrative/api-server.env --stop-service spacetimedb.service --restart-service-after genarrative-api.service --restart-service-after genarrative-external-generation-worker@1.service --restart-service-after genarrative-external-generation-controller.service
|
||||
|
||||
# 备份脚本必须受独立内存上限保护，不能因目录扫描异常拖垮整台 release 主机。
|
||||
Environment=NODE_OPTIONS=--max-old-space-size=768
|
||||
Environment=GENARRATIVE_DATABASE_BACKUP_STOP_MARKER=/var/lib/genarrative/database-backups/.spacetimedb-stopped
|
||||
MemoryHigh=768M
|
||||
MemoryMax=1G
|
||||
OOMPolicy=stop
|
||||
|
||||
# 主进程可能在停库后被 MemoryMax/OOMPolicy 强制终止，JS finally 无法执行；
|
||||
# 仅当备份脚本留下停库 marker 且本次 service 非正常成功时，由 systemd 兜底恢复全部依赖服务。
|
||||
ExecStopPost=/bin/sh -c 'if [ "${SERVICE_RESULT}" != "success" ] && [ -f "${GENARRATIVE_DATABASE_BACKUP_STOP_MARKER}" ]; then systemctl start spacetimedb.service; systemctl restart genarrative-api.service; systemctl restart genarrative-external-generation-worker@1.service; systemctl restart genarrative-external-generation-controller.service; if systemctl is-active --quiet spacetimedb.service && systemctl is-active --quiet genarrative-api.service && systemctl is-active --quiet genarrative-external-generation-worker@1.service && systemctl is-active --quiet genarrative-external-generation-controller.service; then rm -f "${GENARRATIVE_DATABASE_BACKUP_STOP_MARKER}"; fi; fi'
|
||||
|
||||
# 备份需要停止 / 启动 spacetimedb.service，并读取 /stdb、写入 /var/lib/genarrative/database-backups。
|
||||
# 停止 SpacetimeDB 会连带停止 Requires 它的 API / worker / controller，冷备份后必须显式拉起。
|
||||
PrivateTmp=true
|
||||
|
||||
@@ -16,6 +16,14 @@
|
||||
|
||||
---
|
||||
|
||||
## 2026-08-27 外部生成历史采用受控保留清理
|
||||
|
||||
- 背景：`external_generation_job`、`external_generation_job_summary` 与 `external_generation_job_event` 都是持久化表；摘要和 payload 边界收紧后，已确认的终态历史仍会继续占用 SpacetimeDB 常驻内存，且事件审计链会随任务数量增长。
|
||||
- 决策：新增仅 migration operator 可调用的 `prune_external_generation_job_history_and_return`。默认按 `source_module=editor-canvas`、30 天保留期和 `job_id` 游标分批运行；只删除主任务与摘要状态一致、属于 completed / failed / cancelled、摘要已有 `notification_acknowledged_at` 且终态时间达到 cutoff 的任务，并在同一事务内删除该任务的全部事件、摘要和主任务。默认 dry-run，必须固定 dry-run 返回的 cutoff 后再 apply；pending / running、未确认通知、摘要缺失或状态不一致的数据永不删除。其他 source module 必须显式指定并单独评估；资产对象和钱包流水不随任务历史删除；不新增自动定时器或 runtime 清理权限。
|
||||
- 影响范围：`server-rs/crates/spacetime-module/src/external_generation.rs`、外部生成事件 job_id 单列索引、SpacetimeDB 生成 bindings、`scripts/spacetime-maintain-external-generation-jobs.mjs`、架构与生产运维文档。
|
||||
- 验证方式：覆盖终态 / 活跃态 / 已确认与未确认摘要、状态或身份不一致、cutoff 边界测试；运行 SpacetimeDB module tests/check、bindings 生成、schema/encoding/diff 门禁，并在维护窗口先 dry-run 再 apply。
|
||||
- 关联文档：`docs/【后端架构】server-rs与SpacetimeDB数据契约-2026-05-15.md`、`docs/【开发运维】本地开发验证与生产运维-2026-05-15.md`、PR #203。
|
||||
|
||||
## 2026-08-27 SpacetimeDB 工具链统一升级到 2.8.3
|
||||
|
||||
- 背景：SpacetimeDB 2.8.0 引入 TypeScript submodule 与调度延迟观测，2.8.1 修复 v1 WebSocket 订阅移除死锁、TypeScript SDK `array<u8>` 读缓存别名和 Rust string 默认值支持，2.8.2 修复 table accessor 改名自动迁移，2.8.3 修复 scheduled function 从实际执行时间重排导致的长期漂移。仓库若继续锁定 2.7.0，会保留这些已知运行时与 SDK 问题。
|
||||
@@ -1317,8 +1325,8 @@ CI 上 `background_agent_runtime_recovers_stale_running_before_pending_task` 在
|
||||
- OSS 固定恢复入口为 `<prefix>/<database>/latest.json`。CAS 文件和 full/history catalog 保持不可变；latest pointer 只保存最新 full catalog 与已发布 history catalog 的 object key、长度和 SHA，不包含主机绝对路径或文件内容。每次 state 变化先验真全部引用 catalog，再覆盖上传并 HEAD 验真 latest pointer，成功后才落本地 state；history 还必须在 pointer 成功后才允许删除源文件。全新机器可仅凭 bucket、database、prefix 与 OSS 凭据自动下载 pointer 和 full catalog。
|
||||
- dev 带宽不足时，允许把已冻结的 dev 基线经 `10.2.0.10 -> 10.2.4.16` 内网 rsync 到 release 独立 staging，再用 release 出口上传 dev bucket；staging 不得指向 release `/stdb`，不得停止或修改 release 服务，传输凭据必须临时创建并在演练后移除。catalog 不记录 staging 绝对路径，files state 可回传 dev 继续 history。
|
||||
- 恢复边界：恢复时默认从 OSS `latest.json` 自动定位 full catalog，创建目录并按相对路径下载每个对象、逐文件校验长度与 SHA；本地 state 只用于备份续跑，不再是异机恢复前置条件。远程 dev 已完成真实 OSS、清理、重启和异机隔离恢复演练；release timer 与 publish 前备份继续保持原行为。
|
||||
- systemd 接线：主 service 保持 `archive-full`。Server-Provision 新增默认值为 `archive-full` 的 `DATABASE_BACKUP_PROFILE`；dev 或 release 显式选择 `files-history` 时，必须为各自主机指定独立 work-dir，并先用 current release 脚本执行 history dry-run，确认已有 full state 后才安装仓库托管 drop-in，并删除现场手写旧 drop-in。切回默认 profile 必须删除所有 history 覆盖。
|
||||
- 影响范围：`scripts/database-backup-to-oss.mjs`、备份门禁、生产 env 示例、systemd 模板、Server-Provision、SpacetimeDB 运维与恢复流程；release timer 可在独立 baseline 验证后显式选择 profile，publish 前备份是否切换仍需单独决策。
|
||||
- systemd 接线：主 service 保持 `archive-full`。Server-Provision 新增默认值为 `archive-full` 的 `DATABASE_BACKUP_PROFILE`；development 可显式选择 `files-history`，必须指定独立 work-dir 并先用 current release 脚本执行 history dry-run，确认已有 full state 后才安装仓库托管 drop-in；release 拒绝 `files-history`，直到流式 catalog 改造完成，以免大目录扫描再次触发 Node 内存峰值。切回默认 profile 必须删除所有 history 覆盖；备份 unit 同时设置 Node heap 与 systemd memory 上限，避免备份异常拖垮业务主机。
|
||||
- 影响范围：`scripts/database-backup-to-oss.mjs`、备份门禁、生产 env 示例、systemd 模板、Server-Provision、SpacetimeDB 运维与恢复流程；release timer 固定使用 archive-full，publish 前备份是否切换仍需单独决策。
|
||||
- 验证方式：`npm run check:database-backup`、`npm run check:production-ops`、`npm run check:encoding`、`git diff --check`；dev 现场必须完成逐文件 full catalog、重复 full 零 PUT、history dry-run、上传后清理、STDB 重启和按 catalog 隔离恢复 roundtrip。
|
||||
- 关联：<https://github.com/clockworklabs/SpacetimeDB/issues/5542#issuecomment-4981566448>。
|
||||
|
||||
@@ -7763,3 +7771,11 @@ CI 上 `background_agent_runtime_recovers_stale_running_before_pending_task` 在
|
||||
- 决策：`.codex/skills/` 下的 SpacetimeDB 指导收敛为单一 `.codex/skills/genarrative-spacetimedb/SKILL.md`。官方 `spacetimedb` 插件负责通用 concepts、Rust server、CLI、TypeScript client 和 MCP 知识；项目 skill 只保留 Genarrative 的架构边界、schema/migration 门禁、目标 server 安全规则、运行时排障和验证路径。
|
||||
- 路由：涉及 SpacetimeDB 的任务统一先读取项目适配 skill，再按需读取 `spacetimedb:concepts`、`spacetimedb:rust-server`、`spacetimedb:cli`、`spacetimedb:typescript-client` 或 `spacetimedb:mcp`。插件通用示例不得覆盖项目禁止 `maincloud`、禁止人工 `spacetime --root-dir`、显式 server 和后端分层等约束。
|
||||
- 安装：团队环境缺少插件时使用 `codex plugin marketplace add clockworklabs/SpacetimeDB --sparse .agents --sparse codex-plugin` 和 `codex plugin add spacetimedb\@spacetimedb-plugins`；个人配置、缓存和凭据不进入仓库。
|
||||
|
||||
## 2026-08-27 release 内存增长修复与备份 OOM 恢复兜底
|
||||
|
||||
- 决策：外部生成 worker 每轮主动 `try_join_next` 回收已完成 `JoinHandle`，避免持续有队列任务时只归还 semaphore permit 却让 `JoinSet` 句柄集合无界增长；超时脱管任务在 abort 后等待句柄结束，执行许可保持到 work 真正结束或被取消。
|
||||
- 决策：`module-ai` 的阶段终态写入与流式增量统一受文本、结构化 JSON、warning 和全局 retained 工作集上限约束；认证投影恢复对过滤后的 refresh session 重新计数，超过 8192 条直接拒绝启动恢复，避免超限快照灌入内存。
|
||||
- 复审补充：`spacetime-module` 的 AI procedure 复用同一组任务元数据 / payload / 输出 / 结果引用上限，流式文本聚合超限时回滚事务，terminal task 收口后删除 `ai_text_chunk` 明细，避免真实持久化链路绕过内存边界。
|
||||
- 决策：备份脚本停库前写入受保护 `.spacetimedb-stopped` marker，正常恢复完成后清理；systemd 备份 service 通过 `MemoryHigh/MemoryMax/OOMPolicy` 和 `ExecStopPost` 在 Node OOM kill、无法执行 JS finally 时兜底拉起 SpacetimeDB、API、worker、controller，恢复不完整则保留 marker。
|
||||
- 验证：worker/module-auth/module-ai 定向 Rust 测试、database-backup/production-ops/encoding 门禁和 `git diff --check` 必须在提交前通过；release 现场需按 archive-full 重新 provision 并核验旧 files-history drop-in 已删除、四个服务 active。
|
||||
|
||||
@@ -329,6 +329,8 @@ Responses 的终态载荷既是工具调用的恢复源，也是正文的恢复
|
||||
|
||||
- Rust 结构体：`AiTask`
|
||||
- 源码：`server-rs/crates/spacetime-module/src/ai/tasks.rs`
|
||||
- `module-ai` 的进程内热状态不是持久化真相：文本增量按阶段有序聚合并受单阶段 512 KiB 上限约束；terminal task 立即释放增量明细，内存工作集最多保留 1024 个任务。需要长期查询时必须读取 SpacetimeDB 的 `ai_task` / `ai_task_stage` 投影，不得依赖进程重启后仍存在的内存快照。
|
||||
- SpacetimeDB 的 AI 写入 procedure 必须复用同一组任务元数据、payload、文本、结构化输出、warning、失败消息和结果引用上限；流式聚合超过 512 KiB 时在事务内拒绝，terminal task 收口后删除 `ai_text_chunk` 明细，只保留阶段最终快照和结果引用。
|
||||
|
||||
### `ai_task_event`
|
||||
|
||||
@@ -360,13 +362,14 @@ Responses 的终态载荷既是工具调用的恢复源，也是正文的恢复
|
||||
- 源码：`server-rs/crates/spacetime-module/src/external_generation.rs`
|
||||
- 用途：外部生成正式任务列表的轻量投影，按 `job_id` 保存 owner、来源、状态、可选 `phase`、价格、有界错误摘要、通知确认时间、各阶段时间和入队时提取的 `request_prompt`，不包含 request/result payload、worker lease 或 dedupe 内部字段。错误摘要统一拒绝内联媒体并限制为 2048 字符；列表在单次 owner 扫描中同时计数并只保留请求 limit 的固定大小 top-N，不得先收集全量历史再截断。enqueue、claim、renew、phase update、complete、fail 事务同步投影；acknowledge 只更新该轻量表并写审计事件，后续主任务同步必须保留已有确认时间，禁止为了写确认时间加载 / 重写大 payload 行。BFF 的列表、状态和确认只调用 summary procedure；`running + processing` 映射为“正在处理”，其它 running（含旧行 `phase=None`）映射为“正在生成”。历史终态任务由迁移操作员的游标分批 maintenance procedure 在压缩 payload 时同步回填摘要，正式列表不得为兼容旧数据回扫完整主表。
|
||||
- 非阻断告警：摘要字段 `warning_message` 是展示投影，由完成任务的轻量 `result_payload_json.warning.reason` 原样提取，不等同于公开 inline / external v1 的原始结构化诊断字段。complete 和历史 backfill 共用同一构建路径；历史任务按其结果载荷中已写入的 `reason` 快照投影，不为格式升级重写或补前缀。单 job 状态和任务列表 BFF 以 `warning: string` 返回该可直接展示的完整文案，不再返回结构化 code，Web 不得再次补前缀或按字符串推断告警类型。错误与告警摘要都不复制内联媒体并限制为 2048 字符。`phase` 与 `warning_message` 分别表示当前执行阶段和成功降级提示，不得混用；worker / BFF / Web 必须同版本协调发布，不保证滚动混部或旧 Web 缓存下的字符串语义兼容。
|
||||
- 正式读取 procedure 为 `get_external_generation_job_summary_and_return`、`list_external_generation_job_summaries_and_return` 和 `acknowledge_external_generation_job_summaries_and_return`。历史维护 procedure 为 `compact_external_generation_job_payloads_and_return` 与 `backfill_external_generation_job_summaries_and_return`，仅 migration operator 可调用；运维入口统一使用 `npm run spacetime:external-generation:maintain -- ...`，默认 dry-run、单批最多 25 条。B-tree cursor 选择阶段最多反序列化 `limit + 1` 行，apply 再按主键逐条读取选中行；怀疑存在单行异常巨型 JSON 时必须先使用 `--limit 1`。payload 压缩额外固定使用 `source_module = editor-canvas` 的复合 cursor 索引，不得静默改写其它玩法历史任务。
|
||||
- 正式读取 procedure 为 `get_external_generation_job_summary_and_return`、`list_external_generation_job_summaries_and_return` 和 `acknowledge_external_generation_job_summaries_and_return`。历史维护 procedure 为 `compact_external_generation_job_payloads_and_return`、`backfill_external_generation_job_summaries_and_return` 与 `prune_external_generation_job_history_and_return`，仅 migration operator 可调用；运维入口统一使用 `npm run spacetime:external-generation:maintain -- ...`，默认 dry-run、单批最多 25 条。B-tree cursor 选择阶段最多反序列化 `limit + 1` 行，apply 再按主键逐条读取选中行；怀疑存在单行异常巨型 JSON 时必须先使用 `--limit 1`。payload 压缩额外固定使用 `source_module = editor-canvas` 的复合 cursor 索引，不得静默改写其它玩法历史任务。历史清理默认使用 `--prune-history`、`source_module = editor-canvas` 和 30 天保留期；只有主任务与摘要状态一致且属于 completed / failed / cancelled、摘要已有 `notification_acknowledged_at`、终态时间不晚于 cutoff 的记录才是候选。apply 在同一事务内按事件 → 摘要 → 主任务顺序删除，事件不得独立清理；pending / running、未确认通知、摘要缺失或状态不一致的记录永不删除。清理不触碰资产对象或钱包流水，其他 source module 必须显式指定并单独评估。
|
||||
|
||||
### `external_generation_job_event`
|
||||
|
||||
- Rust 结构体：`ExternalGenerationJobEvent`
|
||||
- 源码：`server-rs/crates/spacetime-module/src/external_generation.rs`
|
||||
- 用途：外部生成任务审计事件表，按 `job_id` 和 `owner_user_id` 记录 `enqueued`、`claimed`、`lease_renewed`、`completed`、`failed`、`acknowledged` 等状态转换事实。状态转换只能由 SpacetimeDB procedure 写入，不由前端或 worker 直接改表；该表用于追溯任务生命周期和排障，不替代 `external_generation_job` 当前状态。
|
||||
- 保留策略：事件只会随已确认通知的终态任务由 `prune_external_generation_job_history_and_return` 原子删除，不支持按事件单独清理，以保持任务、摘要和审计链一致。
|
||||
|
||||
### `ai_text_chunk`
|
||||
|
||||
@@ -410,10 +413,12 @@ Responses 的终态载荷既是工具调用的恢复源，也是正文的恢复
|
||||
|
||||
### `auth_store_projection_meta`
|
||||
|
||||
启动投影恢复会对过滤后的 retained refresh session 重新计数；超过 8192 条时直接失败关闭并继续重试，不得把超限快照一次性灌入内存。
|
||||
|
||||
- Rust 结构体：`AuthStoreProjectionMeta`
|
||||
- 源码：`server-rs/crates/spacetime-module/src/auth/tables.rs`
|
||||
|
||||
认证恢复策略：`api-server` 启动时只从 SpacetimeDB 正式认证表（`user_account` / `auth_identity` / `refresh_session`）导出 typed `AuthStoreProjectionView`，再恢复 `module-auth` 的进程内认证工作集；运行中 Bearer `sid` 或 refresh cookie 在本进程工作集内未命中时直接按失效处理，不再从 SpacetimeDB 导出整包认证状态刷新内存，避免旧投影把重复手机号或旧会话重新灌回进程。`module-auth` 只保留内存工作集和 projection 导入 / 导出能力，不再保留 JSON 快照导入 / 导出能力，也不写本地持久化文件；`auth-store.json` / `GENARRATIVE_AUTH_STORE_PATH` 不再是兼容恢复源。认证创建、登录会话、刷新、退出、改密、重置密码、绑定和资料变更等写操作必须在返回客户端前通过 `sync_auth_store_projection` 成功同步 SpacetimeDB 正式认证表；同步失败时接口返回错误，不允许把只存在于当前进程内存的账号或会话当成成功结果。新用户注册奖励、邀请码绑定和登录埋点必须排在认证同步成功之后，避免认证没落库时先写出钱包或邀请关系。若启动恢复阶段 SpacetimeDB 不可连接或超时，`api-server` 会按固定间隔持续重试认证工作集恢复，恢复成功后才开始监听 HTTP，避免一次短超时让进程永久停留在依赖不可用状态。
|
||||
认证恢复策略：`api-server` 启动时只从 SpacetimeDB 正式认证表（`user_account` / `auth_identity` / `refresh_session`）导出 typed `AuthStoreProjectionView`，再恢复 `module-auth` 的进程内认证工作集；运行中 Bearer `sid` 或 refresh cookie 在本进程工作集内未命中时直接按失效处理，不再从 SpacetimeDB 导出整包认证状态刷新内存，避免旧投影把重复手机号或旧会话重新灌回进程。`module-auth` 只保留内存工作集和 projection 导入 / 导出能力，不再保留 JSON 快照导入 / 导出能力，也不写本地持久化文件；`auth-store.json` / `GENARRATIVE_AUTH_STORE_PATH` 不再是兼容恢复源。认证创建、登录会话、刷新、退出、改密、重置密码、绑定和资料变更等写操作必须在返回客户端前通过 `sync_auth_store_projection` 成功同步 SpacetimeDB 正式认证表；同步失败时接口返回错误，不允许把只存在于当前进程内存的账号或会话当成成功结果。新用户注册奖励、邀请码绑定和登录埋点必须排在认证同步成功之后，避免认证没落库时先写出钱包或邀请关系。refresh session 工作集最多保留 8192 条，短信验证码和微信登录 state 分别最多保留 4096 条；写入前清理过期项，达到上限时拒绝新增而不继续膨胀。若启动恢复阶段 SpacetimeDB 不可连接或超时，`api-server` 会按固定间隔持续重试认证工作集恢复，恢复成功后才开始监听 HTTP，避免一次短超时让进程永久停留在依赖不可用状态。
|
||||
|
||||
`auth_store_snapshot` 表和旧 `import_auth_store_snapshot_json` / `export_auth_store_snapshot_from_tables` procedure 已删除。认证投影同步只读写 `user_account`、`auth_identity`、`refresh_session` 和 `auth_store_projection_meta`；`auth_identity` 不再写 `phone_e164`、`display_name`、`avatar_url`，这些账号资料只以 `user_account` 为准。
|
||||
|
||||
@@ -1232,6 +1237,7 @@ RPG 创作入口的配置 ID 是 `rpg`，当前 `visible=true`、`open=true`；
|
||||
|
||||
- Rust 结构体：`RefreshSession`
|
||||
- 源码：`server-rs/crates/spacetime-module/src/auth/tables.rs`
|
||||
- 认证工作集只保留 active 会话以及最近 24 小时内的 revoked / expired 会话；超过宽限期的失效会话在 refresh session 写路径和 projection 导出前从内存索引移除，并随下一次 typed projection 同步从正式表清理。该清理不改变 active 多端登录、单端登出或全端登出语义。
|
||||
|
||||
### `runtime_setting`
|
||||
|
||||
|
||||
File diff suppressed because one or more lines are too long
@@ -34,8 +34,8 @@ pipeline {
|
||||
string(name: 'WEB_LINK', defaultValue: '/srv/genarrative/web', description: 'Nginx 静态站点目录或软链接')
|
||||
string(name: 'API_ENV_FILE', defaultValue: '/etc/genarrative/api-server.env', description: 'api-server 环境文件')
|
||||
string(name: 'API_PORT', defaultValue: '8082', description: 'api-server 本机监听端口')
|
||||
choice(name: 'DATABASE_BACKUP_PROFILE', choices: ['archive-full', 'files-history'], description: '数据库定时备份 profile；默认 archive-full，files-history 仅在指定 work-dir 已有完整 full baseline 后启用')
|
||||
string(name: 'DATABASE_BACKUP_FILES_HISTORY_WORK_DIR', defaultValue: '/var/lib/genarrative/database-backups/files-history', description: 'files-history 的本地 state/catalog 目录；dev/release 必须使用各自已建立 full baseline 的独立目录')
|
||||
choice(name: 'DATABASE_BACKUP_PROFILE', choices: ['archive-full', 'files-history'], description: '数据库定时备份 profile；release 仅允许 archive-full，files-history 仅供 development 在指定 work-dir 已有完整 full baseline 后启用')
|
||||
string(name: 'DATABASE_BACKUP_FILES_HISTORY_WORK_DIR', defaultValue: '/var/lib/genarrative/database-backups/files-history', description: 'development files-history 的本地 state/catalog 目录；必须使用已建立 full baseline 的独立目录')
|
||||
choice(name: 'NGINX_CONFIG_MODE', choices: ['none', 'production-https', 'development-http'], description: 'Nginx 配置模式；开发服无域名时选 development-http，release 正式入口选 production-https')
|
||||
booleanParam(name: 'ENABLE_SERVICES', defaultValue: true, description: '启用并启动 spacetimedb 与 api-server systemd 服务')
|
||||
booleanParam(name: 'ENABLE_OTELCOL', defaultValue: true, description: '安装并启用本机 OpenTelemetry Collector；api-server 模板默认开启 OTLP，如需关闭请在 API_ENV_FILE 中将 GENARRATIVE_OTEL_ENABLED 改为 false')
|
||||
@@ -113,6 +113,9 @@ pipeline {
|
||||
if (!(databaseBackupProfile in ['archive-full', 'files-history'])) {
|
||||
error("DATABASE_BACKUP_PROFILE 只能是 archive-full 或 files-history，当前值: ${params.DATABASE_BACKUP_PROFILE}")
|
||||
}
|
||||
if (params.DEPLOY_TARGET == 'release' && databaseBackupProfile == 'files-history') {
|
||||
error('release 仅允许 archive-full；files-history 会把整棵历史目录加载到 Node 内存，需先完成流式 catalog 改造后才能重新启用。')
|
||||
}
|
||||
def databaseBackupFilesHistoryWorkDir = params.DATABASE_BACKUP_FILES_HISTORY_WORK_DIR?.trim()
|
||||
if (!(databaseBackupFilesHistoryWorkDir ==~ /^\/var\/lib\/genarrative\/database-backups\/[A-Za-z0-9._\/-]+$/) || databaseBackupFilesHistoryWorkDir.contains('..')) {
|
||||
error("DATABASE_BACKUP_FILES_HISTORY_WORK_DIR 必须是 /var/lib/genarrative/database-backups/ 下不含连续点号的绝对路径，当前值: ${params.DATABASE_BACKUP_FILES_HISTORY_WORK_DIR}")
|
||||
|
||||
@@ -50,6 +50,7 @@ async function main() {
|
||||
assertDeferredArchiveDiscoveryIsBoundedAndDeterministic();
|
||||
assertCanonicalQueryAndAuthorizationIncludeMultipartParameters();
|
||||
assertInsufficientSpaceStopsBeforeServiceChanges();
|
||||
assertStopFailureRetainsRecoveryMarker();
|
||||
assertArchiveFailureStillRestoresDependentServices();
|
||||
await assertMultipartUploadRetriesAndVerifiesRemoteLength();
|
||||
await assertUploadBandwidthLimiterSharesBudgetAndPropagatesErrors();
|
||||
@@ -748,6 +749,27 @@ function assertInsufficientSpaceStopsBeforeServiceChanges() {
|
||||
assertFileMissing(fixture.tarLog, '空间不足时不能调用 tar。');
|
||||
}
|
||||
|
||||
function assertStopFailureRetainsRecoveryMarker() {
|
||||
const fixture = createFixture('stop-failure-marker');
|
||||
writeExecutable(
|
||||
path.join(fixture.binDir, 'systemctl'),
|
||||
`#!/usr/bin/env bash
|
||||
printf 'systemctl %s\\n' "$*" >> "${fixture.systemctlLog}"
|
||||
if [ "$1" = stop ]; then
|
||||
exit 9
|
||||
fi
|
||||
exit 0
|
||||
`,
|
||||
);
|
||||
const result = runBackup(fixture, ['--stop-service', 'spacetimedb.service']);
|
||||
|
||||
assertStatus(result, 1, '停止服务失败时备份必须失败。');
|
||||
assertTrue(
|
||||
existsSync(path.join(fixture.workDir, '.spacetimedb-stopped')),
|
||||
'停止服务命令失败时必须保留 marker，供 systemd ExecStopPost 兜底恢复。',
|
||||
);
|
||||
}
|
||||
|
||||
function assertArchiveFailureStillRestoresDependentServices() {
|
||||
const fixture = createFixture('tar-failure');
|
||||
const result = runBackup(fixture, [
|
||||
@@ -776,6 +798,10 @@ function assertArchiveFailureStillRestoresDependentServices() {
|
||||
for (const command of expectedCommands) {
|
||||
assertIncludes(systemctlLog, command, `tar 失败后必须执行: ${command}`);
|
||||
}
|
||||
assertFileMissing(
|
||||
path.join(fixture.workDir, '.spacetimedb-stopped'),
|
||||
'正常执行 finally 恢复全部服务后必须清理停库 marker。',
|
||||
);
|
||||
}
|
||||
|
||||
async function assertMultipartUploadRetriesAndVerifiesRemoteLength() {
|
||||
|
||||
@@ -821,6 +821,31 @@ const checks = [
|
||||
reason:
|
||||
'生产冷备份 service 必须用 node -- 分隔脚本参数，避免 Node 22 抢占业务 --env-file。',
|
||||
},
|
||||
{
|
||||
file: 'deploy/systemd/genarrative-database-backup.service',
|
||||
includes: 'Environment=NODE_OPTIONS=--max-old-space-size=768',
|
||||
reason:
|
||||
'备份 Node 进程必须设置独立 heap 上限，避免目录扫描异常拖垮 release 主机。',
|
||||
},
|
||||
{
|
||||
file: 'deploy/systemd/genarrative-database-backup.service',
|
||||
includes:
|
||||
'Environment=GENARRATIVE_DATABASE_BACKUP_STOP_MARKER=/var/lib/genarrative/database-backups/.spacetimedb-stopped',
|
||||
reason:
|
||||
'备份停库 marker 必须固定在受保护的 release work-dir，供 OOM 后 systemd 兜底恢复服务。',
|
||||
},
|
||||
{
|
||||
file: 'deploy/systemd/genarrative-database-backup.service',
|
||||
includes: 'MemoryMax=1G',
|
||||
reason:
|
||||
'备份 service 必须设置 systemd 内存硬上限，避免异常进程消耗整机内存。',
|
||||
},
|
||||
{
|
||||
file: 'deploy/systemd/genarrative-database-backup.service',
|
||||
includes: 'ExecStopPost=/bin/sh -c',
|
||||
reason:
|
||||
'备份主进程被 OOM kill 后必须由 systemd 兜底恢复停掉的 SpacetimeDB、API、worker 和 controller。',
|
||||
},
|
||||
{
|
||||
file: 'deploy/systemd/genarrative-database-backup.service',
|
||||
excludes: '--storage-format files',
|
||||
@@ -941,10 +966,9 @@ const checks = [
|
||||
},
|
||||
{
|
||||
file: 'jenkins/Jenkinsfile.production-server-provision',
|
||||
excludes:
|
||||
"params.DEPLOY_TARGET == 'release' && databaseBackupProfile == 'files-history'",
|
||||
includes: 'release 仅允许 archive-full；files-history',
|
||||
reason:
|
||||
'release 必须能在显式选择 profile 且 baseline 预检通过后启用 files-history。',
|
||||
'release 必须拒绝 files-history，避免逐文件 catalog 扫描再次触发生产内存峰值。',
|
||||
},
|
||||
{
|
||||
file: 'scripts/database-backup-to-oss.mjs',
|
||||
@@ -954,7 +978,7 @@ const checks = [
|
||||
{
|
||||
file: 'scripts/database-backup-to-oss.mjs',
|
||||
includes:
|
||||
'restoreServicesAfterBackup({stopService, serviceStopped, restartServicesAfter})',
|
||||
'restoreServicesAfterBackup({stopService, serviceStopped, restartServicesAfter, stopMarkerPath})',
|
||||
reason: '生产冷备份打包失败时也必须恢复 SpacetimeDB 及依赖服务。',
|
||||
},
|
||||
{
|
||||
|
||||
@@ -35,6 +35,7 @@ const DEFAULT_LOCAL_DATA_DIR = resolve(REPO_ROOT, 'server-rs/.spacetimedb/local/
|
||||
const DEFAULT_LOCAL_WORK_DIR = resolve(REPO_ROOT, 'server-rs/.data/database-backups');
|
||||
const DEFAULT_PRODUCTION_DATA_DIR = '/stdb';
|
||||
const DEFAULT_PRODUCTION_WORK_DIR = '/var/lib/genarrative/database-backups';
|
||||
const DEFAULT_DATABASE_BACKUP_STOP_MARKER = join(DEFAULT_PRODUCTION_WORK_DIR, '.spacetimedb-stopped');
|
||||
const DEFAULT_SPACE_SAFETY_RATIO = 1.1;
|
||||
const DEFAULT_EXTRA_FREE_BYTES = 512 * 1024 * 1024;
|
||||
const OSS_ALGORITHM = 'OSS4-HMAC-SHA256';
|
||||
@@ -555,7 +556,11 @@ function assertSafeRelativePath(dataDir, absolutePath) {
|
||||
}
|
||||
|
||||
function statFingerprint(absolutePath, rootPath = absolutePath) {
|
||||
const entries = [];
|
||||
// 候选 snapshot 可能包含数十万条目录项；增量更新摘要，避免把每条
|
||||
// fingerprint 字符串同时保存在 entries[] 后再 join，造成一次性内存峰值。
|
||||
const fingerprintHash = createHash('sha256');
|
||||
let isFirstEntry = true;
|
||||
let entryCount = 0;
|
||||
let totalSize = 0n;
|
||||
const visit = (currentPath) => {
|
||||
const stat = lstatSync(currentPath, {bigint: true});
|
||||
@@ -567,7 +572,7 @@ function statFingerprint(absolutePath, rootPath = absolutePath) {
|
||||
if (kind === 'other') {
|
||||
throw new Error(`history 候选只允许普通文件或目录: ${currentPath}`);
|
||||
}
|
||||
entries.push([
|
||||
const entry = [
|
||||
entryPath,
|
||||
kind,
|
||||
stat.dev.toString(),
|
||||
@@ -575,7 +580,13 @@ function statFingerprint(absolutePath, rootPath = absolutePath) {
|
||||
stat.mode.toString(),
|
||||
stat.size.toString(),
|
||||
stat.mtimeNs.toString(),
|
||||
].join('\0'));
|
||||
].join('\0');
|
||||
if (!isFirstEntry) {
|
||||
fingerprintHash.update('\n');
|
||||
}
|
||||
fingerprintHash.update(entry);
|
||||
isFirstEntry = false;
|
||||
entryCount += 1;
|
||||
if (stat.isFile()) {
|
||||
totalSize += stat.size;
|
||||
} else {
|
||||
@@ -586,9 +597,9 @@ function statFingerprint(absolutePath, rootPath = absolutePath) {
|
||||
};
|
||||
visit(rootPath);
|
||||
return {
|
||||
fingerprint: sha256Hex(entries.join('\n')),
|
||||
fingerprint: fingerprintHash.digest('hex'),
|
||||
sizeBytes: totalSize.toString(),
|
||||
entryCount: entries.length,
|
||||
entryCount,
|
||||
};
|
||||
}
|
||||
|
||||
@@ -880,11 +891,37 @@ function collectRestartServicesAfterBackup({args, env}) {
|
||||
return [...new Set(serviceNames.filter(Boolean))];
|
||||
}
|
||||
|
||||
function stopServiceIfNeeded(serviceName) {
|
||||
function databaseBackupStopMarkerPath(workDir) {
|
||||
return resolvePath(firstNonEmpty(
|
||||
process.env.GENARRATIVE_DATABASE_BACKUP_STOP_MARKER,
|
||||
workDir === DEFAULT_PRODUCTION_WORK_DIR
|
||||
? DEFAULT_DATABASE_BACKUP_STOP_MARKER
|
||||
: join(workDir, '.spacetimedb-stopped'),
|
||||
));
|
||||
}
|
||||
|
||||
function writeDatabaseBackupStopMarker(markerPath, serviceName) {
|
||||
atomicWriteJson(markerPath, {
|
||||
serviceName,
|
||||
pid: process.pid,
|
||||
stoppedAt: new Date().toISOString(),
|
||||
});
|
||||
}
|
||||
|
||||
function clearDatabaseBackupStopMarker(markerPath) {
|
||||
if (markerPath) {
|
||||
rmSync(markerPath, {force: true});
|
||||
}
|
||||
}
|
||||
|
||||
function stopServiceIfNeeded(serviceName, stopMarkerPath) {
|
||||
if (!serviceName) {
|
||||
return false;
|
||||
}
|
||||
console.log(`[database-backup] 停止服务以获取冷备份: ${serviceName}`);
|
||||
writeDatabaseBackupStopMarker(stopMarkerPath, serviceName);
|
||||
// stop 命令失败时仍保留 marker：systemd 的 ExecStopPost 需要它判断是否要
|
||||
// 兜底恢复，不能因为当前进程还能捕获异常就抹掉上一次停库证据。
|
||||
runCommand('systemctl', ['stop', serviceName], {stdio: 'inherit'});
|
||||
return true;
|
||||
}
|
||||
@@ -915,7 +952,7 @@ function restartServicesAfterBackup(serviceNames) {
|
||||
}
|
||||
}
|
||||
|
||||
function restoreServicesAfterBackup({stopService, serviceStopped, restartServicesAfter}) {
|
||||
function restoreServicesAfterBackup({stopService, serviceStopped, restartServicesAfter, stopMarkerPath}) {
|
||||
const errors = [];
|
||||
try {
|
||||
startServiceIfNeeded(stopService, serviceStopped);
|
||||
@@ -930,6 +967,7 @@ function restoreServicesAfterBackup({stopService, serviceStopped, restartService
|
||||
if (errors.length > 0) {
|
||||
throw new AggregateError(errors, `恢复冷备份相关服务失败: ${errors.map((error) => error.message).join('; ')}`);
|
||||
}
|
||||
clearDatabaseBackupStopMarker(stopMarkerPath);
|
||||
}
|
||||
|
||||
function createArchive({dataDir, workDir, fileName}) {
|
||||
@@ -1285,14 +1323,17 @@ export async function collectDirectFileEntries({dataDir, candidates = null, obje
|
||||
throw new Error(`files 扫描期间源文件发生变化: ${relativePath}`);
|
||||
}
|
||||
const basePrefix = normalizeObjectPrefix(objectPrefix, database);
|
||||
files.set(relativePath, {
|
||||
const file = {
|
||||
path: relativePath,
|
||||
sizeBytes: Number(after.size),
|
||||
sha256,
|
||||
mode: after.mode,
|
||||
objectKey: `${basePrefix}/files/sha256/${sha256.slice(0, 2)}/${sha256}`,
|
||||
sourceStat: after,
|
||||
});
|
||||
};
|
||||
// 上传前后的 inode/stat 仍用于防止在线扫描漂移，但设为不可枚举，避免
|
||||
// 把仅供本地校验的副本再次写入 catalog 或 result JSON。
|
||||
Object.defineProperty(file, 'sourceStat', {value: after, enumerable: false});
|
||||
files.set(relativePath, file);
|
||||
};
|
||||
|
||||
for (const root of roots.sort((left, right) => left.relativePath.localeCompare(right.relativePath))) {
|
||||
@@ -1309,14 +1350,40 @@ export async function collectDirectFileEntries({dataDir, candidates = null, obje
|
||||
}
|
||||
|
||||
function directCatalogIdentity({mode, baselineCatalogId, rootName, directories, files, symlinks}) {
|
||||
return sha256Hex(JSON.stringify({
|
||||
mode,
|
||||
baselineCatalogId: baselineCatalogId || '',
|
||||
rootName,
|
||||
directories,
|
||||
files: files.map(({path, sizeBytes, sha256, mode, objectKey}) => ({path, sizeBytes, sha256, mode, objectKey})),
|
||||
symlinks,
|
||||
// 不把数十万条文件元数据先拼成一个巨型 JSON 字符串；分段写入 hash
|
||||
// 保持与 JSON.stringify 同样的字段顺序和转义结果，同时把峰值降到单条记录。
|
||||
const hash = createHash('sha256');
|
||||
hash.update('{"mode":');
|
||||
hash.update(JSON.stringify(mode));
|
||||
hash.update(',"baselineCatalogId":');
|
||||
hash.update(JSON.stringify(baselineCatalogId || ''));
|
||||
hash.update(',"rootName":');
|
||||
hash.update(JSON.stringify(rootName));
|
||||
hash.update(',"directories":');
|
||||
updateJsonArrayHash(hash, directories, (directory) => JSON.stringify(directory));
|
||||
hash.update(',"files":');
|
||||
updateJsonArrayHash(hash, files, (file) => JSON.stringify({
|
||||
path: file.path,
|
||||
sizeBytes: file.sizeBytes,
|
||||
sha256: file.sha256,
|
||||
mode: file.mode,
|
||||
objectKey: file.objectKey,
|
||||
}));
|
||||
hash.update(',"symlinks":');
|
||||
updateJsonArrayHash(hash, symlinks, (symlink) => JSON.stringify(symlink));
|
||||
hash.update('}');
|
||||
return hash.digest('hex');
|
||||
}
|
||||
|
||||
function updateJsonArrayHash(hash, values, serialize) {
|
||||
hash.update('[');
|
||||
values.forEach((value, index) => {
|
||||
if (index > 0) {
|
||||
hash.update(',');
|
||||
}
|
||||
hash.update(serialize(value));
|
||||
});
|
||||
hash.update(']');
|
||||
}
|
||||
|
||||
function readDirectFilesState(statePath, {database, bucket}) {
|
||||
@@ -1621,7 +1688,7 @@ export async function runDirectFilesBackup({
|
||||
baselineCatalogId,
|
||||
rootName,
|
||||
directories: collected.directories,
|
||||
files: collected.files.map(({sourceStat: _sourceStat, ...file}) => file),
|
||||
files: collected.files,
|
||||
symlinks: collected.symlinks,
|
||||
};
|
||||
writeManifest({manifestPath: catalogPath, payload: catalog});
|
||||
@@ -3302,12 +3369,13 @@ async function main() {
|
||||
}
|
||||
const stopService = args.stopService || firstNonEmpty(env.GENARRATIVE_DATABASE_BACKUP_STOP_SERVICE);
|
||||
const restartServicesAfter = collectRestartServicesAfterBackup({args, env});
|
||||
const stopMarkerPath = databaseBackupStopMarkerPath(workDir);
|
||||
let serviceStopped = false;
|
||||
let backupError = null;
|
||||
let restoreError = null;
|
||||
try {
|
||||
if (args.mode === 'full' && !args.dryRun) {
|
||||
serviceStopped = stopServiceIfNeeded(stopService);
|
||||
serviceStopped = stopServiceIfNeeded(stopService, stopMarkerPath);
|
||||
}
|
||||
await runDirectFilesBackup({
|
||||
mode: args.mode,
|
||||
@@ -3326,7 +3394,7 @@ async function main() {
|
||||
} finally {
|
||||
try {
|
||||
if (serviceStopped) {
|
||||
restoreServicesAfterBackup({stopService, serviceStopped, restartServicesAfter});
|
||||
restoreServicesAfterBackup({stopService, serviceStopped, restartServicesAfter, stopMarkerPath});
|
||||
} else if (!backupError && args.mode === 'full' && !args.dryRun) {
|
||||
restartServicesAfterBackup(restartServicesAfter);
|
||||
}
|
||||
@@ -3380,16 +3448,17 @@ async function main() {
|
||||
let restoreError = null;
|
||||
const stopService = args.stopService || firstNonEmpty(env.GENARRATIVE_DATABASE_BACKUP_STOP_SERVICE);
|
||||
const restartServicesAfter = collectRestartServicesAfterBackup({args, env});
|
||||
const stopMarkerPath = databaseBackupStopMarkerPath(workDir);
|
||||
try {
|
||||
assertSufficientWorkDirSpace({dataDir, workDir, args, env});
|
||||
serviceStopped = stopServiceIfNeeded(stopService);
|
||||
serviceStopped = stopServiceIfNeeded(stopService, stopMarkerPath);
|
||||
archivePath = createArchive({dataDir, workDir, fileName});
|
||||
} catch (error) {
|
||||
backupError = error;
|
||||
} finally {
|
||||
try {
|
||||
if (serviceStopped) {
|
||||
restoreServicesAfterBackup({stopService, serviceStopped, restartServicesAfter});
|
||||
restoreServicesAfterBackup({stopService, serviceStopped, restartServicesAfter, stopMarkerPath});
|
||||
} else if (!backupError) {
|
||||
restartServicesAfterBackup(restartServicesAfter);
|
||||
}
|
||||
|
||||
@@ -78,6 +78,10 @@ validate_database_backup_profile() {
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
if [[ "${DEPLOY_TARGET}" == "release" && "${DATABASE_BACKUP_PROFILE}" == "files-history" ]]; then
|
||||
echo "[server-provision] release 仅允许 archive-full；files-history 会把整棵历史目录加载到 Node 内存，需先完成流式 catalog 改造后才能重新启用。" >&2
|
||||
exit 1
|
||||
fi
|
||||
if [[ ! "${DATABASE_BACKUP_FILES_HISTORY_WORK_DIR}" =~ ^/var/lib/genarrative/database-backups/[A-Za-z0-9._/-]+$ || "${DATABASE_BACKUP_FILES_HISTORY_WORK_DIR}" == *..* ]]; then
|
||||
echo "[server-provision] DATABASE_BACKUP_FILES_HISTORY_WORK_DIR 必须是 /var/lib/genarrative/database-backups/ 下不含连续点号的绝对路径，当前值: ${DATABASE_BACKUP_FILES_HISTORY_WORK_DIR}" >&2
|
||||
exit 1
|
||||
|
||||
@@ -8,23 +8,29 @@ import {
|
||||
} from './spacetime-migration-common.mjs';
|
||||
|
||||
const MAX_BATCH_SIZE = 25;
|
||||
const DEFAULT_RETENTION_DAYS = 30;
|
||||
const MICROS_PER_DAY = 86_400_000_000;
|
||||
|
||||
function usage() {
|
||||
return `用法：
|
||||
node scripts/spacetime-maintain-external-generation-jobs.mjs --database <name> [选项]
|
||||
|
||||
默认只 dry-run 一批历史终态任务 payload 压缩，不修改数据库。
|
||||
使用 --prune-history 时改为清理已确认通知且超过保留期的历史任务、摘要与事件。
|
||||
|
||||
公共选项：
|
||||
--database <name> 目标数据库（必填，也可用 GENARRATIVE_SPACETIME_DATABASE）
|
||||
--server <name-or-url> spacetime CLI server 名或 URL
|
||||
--server-url <url> 显式 server URL
|
||||
--limit <1-${MAX_BATCH_SIZE}> 单批任务数，默认 10
|
||||
--limit <1-${MAX_BATCH_SIZE}> 单批任务数，默认 10
|
||||
--cursor-job-id <jobId> 从上一批 next_cursor_job_id 继续
|
||||
--apply 执行写入；省略时始终 dry-run
|
||||
--backfill-summaries 改为回填轻量摘要投影
|
||||
--prune-history 改为清理已确认通知的终态历史
|
||||
--owner-user-id <userId> 仅摘要回填可选，限定 owner
|
||||
--completed-before-micros <n> 仅 payload 压缩可选，限定终态完成时间
|
||||
--source-module <module> 仅历史清理可选，默认 editor-canvas
|
||||
--retention-days <n> 仅历史清理可选，默认 ${DEFAULT_RETENTION_DAYS} 天
|
||||
--completed-before-micros <n> 限定终态完成时间；历史清理默认按 retention-days 计算
|
||||
--help 显示帮助
|
||||
|
||||
必须使用已授权 migration operator 的 spacetime CLI 登录态。脚本每次只处理一批；
|
||||
@@ -40,6 +46,9 @@ function parseOptions(argv) {
|
||||
database: process.env.GENARRATIVE_SPACETIME_DATABASE || '',
|
||||
limit: 10,
|
||||
ownerUserId: '',
|
||||
pruneHistory: false,
|
||||
retentionDays: DEFAULT_RETENTION_DAYS,
|
||||
sourceModule: 'editor-canvas',
|
||||
passthrough: [],
|
||||
server: process.env.GENARRATIVE_SPACETIME_SERVER || '',
|
||||
serverUrl: process.env.GENARRATIVE_SPACETIME_SERVER_URL || '',
|
||||
@@ -82,6 +91,15 @@ function parseOptions(argv) {
|
||||
options.apply = true;
|
||||
} else if (arg === '--backfill-summaries') {
|
||||
options.backfillSummaries = true;
|
||||
} else if (arg === '--prune-history') {
|
||||
options.pruneHistory = true;
|
||||
} else if (arg === '--source-module') {
|
||||
options.sourceModule = readValue(arg).trim();
|
||||
if (!options.sourceModule) {
|
||||
throw new Error('--source-module 不能为空。');
|
||||
}
|
||||
} else if (arg === '--retention-days') {
|
||||
options.retentionDays = parsePositiveInteger(readValue(arg), arg);
|
||||
} else if (arg === '--help' || arg === '-h') {
|
||||
options.help = true;
|
||||
} else {
|
||||
@@ -95,12 +113,49 @@ function parseOptions(argv) {
|
||||
if (options.ownerUserId && !options.backfillSummaries) {
|
||||
throw new Error('--owner-user-id 只能与 --backfill-summaries 一起使用。');
|
||||
}
|
||||
if (options.backfillSummaries && options.pruneHistory) {
|
||||
throw new Error('--backfill-summaries 与 --prune-history 不能同时使用。');
|
||||
}
|
||||
if (options.sourceModule !== 'editor-canvas' && !options.pruneHistory) {
|
||||
throw new Error('--source-module 只能与 --prune-history 一起使用。');
|
||||
}
|
||||
if (
|
||||
options.retentionDays !== DEFAULT_RETENTION_DAYS &&
|
||||
!options.pruneHistory
|
||||
) {
|
||||
throw new Error('--retention-days 只能与 --prune-history 一起使用。');
|
||||
}
|
||||
if (options.completedBeforeMicros !== null && options.backfillSummaries) {
|
||||
throw new Error('--completed-before-micros 不能用于摘要回填。');
|
||||
}
|
||||
if (
|
||||
options.completedBeforeMicros !== null &&
|
||||
options.pruneHistory &&
|
||||
options.retentionDays !== DEFAULT_RETENTION_DAYS
|
||||
) {
|
||||
throw new Error(
|
||||
'--completed-before-micros 与 --retention-days 不能同时使用。',
|
||||
);
|
||||
}
|
||||
return options;
|
||||
}
|
||||
|
||||
function resolveRetentionCutoffMicros(options) {
|
||||
if (!options.pruneHistory) {
|
||||
return options.completedBeforeMicros;
|
||||
}
|
||||
if (options.completedBeforeMicros !== null) {
|
||||
return options.completedBeforeMicros;
|
||||
}
|
||||
const cutoff = Date.now() * 1000 - options.retentionDays * MICROS_PER_DAY;
|
||||
if (!Number.isSafeInteger(cutoff)) {
|
||||
throw new Error(
|
||||
'--retention-days 计算出的 completed_before_micros 超出安全整数范围。',
|
||||
);
|
||||
}
|
||||
return cutoff;
|
||||
}
|
||||
|
||||
try {
|
||||
const options = parseOptions(process.argv.slice(2));
|
||||
if (options.help) {
|
||||
@@ -113,24 +168,36 @@ try {
|
||||
);
|
||||
}
|
||||
|
||||
const procedureName = options.backfillSummaries
|
||||
? 'backfill_external_generation_job_summaries_and_return'
|
||||
: 'compact_external_generation_job_payloads_and_return';
|
||||
const input = options.backfillSummaries
|
||||
const completedBeforeMicros = resolveRetentionCutoffMicros(options);
|
||||
|
||||
const procedureName = options.pruneHistory
|
||||
? 'prune_external_generation_job_history_and_return'
|
||||
: options.backfillSummaries
|
||||
? 'backfill_external_generation_job_summaries_and_return'
|
||||
: 'compact_external_generation_job_payloads_and_return';
|
||||
const input = options.pruneHistory
|
||||
? {
|
||||
owner_user_id: encodeSpacetimeCliOption(options.ownerUserId || null),
|
||||
source_module: options.sourceModule,
|
||||
limit: options.limit,
|
||||
cursor_job_id: encodeSpacetimeCliOption(options.cursorJobId || null),
|
||||
completed_before_micros: completedBeforeMicros,
|
||||
dry_run: !options.apply,
|
||||
}
|
||||
: {
|
||||
dry_run: !options.apply,
|
||||
limit: options.limit,
|
||||
cursor_job_id: encodeSpacetimeCliOption(options.cursorJobId || null),
|
||||
completed_before_micros: encodeSpacetimeCliOption(
|
||||
options.completedBeforeMicros,
|
||||
),
|
||||
};
|
||||
: options.backfillSummaries
|
||||
? {
|
||||
owner_user_id: encodeSpacetimeCliOption(options.ownerUserId || null),
|
||||
limit: options.limit,
|
||||
cursor_job_id: encodeSpacetimeCliOption(options.cursorJobId || null),
|
||||
dry_run: !options.apply,
|
||||
}
|
||||
: {
|
||||
dry_run: !options.apply,
|
||||
limit: options.limit,
|
||||
cursor_job_id: encodeSpacetimeCliOption(options.cursorJobId || null),
|
||||
completed_before_micros: encodeSpacetimeCliOption(
|
||||
completedBeforeMicros,
|
||||
),
|
||||
};
|
||||
const result = await callSpacetimeProcedureViaCli(
|
||||
options,
|
||||
procedureName,
|
||||
@@ -138,10 +205,29 @@ try {
|
||||
);
|
||||
ensureProcedureOk(result);
|
||||
|
||||
console.log(JSON.stringify({ procedure: procedureName, ...result }, null, 2));
|
||||
const pendingApplyCount = options.backfillSummaries
|
||||
? Number(result.selected_count ?? 0)
|
||||
: Number(result.matched_count ?? 0);
|
||||
console.log(
|
||||
JSON.stringify(
|
||||
{
|
||||
procedure: procedureName,
|
||||
...(options.pruneHistory
|
||||
? {
|
||||
source_module: options.sourceModule,
|
||||
completed_before_micros: completedBeforeMicros,
|
||||
...(options.completedBeforeMicros === null
|
||||
? { retention_days: options.retentionDays }
|
||||
: {}),
|
||||
}
|
||||
: {}),
|
||||
...result,
|
||||
},
|
||||
null,
|
||||
2,
|
||||
),
|
||||
);
|
||||
const pendingApplyCount =
|
||||
options.pruneHistory || options.backfillSummaries
|
||||
? Number(result.selected_count ?? 0)
|
||||
: Number(result.matched_count ?? 0);
|
||||
if (result.has_more && options.apply) {
|
||||
console.log(
|
||||
`仍有后续批次；下一次追加 --cursor-job-id ${result.next_cursor_job_id ?? '<missing>'}。`,
|
||||
@@ -150,8 +236,11 @@ try {
|
||||
const currentCursor = options.cursorJobId
|
||||
? `保留 --cursor-job-id ${options.cursorJobId}`
|
||||
: '仍从首批开始';
|
||||
const cutoffHint = options.pruneHistory
|
||||
? `并固定 --completed-before-micros ${completedBeforeMicros}`
|
||||
: '';
|
||||
console.log(
|
||||
`当前仅 dry-run；请${currentCursor}并追加 --apply 重跑同一批。apply 成功后再使用其 next_cursor_job_id 进入下一批。`,
|
||||
`当前仅 dry-run；请${currentCursor}${cutoffHint}并追加 --apply 重跑同一批。apply 成功后再使用其 next_cursor_job_id 进入下一批。`,
|
||||
);
|
||||
}
|
||||
} catch (error) {
|
||||
|
||||
@@ -14,6 +14,7 @@ use spacetime_client::{
|
||||
ExternalGenerationJobRenewLeaseRecordInput, ExternalGenerationQueueWakeSubscription,
|
||||
};
|
||||
use tokio::{
|
||||
sync::{OwnedSemaphorePermit, Semaphore},
|
||||
task::{JoinHandle, JoinSet},
|
||||
time::sleep,
|
||||
};
|
||||
@@ -92,6 +93,10 @@ pub(crate) async fn run_external_generation_worker(state: AppState) -> Result<()
|
||||
let concurrency = state.config.external_generation_worker_concurrency.max(1);
|
||||
let poll_interval = state.config.external_generation_worker_poll_interval;
|
||||
let lease = state.config.external_generation_worker_lease;
|
||||
// 超时任务不能立即取消（在途 procedure 仍可能写回），因此执行容量必须同时
|
||||
// 约束 active 与 detached work；否则每次超时都会释放 tasks 槽位，实际内存占用
|
||||
// 会超过配置并发。
|
||||
let work_slots = std::sync::Arc::new(Semaphore::new(concurrency));
|
||||
let mut tasks = JoinSet::new();
|
||||
let mut shutdown = external_generation_worker_shutdown_signal();
|
||||
let mut queue_wake = None;
|
||||
@@ -113,16 +118,29 @@ pub(crate) async fn run_external_generation_worker(state: AppState) -> Result<()
|
||||
);
|
||||
|
||||
loop {
|
||||
// 持续有队列任务时不会进入等待分支，因此必须在每轮主动回收已完成的
|
||||
// JoinHandle；否则 permit 虽已归还，JoinSet 仍会保留每个历史任务的句柄。
|
||||
reap_finished_external_generation_worker_tasks(&mut tasks);
|
||||
ensure_external_generation_queue_wake_subscription(&state, &mut queue_wake).await;
|
||||
|
||||
while tasks.len() >= concurrency {
|
||||
if await_worker_task_or_shutdown(&mut tasks, &mut shutdown).await {
|
||||
drain_external_generation_worker_tasks(&mut tasks).await;
|
||||
return Ok(());
|
||||
while work_slots.available_permits() == 0 {
|
||||
tokio::select! {
|
||||
_ = shutdown.as_mut() => {
|
||||
drain_external_generation_worker_tasks(&mut tasks).await;
|
||||
return Ok(());
|
||||
}
|
||||
permit = work_slots.clone().acquire_owned() => {
|
||||
if permit.is_err() {
|
||||
drain_external_generation_worker_tasks(&mut tasks).await;
|
||||
return Ok(());
|
||||
}
|
||||
// 只用 acquire 作为容量变化唤醒信号，许可立即归还；真正领取任务
|
||||
// 时在下方按返回的 job 数量逐个 try_acquire。
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let available = concurrency.saturating_sub(tasks.len()).max(1);
|
||||
let available = work_slots.available_permits().max(1);
|
||||
let now_micros = current_utc_micros();
|
||||
let lease_expires_at_micros = now_micros.saturating_add(duration_micros_i64(lease));
|
||||
|
||||
@@ -178,9 +196,13 @@ pub(crate) async fn run_external_generation_worker(state: AppState) -> Result<()
|
||||
for job in jobs {
|
||||
let state = state.clone();
|
||||
let worker_id = worker_id.clone();
|
||||
let permit = work_slots
|
||||
.clone()
|
||||
.try_acquire_owned()
|
||||
.expect("claimed job must have an execution capacity permit");
|
||||
tasks.spawn(async move {
|
||||
if let Err(error) =
|
||||
process_external_generation_job(state, worker_id, lease, job).await
|
||||
process_external_generation_job(state, worker_id, lease, job, permit).await
|
||||
{
|
||||
error!(error = %error, "external generation worker 执行任务失败");
|
||||
}
|
||||
@@ -255,13 +277,11 @@ async fn await_worker_task(tasks: &mut JoinSet<()>) {
|
||||
}
|
||||
}
|
||||
|
||||
async fn await_worker_task_or_shutdown(
|
||||
tasks: &mut JoinSet<()>,
|
||||
shutdown: &mut ExternalGenerationShutdownSignal,
|
||||
) -> bool {
|
||||
tokio::select! {
|
||||
_ = shutdown.as_mut() => true,
|
||||
_ = await_worker_task(tasks) => false,
|
||||
fn reap_finished_external_generation_worker_tasks(tasks: &mut JoinSet<()>) {
|
||||
while let Some(result) = tasks.try_join_next() {
|
||||
if let Err(error) = result {
|
||||
error!(error = %error, "external generation worker 子任务 panic");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -326,6 +346,7 @@ async fn process_external_generation_job(
|
||||
worker_id: String,
|
||||
lease: Duration,
|
||||
job: ExternalGenerationJobRecord,
|
||||
permit: OwnedSemaphorePermit,
|
||||
) -> Result<(), String> {
|
||||
let heartbeat_interval = external_generation_worker_heartbeat_interval(lease);
|
||||
let job_timeout = external_generation_worker_job_timeout(&state.config, job.job_kind.as_str());
|
||||
@@ -377,13 +398,14 @@ async fn process_external_generation_job(
|
||||
job_id = %job.job_id,
|
||||
job_kind = %job.job_kind,
|
||||
timeout_seconds = job_timeout.as_secs(),
|
||||
"external generation worker 任务超过执行预算，停止续租并释放 worker 槽位，在途执行交由租约仲裁"
|
||||
"external generation worker 任务超过执行预算，停止续租并保留 worker 槽位，在途执行交由租约仲裁"
|
||||
);
|
||||
detach_external_generation_work_until_lease_expiry(
|
||||
work_handle,
|
||||
&job,
|
||||
lease,
|
||||
"任务超过执行预算",
|
||||
Some(permit),
|
||||
);
|
||||
Err(message)
|
||||
}
|
||||
@@ -393,6 +415,7 @@ async fn process_external_generation_job(
|
||||
&job,
|
||||
lease,
|
||||
"任务租约续期失败",
|
||||
Some(permit),
|
||||
);
|
||||
Err(error)
|
||||
}
|
||||
@@ -417,6 +440,7 @@ fn detach_external_generation_work_until_lease_expiry(
|
||||
job: &ExternalGenerationJobRecord,
|
||||
lease: Duration,
|
||||
reason: &'static str,
|
||||
permit: Option<OwnedSemaphorePermit>,
|
||||
) {
|
||||
let job_id = job.job_id.clone();
|
||||
let job_kind = job.job_kind.clone();
|
||||
@@ -445,6 +469,9 @@ fn detach_external_generation_work_until_lease_expiry(
|
||||
),
|
||||
Err(_) => {
|
||||
work_handle.abort();
|
||||
// 仅调用 abort 不会从 JoinHandle/JoinSet 中消费完成结果；等待被取消
|
||||
// 的 handle，确保 permit 与任务句柄在同一生命周期内一起释放。
|
||||
let _ = work_handle.await;
|
||||
warn!(
|
||||
job_id = %job_id,
|
||||
job_kind = %job_kind,
|
||||
@@ -454,6 +481,9 @@ fn detach_external_generation_work_until_lease_expiry(
|
||||
);
|
||||
}
|
||||
}
|
||||
// 保持执行许可直到 work 真正结束或被取消，避免超时任务脱管后继续
|
||||
// 累积图片/音频响应占用。
|
||||
drop(permit);
|
||||
});
|
||||
}
|
||||
|
||||
@@ -1972,6 +2002,7 @@ mod tests {
|
||||
&job,
|
||||
Duration::from_millis(200),
|
||||
"任务超过执行预算",
|
||||
None,
|
||||
);
|
||||
|
||||
tokio::time::sleep(Duration::from_millis(100)).await;
|
||||
@@ -1981,6 +2012,61 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn worker_detached_work_keeps_execution_slot_until_finished() {
|
||||
let slots = std::sync::Arc::new(tokio::sync::Semaphore::new(1));
|
||||
let permit = slots
|
||||
.clone()
|
||||
.acquire_owned()
|
||||
.await
|
||||
.expect("the only execution slot should be available");
|
||||
let work_handle = tokio::spawn(async {
|
||||
tokio::time::sleep(Duration::from_millis(20)).await;
|
||||
Ok(())
|
||||
});
|
||||
let job = external_generation_job_record_fixture(Some("lease-1"));
|
||||
|
||||
detach_external_generation_work_until_lease_expiry(
|
||||
work_handle,
|
||||
&job,
|
||||
Duration::from_millis(200),
|
||||
"任务超过执行预算",
|
||||
Some(permit),
|
||||
);
|
||||
|
||||
assert!(
|
||||
slots.try_acquire().is_err(),
|
||||
"脱管 work 完成前不得重新领取执行容量"
|
||||
);
|
||||
tokio::time::sleep(Duration::from_millis(100)).await;
|
||||
assert!(
|
||||
slots.try_acquire().is_ok(),
|
||||
"脱管 work 完成后应归还执行容量"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn worker_reaps_completed_tasks_while_queue_remains_busy() {
|
||||
let mut tasks = JoinSet::new();
|
||||
let completed = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
|
||||
const TASK_COUNT: usize = 128;
|
||||
for _ in 0..TASK_COUNT {
|
||||
let completed = completed.clone();
|
||||
tasks.spawn(async move {
|
||||
completed.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
|
||||
});
|
||||
}
|
||||
|
||||
while completed.load(std::sync::atomic::Ordering::SeqCst) < TASK_COUNT {
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
assert_eq!(tasks.len(), TASK_COUNT);
|
||||
|
||||
reap_finished_external_generation_worker_tasks(&mut tasks);
|
||||
|
||||
assert!(tasks.is_empty(), "已完成任务的 JoinHandle 应在每轮被回收");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn worker_detached_work_is_aborted_after_lease_arbitration_window() {
|
||||
let connection = std::sync::Arc::new(tokio::sync::Semaphore::new(1));
|
||||
@@ -1999,6 +2085,7 @@ mod tests {
|
||||
&job,
|
||||
Duration::from_millis(10),
|
||||
"任务超过执行预算",
|
||||
None,
|
||||
);
|
||||
|
||||
let reacquired =
|
||||
|
||||
@@ -28,7 +28,7 @@ impl AiTaskService {
|
||||
validate_task_create_input(&input).map_err(AiTaskServiceError::Field)?;
|
||||
|
||||
let snapshot = AiTaskSnapshot {
|
||||
task_id: input.task_id.clone(),
|
||||
task_id: normalize_required_string(input.task_id).unwrap_or_default(),
|
||||
task_kind: input.task_kind,
|
||||
owner_user_id: normalize_required_string(input.owner_user_id).unwrap_or_default(),
|
||||
request_label: normalize_required_string(input.request_label).unwrap_or_default(),
|
||||
|
||||
@@ -1,10 +1,12 @@
|
||||
use std::{
|
||||
collections::HashMap,
|
||||
collections::{BTreeMap, HashMap},
|
||||
sync::{Arc, Mutex},
|
||||
};
|
||||
|
||||
use crate::{
|
||||
AiTaskServiceError, AiTaskSnapshot, AiTaskStageStatus, AiTaskStatus, AiTextChunkSnapshot,
|
||||
MAX_AI_TASK_RETAINED_OUTPUT_BYTES, MAX_AI_TASK_RETAINED_TASKS, MAX_AI_TASK_TEXT_OUTPUT_BYTES,
|
||||
validate_ai_task_snapshot_memory_limits,
|
||||
};
|
||||
|
||||
use super::ensure_task_is_not_terminal;
|
||||
@@ -17,7 +19,9 @@ pub struct InMemoryAiTaskStore {
|
||||
#[derive(Debug, Default)]
|
||||
struct InMemoryAiTaskStoreState {
|
||||
tasks: HashMap<String, AiTaskSnapshot>,
|
||||
text_chunks: HashMap<String, Vec<AiTextChunkSnapshot>>,
|
||||
// Keep only the ordered deltas needed to handle an out-of-order chunk.
|
||||
// Completed tasks drop this map immediately; it is not a second durable log.
|
||||
text_chunks: HashMap<String, HashMap<crate::AiTaskStageKind, BTreeMap<u32, String>>>,
|
||||
}
|
||||
|
||||
impl InMemoryAiTaskStore {
|
||||
@@ -34,7 +38,48 @@ impl InMemoryAiTaskStore {
|
||||
return Err(AiTaskServiceError::TaskAlreadyExists);
|
||||
}
|
||||
|
||||
state.text_chunks.insert(task.task_id.clone(), Vec::new());
|
||||
validate_task_memory_limits(&task)?;
|
||||
|
||||
let oldest_terminal = if state.tasks.len() >= MAX_AI_TASK_RETAINED_TASKS {
|
||||
let oldest_terminal = state
|
||||
.tasks
|
||||
.values()
|
||||
.filter(|value| value.status.is_terminal())
|
||||
.min_by_key(|value| value.completed_at_micros.or(Some(value.updated_at_micros)))
|
||||
.map(|value| value.task_id.clone());
|
||||
if oldest_terminal.is_none() {
|
||||
return Err(AiTaskServiceError::Store(
|
||||
"AI 任务仓储已达到内存容量上限".to_string(),
|
||||
));
|
||||
}
|
||||
oldest_terminal
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let retained_output_bytes = retained_output_bytes(&state)
|
||||
.saturating_sub(
|
||||
oldest_terminal
|
||||
.as_deref()
|
||||
.and_then(|task_id| state.tasks.get(task_id))
|
||||
.map(task_output_bytes)
|
||||
.unwrap_or_default(),
|
||||
)
|
||||
.saturating_add(task_output_bytes(&task));
|
||||
if retained_output_bytes > MAX_AI_TASK_RETAINED_OUTPUT_BYTES {
|
||||
return Err(AiTaskServiceError::Store(
|
||||
"AI 任务仓储输出工作集超过内存上限".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
if let Some(task_id) = oldest_terminal {
|
||||
state.tasks.remove(&task_id);
|
||||
state.text_chunks.remove(&task_id);
|
||||
}
|
||||
|
||||
state
|
||||
.text_chunks
|
||||
.insert(task.task_id.clone(), HashMap::new());
|
||||
state.tasks.insert(task.task_id.clone(), task.clone());
|
||||
Ok(task)
|
||||
}
|
||||
@@ -51,12 +96,37 @@ impl InMemoryAiTaskStore {
|
||||
.inner
|
||||
.lock()
|
||||
.map_err(|_| AiTaskServiceError::Store("AI 任务仓储锁已中毒".to_string()))?;
|
||||
let task = state
|
||||
.tasks
|
||||
.get_mut(task_id.trim())
|
||||
.ok_or(AiTaskServiceError::TaskNotFound)?;
|
||||
apply(task)?;
|
||||
Ok(task.clone())
|
||||
let (previous_task, snapshot) = {
|
||||
let task = state
|
||||
.tasks
|
||||
.get_mut(task_id.trim())
|
||||
.ok_or(AiTaskServiceError::TaskNotFound)?;
|
||||
let previous_task = task.clone();
|
||||
if let Err(error) = apply(task) {
|
||||
*task = previous_task;
|
||||
return Err(error);
|
||||
}
|
||||
(previous_task, task.clone())
|
||||
};
|
||||
if let Err(error) = validate_task_memory_limits(&snapshot) {
|
||||
state
|
||||
.tasks
|
||||
.insert(task_id.trim().to_string(), previous_task);
|
||||
return Err(error);
|
||||
}
|
||||
let retained_output_bytes = retained_output_bytes(&state);
|
||||
if retained_output_bytes > MAX_AI_TASK_RETAINED_OUTPUT_BYTES {
|
||||
state
|
||||
.tasks
|
||||
.insert(task_id.trim().to_string(), previous_task);
|
||||
return Err(AiTaskServiceError::Store(
|
||||
"AI 任务仓储输出工作集超过内存上限".to_string(),
|
||||
));
|
||||
}
|
||||
if snapshot.status.is_terminal() {
|
||||
state.text_chunks.remove(task_id.trim());
|
||||
}
|
||||
Ok(snapshot)
|
||||
}
|
||||
|
||||
pub(super) fn append_text_chunk(
|
||||
@@ -67,13 +137,75 @@ impl InMemoryAiTaskStore {
|
||||
.inner
|
||||
.lock()
|
||||
.map_err(|_| AiTaskServiceError::Store("AI 任务仓储锁已中毒".to_string()))?;
|
||||
{
|
||||
if chunk.delta_text.len() > MAX_AI_TASK_TEXT_OUTPUT_BYTES {
|
||||
return Err(AiTaskServiceError::Store(
|
||||
"AI 任务文本输出超过内存上限".to_string(),
|
||||
));
|
||||
}
|
||||
let (previous_stage_output_bytes, previous_latest_output_bytes, previous_task) = {
|
||||
let task = state
|
||||
.tasks
|
||||
.get(&chunk.task_id)
|
||||
.ok_or(AiTaskServiceError::TaskNotFound)?;
|
||||
ensure_task_is_not_terminal(task.status)?;
|
||||
let stage = task
|
||||
.stages
|
||||
.iter()
|
||||
.find(|stage| stage.stage_kind == chunk.stage_kind)
|
||||
.ok_or(AiTaskServiceError::StageNotFound)?;
|
||||
(
|
||||
stage.text_output.as_ref().map_or(0, String::len),
|
||||
task.latest_text_output.as_ref().map_or(0, String::len),
|
||||
task.clone(),
|
||||
)
|
||||
};
|
||||
|
||||
let (previous_chunk, aggregated_bytes, aggregated_text) = {
|
||||
let chunks = state
|
||||
.text_chunks
|
||||
.get_mut(&chunk.task_id)
|
||||
.ok_or(AiTaskServiceError::TaskNotFound)?;
|
||||
let stage_chunks = chunks.entry(chunk.stage_kind).or_default();
|
||||
let previous_chunk = stage_chunks.insert(chunk.sequence, chunk.delta_text.clone());
|
||||
let aggregated_bytes = stage_chunks
|
||||
.values()
|
||||
.fold(0_usize, |total, delta| total.saturating_add(delta.len()));
|
||||
let mut aggregated_text = String::with_capacity(aggregated_bytes);
|
||||
for delta in stage_chunks.values() {
|
||||
aggregated_text.push_str(delta);
|
||||
}
|
||||
(previous_chunk, aggregated_bytes, aggregated_text)
|
||||
};
|
||||
if aggregated_bytes > MAX_AI_TASK_TEXT_OUTPUT_BYTES {
|
||||
rollback_text_chunk(&mut state, &chunk, previous_chunk);
|
||||
return Err(AiTaskServiceError::Store(
|
||||
"AI 任务文本输出超过内存上限".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let projected_retained_output_bytes = retained_output_bytes(&state)
|
||||
.saturating_sub(previous_stage_output_bytes)
|
||||
.saturating_sub(previous_latest_output_bytes)
|
||||
.saturating_add(aggregated_bytes.saturating_mul(2));
|
||||
if projected_retained_output_bytes > MAX_AI_TASK_RETAINED_OUTPUT_BYTES {
|
||||
rollback_text_chunk(&mut state, &chunk, previous_chunk);
|
||||
return Err(AiTaskServiceError::Store(
|
||||
"AI 任务仓储输出工作集超过内存上限".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let normalized_output = if aggregated_text.trim().is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(aggregated_text)
|
||||
};
|
||||
|
||||
let snapshot = {
|
||||
let task = state
|
||||
.tasks
|
||||
.get_mut(&chunk.task_id)
|
||||
.ok_or(AiTaskServiceError::TaskNotFound)?;
|
||||
ensure_task_is_not_terminal(task.status)?;
|
||||
|
||||
let stage = task
|
||||
.stages
|
||||
.iter_mut()
|
||||
@@ -83,45 +215,23 @@ impl InMemoryAiTaskStore {
|
||||
stage.status = AiTaskStageStatus::Running;
|
||||
stage.started_at_micros = Some(chunk.created_at_micros);
|
||||
}
|
||||
|
||||
task.status = AiTaskStatus::Running;
|
||||
task.started_at_micros
|
||||
.get_or_insert(chunk.created_at_micros);
|
||||
}
|
||||
|
||||
let chunks = state
|
||||
.text_chunks
|
||||
.get_mut(&chunk.task_id)
|
||||
.ok_or(AiTaskServiceError::TaskNotFound)?;
|
||||
chunks.push(chunk.clone());
|
||||
chunks.sort_by_key(|value| value.sequence);
|
||||
|
||||
let aggregated_text = chunks
|
||||
.iter()
|
||||
.filter(|value| value.stage_kind == chunk.stage_kind)
|
||||
.map(|value| value.delta_text.as_str())
|
||||
.collect::<Vec<_>>()
|
||||
.join("");
|
||||
let normalized_output = if aggregated_text.trim().is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(aggregated_text)
|
||||
stage.text_output = normalized_output.clone();
|
||||
task.latest_text_output = normalized_output;
|
||||
task.updated_at_micros = chunk.created_at_micros;
|
||||
task.version += 1;
|
||||
task.clone()
|
||||
};
|
||||
|
||||
let task = state
|
||||
.tasks
|
||||
.get_mut(&chunk.task_id)
|
||||
.ok_or(AiTaskServiceError::TaskNotFound)?;
|
||||
let stage = task
|
||||
.stages
|
||||
.iter_mut()
|
||||
.find(|stage| stage.stage_kind == chunk.stage_kind)
|
||||
.ok_or(AiTaskServiceError::StageNotFound)?;
|
||||
stage.text_output = normalized_output.clone();
|
||||
task.latest_text_output = normalized_output;
|
||||
task.updated_at_micros = chunk.created_at_micros;
|
||||
task.version += 1;
|
||||
Ok(task.clone())
|
||||
if let Err(error) = validate_task_memory_limits(&snapshot)
|
||||
.and_then(|_| validate_retained_output_bytes(&state))
|
||||
{
|
||||
state.tasks.insert(chunk.task_id.clone(), previous_task);
|
||||
rollback_text_chunk(&mut state, &chunk, previous_chunk);
|
||||
return Err(error);
|
||||
}
|
||||
Ok(snapshot)
|
||||
}
|
||||
|
||||
pub(super) fn get_task(&self, task_id: &str) -> Result<AiTaskSnapshot, AiTaskServiceError> {
|
||||
@@ -136,3 +246,109 @@ impl InMemoryAiTaskStore {
|
||||
.ok_or(AiTaskServiceError::TaskNotFound)
|
||||
}
|
||||
}
|
||||
|
||||
fn rollback_text_chunk(
|
||||
state: &mut InMemoryAiTaskStoreState,
|
||||
chunk: &AiTextChunkSnapshot,
|
||||
previous_chunk: Option<String>,
|
||||
) {
|
||||
if let Some(stage_chunks) = state
|
||||
.text_chunks
|
||||
.get_mut(&chunk.task_id)
|
||||
.and_then(|chunks| chunks.get_mut(&chunk.stage_kind))
|
||||
{
|
||||
if let Some(previous_chunk) = previous_chunk {
|
||||
stage_chunks.insert(chunk.sequence, previous_chunk);
|
||||
} else {
|
||||
stage_chunks.remove(&chunk.sequence);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn retained_output_bytes(state: &InMemoryAiTaskStoreState) -> usize {
|
||||
let snapshot_bytes = state.tasks.values().fold(0_usize, |total, task| {
|
||||
total.saturating_add(task_output_bytes(task))
|
||||
});
|
||||
state
|
||||
.text_chunks
|
||||
.values()
|
||||
.fold(snapshot_bytes, |total, stages| {
|
||||
stages.values().fold(total, |stage_total, chunks| {
|
||||
chunks.values().fold(stage_total, |chunk_total, delta| {
|
||||
chunk_total.saturating_add(delta.len())
|
||||
})
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
fn validate_retained_output_bytes(
|
||||
state: &InMemoryAiTaskStoreState,
|
||||
) -> Result<(), AiTaskServiceError> {
|
||||
if retained_output_bytes(state) > MAX_AI_TASK_RETAINED_OUTPUT_BYTES {
|
||||
return Err(AiTaskServiceError::Store(
|
||||
"AI 任务仓储输出工作集超过内存上限".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn task_output_bytes(task: &AiTaskSnapshot) -> usize {
|
||||
let task_metadata = task
|
||||
.task_id
|
||||
.len()
|
||||
.saturating_add(task.owner_user_id.len())
|
||||
.saturating_add(task.request_label.len())
|
||||
.saturating_add(task.source_module.len())
|
||||
.saturating_add(task.source_entity_id.as_ref().map_or(0, String::len))
|
||||
.saturating_add(task.stages.iter().fold(0_usize, |total, stage| {
|
||||
total
|
||||
.saturating_add(stage.label.len())
|
||||
.saturating_add(stage.detail.len())
|
||||
}));
|
||||
let request_payload = task.request_payload_json.as_ref().map_or(0, String::len);
|
||||
let failure_message = task.failure_message.as_ref().map_or(0, String::len);
|
||||
let result_references = task
|
||||
.result_references
|
||||
.iter()
|
||||
.fold(0_usize, |total, reference| {
|
||||
total
|
||||
.saturating_add(reference.result_ref_id.len())
|
||||
.saturating_add(reference.task_id.len())
|
||||
.saturating_add(reference.reference_id.len())
|
||||
.saturating_add(reference.label.as_ref().map_or(0, String::len))
|
||||
});
|
||||
let latest_text = task.latest_text_output.as_ref().map_or(0, String::len);
|
||||
let latest_structured = task
|
||||
.latest_structured_payload_json
|
||||
.as_ref()
|
||||
.map_or(0, String::len);
|
||||
let stage_bytes = task.stages.iter().fold(0_usize, |total, stage| {
|
||||
let text = stage.text_output.as_ref().map_or(0, String::len);
|
||||
let structured = stage
|
||||
.structured_payload_json
|
||||
.as_ref()
|
||||
.map_or(0, String::len);
|
||||
let warnings = stage
|
||||
.warning_messages
|
||||
.iter()
|
||||
.fold(0_usize, |warning_total, warning| {
|
||||
warning_total.saturating_add(warning.len())
|
||||
});
|
||||
total
|
||||
.saturating_add(text)
|
||||
.saturating_add(structured)
|
||||
.saturating_add(warnings)
|
||||
});
|
||||
task_metadata
|
||||
.saturating_add(request_payload)
|
||||
.saturating_add(failure_message)
|
||||
.saturating_add(result_references)
|
||||
.saturating_add(latest_text)
|
||||
.saturating_add(latest_structured)
|
||||
.saturating_add(stage_bytes)
|
||||
}
|
||||
|
||||
fn validate_task_memory_limits(task: &AiTaskSnapshot) -> Result<(), AiTaskServiceError> {
|
||||
validate_ai_task_snapshot_memory_limits(task)
|
||||
.map_err(|message| AiTaskServiceError::Store(message.to_string()))
|
||||
}
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
mod ids;
|
||||
mod limits;
|
||||
mod stages;
|
||||
mod types;
|
||||
|
||||
@@ -8,6 +9,16 @@ pub use ids::{
|
||||
generate_ai_task_stage_id, generate_ai_text_chunk_id, normalize_optional_text,
|
||||
normalize_string_list,
|
||||
};
|
||||
pub use limits::{
|
||||
MAX_AI_TASK_FAILURE_MESSAGE_BYTES, MAX_AI_TASK_ID_BYTES, MAX_AI_TASK_OWNER_USER_ID_BYTES,
|
||||
MAX_AI_TASK_REFERENCE_ID_BYTES, MAX_AI_TASK_REFERENCE_LABEL_BYTES,
|
||||
MAX_AI_TASK_REQUEST_LABEL_BYTES, MAX_AI_TASK_REQUEST_PAYLOAD_BYTES,
|
||||
MAX_AI_TASK_RESULT_REFERENCES, MAX_AI_TASK_RETAINED_OUTPUT_BYTES, MAX_AI_TASK_RETAINED_TASKS,
|
||||
MAX_AI_TASK_SOURCE_ENTITY_ID_BYTES, MAX_AI_TASK_SOURCE_MODULE_BYTES,
|
||||
MAX_AI_TASK_STAGE_DETAIL_BYTES, MAX_AI_TASK_STAGE_LABEL_BYTES,
|
||||
MAX_AI_TASK_STRUCTURED_OUTPUT_BYTES, MAX_AI_TASK_TEXT_OUTPUT_BYTES, MAX_AI_TASK_WARNING_BYTES,
|
||||
validate_ai_task_snapshot_memory_limits,
|
||||
};
|
||||
pub use types::{
|
||||
AiResultReferenceKind, AiResultReferenceSnapshot, AiTaskKind, AiTaskSnapshot,
|
||||
AiTaskStageBlueprint, AiTaskStageKind, AiTaskStageSnapshot, AiTaskStageStatus, AiTaskStatus,
|
||||
|
||||
@@ -0,0 +1,115 @@
|
||||
use super::types::AiTaskSnapshot;
|
||||
|
||||
pub const MAX_AI_TASK_RETAINED_TASKS: usize = 1024;
|
||||
pub const MAX_AI_TASK_ID_BYTES: usize = 256;
|
||||
pub const MAX_AI_TASK_OWNER_USER_ID_BYTES: usize = 256;
|
||||
pub const MAX_AI_TASK_REQUEST_LABEL_BYTES: usize = 4 * 1024;
|
||||
pub const MAX_AI_TASK_SOURCE_MODULE_BYTES: usize = 256;
|
||||
pub const MAX_AI_TASK_SOURCE_ENTITY_ID_BYTES: usize = 512;
|
||||
pub const MAX_AI_TASK_STAGE_LABEL_BYTES: usize = 4 * 1024;
|
||||
pub const MAX_AI_TASK_STAGE_DETAIL_BYTES: usize = 8 * 1024;
|
||||
pub const MAX_AI_TASK_TEXT_OUTPUT_BYTES: usize = 512 * 1024;
|
||||
pub const MAX_AI_TASK_STRUCTURED_OUTPUT_BYTES: usize = 512 * 1024;
|
||||
pub const MAX_AI_TASK_WARNING_BYTES: usize = 64 * 1024;
|
||||
pub const MAX_AI_TASK_REQUEST_PAYLOAD_BYTES: usize = 512 * 1024;
|
||||
pub const MAX_AI_TASK_FAILURE_MESSAGE_BYTES: usize = 64 * 1024;
|
||||
pub const MAX_AI_TASK_RESULT_REFERENCES: usize = 64;
|
||||
pub const MAX_AI_TASK_REFERENCE_ID_BYTES: usize = 512;
|
||||
pub const MAX_AI_TASK_REFERENCE_LABEL_BYTES: usize = 2 * 1024;
|
||||
pub const MAX_AI_TASK_RETAINED_OUTPUT_BYTES: usize = 64 * 1024 * 1024;
|
||||
|
||||
pub fn validate_ai_task_snapshot_memory_limits(task: &AiTaskSnapshot) -> Result<(), &'static str> {
|
||||
if task.task_id.len() > MAX_AI_TASK_ID_BYTES {
|
||||
return Err("AI 任务 ID 超过内存上限");
|
||||
}
|
||||
if task.owner_user_id.len() > MAX_AI_TASK_OWNER_USER_ID_BYTES {
|
||||
return Err("AI 任务用户 ID 超过内存上限");
|
||||
}
|
||||
if task.request_label.len() > MAX_AI_TASK_REQUEST_LABEL_BYTES {
|
||||
return Err("AI 任务请求标签超过内存上限");
|
||||
}
|
||||
if task.source_module.len() > MAX_AI_TASK_SOURCE_MODULE_BYTES {
|
||||
return Err("AI 任务来源模块超过内存上限");
|
||||
}
|
||||
if task
|
||||
.source_entity_id
|
||||
.as_ref()
|
||||
.is_some_and(|entity_id| entity_id.len() > MAX_AI_TASK_SOURCE_ENTITY_ID_BYTES)
|
||||
{
|
||||
return Err("AI 任务来源实体 ID 超过内存上限");
|
||||
}
|
||||
if task.stages.iter().any(|stage| {
|
||||
stage.label.len() > MAX_AI_TASK_STAGE_LABEL_BYTES
|
||||
|| stage.detail.len() > MAX_AI_TASK_STAGE_DETAIL_BYTES
|
||||
}) {
|
||||
return Err("AI 任务阶段元数据超过内存上限");
|
||||
}
|
||||
if task
|
||||
.request_payload_json
|
||||
.as_ref()
|
||||
.is_some_and(|payload| payload.len() > MAX_AI_TASK_REQUEST_PAYLOAD_BYTES)
|
||||
{
|
||||
return Err("AI 任务请求 payload 超过内存上限");
|
||||
}
|
||||
if task
|
||||
.failure_message
|
||||
.as_ref()
|
||||
.is_some_and(|message| message.len() > MAX_AI_TASK_FAILURE_MESSAGE_BYTES)
|
||||
{
|
||||
return Err("AI 任务失败消息超过内存上限");
|
||||
}
|
||||
if task.stages.iter().any(|stage| {
|
||||
stage
|
||||
.text_output
|
||||
.as_ref()
|
||||
.is_some_and(|text| text.len() > MAX_AI_TASK_TEXT_OUTPUT_BYTES)
|
||||
}) || task
|
||||
.latest_text_output
|
||||
.as_ref()
|
||||
.is_some_and(|text| text.len() > MAX_AI_TASK_TEXT_OUTPUT_BYTES)
|
||||
{
|
||||
return Err("AI 任务文本输出超过内存上限");
|
||||
}
|
||||
if task.stages.iter().any(|stage| {
|
||||
stage
|
||||
.structured_payload_json
|
||||
.as_ref()
|
||||
.is_some_and(|payload| payload.len() > MAX_AI_TASK_STRUCTURED_OUTPUT_BYTES)
|
||||
}) || task
|
||||
.latest_structured_payload_json
|
||||
.as_ref()
|
||||
.is_some_and(|payload| payload.len() > MAX_AI_TASK_STRUCTURED_OUTPUT_BYTES)
|
||||
{
|
||||
return Err("AI 任务结构化输出超过内存上限");
|
||||
}
|
||||
if task.stages.iter().any(|stage| {
|
||||
stage
|
||||
.warning_messages
|
||||
.iter()
|
||||
.fold(0_usize, |total, warning| {
|
||||
total.saturating_add(warning.len())
|
||||
})
|
||||
> MAX_AI_TASK_WARNING_BYTES
|
||||
}) {
|
||||
return Err("AI 任务 warning 输出超过内存上限");
|
||||
}
|
||||
if task.result_references.len() > MAX_AI_TASK_RESULT_REFERENCES {
|
||||
return Err("AI 任务结果引用数量超过内存上限");
|
||||
}
|
||||
if task
|
||||
.result_references
|
||||
.iter()
|
||||
.any(|reference| reference.reference_id.len() > MAX_AI_TASK_REFERENCE_ID_BYTES)
|
||||
{
|
||||
return Err("AI 任务结果引用 ID 超过内存上限");
|
||||
}
|
||||
if task.result_references.iter().any(|reference| {
|
||||
reference
|
||||
.label
|
||||
.as_ref()
|
||||
.is_some_and(|label| label.len() > MAX_AI_TASK_REFERENCE_LABEL_BYTES)
|
||||
}) {
|
||||
return Err("AI 任务结果引用标签超过内存上限");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
@@ -14,9 +14,17 @@ pub use domain::{
|
||||
AI_RESULT_REF_ID_PREFIX, AI_TASK_ID_PREFIX, AI_TASK_STAGE_ID_PREFIX, AI_TEXT_CHUNK_ID_PREFIX,
|
||||
AiResultReferenceKind, AiResultReferenceSnapshot, AiTaskKind, AiTaskSnapshot,
|
||||
AiTaskStageBlueprint, AiTaskStageKind, AiTaskStageSnapshot, AiTaskStageStatus, AiTaskStatus,
|
||||
AiTextChunkSnapshot, INITIAL_AI_TASK_VERSION, generate_ai_result_ref_id, generate_ai_task_id,
|
||||
generate_ai_task_stage_id, generate_ai_text_chunk_id, normalize_optional_text,
|
||||
normalize_string_list,
|
||||
AiTextChunkSnapshot, INITIAL_AI_TASK_VERSION, MAX_AI_TASK_FAILURE_MESSAGE_BYTES,
|
||||
MAX_AI_TASK_ID_BYTES, MAX_AI_TASK_OWNER_USER_ID_BYTES, MAX_AI_TASK_REFERENCE_ID_BYTES,
|
||||
MAX_AI_TASK_REFERENCE_LABEL_BYTES, MAX_AI_TASK_REQUEST_LABEL_BYTES,
|
||||
MAX_AI_TASK_REQUEST_PAYLOAD_BYTES, MAX_AI_TASK_RESULT_REFERENCES,
|
||||
MAX_AI_TASK_RETAINED_OUTPUT_BYTES, MAX_AI_TASK_RETAINED_TASKS,
|
||||
MAX_AI_TASK_SOURCE_ENTITY_ID_BYTES, MAX_AI_TASK_SOURCE_MODULE_BYTES,
|
||||
MAX_AI_TASK_STAGE_DETAIL_BYTES, MAX_AI_TASK_STAGE_LABEL_BYTES,
|
||||
MAX_AI_TASK_STRUCTURED_OUTPUT_BYTES, MAX_AI_TASK_TEXT_OUTPUT_BYTES, MAX_AI_TASK_WARNING_BYTES,
|
||||
generate_ai_result_ref_id, generate_ai_task_id, generate_ai_task_stage_id,
|
||||
generate_ai_text_chunk_id, normalize_optional_text, normalize_string_list,
|
||||
validate_ai_task_snapshot_memory_limits,
|
||||
};
|
||||
pub use errors::{AiTaskFieldError, AiTaskServiceError};
|
||||
pub use events::AiTaskDomainEvent;
|
||||
|
||||
@@ -43,6 +43,32 @@ fn create_task_rejects_duplicate_stage_blueprints() {
|
||||
assert_eq!(error, AiTaskFieldError::DuplicateStageBlueprint);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn create_task_rejects_oversized_request_payload() {
|
||||
let service = build_service();
|
||||
let mut input = build_create_input(AiTaskKind::StoryGeneration);
|
||||
input.request_payload_json = Some("x".repeat(MAX_AI_TASK_REQUEST_PAYLOAD_BYTES + 1));
|
||||
|
||||
let error = service
|
||||
.create_task(input)
|
||||
.expect_err("request payload over the memory cap should fail");
|
||||
assert!(
|
||||
matches!(error, AiTaskServiceError::Store(message) if message.contains("请求 payload"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn create_task_rejects_oversized_request_metadata() {
|
||||
let service = build_service();
|
||||
let mut input = build_create_input(AiTaskKind::StoryGeneration);
|
||||
input.request_label = "x".repeat(MAX_AI_TASK_REQUEST_LABEL_BYTES + 1);
|
||||
|
||||
let error = service
|
||||
.create_task(input)
|
||||
.expect_err("request metadata over the memory cap should fail");
|
||||
assert!(matches!(error, AiTaskServiceError::Store(message) if message.contains("请求标签")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn generate_ai_task_stage_id_contains_task_and_stage_slug() {
|
||||
let stage_id = generate_ai_task_stage_id("aitask_demo", AiTaskStageKind::NormalizeResult);
|
||||
@@ -112,6 +138,47 @@ fn append_text_chunk_aggregates_stream_output_by_stage() {
|
||||
assert_eq!(second_chunk.sequence, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn append_text_chunk_rejects_output_over_stage_memory_limit_without_mutating_task() {
|
||||
let service = build_service();
|
||||
let task = service
|
||||
.create_task(build_create_input(AiTaskKind::CharacterChat))
|
||||
.expect("task should create");
|
||||
let max_output = "a".repeat(512 * 1024);
|
||||
|
||||
let (updated, _) = service
|
||||
.append_text_chunk(
|
||||
&task.task_id,
|
||||
AiTaskStageKind::RequestModel,
|
||||
1,
|
||||
max_output.clone(),
|
||||
task.created_at_micros + 1,
|
||||
)
|
||||
.expect("the stage limit itself should be accepted");
|
||||
assert_eq!(
|
||||
updated.latest_text_output.as_deref().map(str::len),
|
||||
Some(max_output.len())
|
||||
);
|
||||
|
||||
let error = service
|
||||
.append_text_chunk(
|
||||
&task.task_id,
|
||||
AiTaskStageKind::RequestModel,
|
||||
2,
|
||||
"b".to_string(),
|
||||
task.created_at_micros + 2,
|
||||
)
|
||||
.expect_err("output beyond the stage limit should fail");
|
||||
assert!(matches!(error, AiTaskServiceError::Store(_)));
|
||||
let after_rejection = service
|
||||
.get_task(&task.task_id)
|
||||
.expect("task should remain readable");
|
||||
assert_eq!(
|
||||
after_rejection.latest_text_output.as_deref().map(str::len),
|
||||
Some(max_output.len())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn complete_stage_updates_latest_outputs() {
|
||||
let service = build_service();
|
||||
@@ -147,6 +214,150 @@ fn complete_stage_updates_latest_outputs() {
|
||||
assert_eq!(stage.warning_messages, vec!["使用了 fallback 选项池"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn complete_stage_rejects_oversized_text_output_without_mutating_task() {
|
||||
let service = build_service();
|
||||
let task = service
|
||||
.create_task(build_create_input(AiTaskKind::StoryGeneration))
|
||||
.expect("task should create");
|
||||
let oversized = "x".repeat(512 * 1024 + 1);
|
||||
|
||||
let error = service
|
||||
.complete_stage(AiStageCompletionInput {
|
||||
task_id: task.task_id.clone(),
|
||||
stage_kind: AiTaskStageKind::NormalizeResult,
|
||||
text_output: Some(oversized),
|
||||
structured_payload_json: None,
|
||||
warning_messages: Vec::new(),
|
||||
completed_at_micros: task.created_at_micros + 1,
|
||||
})
|
||||
.expect_err("text output over the per-stage cap should fail");
|
||||
assert!(matches!(error, AiTaskServiceError::Store(message) if message.contains("文本输出")));
|
||||
|
||||
let unchanged = service
|
||||
.get_task(&task.task_id)
|
||||
.expect("task should remain readable");
|
||||
let stage = unchanged
|
||||
.stages
|
||||
.iter()
|
||||
.find(|stage| stage.stage_kind == AiTaskStageKind::NormalizeResult)
|
||||
.expect("normalize stage should exist");
|
||||
assert_eq!(stage.status, AiTaskStageStatus::Pending);
|
||||
assert!(stage.text_output.is_none());
|
||||
assert!(unchanged.latest_text_output.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn complete_stage_rejects_oversized_structured_output_without_mutating_task() {
|
||||
let service = build_service();
|
||||
let task = service
|
||||
.create_task(build_create_input(AiTaskKind::StoryGeneration))
|
||||
.expect("task should create");
|
||||
let oversized = "x".repeat(512 * 1024 + 1);
|
||||
|
||||
let error = service
|
||||
.complete_stage(AiStageCompletionInput {
|
||||
task_id: task.task_id.clone(),
|
||||
stage_kind: AiTaskStageKind::NormalizeResult,
|
||||
text_output: None,
|
||||
structured_payload_json: Some(oversized),
|
||||
warning_messages: Vec::new(),
|
||||
completed_at_micros: task.created_at_micros + 1,
|
||||
})
|
||||
.expect_err("structured output over the per-stage cap should fail");
|
||||
assert!(matches!(error, AiTaskServiceError::Store(message) if message.contains("结构化输出")));
|
||||
|
||||
let unchanged = service
|
||||
.get_task(&task.task_id)
|
||||
.expect("task should remain readable");
|
||||
let stage = unchanged
|
||||
.stages
|
||||
.iter()
|
||||
.find(|stage| stage.stage_kind == AiTaskStageKind::NormalizeResult)
|
||||
.expect("normalize stage should exist");
|
||||
assert_eq!(stage.status, AiTaskStageStatus::Pending);
|
||||
assert!(stage.structured_payload_json.is_none());
|
||||
assert!(unchanged.latest_structured_payload_json.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn complete_stage_rejects_oversized_warning_output_without_mutating_task() {
|
||||
let service = build_service();
|
||||
let task = service
|
||||
.create_task(build_create_input(AiTaskKind::StoryGeneration))
|
||||
.expect("task should create");
|
||||
let oversized_warning = "w".repeat(64 * 1024 + 1);
|
||||
|
||||
let error = service
|
||||
.complete_stage(AiStageCompletionInput {
|
||||
task_id: task.task_id.clone(),
|
||||
stage_kind: AiTaskStageKind::NormalizeResult,
|
||||
text_output: None,
|
||||
structured_payload_json: None,
|
||||
warning_messages: vec![oversized_warning],
|
||||
completed_at_micros: task.created_at_micros + 1,
|
||||
})
|
||||
.expect_err("warning output over the per-stage cap should fail");
|
||||
assert!(matches!(error, AiTaskServiceError::Store(message) if message.contains("warning")));
|
||||
|
||||
let unchanged = service
|
||||
.get_task(&task.task_id)
|
||||
.expect("task should remain readable");
|
||||
let stage = unchanged
|
||||
.stages
|
||||
.iter()
|
||||
.find(|stage| stage.stage_kind == AiTaskStageKind::NormalizeResult)
|
||||
.expect("normalize stage should exist");
|
||||
assert_eq!(stage.status, AiTaskStageStatus::Pending);
|
||||
assert!(stage.warning_messages.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn complete_stage_enforces_global_retained_output_cap() {
|
||||
let service = build_service();
|
||||
let structured_payload = "x".repeat(512 * 1024);
|
||||
for index in 0..63 {
|
||||
let task = service
|
||||
.create_task(AiTaskCreateInput {
|
||||
task_id: format!("task-structured-cap-{index}"),
|
||||
..build_create_input(AiTaskKind::StoryGeneration)
|
||||
})
|
||||
.expect("task should create");
|
||||
service
|
||||
.complete_stage(AiStageCompletionInput {
|
||||
task_id: task.task_id,
|
||||
stage_kind: AiTaskStageKind::NormalizeResult,
|
||||
text_output: None,
|
||||
structured_payload_json: Some(structured_payload.clone()),
|
||||
warning_messages: Vec::new(),
|
||||
completed_at_micros: task.created_at_micros + 1,
|
||||
})
|
||||
.expect("63 MiB retained output should remain within the global cap");
|
||||
}
|
||||
|
||||
let task = service
|
||||
.create_task(AiTaskCreateInput {
|
||||
task_id: "task-structured-cap-overflow".to_string(),
|
||||
..build_create_input(AiTaskKind::StoryGeneration)
|
||||
})
|
||||
.expect("the overflow candidate task itself should create");
|
||||
let error = service
|
||||
.complete_stage(AiStageCompletionInput {
|
||||
task_id: task.task_id.clone(),
|
||||
stage_kind: AiTaskStageKind::NormalizeResult,
|
||||
text_output: None,
|
||||
structured_payload_json: Some(structured_payload),
|
||||
warning_messages: Vec::new(),
|
||||
completed_at_micros: task.created_at_micros + 1,
|
||||
})
|
||||
.expect_err("global retained output cap should reject the overflow");
|
||||
assert!(matches!(error, AiTaskServiceError::Store(message) if message.contains("工作集")));
|
||||
let unchanged = service
|
||||
.get_task(&task.task_id)
|
||||
.expect("overflow task should remain readable");
|
||||
assert!(unchanged.latest_structured_payload_json.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn attach_result_reference_appends_binding() {
|
||||
let service = build_service();
|
||||
@@ -172,6 +383,47 @@ fn attach_result_reference_appends_binding() {
|
||||
assert_eq!(updated.result_references[0].reference_id, "profile_001");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn attach_result_reference_rejects_unbounded_reference_growth() {
|
||||
let service = build_service();
|
||||
let task = service
|
||||
.create_task(build_create_input(AiTaskKind::CustomWorldGeneration))
|
||||
.expect("task should create");
|
||||
|
||||
for index in 0..MAX_AI_TASK_RESULT_REFERENCES {
|
||||
service
|
||||
.attach_result_reference(
|
||||
&task.task_id,
|
||||
AiResultReferenceKind::CustomWorldProfile,
|
||||
format!("profile_{index}"),
|
||||
None,
|
||||
task.created_at_micros + index as i64 + 1,
|
||||
)
|
||||
.expect("references within the cap should attach");
|
||||
}
|
||||
|
||||
let error = service
|
||||
.attach_result_reference(
|
||||
&task.task_id,
|
||||
AiResultReferenceKind::CustomWorldProfile,
|
||||
"profile_overflow".to_string(),
|
||||
None,
|
||||
task.created_at_micros + MAX_AI_TASK_RESULT_REFERENCES as i64 + 1,
|
||||
)
|
||||
.expect_err("references over the cap should fail");
|
||||
assert!(
|
||||
matches!(error, AiTaskServiceError::Store(message) if message.contains("结果引用数量"))
|
||||
);
|
||||
|
||||
let unchanged = service
|
||||
.get_task(&task.task_id)
|
||||
.expect("task should remain readable");
|
||||
assert_eq!(
|
||||
unchanged.result_references.len(),
|
||||
MAX_AI_TASK_RESULT_REFERENCES
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn fail_and_cancel_task_move_into_terminal_states() {
|
||||
let service = build_service();
|
||||
|
||||
@@ -28,6 +28,11 @@ use shared_kernel::{
|
||||
use time::{Duration, OffsetDateTime};
|
||||
use tracing::{info, warn};
|
||||
|
||||
const REFRESH_SESSION_STALE_RETENTION: Duration = Duration::days(1);
|
||||
const MAX_REFRESH_SESSIONS: usize = 8_192;
|
||||
const MAX_PHONE_CODES: usize = 4_096;
|
||||
const MAX_WECHAT_STATES: usize = 4_096;
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct InMemoryAuthStore {
|
||||
inner: Arc<Mutex<InMemoryAuthStoreState>>,
|
||||
@@ -364,6 +369,7 @@ impl RefreshSessionService {
|
||||
input: CreateRefreshSessionInput,
|
||||
now: OffsetDateTime,
|
||||
) -> Result<CreateRefreshSessionResult, RefreshSessionError> {
|
||||
self.store.prune_stale_sessions(now)?;
|
||||
self.store
|
||||
.find_by_user_id(&input.user_id)
|
||||
.map_err(map_password_store_error)?
|
||||
@@ -400,6 +406,7 @@ impl RefreshSessionService {
|
||||
input: RotateRefreshSessionInput,
|
||||
now: OffsetDateTime,
|
||||
) -> Result<RotateRefreshSessionResult, RefreshSessionError> {
|
||||
self.store.prune_stale_sessions(now)?;
|
||||
let Some(refresh_token_hash) = normalize_required_string(&input.refresh_token_hash) else {
|
||||
return Err(RefreshSessionError::MissingToken);
|
||||
};
|
||||
@@ -454,6 +461,7 @@ impl RefreshSessionService {
|
||||
user_id: &str,
|
||||
now: OffsetDateTime,
|
||||
) -> Result<ListActiveRefreshSessionsResult, RefreshSessionError> {
|
||||
self.store.prune_stale_sessions(now)?;
|
||||
self.store
|
||||
.find_by_user_id(user_id)
|
||||
.map_err(map_password_store_error)?
|
||||
@@ -468,6 +476,7 @@ impl RefreshSessionService {
|
||||
input: RevokeRefreshSessionByUserInput,
|
||||
now: OffsetDateTime,
|
||||
) -> Result<RevokeRefreshSessionResult, RefreshSessionError> {
|
||||
self.store.prune_stale_sessions(now)?;
|
||||
self.store
|
||||
.find_by_user_id(&input.user_id)
|
||||
.map_err(map_password_store_error)?
|
||||
@@ -492,6 +501,7 @@ impl RefreshSessionService {
|
||||
session_id: &str,
|
||||
now: OffsetDateTime,
|
||||
) -> Result<bool, RefreshSessionError> {
|
||||
self.store.prune_stale_sessions(now)?;
|
||||
self.store
|
||||
.is_session_active_for_user(user_id, session_id.trim(), now)
|
||||
}
|
||||
@@ -511,6 +521,7 @@ impl PhoneAuthService {
|
||||
input: SendPhoneCodeInput,
|
||||
now: OffsetDateTime,
|
||||
) -> Result<SendPhoneCodeResult, PhoneAuthError> {
|
||||
self.store.prune_expired_phone_codes(now)?;
|
||||
let scene = input.scene.clone();
|
||||
validate_mainland_china_country_code(input.country_code.as_deref())?;
|
||||
let normalized_phone = normalize_mainland_china_phone_number(&input.pure_phone_number)?;
|
||||
@@ -525,6 +536,8 @@ impl PhoneAuthService {
|
||||
);
|
||||
self.store
|
||||
.ensure_phone_code_not_cooling_down(&normalized_phone.e164, &scene, now)?;
|
||||
self.store
|
||||
.ensure_phone_code_capacity(&normalized_phone.e164, &scene)?;
|
||||
let expires_at = now
|
||||
.checked_add(Duration::minutes(SMS_CODE_TTL_MINUTES))
|
||||
.ok_or_else(|| PhoneAuthError::Store("短信验证码过期时间计算溢出".to_string()))?;
|
||||
@@ -787,6 +800,7 @@ impl WechatAuthStateService {
|
||||
input: CreateWechatAuthStateInput,
|
||||
now: OffsetDateTime,
|
||||
) -> Result<CreateWechatAuthStateResult, WechatAuthError> {
|
||||
self.store.prune_wechat_states(now)?;
|
||||
let created_at = format_rfc3339(now).map_err(|message| {
|
||||
WechatAuthError::Store(format!("微信 state 时间格式化失败：{message}"))
|
||||
})?;
|
||||
@@ -1063,10 +1077,25 @@ impl InMemoryAuthStoreState {
|
||||
}
|
||||
}
|
||||
|
||||
let now = OffsetDateTime::now_utc();
|
||||
let mut retained_refresh_session_count = 0_usize;
|
||||
for session in view.refresh_sessions {
|
||||
if !existing_user_ids.contains(&session.user_id) {
|
||||
continue;
|
||||
}
|
||||
if should_prune_refresh_session_fields(
|
||||
&session.expires_at,
|
||||
session.revoked_at.as_deref(),
|
||||
now,
|
||||
) {
|
||||
continue;
|
||||
}
|
||||
retained_refresh_session_count += 1;
|
||||
if retained_refresh_session_count > MAX_REFRESH_SESSIONS {
|
||||
return Err(format!(
|
||||
"认证投影中的 refresh session 数量超过内存上限（最多 {MAX_REFRESH_SESSIONS} 条）"
|
||||
));
|
||||
}
|
||||
let client_info =
|
||||
serde_json::from_str::<RefreshSessionClientInfo>(&session.client_info_json)
|
||||
.map_err(|error| format!("解析 refresh session 客户端信息失败：{error}"))?;
|
||||
@@ -1185,6 +1214,8 @@ impl InMemoryAuthStore {
|
||||
&self,
|
||||
updated_at_micros: i64,
|
||||
) -> Result<AuthStoreProjectionView, String> {
|
||||
self.prune_stale_sessions(OffsetDateTime::now_utc())
|
||||
.map_err(|error| error.to_string())?;
|
||||
let state = self
|
||||
.inner
|
||||
.lock()
|
||||
@@ -1262,6 +1293,38 @@ impl InMemoryAuthStore {
|
||||
})
|
||||
}
|
||||
|
||||
fn prune_stale_sessions(&self, now: OffsetDateTime) -> Result<(), RefreshSessionError> {
|
||||
let mut state = self
|
||||
.inner
|
||||
.lock()
|
||||
.map_err(|_| RefreshSessionError::Store("会话仓储锁已中毒".to_string()))?;
|
||||
let stale_session_ids = state
|
||||
.sessions_by_id
|
||||
.iter()
|
||||
.filter(|(_, stored)| should_prune_refresh_session(&stored.session, now))
|
||||
.map(|(session_id, _)| session_id.clone())
|
||||
.collect::<Vec<_>>();
|
||||
if stale_session_ids.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
for session_id in stale_session_ids {
|
||||
let Some(stored) = state.sessions_by_id.remove(&session_id) else {
|
||||
continue;
|
||||
};
|
||||
if state
|
||||
.session_id_by_refresh_token_hash
|
||||
.get(&stored.session.refresh_token_hash)
|
||||
.is_some_and(|mapped_id| mapped_id == &session_id)
|
||||
{
|
||||
state
|
||||
.session_id_by_refresh_token_hash
|
||||
.remove(&stored.session.refresh_token_hash);
|
||||
}
|
||||
}
|
||||
self.persist_refresh_state(&state)
|
||||
}
|
||||
|
||||
fn persist_state(&self, state: &InMemoryAuthStoreState) -> Result<(), String> {
|
||||
let _ = state;
|
||||
Ok(())
|
||||
@@ -1894,6 +1957,11 @@ impl InMemoryAuthStore {
|
||||
"refresh token hash 已存在，无法重复创建会话".to_string(),
|
||||
));
|
||||
}
|
||||
if state.sessions_by_id.len() >= MAX_REFRESH_SESSIONS {
|
||||
return Err(RefreshSessionError::Store(
|
||||
"refresh session 内存容量已达到上限".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
state.session_id_by_refresh_token_hash.insert(
|
||||
session.refresh_token_hash.clone(),
|
||||
@@ -1918,10 +1986,41 @@ impl InMemoryAuthStore {
|
||||
.map_err(|_| PhoneAuthError::Store("短信验证码仓储锁已中毒".to_string()))?;
|
||||
// 手机号和业务场景共同决定同一份验证码快照，重复发送时直接覆盖旧值。
|
||||
let key = build_phone_code_key(&code.phone_number, &code.scene);
|
||||
if !state.phone_codes_by_key.contains_key(&key)
|
||||
&& state.phone_codes_by_key.len() >= MAX_PHONE_CODES
|
||||
{
|
||||
return Err(PhoneAuthError::Store(
|
||||
"短信验证码内存容量已达到上限，请稍后重试".to_string(),
|
||||
));
|
||||
}
|
||||
state.phone_codes_by_key.insert(key, code);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn prune_expired_phone_codes(&self, now: OffsetDateTime) -> Result<(), PhoneAuthError> {
|
||||
let mut state = self
|
||||
.inner
|
||||
.lock()
|
||||
.map_err(|_| PhoneAuthError::Store("短信验证码仓储锁已中毒".to_string()))?;
|
||||
let expired_keys = state
|
||||
.phone_codes_by_key
|
||||
.iter()
|
||||
.filter_map(|(key, stored)| {
|
||||
OffsetDateTime::parse(
|
||||
&stored.expires_at,
|
||||
&time::format_description::well_known::Rfc3339,
|
||||
)
|
||||
.ok()
|
||||
.filter(|expires_at| *expires_at <= now)
|
||||
.map(|_| key.clone())
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
for key in expired_keys {
|
||||
state.phone_codes_by_key.remove(&key);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn ensure_phone_code_not_cooling_down(
|
||||
&self,
|
||||
phone_number: &str,
|
||||
@@ -1961,6 +2060,26 @@ impl InMemoryAuthStore {
|
||||
})
|
||||
}
|
||||
|
||||
fn ensure_phone_code_capacity(
|
||||
&self,
|
||||
phone_number: &str,
|
||||
scene: &PhoneAuthScene,
|
||||
) -> Result<(), PhoneAuthError> {
|
||||
let state = self
|
||||
.inner
|
||||
.lock()
|
||||
.map_err(|_| PhoneAuthError::Store("短信验证码仓储锁已中毒".to_string()))?;
|
||||
let key = build_phone_code_key(phone_number, scene);
|
||||
if state.phone_codes_by_key.contains_key(&key)
|
||||
|| state.phone_codes_by_key.len() < MAX_PHONE_CODES
|
||||
{
|
||||
return Ok(());
|
||||
}
|
||||
Err(PhoneAuthError::Store(
|
||||
"短信验证码内存容量已达到上限，请稍后重试".to_string(),
|
||||
))
|
||||
}
|
||||
|
||||
fn get_active_phone_code(
|
||||
&self,
|
||||
phone_number: &str,
|
||||
@@ -2041,6 +2160,11 @@ impl InMemoryAuthStore {
|
||||
{
|
||||
return Err(WechatAuthError::Store("微信 state 已存在".to_string()));
|
||||
}
|
||||
if state.wechat_states_by_token.len() >= MAX_WECHAT_STATES {
|
||||
return Err(WechatAuthError::Store(
|
||||
"微信登录 state 内存容量已达到上限，请稍后重试".to_string(),
|
||||
));
|
||||
}
|
||||
state.wechat_states_by_token.insert(
|
||||
state_record.state_token.clone(),
|
||||
StoredWechatAuthState {
|
||||
@@ -2405,6 +2529,33 @@ impl InMemoryAuthStore {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn prune_wechat_states(&self, now: OffsetDateTime) -> Result<(), WechatAuthError> {
|
||||
let mut state = self
|
||||
.inner
|
||||
.lock()
|
||||
.map_err(|_| WechatAuthError::Store("微信 state 仓储锁已中毒".to_string()))?;
|
||||
let stale_tokens = state
|
||||
.wechat_states_by_token
|
||||
.iter()
|
||||
.filter_map(|(token, stored)| {
|
||||
if stored.state.consumed_at.is_some() {
|
||||
return Some(token.clone());
|
||||
}
|
||||
OffsetDateTime::parse(
|
||||
&stored.state.expires_at,
|
||||
&time::format_description::well_known::Rfc3339,
|
||||
)
|
||||
.ok()
|
||||
.filter(|expires_at| *expires_at <= now)
|
||||
.map(|_| token.clone())
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
for token in stale_tokens {
|
||||
state.wechat_states_by_token.remove(&token);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn revoke_session_by_user_and_session_id(
|
||||
&self,
|
||||
user_id: &str,
|
||||
@@ -2568,6 +2719,25 @@ impl InMemoryAuthStore {
|
||||
}
|
||||
}
|
||||
|
||||
fn should_prune_refresh_session(session: &RefreshSessionRecord, now: OffsetDateTime) -> bool {
|
||||
should_prune_refresh_session_fields(&session.expires_at, session.revoked_at.as_deref(), now)
|
||||
}
|
||||
|
||||
fn should_prune_refresh_session_fields(
|
||||
expires_at: &str,
|
||||
revoked_at: Option<&str>,
|
||||
now: OffsetDateTime,
|
||||
) -> bool {
|
||||
let stale_before = now.saturating_sub(REFRESH_SESSION_STALE_RETENTION);
|
||||
if let Some(revoked_at) = revoked_at {
|
||||
return OffsetDateTime::parse(revoked_at, &time::format_description::well_known::Rfc3339)
|
||||
.is_ok_and(|timestamp| timestamp <= stale_before);
|
||||
}
|
||||
|
||||
OffsetDateTime::parse(expires_at, &time::format_description::well_known::Rfc3339)
|
||||
.is_ok_and(|timestamp| timestamp <= stale_before)
|
||||
}
|
||||
|
||||
fn map_sms_provider_error_to_phone_error(error: SmsProviderError) -> PhoneAuthError {
|
||||
match error {
|
||||
SmsProviderError::InvalidVerifyCode => PhoneAuthError::InvalidVerifyCode,
|
||||
@@ -4012,6 +4182,105 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn stale_refresh_sessions_are_pruned_from_both_indexes() {
|
||||
let store = build_store();
|
||||
let refresh_service = build_refresh_service(store.clone());
|
||||
let user = create_phone_login_user(store.clone(), "13800138008").await;
|
||||
let now = OffsetDateTime::now_utc();
|
||||
|
||||
refresh_service
|
||||
.create_session(
|
||||
CreateRefreshSessionInput {
|
||||
user_id: user.id.clone(),
|
||||
refresh_token_hash: hash_refresh_session_token("stale-revoked"),
|
||||
issued_by_provider: AuthLoginMethod::Password,
|
||||
client_info: build_client_info(),
|
||||
},
|
||||
now - Duration::days(2),
|
||||
)
|
||||
.expect("stale session should create");
|
||||
store
|
||||
.revoke_session_by_refresh_token_hash(
|
||||
&hash_refresh_session_token("stale-revoked"),
|
||||
now - Duration::days(2),
|
||||
)
|
||||
.expect("stale session should revoke");
|
||||
|
||||
refresh_service
|
||||
.create_session(
|
||||
CreateRefreshSessionInput {
|
||||
user_id: user.id.clone(),
|
||||
refresh_token_hash: hash_refresh_session_token("recent-revoked"),
|
||||
issued_by_provider: AuthLoginMethod::Password,
|
||||
client_info: build_client_info(),
|
||||
},
|
||||
now,
|
||||
)
|
||||
.expect("recent session should create");
|
||||
store
|
||||
.revoke_session_by_refresh_token_hash(
|
||||
&hash_refresh_session_token("recent-revoked"),
|
||||
now,
|
||||
)
|
||||
.expect("recent session should revoke");
|
||||
|
||||
let projection = store
|
||||
.export_projection_view(now.unix_timestamp())
|
||||
.expect("projection export should prune stale sessions");
|
||||
assert_eq!(projection.refresh_sessions.len(), 1);
|
||||
assert_eq!(
|
||||
projection.refresh_sessions[0].refresh_token_hash,
|
||||
hash_refresh_session_token("recent-revoked")
|
||||
);
|
||||
|
||||
let stale_error = refresh_service
|
||||
.rotate_session(
|
||||
RotateRefreshSessionInput {
|
||||
refresh_token_hash: hash_refresh_session_token("stale-revoked"),
|
||||
next_refresh_token_hash: hash_refresh_session_token("stale-next"),
|
||||
},
|
||||
now,
|
||||
)
|
||||
.expect_err("pruned session should no longer be indexed");
|
||||
assert_eq!(stale_error, RefreshSessionError::SessionNotFound);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn projection_restore_rejects_too_many_retained_refresh_sessions() {
|
||||
let client_info_json =
|
||||
serde_json::to_string(&build_client_info()).expect("client info should serialize");
|
||||
let refresh_sessions = (0..=MAX_REFRESH_SESSIONS)
|
||||
.map(|index| AuthStoreProjectionRefreshSession {
|
||||
session_id: format!("session-{index}"),
|
||||
user_id: "user_projection_cap".to_string(),
|
||||
refresh_token_hash: format!("hash-{index}"),
|
||||
issued_by_provider: "password".to_string(),
|
||||
client_info_json: client_info_json.clone(),
|
||||
expires_at: "2999-01-01T00:00:00Z".to_string(),
|
||||
revoked_at: None,
|
||||
created_at: "2026-01-01T00:00:00Z".to_string(),
|
||||
updated_at: "2026-01-01T00:00:00Z".to_string(),
|
||||
last_seen_at: "2026-01-01T00:00:00Z".to_string(),
|
||||
})
|
||||
.collect();
|
||||
|
||||
let error = InMemoryAuthStore::from_projection_view(AuthStoreProjectionView {
|
||||
updated_at_micros: 1,
|
||||
users: vec![projection_user(
|
||||
"user_projection_cap",
|
||||
"projection_cap",
|
||||
None,
|
||||
)],
|
||||
identities: vec![],
|
||||
refresh_sessions,
|
||||
})
|
||||
.expect_err("projection restore must enforce the refresh session cap");
|
||||
|
||||
assert!(error.contains("refresh session"));
|
||||
assert!(error.contains(&MAX_REFRESH_SESSIONS.to_string()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn wechat_login_hits_existing_user_by_union_id_before_openid() {
|
||||
let store = build_store();
|
||||
|
||||
@@ -409,6 +409,8 @@ pub mod external_generation_job_procedure_result_type;
|
||||
pub mod external_generation_job_renew_lease_input_type;
|
||||
pub mod external_generation_job_result_procedure_result_type;
|
||||
pub mod external_generation_job_result_snapshot_type;
|
||||
pub mod external_generation_job_retention_input_type;
|
||||
pub mod external_generation_job_retention_procedure_result_type;
|
||||
pub mod external_generation_job_snapshot_type;
|
||||
pub mod external_generation_job_summary_backfill_input_type;
|
||||
pub mod external_generation_job_summary_backfill_procedure_result_type;
|
||||
@@ -570,6 +572,7 @@ pub mod profile_wallet_ledger_table;
|
||||
pub mod profile_wallet_ledger_type;
|
||||
pub mod profile_wallet_manual_restriction_table;
|
||||
pub mod profile_wallet_manual_restriction_type;
|
||||
pub mod prune_external_generation_job_history_and_return_procedure;
|
||||
pub mod public_work_like_table;
|
||||
pub mod public_work_like_type;
|
||||
pub mod public_work_play_daily_stat_table;
|
||||
@@ -1268,6 +1271,8 @@ pub use external_generation_job_procedure_result_type::ExternalGenerationJobProc
|
||||
pub use external_generation_job_renew_lease_input_type::ExternalGenerationJobRenewLeaseInput;
|
||||
pub use external_generation_job_result_procedure_result_type::ExternalGenerationJobResultProcedureResult;
|
||||
pub use external_generation_job_result_snapshot_type::ExternalGenerationJobResultSnapshot;
|
||||
pub use external_generation_job_retention_input_type::ExternalGenerationJobRetentionInput;
|
||||
pub use external_generation_job_retention_procedure_result_type::ExternalGenerationJobRetentionProcedureResult;
|
||||
pub use external_generation_job_snapshot_type::ExternalGenerationJobSnapshot;
|
||||
pub use external_generation_job_summary_backfill_input_type::ExternalGenerationJobSummaryBackfillInput;
|
||||
pub use external_generation_job_summary_backfill_procedure_result_type::ExternalGenerationJobSummaryBackfillProcedureResult;
|
||||
@@ -1429,6 +1434,7 @@ pub use profile_wallet_ledger_table::*;
|
||||
pub use profile_wallet_ledger_type::ProfileWalletLedger;
|
||||
pub use profile_wallet_manual_restriction_table::*;
|
||||
pub use profile_wallet_manual_restriction_type::ProfileWalletManualRestriction;
|
||||
pub use prune_external_generation_job_history_and_return_procedure::prune_external_generation_job_history_and_return;
|
||||
pub use public_work_like_table::*;
|
||||
pub use public_work_like_type::PublicWorkLike;
|
||||
pub use public_work_play_daily_stat_table::*;
|
||||
|
||||
+2
@@ -56,6 +56,7 @@ impl __sdk::__query_builder::HasCols for ExternalGenerationJobEvent {
|
||||
/// Provides typed access to indexed columns for query building.
|
||||
pub struct ExternalGenerationJobEventIxCols {
|
||||
pub event_id: __sdk::__query_builder::IxCol<ExternalGenerationJobEvent, String>,
|
||||
pub job_id: __sdk::__query_builder::IxCol<ExternalGenerationJobEvent, String>,
|
||||
}
|
||||
|
||||
impl __sdk::__query_builder::HasIxCols for ExternalGenerationJobEvent {
|
||||
@@ -63,6 +64,7 @@ impl __sdk::__query_builder::HasIxCols for ExternalGenerationJobEvent {
|
||||
fn ix_cols(table_name: &'static str) -> Self::IxCols {
|
||||
ExternalGenerationJobEventIxCols {
|
||||
event_id: __sdk::__query_builder::IxCol::new(table_name, "event_id"),
|
||||
job_id: __sdk::__query_builder::IxCol::new(table_name, "job_id"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+19
@@ -0,0 +1,19 @@
|
||||
// THIS FILE IS AUTOMATICALLY GENERATED BY SPACETIMEDB. EDITS TO THIS FILE
|
||||
// WILL NOT BE SAVED. MODIFY TABLES IN YOUR MODULE SOURCE CODE INSTEAD.
|
||||
|
||||
#![allow(unused, clippy::all)]
|
||||
use spacetimedb_sdk::__codegen::{self as __sdk, __lib, __sats, __ws};
|
||||
|
||||
#[derive(__lib::ser::Serialize, __lib::de::Deserialize, Clone, PartialEq, Debug)]
|
||||
#[sats(crate = __lib)]
|
||||
pub struct ExternalGenerationJobRetentionInput {
|
||||
pub source_module: String,
|
||||
pub limit: u32,
|
||||
pub cursor_job_id: Option<String>,
|
||||
pub completed_before_micros: i64,
|
||||
pub dry_run: bool,
|
||||
}
|
||||
|
||||
impl __sdk::InModule for ExternalGenerationJobRetentionInput {
|
||||
type Module = super::RemoteModule;
|
||||
}
|
||||
+24
@@ -0,0 +1,24 @@
|
||||
// THIS FILE IS AUTOMATICALLY GENERATED BY SPACETIMEDB. EDITS TO THIS FILE
|
||||
// WILL NOT BE SAVED. MODIFY TABLES IN YOUR MODULE SOURCE CODE INSTEAD.
|
||||
|
||||
#![allow(unused, clippy::all)]
|
||||
use spacetimedb_sdk::__codegen::{self as __sdk, __lib, __sats, __ws};
|
||||
|
||||
#[derive(__lib::ser::Serialize, __lib::de::Deserialize, Clone, PartialEq, Debug)]
|
||||
#[sats(crate = __lib)]
|
||||
pub struct ExternalGenerationJobRetentionProcedureResult {
|
||||
pub ok: bool,
|
||||
pub dry_run: bool,
|
||||
pub scanned_count: u64,
|
||||
pub selected_count: u32,
|
||||
pub deleted_job_count: u32,
|
||||
pub deleted_summary_count: u32,
|
||||
pub deleted_event_count: u32,
|
||||
pub next_cursor_job_id: Option<String>,
|
||||
pub has_more: bool,
|
||||
pub error_message: Option<String>,
|
||||
}
|
||||
|
||||
impl __sdk::InModule for ExternalGenerationJobRetentionProcedureResult {
|
||||
type Module = super::RemoteModule;
|
||||
}
|
||||
+62
@@ -0,0 +1,62 @@
|
||||
// THIS FILE IS AUTOMATICALLY GENERATED BY SPACETIMEDB. EDITS TO THIS FILE
|
||||
// WILL NOT BE SAVED. MODIFY TABLES IN YOUR MODULE SOURCE CODE INSTEAD.
|
||||
|
||||
#![allow(unused, clippy::all)]
|
||||
use spacetimedb_sdk::__codegen::{self as __sdk, __lib, __sats, __ws};
|
||||
|
||||
use super::external_generation_job_retention_input_type::ExternalGenerationJobRetentionInput;
|
||||
use super::external_generation_job_retention_procedure_result_type::ExternalGenerationJobRetentionProcedureResult;
|
||||
|
||||
#[derive(__lib::ser::Serialize, __lib::de::Deserialize, Clone, PartialEq, Debug)]
|
||||
#[sats(crate = __lib)]
|
||||
struct PruneExternalGenerationJobHistoryAndReturnArgs {
|
||||
pub input: ExternalGenerationJobRetentionInput,
|
||||
}
|
||||
|
||||
impl __sdk::InModule for PruneExternalGenerationJobHistoryAndReturnArgs {
|
||||
type Module = super::RemoteModule;
|
||||
}
|
||||
|
||||
#[allow(non_camel_case_types)]
|
||||
/// Extension trait for access to the procedure `prune_external_generation_job_history_and_return`.
|
||||
///
|
||||
/// Implemented for [`super::RemoteProcedures`].
|
||||
pub trait prune_external_generation_job_history_and_return {
|
||||
fn prune_external_generation_job_history_and_return(
|
||||
&self,
|
||||
input: ExternalGenerationJobRetentionInput,
|
||||
) {
|
||||
self.prune_external_generation_job_history_and_return_then(input, |_, _| {});
|
||||
}
|
||||
|
||||
fn prune_external_generation_job_history_and_return_then(
|
||||
&self,
|
||||
input: ExternalGenerationJobRetentionInput,
|
||||
|
||||
__callback: impl FnOnce(
|
||||
&super::ProcedureEventContext,
|
||||
Result<ExternalGenerationJobRetentionProcedureResult, __sdk::InternalError>,
|
||||
) + Send
|
||||
+ 'static,
|
||||
);
|
||||
}
|
||||
|
||||
impl prune_external_generation_job_history_and_return for super::RemoteProcedures {
|
||||
fn prune_external_generation_job_history_and_return_then(
|
||||
&self,
|
||||
input: ExternalGenerationJobRetentionInput,
|
||||
|
||||
__callback: impl FnOnce(
|
||||
&super::ProcedureEventContext,
|
||||
Result<ExternalGenerationJobRetentionProcedureResult, __sdk::InternalError>,
|
||||
) + Send
|
||||
+ 'static,
|
||||
) {
|
||||
self.imp
|
||||
.invoke_procedure_with_callback::<_, ExternalGenerationJobRetentionProcedureResult>(
|
||||
"prune_external_generation_job_history_and_return",
|
||||
PruneExternalGenerationJobHistoryAndReturnArgs { input },
|
||||
__callback,
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -139,17 +139,6 @@ pub(crate) fn build_ai_text_chunk_row_id(snapshot: &AiTextChunkSnapshot) -> Stri
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn build_ai_text_chunk_snapshot_from_row(row: &AiTextChunk) -> AiTextChunkSnapshot {
|
||||
AiTextChunkSnapshot {
|
||||
chunk_id: row.chunk_id.clone(),
|
||||
task_id: row.task_id.clone(),
|
||||
stage_kind: row.stage_kind,
|
||||
sequence: row.sequence,
|
||||
delta_text: row.delta_text.clone(),
|
||||
created_at_micros: row.created_at.to_micros_since_unix_epoch(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn build_ai_result_reference_row(
|
||||
snapshot: &AiResultReferenceSnapshot,
|
||||
) -> AiResultReference {
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
use crate::*;
|
||||
use module_ai::{
|
||||
generate_ai_result_ref_id, generate_ai_text_chunk_id, normalize_optional_text,
|
||||
normalize_string_list,
|
||||
MAX_AI_TASK_TEXT_OUTPUT_BYTES, generate_ai_result_ref_id, generate_ai_text_chunk_id,
|
||||
normalize_optional_text, normalize_string_list, validate_ai_task_snapshot_memory_limits,
|
||||
};
|
||||
|
||||
#[spacetimedb::table(
|
||||
@@ -178,6 +178,9 @@ pub(crate) fn append_ai_text_chunk_tx(
|
||||
if input.sequence == 0 {
|
||||
return Err("ai_text_chunk.sequence 必须大于 0".to_string());
|
||||
}
|
||||
if input.delta_text.trim().len() > MAX_AI_TASK_TEXT_OUTPUT_BYTES {
|
||||
return Err("AI 任务文本输出超过内存上限".to_string());
|
||||
}
|
||||
|
||||
let mut snapshot = get_ai_task_snapshot_tx(ctx, &input.task_id)?;
|
||||
ensure_ai_task_can_transition(snapshot.status)?;
|
||||
@@ -200,7 +203,7 @@ pub(crate) fn append_ai_text_chunk_tx(
|
||||
.ai_text_chunk()
|
||||
.insert(build_ai_text_chunk_row(&chunk));
|
||||
|
||||
let aggregated_text = collect_ai_stage_text_output(ctx, &chunk.task_id, chunk.stage_kind);
|
||||
let aggregated_text = collect_ai_stage_text_output(ctx, &chunk.task_id, chunk.stage_kind)?;
|
||||
|
||||
snapshot.status = AiTaskStatus::Running;
|
||||
if snapshot.started_at_micros.is_none() {
|
||||
@@ -215,6 +218,7 @@ pub(crate) fn append_ai_text_chunk_tx(
|
||||
snapshot.updated_at_micros = input.created_at_micros;
|
||||
snapshot.version += 1;
|
||||
|
||||
validate_ai_task_snapshot_memory_limits(&snapshot).map_err(str::to_string)?;
|
||||
persist_ai_task_snapshot(ctx, &snapshot)?;
|
||||
emit_ai_task_event(
|
||||
ctx,
|
||||
@@ -252,6 +256,7 @@ pub(crate) fn complete_ai_stage_tx(
|
||||
snapshot.updated_at_micros = input.completed_at_micros;
|
||||
snapshot.version += 1;
|
||||
|
||||
validate_ai_task_snapshot_memory_limits(&snapshot).map_err(str::to_string)?;
|
||||
persist_ai_task_snapshot(ctx, &snapshot)?;
|
||||
emit_ai_task_event(
|
||||
ctx,
|
||||
@@ -285,26 +290,27 @@ pub(crate) fn attach_ai_result_reference_tx(
|
||||
label: normalize_optional_text(input.label),
|
||||
created_at_micros: input.created_at_micros,
|
||||
};
|
||||
ctx.db
|
||||
.ai_result_reference()
|
||||
.insert(build_ai_result_reference_row(&reference));
|
||||
|
||||
snapshot.result_references.push(reference);
|
||||
snapshot.updated_at_micros = input.created_at_micros;
|
||||
snapshot.version += 1;
|
||||
|
||||
persist_ai_task_snapshot(ctx, &snapshot)?;
|
||||
validate_ai_task_snapshot_memory_limits(&snapshot).map_err(str::to_string)?;
|
||||
let reference = snapshot
|
||||
.result_references
|
||||
.last()
|
||||
.cloned()
|
||||
.ok_or_else(|| "ai_result_reference 写入后缺少快照".to_string())?;
|
||||
ctx.db
|
||||
.ai_result_reference()
|
||||
.insert(build_ai_result_reference_row(&reference));
|
||||
persist_ai_task_snapshot(ctx, &snapshot)?;
|
||||
emit_ai_task_event(
|
||||
ctx,
|
||||
&snapshot,
|
||||
AiTaskEventKind::ResultReferenceAttached,
|
||||
None,
|
||||
None,
|
||||
Some(build_ai_result_reference_row_id(reference)),
|
||||
Some(build_ai_result_reference_row_id(&reference)),
|
||||
input.created_at_micros,
|
||||
);
|
||||
Ok(snapshot)
|
||||
@@ -333,29 +339,48 @@ pub(crate) fn replace_ai_task_stages(
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn delete_ai_text_chunks_for_task(ctx: &ReducerContext, task_id: &str) {
|
||||
let chunk_row_ids = ctx
|
||||
.db
|
||||
.ai_text_chunk()
|
||||
.by_ai_text_chunk_task_id()
|
||||
.filter(task_id)
|
||||
.map(|row| row.text_chunk_row_id.clone())
|
||||
.collect::<Vec<_>>();
|
||||
for row_id in chunk_row_ids {
|
||||
ctx.db.ai_text_chunk().text_chunk_row_id().delete(&row_id);
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn collect_ai_stage_text_output(
|
||||
ctx: &ReducerContext,
|
||||
task_id: &str,
|
||||
stage_kind: AiTaskStageKind,
|
||||
) -> Option<String> {
|
||||
let mut chunks = ctx
|
||||
) -> Result<Option<String>, String> {
|
||||
let mut chunks = Vec::new();
|
||||
let mut aggregated_bytes = 0_usize;
|
||||
for row in ctx
|
||||
.db
|
||||
.ai_text_chunk()
|
||||
.by_ai_text_chunk_task_id()
|
||||
.filter(task_id)
|
||||
.filter(|row| row.task_id == task_id && row.stage_kind == stage_kind)
|
||||
.map(|row| build_ai_text_chunk_snapshot_from_row(&row))
|
||||
.collect::<Vec<_>>();
|
||||
chunks.sort_by_key(|chunk| chunk.sequence);
|
||||
{
|
||||
aggregated_bytes = aggregated_bytes.saturating_add(row.delta_text.len());
|
||||
if aggregated_bytes > MAX_AI_TASK_TEXT_OUTPUT_BYTES {
|
||||
return Err("AI 任务文本输出超过内存上限".to_string());
|
||||
}
|
||||
chunks.push((row.sequence, row.delta_text.clone()));
|
||||
}
|
||||
chunks.sort_by_key(|(sequence, _)| *sequence);
|
||||
|
||||
let aggregated = chunks
|
||||
.into_iter()
|
||||
.map(|chunk| chunk.delta_text)
|
||||
.collect::<Vec<_>>()
|
||||
.join("");
|
||||
let mut aggregated = String::with_capacity(aggregated_bytes);
|
||||
for (_, delta) in chunks {
|
||||
aggregated.push_str(&delta);
|
||||
}
|
||||
if aggregated.trim().is_empty() {
|
||||
None
|
||||
Ok(None)
|
||||
} else {
|
||||
Some(aggregated)
|
||||
Ok(Some(aggregated))
|
||||
}
|
||||
}
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user