perf(storage): optimize database hot paths
This commit is contained in:
parent
531ad51015
commit
6b37ce141f
|
|
@ -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。
|
||||
|
|
|
|||
|
|
@ -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。
|
||||
|
|
|
|||
2
RUN.md
2
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 已设):
|
||||
|
||||
|
|
|
|||
|
|
@ -1,3 +1,3 @@
|
|||
# zcbot 版本号单一事实源:web/app.py 的 FastAPI version、/healthz 返回、前端展示都引这里。
|
||||
# 改版本只动这一行。
|
||||
__version__ = "0.64.1"
|
||||
__version__ = "0.64.2"
|
||||
|
|
|
|||
|
|
@ -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 = [
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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(),
|
||||
]
|
||||
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
@ -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__":
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
Loading…
Reference in New Issue