From 6b37ce141f6a7508ed58064280b6d9b79368c0bf Mon Sep 17 00:00:00 2001 From: caoqianming Date: Wed, 12 Aug 2026 13:54:40 +0800 Subject: [PATCH] perf(storage): optimize database hot paths --- DESIGN.md | 2 + PROGRESS.md | 4 +- RUN.md | 2 + core/__init__.py | 2 +- core/artifact_lifecycle.py | 14 +++- core/session.py | 17 ++-- core/storage/message_index.py | 35 ++++++++ core/storage/models.py | 30 ++++++- core/storage/utils.py | 25 ++---- core/working_dirs.py | 2 + .../20260812_1800_0029_storage_hot_paths.py | 79 +++++++++++++++++++ tests/test_artifacts.py | 9 ++- tests/test_message_index.py | 39 +++++++++ tests/test_message_payload.py | 10 ++- tests/test_rename_working_dir.py | 3 +- tests/test_run_lifecycle.py | 2 +- tests/test_storage_migration.py | 37 +++++++++ tests/test_web_common.py | 15 +++- web/common.py | 3 + web/routers/messages.py | 1 + web/run_lifecycle.py | 8 +- 21 files changed, 296 insertions(+), 43 deletions(-) create mode 100644 core/storage/message_index.py create mode 100644 db/migrations/versions/20260812_1800_0029_storage_hot_paths.py create mode 100644 tests/test_message_index.py create mode 100644 tests/test_storage_migration.py diff --git a/DESIGN.md b/DESIGN.md index afabb28..d8c5e17 100644 --- a/DESIGN.md +++ b/DESIGN.md @@ -217,6 +217,7 @@ tasks(task_id pk, user_id fk, name NOT NULL, auto_title_pending default false, channel default 'web', -- 0013 渠道来源,仅 INSERT 写定 run_status default 'idle', run_error, -- 0004 合并 runs 表 run_owner, -- 0020 蓝绿实例归属,reaper 只收自己色;单实例 NULL + next_message_idx, -- 0029 task 行锁下原子分配 messages.idx scheduled_job_id, -- 0017 定时任务归属,普通列表过滤 context_base_idx, -- 0019 §8.8 软重置窗口起点 deleted_at, -- 0010 软删 @@ -232,6 +233,7 @@ scheduled_jobs(§8.5) channel_bindings(§8.7,判别列+JSONB) - working_dir 存相对 ROOT posix 串,读写统一过 `core/paths.py`;入口 `validate_task_name` 拒空/`/\NUL`/`.` 起头。 - `auto_title_pending`(0023)只是一轮 UI 命名闸，不是 task 状态机；旧创建入口/存量行恒 false，快速入口首发后消费，人工改名优先清闸。 - **0004 简化**:runs 表只写不读、run_id 单活 run 下全冗余 → 合并 `run_status/run_error` 入 tasks。**0006**:`tasks.model_profile` 为 source-of-truth(PATCH 切、下条 send 生效);usage_events 重建 v2 多态形态,统计 source-of-truth;tasks 三列保留作粗概览。run_status 终态:ok 收回 idle,error(出错)与 cancelled(用户停止)是持久终态 —— 前端 `renderPersistedRunTerminal` 据此在每次重渲后补持久卡(扛过收尾 loadMessages 整屏重建),刷新/切任务仍在;下次起新 run(post_message 写 running)覆盖清掉。 +- **0029 消息序号**:`tasks.next_message_idx` 在 task 行锁下统一分配 `messages.idx`，Web、agent 与渠道追加不再各自维护序号或依赖冲突重试；分配时仍与 `max(idx)` 校准，允许蓝绿发布窗口内旧实例继续写入。清空消息与计数器在同一事务归零。 - **No-subtask**:同 user 下前缀互含即拒(归一 posix 后 Python 端比对);同 working_dir 允许。 - **文件面板先备料**:user_root 与普通目录都可显式新建直接子目录；创建成功后前端进入该目录，用户可先上传/选入资料，再把顶层空目录选作新对话 working_dir。目录 leaf 复用 `validate_task_name`，不允许借 UI 创建点目录或路径式名称。 - **DB-aware service 是顶层 working_dir mutation 的唯一原语,DB-FS 一致性服务端内化**:文件面板 rename 与对话 `rename_working_dir` 共用同一服务(事务锁关联 task、running→409、DB UPDATE 先于 FS)；对话工具只登记本轮内存动作，等 agent 正常回复完、当前 task 退出 running 后执行，避免 executor/system prompt 仍握旧 cwd。服务在收尾前退出时动作丢失但目录不变，不引持久队列或 migration。delete 仍仅走 files API，被 task 引用时 409。 diff --git a/PROGRESS.md b/PROGRESS.md index f14ce10..769c246 100644 --- a/PROGRESS.md +++ b/PROGRESS.md @@ -2,7 +2,7 @@ > 配合 `DESIGN.md`。本文件只记 phase 状态、决策偏差、文件量、下一步。每条 1-2 句:做了啥 + 关键判断;细节查 `git log` / `git diff` / `DESIGN §7.9`。 -最后更新:2026-08-12(PDF/PPT 连续预览滚轮与触摸修复,bump 0.64.1) +最后更新:2026-08-12(数据库热路径与消息序号优化,bump 0.64.2) --- @@ -23,6 +23,8 @@ ### 2026-08-12 +- **08-12 / 0.64.2 / 数据库热路径 + 消息序号原子化**:新增 0029 migration 平滑回填 `tasks.next_message_idx`，Web、agent 与渠道消息统一在 task 行锁下分配序号，并以实际 `max(idx)` 兼容蓝绿旧实例并行；任务列表、调度扫描/历史和 usage 时间序列补部分/BRIN 索引，列表追加 UUID 稳定排序，artifact 目录操作只查询并锁定目标子树。完整 532 项 unittest 全绿（17 skip），0029 PostgreSQL DDL 编译、Alembic 单 head、Ruff 致命规则及 diff 检查通过；未配置显式测试库，DB 集成测试安全跳过，未连接或写入生产 DB。 + - **08-12 / 0.64.1 / PDF/PPT 连续预览滚轮与触摸修复**:外层预览的防滚动穿透逻辑不再拦截来自 PDF 内层 viewport 的普通 wheel 事件，保留图片 Ctrl+滚轮缩放和其他预览的边界隔离；PDF viewport 同步声明触摸平移与 WebView/iOS 惯性滚动，使鼠标滚轮、触控板和手机上下滑动都直接驱动连续页列表。Node 前端预览 11 项、JavaScript 语法及 diff 检查通过；无 schema、migration、HTTP API、依赖或运行方式变化，未连接生产 DB。 - **08-12 / 0.64.0 / artifact 稳定身份 + 隐藏回收**:新增 0028 `artifacts` 生命周期表并回填存量结构化引用；新发布消息写带 `artifact_id` 的 v2 引用，移动/重命名保持身份，复制创建独立身份并记录来源，历史卡片按身份解析最新路径。删除已发布产物时移动到用户隐藏目录 `.zcbot_artifact_trash/` 并软删记录，普通文件仍物理删除；Python 35 项（测试库门控 1 skip）、Node 前端 11 项、mypy、Alembic 单 head、编译、Ruff 致命规则及 diff 检查通过，未连接或写入生产 DB。 diff --git a/RUN.md b/RUN.md index 018fecc..794c5fa 100644 --- a/RUN.md +++ b/RUN.md @@ -750,6 +750,8 @@ systemctl restart zcbot .venv/bin/python main.py db upgrade head ``` +0.64.2 部署必须先升级到 migration `0029`；它会回填消息序号计数器并创建任务、调度和用量热路径索引，不删除或重置存量数据。 + 启用 docker backend 重启 web(确保 `.env` 有 `ZCBOT_SANDBOX_BACKEND=docker` 或 systemd unit 已设): diff --git a/core/__init__.py b/core/__init__.py index cd0184d..b4017d7 100644 --- a/core/__init__.py +++ b/core/__init__.py @@ -1,3 +1,3 @@ # zcbot 版本号单一事实源:web/app.py 的 FastAPI version、/healthz 返回、前端展示都引这里。 # 改版本只动这一行。 -__version__ = "0.64.1" +__version__ = "0.64.2" diff --git a/core/artifact_lifecycle.py b/core/artifact_lifecycle.py index 94d4103..2bf6b16 100644 --- a/core/artifact_lifecycle.py +++ b/core/artifact_lifecycle.py @@ -9,7 +9,7 @@ from datetime import datetime, timezone from pathlib import Path from uuid import UUID, uuid4 -from sqlalchemy import select +from sqlalchemy import or_, select from sqlalchemy.dialects.postgresql import insert as pg_insert from .artifacts import ARTIFACT_TRASH_DIR, ArtifactRef @@ -29,6 +29,15 @@ def _hash_file(path: Path) -> str: return digest.hexdigest() +def artifact_subtree_clause(path: str): + """精确匹配文件/目录自身及其子树，并转义 LIKE 元字符。""" + escaped = path.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_") + return or_( + Artifact.current_path == path, + Artifact.current_path.like(escaped + "/%", escape="\\"), + ) + + def register_published_artifacts( *, user_id: UUID, @@ -98,6 +107,7 @@ def rename_active_artifacts( select(Artifact).where( Artifact.user_id == user_id, Artifact.status == "active", + artifact_subtree_clause(old_rel), ) ).scalars().all() changed = 0 @@ -129,6 +139,7 @@ def copy_active_artifacts( select(Artifact).where( Artifact.user_id == user_id, Artifact.status == "active", + artifact_subtree_clause(source_rel), ) ).scalars().all() created = 0 @@ -176,6 +187,7 @@ def trash_active_artifacts( select(Artifact).where( Artifact.user_id == user_id, Artifact.status == "active", + artifact_subtree_clause(source_rel), ).with_for_update() ).scalars().all() matches = [ diff --git a/core/session.py b/core/session.py index 2ca198c..963a9be 100644 --- a/core/session.py +++ b/core/session.py @@ -10,16 +10,15 @@ system prompt **不入库** —— 每次 build_agent 重建拼到 messages[0]( """ from __future__ import annotations -from pathlib import Path from typing import Any, Dict, List, Optional from uuid import UUID -from sqlalchemy import delete, func, select +from sqlalchemy import delete, func, select, update -from .storage import session_scope -from .storage.models import Message, Task, sanitize_jsonb_nul -from .file_store import atomic_write_text from .markdown_guard import normalize_markdown_fences +from .storage import session_scope +from .storage.message_index import allocate_message_idx +from .storage.models import Message, Task, sanitize_jsonb_nul def _to_dict(msg: Any) -> Any: @@ -97,16 +96,17 @@ class Session: return None with session_scope() as s: + idx = allocate_message_idx(s, self.task_id) row = Message( task_id=self.task_id, - idx=self._db_idx, + idx=idx, payload=msg_dict, artifact_refs=artifact_refs, ) s.add(row) s.flush() # 触发 INSERT 拿到 server-default 生成的 message_id msg_id = row.message_id - self._db_idx += 1 + self._db_idx = idx + 1 return msg_id def last_measured_usage(self) -> Optional[tuple]: @@ -182,6 +182,9 @@ class Session: self._n_head = 0 with session_scope() as s: s.execute(delete(Message).where(Message.task_id == self.task_id)) + s.execute( + update(Task).where(Task.task_id == self.task_id).values(next_message_idx=0) + ) self._db_idx = 0 self._base_idx = 0 diff --git a/core/storage/message_index.py b/core/storage/message_index.py new file mode 100644 index 0000000..4c29e0c --- /dev/null +++ b/core/storage/message_index.py @@ -0,0 +1,35 @@ +"""同一 task 内消息序号的统一事务分配入口。""" +from __future__ import annotations + +from typing import Optional +from uuid import UUID + +from sqlalchemy import func, select +from sqlalchemy.orm import Session as OrmSession + +from .models import Message, Task + + +def allocate_message_idx( + session: OrmSession, + task_id: UUID, + *, + locked_task: Optional[Task] = None, +) -> int: + """锁定 task 行并返回唯一递增 idx；max 校准兼容蓝绿旧实例。""" + task = locked_task + if task is None: + task = session.execute( + select(Task).where(Task.task_id == task_id).with_for_update() + ).scalar_one() + elif task.task_id != task_id: + raise ValueError("locked_task 与 task_id 不一致") + + max_idx = session.execute( + select(func.max(Message.idx)).where(Message.task_id == task_id) + ).scalar_one() + persisted_next = int(getattr(task, "next_message_idx", 0) or 0) + observed_next = -1 if max_idx is None else int(max_idx) + next_idx = max(persisted_next, observed_next + 1) + task.next_message_idx = next_idx + 1 + return next_idx diff --git a/core/storage/models.py b/core/storage/models.py index 6ed9974..a075f56 100644 --- a/core/storage/models.py +++ b/core/storage/models.py @@ -33,7 +33,8 @@ from sqlalchemy import ( func, text, ) -from sqlalchemy.dialects.postgresql import JSONB, UUID as PG_UUID +from sqlalchemy.dialects.postgresql import JSONB +from sqlalchemy.dialects.postgresql import UUID as PG_UUID from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, validates @@ -85,6 +86,16 @@ class User(Base): class Task(Base): __tablename__ = "tasks" + __table_args__ = ( + Index( + "ix_tasks_user_active_updated", "user_id", "updated_at", "task_id", + postgresql_where=text("deleted_at IS NULL AND scheduled_job_id IS NULL"), + ), + Index( + "ix_tasks_schedule_created", "scheduled_job_id", "created_at", "task_id", + postgresql_where=text("deleted_at IS NULL AND scheduled_job_id IS NOT NULL"), + ), + ) task_id: Mapped[UUID] = mapped_column(PG_UUID(as_uuid=True), primary_key=True, default=uuid4) user_id: Mapped[UUID] = mapped_column( @@ -130,6 +141,10 @@ class Task(Base): # (blue/green),启动 reaper 只收自己实例的孤儿。单实例部署恒 NULL,行为不变。 # 只在 running/cancelling 期间有意义,run 收尾不清(reaper 只看活跃行)。 run_owner: Mapped[Optional[str]] = mapped_column(Text, nullable=True) + # 0029:同一 task 的下一条消息序号；所有消息入口以 task 行锁串行分配。 + next_message_idx: Mapped[int] = mapped_column( + Integer, nullable=False, default=0, server_default="0" + ) # 喂给模型的上下文窗口起点(0019,channel 长会话软重置)。Session.load 只把 idx >= # context_base_idx 的消息装进 LLM 上下文;之前的历史仍全量留 messages 表(web 翻得到)。 # web 普通任务恒 0 = 喂全量;channel 入站按 gap / 「新话题」推进。详 DESIGN §8.7。 @@ -255,6 +270,9 @@ class UsageEvent(Base): 按 user 聚合的统计 query 走 (user_id, created_at) 索引,不 JOIN tasks 表。 """ __tablename__ = "usage_events" + __table_args__ = ( + Index("ix_usage_created_brin", "created_at", postgresql_using="brin"), + ) event_id: Mapped[UUID] = mapped_column(PG_UUID(as_uuid=True), primary_key=True, default=uuid4) user_id: Mapped[UUID] = mapped_column( @@ -321,6 +339,16 @@ class ScheduledJob(Base): """ __tablename__ = "scheduled_jobs" + __table_args__ = ( + Index( + "ix_scheduled_jobs_due_active", "next_run_at", + postgresql_where=text("enabled IS TRUE AND deleted_at IS NULL"), + ), + Index( + "ix_scheduled_jobs_user_active_created", "user_id", "created_at", + postgresql_where=text("deleted_at IS NULL"), + ), + ) job_id: Mapped[UUID] = mapped_column(PG_UUID(as_uuid=True), primary_key=True, default=uuid4) user_id: Mapped[UUID] = mapped_column( diff --git a/core/storage/utils.py b/core/storage/utils.py index 29f1867..7454404 100644 --- a/core/storage/utils.py +++ b/core/storage/utils.py @@ -6,9 +6,9 @@ from uuid import UUID from sqlalchemy import func, select, update from sqlalchemy.dialects.postgresql import insert -from sqlalchemy.exc import IntegrityError from .engine import session_scope +from .message_index import allocate_message_idx from .models import Message, Task @@ -64,26 +64,15 @@ def ensure_local_task_row( def append_channel_message( task_id: UUID, content: str, *, role: str = "assistant", kind: Optional[str] = None ) -> None: - """往 task 追加一条非 agent-run 产生的消息(push 出站记录等)。原子算 idx - (SELECT max(idx)+1)+INSERT;撞 uq_messages_task_idx(与入站 agent run 并发 - append)→ 重试。payload 形态同 Session.append 的 {role, content};不设 + """往 task 追加一条非 agent-run 产生的消息(push 出站记录等)。 + + 与 agent/Web 入口共用 task 行锁下的序号分配。payload 形态同 Session.append；不设 model_profile / tokens_*(非模型产出,usage 不计)。kind 写 messages.kind 列 (独立列,不进 payload):"push" 标记 push 记录,extract_last_assistant_text 据此跳过。""" payload = {"role": role, "content": content} - last_err: Optional[Exception] = None - for _ in range(3): - try: - with session_scope() as s: - max_idx = s.execute( - select(func.max(Message.idx)).where(Message.task_id == task_id) - ).scalar() - next_idx = (max_idx if max_idx is not None else -1) + 1 - s.add(Message(task_id=task_id, idx=next_idx, payload=payload, kind=kind)) - return - except IntegrityError as e: - last_err = e - continue - raise RuntimeError(f"append_channel_message: idx 冲突重试耗尽: {last_err}") + with session_scope() as s: + next_idx = allocate_message_idx(s, task_id) + s.add(Message(task_id=task_id, idx=next_idx, payload=payload, kind=kind)) def upsert_task( diff --git a/core/working_dirs.py b/core/working_dirs.py index 75e4770..37e71b0 100644 --- a/core/working_dirs.py +++ b/core/working_dirs.py @@ -11,6 +11,7 @@ from uuid import UUID from sqlalchemy import select, update +from .artifact_lifecycle import artifact_subtree_clause from .paths import to_db_path from .storage import NoSubtaskError, check_no_subtask, session_scope from .storage.models import Artifact, Task @@ -100,6 +101,7 @@ def rename_working_dir( select(Artifact).where( Artifact.user_id == user_id, Artifact.status == "active", + artifact_subtree_clause(old_rel), ) ).scalars().all() for artifact in artifacts: diff --git a/db/migrations/versions/20260812_1800_0029_storage_hot_paths.py b/db/migrations/versions/20260812_1800_0029_storage_hot_paths.py new file mode 100644 index 0000000..2424a92 --- /dev/null +++ b/db/migrations/versions/20260812_1800_0029_storage_hot_paths.py @@ -0,0 +1,79 @@ +"""Add atomic message indexes and storage hot-path indexes. + +Revision ID: 0029 +Revises: 0028 +Create Date: 2026-08-12 +""" +from typing import Sequence, Union + +import sqlalchemy as sa +from alembic import op + + +revision: str = "0029" +down_revision: Union[str, None] = "0028" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.add_column( + "tasks", + sa.Column("next_message_idx", sa.Integer(), server_default="0", nullable=False), + ) + op.execute( + """ + UPDATE tasks AS t + SET next_message_idx = COALESCE(m.next_idx, 0) + FROM ( + SELECT task_id, MAX(idx) + 1 AS next_idx + FROM messages + GROUP BY task_id + ) AS m + WHERE m.task_id = t.task_id + """ + ) + + op.create_index( + "ix_tasks_user_active_updated", + "tasks", + ["user_id", "updated_at", "task_id"], + postgresql_where=sa.text("deleted_at IS NULL AND scheduled_job_id IS NULL"), + ) + op.create_index( + "ix_tasks_schedule_created", + "tasks", + ["scheduled_job_id", "created_at", "task_id"], + postgresql_where=sa.text( + "deleted_at IS NULL AND scheduled_job_id IS NOT NULL" + ), + ) + op.create_index( + "ix_scheduled_jobs_due_active", + "scheduled_jobs", + ["next_run_at"], + postgresql_where=sa.text("enabled IS TRUE AND deleted_at IS NULL"), + ) + op.create_index( + "ix_scheduled_jobs_user_active_created", + "scheduled_jobs", + ["user_id", "created_at"], + postgresql_where=sa.text("deleted_at IS NULL"), + ) + op.create_index( + "ix_usage_created_brin", + "usage_events", + ["created_at"], + postgresql_using="brin", + ) + + +def downgrade() -> None: + op.drop_index("ix_usage_created_brin", table_name="usage_events") + op.drop_index( + "ix_scheduled_jobs_user_active_created", table_name="scheduled_jobs" + ) + op.drop_index("ix_scheduled_jobs_due_active", table_name="scheduled_jobs") + op.drop_index("ix_tasks_schedule_created", table_name="tasks") + op.drop_index("ix_tasks_user_active_updated", table_name="tasks") + op.drop_column("tasks", "next_message_idx") diff --git a/tests/test_artifacts.py b/tests/test_artifacts.py index d8fa02f..13d349e 100644 --- a/tests/test_artifacts.py +++ b/tests/test_artifacts.py @@ -6,15 +6,15 @@ from types import SimpleNamespace from unittest.mock import MagicMock, patch from uuid import uuid4 +from core.artifact_lifecycle import artifact_subtree_clause, trash_active_artifacts from core.artifacts import ( - ArtifactRef, ArtifactPathError, + ArtifactRef, ToolExecutionResult, resolve_artifact_path, ) from core.executor import ExecCtx from core.executor_host import HostExecutor -from core.artifact_lifecycle import trash_active_artifacts from tools.publish_artifacts import PublishArtifactsTool from web.routers.files import _task_file_target @@ -86,6 +86,11 @@ class ArtifactPathTests(unittest.TestCase): self.assertTrue(trashed.is_file()) self.assertEqual(trashed.read_bytes(), b"pdf") + def test_subtree_query_escapes_like_wildcards(self) -> None: + clause = artifact_subtree_clause("项目_100%/报告") + compiled = clause.compile().params + self.assertIn("项目\\_100\\%/报告/%", compiled.values()) + def test_explicit_dot_slash_disambiguates_same_named_subdirectory(self) -> None: nested = self.wd / "技术讨论" / "nested.html" nested.parent.mkdir() diff --git a/tests/test_message_index.py b/tests/test_message_index.py new file mode 100644 index 0000000..e990352 --- /dev/null +++ b/tests/test_message_index.py @@ -0,0 +1,39 @@ +import unittest +from types import SimpleNamespace +from unittest.mock import MagicMock +from uuid import uuid4 + +from core.storage.message_index import allocate_message_idx + + +class MessageIndexTests(unittest.TestCase): + def test_allocator_uses_persisted_counter_when_ahead(self) -> None: + tid = uuid4() + task = SimpleNamespace(task_id=tid, next_message_idx=8) + session = MagicMock() + session.execute.return_value.scalar_one.return_value = 5 + + self.assertEqual(allocate_message_idx(session, tid, locked_task=task), 8) + self.assertEqual(task.next_message_idx, 9) + + def test_allocator_catches_up_with_legacy_writer(self) -> None: + tid = uuid4() + task = SimpleNamespace(task_id=tid, next_message_idx=3) + session = MagicMock() + session.execute.return_value.scalar_one.return_value = 6 + + self.assertEqual(allocate_message_idx(session, tid, locked_task=task), 7) + self.assertEqual(task.next_message_idx, 8) + + def test_empty_task_starts_at_zero(self) -> None: + tid = uuid4() + task = SimpleNamespace(task_id=tid, next_message_idx=0) + session = MagicMock() + session.execute.return_value.scalar_one.return_value = None + + self.assertEqual(allocate_message_idx(session, tid, locked_task=task), 0) + self.assertEqual(task.next_message_idx, 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_message_payload.py b/tests/test_message_payload.py index ec2c008..3a9a602 100644 --- a/tests/test_message_payload.py +++ b/tests/test_message_payload.py @@ -56,7 +56,10 @@ class MessagePayloadSanitizationTests(unittest.TestCase): yield fake_db session = Session(task_id=uuid4()) - with patch("core.session.session_scope", side_effect=fake_session_scope): + with ( + patch("core.session.session_scope", side_effect=fake_session_scope), + patch("core.session.allocate_message_idx", return_value=0), + ): session.append({"role": "tool", "content": "before\x00after"}) expected = {"role": "tool", "content": "beforeafter"} @@ -81,7 +84,10 @@ class MessagePayloadSanitizationTests(unittest.TestCase): broken = "```markdown\n```mermaid\nA --> B\n```\n```\n正文\n" session = Session(task_id=uuid4()) - with patch("core.session.session_scope", side_effect=fake_session_scope): + with ( + patch("core.session.session_scope", side_effect=fake_session_scope), + patch("core.session.allocate_message_idx", return_value=0), + ): session.append({"role": "assistant", "content": broken}) stored = session.messages[0]["content"] diff --git a/tests/test_rename_working_dir.py b/tests/test_rename_working_dir.py index 033edf5..f59dc9f 100644 --- a/tests/test_rename_working_dir.py +++ b/tests/test_rename_working_dir.py @@ -189,6 +189,7 @@ class WorkingDirServiceTests(unittest.TestCase): session.execute.side_effect = [ SimpleNamespace(all=lambda: rows), MagicMock(), + SimpleNamespace(scalars=lambda: SimpleNamespace(all=lambda: [])), ] @contextmanager @@ -214,7 +215,7 @@ class WorkingDirServiceTests(unittest.TestCase): self.assertEqual(result.tasks_updated, 2) self.assertFalse(old.exists()) self.assertEqual((new / "产物.txt").read_text(encoding="utf-8"), "ok") - self.assertEqual(session.execute.call_count, 2) + self.assertEqual(session.execute.call_count, 3) def test_active_associated_task_blocks_before_fs_change(self) -> None: uid = uuid4() diff --git a/tests/test_run_lifecycle.py b/tests/test_run_lifecycle.py index 5a70c9a..d34afd3 100644 --- a/tests/test_run_lifecycle.py +++ b/tests/test_run_lifecycle.py @@ -23,7 +23,7 @@ class ClaimRunTests(unittest.TestCase): session = MagicMock() session.execute.side_effect = [ SimpleNamespace(scalar_one_or_none=lambda: task), - SimpleNamespace(scalar_one=lambda: 7), + SimpleNamespace(scalar_one=lambda: 6), MagicMock(), ] diff --git a/tests/test_storage_migration.py b/tests/test_storage_migration.py new file mode 100644 index 0000000..a90e79c --- /dev/null +++ b/tests/test_storage_migration.py @@ -0,0 +1,37 @@ +from __future__ import annotations + +import importlib +import unittest +from unittest.mock import patch + +from alembic.migration import MigrationContext +from alembic.operations import Operations +from sqlalchemy import create_mock_engine +from sqlalchemy.dialects import postgresql + + +class StorageMigrationTests(unittest.TestCase): + def test_0029_upgrade_compiles_as_postgresql_ddl(self) -> None: + statements: list[str] = [] + + def capture(sql, *multiparams, **params): + statements.append(str(sql.compile(dialect=postgresql.dialect()))) + + engine = create_mock_engine("postgresql+psycopg://", capture) + operations = Operations(MigrationContext.configure(engine.connect())) + migration = importlib.import_module( + "db.migrations.versions.20260812_1800_0029_storage_hot_paths" + ) + with patch.object(migration, "op", operations): + migration.upgrade() + + rendered = "\n".join(statements) + self.assertIn("next_message_idx", rendered) + self.assertIn("MAX(idx) + 1", rendered) + self.assertIn("ix_tasks_user_active_updated", rendered) + self.assertIn("ix_scheduled_jobs_due_active", rendered) + self.assertIn("USING brin", rendered) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_web_common.py b/tests/test_web_common.py index 6c3d431..f33b598 100644 --- a/tests/test_web_common.py +++ b/tests/test_web_common.py @@ -7,17 +7,26 @@ class TaskOrderingTests(unittest.TestCase): def test_default_ordering_uses_latest_update_first(self): ordering = parse_ordering(None) - self.assertEqual([str(clause) for clause in ordering], ["tasks.updated_at DESC"]) + self.assertEqual( + [str(clause) for clause in ordering], + ["tasks.updated_at DESC", "tasks.task_id DESC"], + ) def test_invalid_ordering_falls_back_to_latest_update_first(self): ordering = parse_ordering("not_a_field") - self.assertEqual([str(clause) for clause in ordering], ["tasks.updated_at DESC"]) + self.assertEqual( + [str(clause) for clause in ordering], + ["tasks.updated_at DESC", "tasks.task_id DESC"], + ) def test_explicit_created_at_ordering_remains_supported(self): ordering = parse_ordering("-created_at") - self.assertEqual([str(clause) for clause in ordering], ["tasks.created_at DESC"]) + self.assertEqual( + [str(clause) for clause in ordering], + ["tasks.created_at DESC", "tasks.task_id DESC"], + ) if __name__ == "__main__": diff --git a/web/common.py b/web/common.py index cad4391..33e4f8f 100644 --- a/web/common.py +++ b/web/common.py @@ -76,6 +76,9 @@ def parse_ordering(s: Optional[str]) -> list: if not cols: # 用户传了全无效字段 → fallback 默认 cols = [Task.updated_at.desc()] + # 所有允许字段都可能重复；UUID 兜底保证 OFFSET 分页在同一快照内稳定。 + last_desc = str(cols[-1]).upper().endswith(" DESC") + cols.append(Task.task_id.desc() if last_desc else Task.task_id.asc()) return cols diff --git a/web/routers/messages.py b/web/routers/messages.py index 8b4e5d4..321313e 100644 --- a/web/routers/messages.py +++ b/web/routers/messages.py @@ -408,6 +408,7 @@ def register_message_routes(app, *, require_user) -> None: "cost_cny": 0, "run_status": "idle", "run_error": None, + "next_message_idx": 0, # 全删后 idx 从 0 重起,base 必须归零否则 load 窗口起点悬空(0019); # 摘要一并清,否则清空后的对话还会被注入旧前情摘要(0021) "context_base_idx": 0, diff --git a/web/run_lifecycle.py b/web/run_lifecycle.py index c655260..1eee557 100644 --- a/web/run_lifecycle.py +++ b/web/run_lifecycle.py @@ -10,9 +10,10 @@ from dataclasses import dataclass, field from typing import Any, Callable, Optional from uuid import UUID -from sqlalchemy import func, select, update +from sqlalchemy import select, update from core.storage import session_scope +from core.storage.message_index import allocate_message_idx from core.storage.models import Message, Task from .broker import broker @@ -76,10 +77,7 @@ def claim_run_with_message( extra_values, metadata = prepare(s, task) # task 行锁串行化同一 task 的 idx 分配，无需额外序列表或 advisory lock。 - next_idx = s.execute( - select(func.coalesce(func.max(Message.idx), -1) + 1) - .where(Message.task_id == task_id) - ).scalar_one() + next_idx = allocate_message_idx(s, task_id, locked_task=task) s.add(Message( task_id=task_id, idx=int(next_idx),