From 85796f649293e5f205223cc28b57f3a688f1df78 Mon Sep 17 00:00:00 2001 From: yu1128-L <2241397554@qq.com> Date: Thu, 8 Oct 2026 20:29:15 +0800 Subject: [PATCH] fix(events): preserve thread seq watermark across run deletion (#6477) * fix(events): preserve thread seq watermark across run deletion * fix(events): clean up watermarks for empty threads * perf(events): avoid duplicate fence on run deletion delete_by_run already reaches _max_seq_for_thread, which acquires the transaction-scoped thread mutation fence. Keep the count and delete under that same transaction lock without taking the Postgres advisory lock twice, and pin the single acquisition in the existing test. --------- Co-authored-by: yu1128-L <296164533+yu1128-L@users.noreply.github.com> Co-authored-by: Willem Jiang --- README.md | 2 + .../deerflow/persistence/migrations/AGENTS.md | 2 +- .../versions/0034_run_event_seq_watermark.py | 39 ++++ .../deerflow/persistence/models/__init__.py | 3 +- .../deerflow/persistence/models/run_event.py | 11 +- .../harness/deerflow/runtime/AGENTS.md | 6 + .../deerflow/runtime/events/store/db.py | 39 +++- .../test_db_event_store_lock_lifecycle.py | 6 + ...migration_0027_0028_notification_outbox.py | 2 +- ...st_migration_0029_scheduler_agent_tasks.py | 2 +- ...migration_0032_activity_and_task_events.py | 6 +- ..._migration_0034_run_event_seq_watermark.py | 125 ++++++++++++ backend/tests/test_run_event_store.py | 186 +++++++++++++++++- 13 files changed, 410 insertions(+), 19 deletions(-) create mode 100644 backend/packages/harness/deerflow/persistence/migrations/versions/0034_run_event_seq_watermark.py create mode 100644 backend/tests/test_migration_0034_run_event_seq_watermark.py diff --git a/README.md b/README.md index b1278da62f..2ed78d765e 100644 --- a/README.md +++ b/README.md @@ -2376,6 +2376,8 @@ editing the saved agent configuration to refresh the selection. DeerFlow can be used as an embedded Python library without running the full HTTP services. The `DeerFlowClient` provides direct in-process access to all agent and Gateway capabilities, returning the same response schemas as the HTTP Gateway API. The HTTP Gateway also exposes `DELETE /api/threads/{thread_id}` to remove DeerFlow-managed local thread data after the LangGraph thread itself has been deleted: +For database-backed run events, deleting a run preserves its thread's sequence watermark. Thread deletion removes that watermark once no events remain, allowing a recreated thread to restart at sequence 1. Owner-scoped deletion preserves the watermark when another owner's events remain. + Thread IDs may be supplied by callers and do not have to be UUIDs. Explicit IDs must contain 1–64 ASCII letters, digits, hyphens, or underscores (`^[A-Za-z0-9_-]{1,64}$`). DeerFlow generates a UUID only when `thread_id` is diff --git a/backend/packages/harness/deerflow/persistence/migrations/AGENTS.md b/backend/packages/harness/deerflow/persistence/migrations/AGENTS.md index ba22604181..8e7f49e32a 100644 --- a/backend/packages/harness/deerflow/persistence/migrations/AGENTS.md +++ b/backend/packages/harness/deerflow/persistence/migrations/AGENTS.md @@ -28,7 +28,7 @@ The empty-DB path keeps using `create_all` because `Base.metadata` is the only a `0026_mcp_task_lease_tokens` → `0027_notification_deliveries` → `0028_parked_attempts` → `0029_scheduler_agent_tasks` → `0030_notification_claim_tokens` → `0031_scheduled_streak_boundary` → -`0032_activity_and_task_events` → `0033_batch_result_artifact` +`0032_activity_and_task_events` → `0033_batch_result_artifact` → `0034_run_event_seq_watermark` (current head). The preference revision adds a separate owner/key table with a cascading users foreign key and does not alter users; the project-documents revision adds a new owner-scoped diff --git a/backend/packages/harness/deerflow/persistence/migrations/versions/0034_run_event_seq_watermark.py b/backend/packages/harness/deerflow/persistence/migrations/versions/0034_run_event_seq_watermark.py new file mode 100644 index 0000000000..d33514dc03 --- /dev/null +++ b/backend/packages/harness/deerflow/persistence/migrations/versions/0034_run_event_seq_watermark.py @@ -0,0 +1,39 @@ +"""Persist the per-thread run-event sequence watermark.""" + +from __future__ import annotations + +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op + +revision: str = "0034_run_event_seq_watermark" +down_revision: str | Sequence[str] | None = "0033_batch_result_artifact" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + +_TABLE = "run_event_thread_seq" + + +def _has_table() -> bool: + return sa.inspect(op.get_bind()).has_table(_TABLE) + + +def upgrade() -> None: + if _has_table(): + return + + op.create_table( + _TABLE, + sa.Column("thread_id", sa.String(64), nullable=False), + sa.Column("seq", sa.BigInteger(), nullable=False, server_default=sa.text("0")), + sa.PrimaryKeyConstraint("thread_id", name="pk_run_event_thread_seq"), + ) + # Existing rows predate the durable watermark; seed it so a subsequent + # single-run deletion cannot lower the thread's allocation floor. + op.execute(sa.text(f"INSERT INTO {_TABLE} (thread_id, seq) SELECT thread_id, MAX(seq) FROM run_events GROUP BY thread_id")) + + +def downgrade() -> None: + if _has_table(): + op.drop_table(_TABLE) diff --git a/backend/packages/harness/deerflow/persistence/models/__init__.py b/backend/packages/harness/deerflow/persistence/models/__init__.py index 1846140092..825ee5db59 100644 --- a/backend/packages/harness/deerflow/persistence/models/__init__.py +++ b/backend/packages/harness/deerflow/persistence/models/__init__.py @@ -24,7 +24,7 @@ from deerflow.persistence.channel_connections.model import ( from deerflow.persistence.feedback.model import FeedbackRow from deerflow.persistence.managed_subagents.model import ManagedSubagentRow from deerflow.persistence.mcp_tasks.model import McpTaskRow -from deerflow.persistence.models.run_event import RunEventRow +from deerflow.persistence.models.run_event import RunEventRow, RunEventThreadSeqRow from deerflow.persistence.notification_deliveries.model import NotificationDeliveryRow from deerflow.persistence.personal_access_tokens.model import PersonalAccessTokenRow from deerflow.persistence.projects.model import ProjectDocumentRow, ProjectRow @@ -52,6 +52,7 @@ __all__ = [ "ProjectDocumentRow", "ProjectRow", "RunEventRow", + "RunEventThreadSeqRow", "RunChangeClockRow", "RunRow", "ScheduledTaskEventRow", diff --git a/backend/packages/harness/deerflow/persistence/models/run_event.py b/backend/packages/harness/deerflow/persistence/models/run_event.py index e0cc519744..479602858f 100644 --- a/backend/packages/harness/deerflow/persistence/models/run_event.py +++ b/backend/packages/harness/deerflow/persistence/models/run_event.py @@ -4,7 +4,7 @@ from __future__ import annotations from datetime import UTC, datetime -from sqlalchemy import JSON, DateTime, Index, String, Text, UniqueConstraint +from sqlalchemy import JSON, BigInteger, DateTime, Index, String, Text, UniqueConstraint, text from sqlalchemy.orm import Mapped, mapped_column from deerflow.constants import RUN_EVENT_CATEGORY_MAX_LENGTH, RUN_EVENT_TYPE_MAX_LENGTH @@ -34,3 +34,12 @@ class RunEventRow(Base): Index("ix_events_thread_cat_seq", "thread_id", "category", "seq"), Index("ix_events_run", "thread_id", "run_id", "seq"), ) + + +class RunEventThreadSeqRow(Base): + """Durable per-thread high-water mark for run-event sequence allocation.""" + + __tablename__ = "run_event_thread_seq" + + thread_id: Mapped[str] = mapped_column(String(64), primary_key=True) + seq: Mapped[int] = mapped_column(BigInteger, nullable=False, default=0, server_default=text("0")) diff --git a/backend/packages/harness/deerflow/runtime/AGENTS.md b/backend/packages/harness/deerflow/runtime/AGENTS.md index 2e3728c760..baccfa8ea6 100644 --- a/backend/packages/harness/deerflow/runtime/AGENTS.md +++ b/backend/packages/harness/deerflow/runtime/AGENTS.md @@ -157,6 +157,12 @@ checkpoint-write admission boundary must repeat the complete audit after admission; a pre-admission exact hit can be superseded by a later event just as a pre-admission miss can become an exact hit. +**DB run-event sequence watermark:** Run deletion retains `run_event_thread_seq`. +Thread deletion removes it in the same mutation-fenced transaction only when no +events remain for the thread, including owner-scoped and zero-row deletions. +Check all owners and categories; surviving rows retain the allocation floor. +`tests/test_run_event_store.py` pins cleanup, recreation and cursor visibility. + **Changed-run discovery:** Use the durable `(change_seq, run_id)` cursor and repeat history audits after admission. Details: `backend/docs/runtime-guidance-details.md`. **Terminal run cleanup:** Close streams, journals, and graph references even on cancellation. Details: `backend/docs/runtime-guidance-details.md`. diff --git a/backend/packages/harness/deerflow/runtime/events/store/db.py b/backend/packages/harness/deerflow/runtime/events/store/db.py index 5abded179a..199639142a 100644 --- a/backend/packages/harness/deerflow/runtime/events/store/db.py +++ b/backend/packages/harness/deerflow/runtime/events/store/db.py @@ -17,7 +17,7 @@ from typing import Any from sqlalchemy import delete, func, or_, select, text from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker -from deerflow.persistence.models.run_event import RunEventRow +from deerflow.persistence.models.run_event import RunEventRow, RunEventThreadSeqRow from deerflow.runtime.events.message_identity import message_identity from deerflow.runtime.events.store.base import RunEventStore from deerflow.runtime.user_context import AUTO, _AutoSentinel, get_current_user, resolve_user_id @@ -168,11 +168,12 @@ class DbRunEventStore(RunEventStore): @staticmethod async def _max_seq_for_thread(session: AsyncSession, thread_id: str) -> int | None: - """Return the current max seq while serializing writers per thread. + """Return the current seq high-water mark while serializing writers. - Takes the shared thread mutation fence before reading the aggregate, so + The row aggregate is the source of truth for existing data, while the + durable watermark preserves sequence values after a run is deleted. + Takes the shared thread mutation fence before reading either value, so the read is ordered against every other mutation of the same thread. - Other dialects keep the existing row-locking statement. """ await DbRunEventStore._acquire_thread_mutation_fence(session, thread_id) @@ -181,9 +182,24 @@ class DbRunEventStore(RunEventStore): dialect_name = bind.dialect.name if bind is not None else "" if dialect_name == "postgresql": - return await session.scalar(stmt) + row_max = await session.scalar(stmt) + else: + row_max = await session.scalar(stmt.with_for_update()) - return await session.scalar(stmt.with_for_update()) + watermark = await session.scalar(select(RunEventThreadSeqRow.seq).where(RunEventThreadSeqRow.thread_id == thread_id)) + values = [value for value in (row_max, watermark) if value is not None] + return max(values) if values else None + + @staticmethod + async def _raise_seq_watermark(session: AsyncSession, thread_id: str, seq: int) -> None: + """Persist *seq* as the thread's new minimum allocation floor.""" + if seq <= 0: + return + watermark = await session.get(RunEventThreadSeqRow, thread_id) + if watermark is None: + session.add(RunEventThreadSeqRow(thread_id=thread_id, seq=seq)) + elif watermark.seq < seq: + watermark.seq = seq async def put(self, *, thread_id, run_id, event_type, category, content="", metadata=None, created_at=None): # noqa: D401 """Write a single event — low-frequency path only. @@ -214,6 +230,7 @@ class DbRunEventStore(RunEventStore): created_at=datetime.fromisoformat(created_at) if created_at else datetime.now(UTC), ) session.add(row) + await self._raise_seq_watermark(session, thread_id, seq) return self._row_to_dict(row) async def put_batch(self, events): @@ -251,6 +268,7 @@ class DbRunEventStore(RunEventStore): ) session.add(row) rows.append(row) + await self._raise_seq_watermark(session, thread_id, seq) return [self._row_to_dict(r) for r in rows] async def put_if_absent( @@ -304,6 +322,7 @@ class DbRunEventStore(RunEventStore): created_at=datetime.fromisoformat(created_at) if created_at else datetime.now(UTC), ) session.add(row) + await self._raise_seq_watermark(session, thread_id, row.seq) return self._row_to_dict(row), True async def list_messages( @@ -529,6 +548,11 @@ class DbRunEventStore(RunEventStore): count = await session.scalar(count_stmt) or 0 if count > 0: await session.execute(delete(RunEventRow).where(*count_conditions)) + # Owner-scoped deletion is the Gateway's normal path. + # Reset only an empty thread: surviving events from any + # owner still need the allocation floor for deleted runs. + remaining_events = select(RunEventRow.id).where(RunEventRow.thread_id == thread_id).exists() + await session.execute(delete(RunEventThreadSeqRow).where(RunEventThreadSeqRow.thread_id == thread_id, ~remaining_events)) # Retire the live-thread pin, but never remove the weak registry # entry directly. asyncio.Lock.release() clears ``locked()`` before # a queued waiter resumes, so an unlocked check can observe the @@ -555,7 +579,7 @@ class DbRunEventStore(RunEventStore): async with self._get_write_lock(thread_id): async with self._sf() as session: async with session.begin(): - await self._acquire_thread_mutation_fence(session, thread_id) + max_seq = await self._max_seq_for_thread(session, thread_id) count_conditions = [RunEventRow.thread_id == thread_id, RunEventRow.run_id == run_id] if resolved_user_id is not None: count_conditions.append(RunEventRow.user_id == resolved_user_id) @@ -563,4 +587,5 @@ class DbRunEventStore(RunEventStore): count = await session.scalar(count_stmt) or 0 if count > 0: await session.execute(delete(RunEventRow).where(*count_conditions)) + await self._raise_seq_watermark(session, thread_id, max_seq or 0) return count diff --git a/backend/tests/test_db_event_store_lock_lifecycle.py b/backend/tests/test_db_event_store_lock_lifecycle.py index 34ccdc90de..db12462c60 100644 --- a/backend/tests/test_db_event_store_lock_lifecycle.py +++ b/backend/tests/test_db_event_store_lock_lifecycle.py @@ -41,6 +41,12 @@ class _PausedDeleteSession: async def execute(self, _stmt, _params=None): return None + async def get(self, _model, _thread_id): + return None + + def add(self, _obj) -> None: + return None + async def commit(self) -> None: return None diff --git a/backend/tests/test_migration_0027_0028_notification_outbox.py b/backend/tests/test_migration_0027_0028_notification_outbox.py index 66a9e07210..e7b72b4465 100644 --- a/backend/tests/test_migration_0027_0028_notification_outbox.py +++ b/backend/tests/test_migration_0027_0028_notification_outbox.py @@ -23,7 +23,7 @@ OUTBOX = "0027_notification_deliveries" PARKED = "0028_parked_attempts" SCHEDULER = "0029_scheduler_agent_tasks" CLAIM_TOKENS = "0030_notification_claim_tokens" -HEAD = "0033_batch_result_artifact" +HEAD = "0034_run_event_seq_watermark" PREVIOUS = "0026_mcp_task_lease_tokens" TABLE = "notification_deliveries" diff --git a/backend/tests/test_migration_0029_scheduler_agent_tasks.py b/backend/tests/test_migration_0029_scheduler_agent_tasks.py index 70fe31f296..43673413b1 100644 --- a/backend/tests/test_migration_0029_scheduler_agent_tasks.py +++ b/backend/tests/test_migration_0029_scheduler_agent_tasks.py @@ -23,7 +23,7 @@ from deerflow.persistence.postgres_schema import build_asyncpg_connect_args from deerflow.persistence.scheduled_tasks import ScheduledTaskRepository REVISION = "0029_scheduler_agent_tasks" -CURRENT_HEAD = "0033_batch_result_artifact" +CURRENT_HEAD = "0034_run_event_seq_watermark" NEXT = "0030_notification_claim_tokens" PREVIOUS = "0028_parked_attempts" TASK_FIELDS = {"origin_thread_id", "goal_objective", "max_runs", "end_at", "standing_notes"} diff --git a/backend/tests/test_migration_0032_activity_and_task_events.py b/backend/tests/test_migration_0032_activity_and_task_events.py index 3a57e5d691..c6644be559 100644 --- a/backend/tests/test_migration_0032_activity_and_task_events.py +++ b/backend/tests/test_migration_0032_activity_and_task_events.py @@ -17,7 +17,8 @@ from deerflow.persistence import bootstrap from deerflow.persistence.postgres_schema import build_asyncpg_connect_args REVISION = "0032_activity_and_task_events" -CURRENT_HEAD = "0033_batch_result_artifact" +NEXT = "0033_batch_result_artifact" +CURRENT_HEAD = "0034_run_event_seq_watermark" PREVIOUS = "0031_scheduled_streak_boundary" TABLES = {"thread_read_markers", "thread_read_versions", "scheduled_task_events"} INDEXES = { @@ -32,7 +33,8 @@ async def test_0032_remains_in_the_single_migration_chain(): script = ScriptDirectory(str(bootstrap._MIGRATIONS_DIR)) assert script.get_heads() == [CURRENT_HEAD] assert script.get_revision(REVISION).down_revision == PREVIOUS - assert script.get_revision(CURRENT_HEAD).down_revision == REVISION + assert script.get_revision(NEXT).down_revision == REVISION + assert script.get_revision(CURRENT_HEAD).down_revision == NEXT # alembic_version.version_num is VARCHAR(32). assert len(REVISION) <= 32 diff --git a/backend/tests/test_migration_0034_run_event_seq_watermark.py b/backend/tests/test_migration_0034_run_event_seq_watermark.py new file mode 100644 index 0000000000..1e2f9adfda --- /dev/null +++ b/backend/tests/test_migration_0034_run_event_seq_watermark.py @@ -0,0 +1,125 @@ +"""Migration 0034: durable run-event sequence watermark.""" + +from __future__ import annotations + +import asyncio +import os +import uuid +from datetime import UTC, datetime +from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit + +import pytest +import sqlalchemy as sa +from alembic import command +from alembic.script import ScriptDirectory +from sqlalchemy.ext.asyncio import create_async_engine + +from deerflow.persistence import bootstrap +from deerflow.persistence.postgres_schema import build_asyncpg_connect_args + +REVISION = "0034_run_event_seq_watermark" +PREVIOUS = "0033_batch_result_artifact" +TABLE = "run_event_thread_seq" +pytestmark = pytest.mark.asyncio + + +async def test_0034_is_the_single_migration_head(): + script = ScriptDirectory(str(bootstrap._MIGRATIONS_DIR)) + assert script.get_heads() == [REVISION] + assert script.get_revision(REVISION).down_revision == PREVIOUS + assert len(REVISION) <= 32 + + +def _engine(tmp_path, backend): + if backend == "sqlite": + return create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'migration.db'}"), None + uri = os.environ.get("TEST_POSTGRES_URI") + if not uri: + pytest.skip("requires TEST_POSTGRES_URI (real Postgres migration 0034)") + parts = urlsplit(uri) + scheme = "postgresql+asyncpg" if parts.scheme in {"postgres", "postgresql"} else parts.scheme + query = urlencode([(key, value) for key, value in parse_qsl(parts.query, keep_blank_values=True) if key not in {"sslmode", "channel_binding"}]) + schema = f"run_event_seq_{uuid.uuid4().hex}" + return create_async_engine(urlunsplit(parts._replace(scheme=scheme, query=query)), connect_args=build_asyncpg_connect_args(schema)), schema + + +@pytest.mark.parametrize("backend", ["sqlite", "postgres"]) +async def test_0034_backfills_and_round_trips(tmp_path, backend): + engine, schema = _engine(tmp_path, backend) + cfg = bootstrap._get_alembic_config(engine, postgres_schema=schema or "") + try: + if schema: + async with engine.begin() as conn: + await conn.execute(sa.text(f'CREATE SCHEMA "{schema}"')) + await asyncio.to_thread(command.upgrade, cfg, PREVIOUS) + + run_events = sa.table( + "run_events", + sa.column("thread_id", sa.String(64)), + sa.column("run_id", sa.String(64)), + sa.column("user_id", sa.String(64)), + sa.column("event_type", sa.String(32)), + sa.column("category", sa.String(16)), + sa.column("content", sa.Text()), + sa.column("event_metadata", sa.JSON()), + sa.column("seq", sa.Integer()), + sa.column("created_at", sa.DateTime(timezone=True)), + ) + async with engine.begin() as conn: + await conn.execute( + sa.insert(run_events), + [ + { + "thread_id": "t1", + "run_id": "r1", + "user_id": None, + "event_type": "human_message", + "category": "message", + "content": "", + "event_metadata": {}, + "seq": 1, + "created_at": datetime.now(UTC), + }, + { + "thread_id": "t1", + "run_id": "r2", + "user_id": None, + "event_type": "human_message", + "category": "message", + "content": "", + "event_metadata": {}, + "seq": 2, + "created_at": datetime.now(UTC), + }, + { + "thread_id": "t2", + "run_id": "r1", + "user_id": None, + "event_type": "human_message", + "category": "message", + "content": "", + "event_metadata": {}, + "seq": 5, + "created_at": datetime.now(UTC), + }, + ], + ) + + await asyncio.to_thread(command.upgrade, cfg, REVISION) + async with engine.connect() as conn: + assert TABLE in await conn.run_sync(lambda sync: set(sa.inspect(sync).get_table_names())) + rows = (await conn.execute(sa.text(f"SELECT thread_id, seq FROM {TABLE} ORDER BY thread_id"))).all() + assert [tuple(row) for row in rows] == [("t1", 2), ("t2", 5)] + + await asyncio.to_thread(command.downgrade, cfg, PREVIOUS) + async with engine.connect() as conn: + assert TABLE not in await conn.run_sync(lambda sync: set(sa.inspect(sync).get_table_names())) + + await asyncio.to_thread(command.upgrade, cfg, REVISION) + async with engine.connect() as conn: + assert TABLE in await conn.run_sync(lambda sync: set(sa.inspect(sync).get_table_names())) + finally: + if schema: + async with engine.begin() as conn: + await conn.execute(sa.text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE')) + await engine.dispose() diff --git a/backend/tests/test_run_event_store.py b/backend/tests/test_run_event_store.py index c7b749db6b..dcb4d660e9 100644 --- a/backend/tests/test_run_event_store.py +++ b/backend/tests/test_run_event_store.py @@ -492,6 +492,54 @@ class TestDelete: assert len(messages) == 1 assert messages[0]["run_id"] == "r1" + @pytest.mark.anyio + @pytest.mark.parametrize("backend", ["memory", "jsonl", "db"]) + async def test_delete_by_run_preserves_seq_watermark(self, tmp_path, backend): + from deerflow.persistence.engine import close_engine, get_session_factory, init_engine + from deerflow.runtime.events.store.db import DbRunEventStore + from deerflow.runtime.events.store.jsonl import JsonlRunEventStore + + if backend == "db": + await init_engine("sqlite", url=f"sqlite+aiosqlite:///{tmp_path / 'test.db'}", sqlite_dir=str(tmp_path)) + event_store = DbRunEventStore(get_session_factory()) + elif backend == "jsonl": + event_store = JsonlRunEventStore(base_dir=tmp_path / "jsonl") + else: + event_store = MemoryRunEventStore() + + try: + r1 = await event_store.put( + thread_id="t1", + run_id="r1", + event_type="human_message", + category="message", + content="one", + ) + r2 = await event_store.put( + thread_id="t1", + run_id="r2", + event_type="human_message", + category="message", + content="two", + ) + assert [r1["seq"], r2["seq"]] == [1, 2] + + assert await event_store.delete_by_run("t1", "r2", user_id=None) == 1 + + r3 = await event_store.put( + thread_id="t1", + run_id="r3", + event_type="human_message", + category="message", + content="three", + ) + assert r3["seq"] > r2["seq"] + messages = await event_store.list_messages("t1", limit=50, after_seq=r2["seq"], user_id=None) + assert [message["seq"] for message in messages] == [r3["seq"]] + finally: + if backend == "db": + await close_engine() + @pytest.mark.anyio async def test_delete_by_thread_accepts_owner_scope(self, store): """Every backend accepts the owner scope the Gateway passes (#2803 wiring). @@ -564,7 +612,7 @@ class TestDbRunEventStore: def __init__(self): self.dialect = postgresql.dialect() self.execute_calls = [] - self.scalar_stmt = None + self.scalar_stmts = [] def get_bind(self): return self @@ -573,7 +621,7 @@ class TestDbRunEventStore: self.execute_calls.append((stmt, params)) async def scalar(self, stmt): - self.scalar_stmt = stmt + self.scalar_stmts.append(stmt) return 41 session = FakeSession() @@ -584,8 +632,9 @@ class TestDbRunEventStore: assert session.execute_calls assert session.execute_calls[0][1] == {"thread_id": "thread-1"} assert "pg_advisory_xact_lock" in str(session.execute_calls[0][0]) - compiled = str(session.scalar_stmt.compile(dialect=postgresql.dialect())) - assert "FOR UPDATE" not in compiled + compiled = [str(stmt.compile(dialect=postgresql.dialect())) for stmt in session.scalar_stmts] + assert compiled + assert all("FOR UPDATE" not in stmt for stmt in compiled) @pytest.mark.anyio async def test_delete_by_thread_takes_postgres_advisory_lock(self): @@ -659,6 +708,12 @@ class TestDbRunEventStore: async def scalar(self, _stmt): return 2 + async def get(self, _model, _thread_id): + return None + + def add(self, _obj) -> None: + return None + async def commit(self) -> None: return None @@ -670,6 +725,7 @@ class TestDbRunEventStore: assert session.execute_calls assert "pg_advisory_xact_lock" in str(session.execute_calls[0][0]) assert session.execute_calls[0][1] == {"thread_id": "thread-1"} + assert sum("pg_advisory_xact_lock" in str(call[0]) for call in session.execute_calls) == 1 @pytest.mark.anyio async def test_basic_crud(self, tmp_path): @@ -948,6 +1004,126 @@ class TestDbRunEventStore: await close_engine() + @pytest.mark.anyio + @pytest.mark.parametrize("user_id", ["owner-a", None], ids=["owner-scoped", "unscoped"]) + @pytest.mark.parametrize("delete_run_first", [False, True], ids=["with-events", "already-empty"]) + async def test_delete_by_thread_clears_watermark_when_empty(self, tmp_path, user_id, delete_run_first): + from types import SimpleNamespace + + from deerflow.persistence.engine import close_engine, get_session_factory, init_engine + from deerflow.persistence.models.run_event import RunEventThreadSeqRow + from deerflow.runtime.events.store.db import DbRunEventStore + from deerflow.runtime.user_context import reset_current_user, set_current_user + + await init_engine("sqlite", url=f"sqlite+aiosqlite:///{tmp_path / 'test.db'}", sqlite_dir=str(tmp_path)) + owner_token = set_current_user(SimpleNamespace(id="owner-a")) + try: + session_factory = get_session_factory() + s = DbRunEventStore(session_factory) + for _ in range(2): + await s.put(thread_id="t1", run_id="r1", event_type="human_message", category="message") + await s.put(thread_id="t2", run_id="r2", event_type="human_message", category="message") + if delete_run_first: + assert await s.delete_by_run("t1", "r1", user_id=user_id) == 2 + async with session_factory() as session: + assert (await session.get(RunEventThreadSeqRow, "t1")).seq == 2 + + assert await s.delete_by_thread("t1", user_id=user_id) == (0 if delete_run_first else 2) + + assert await s.list_messages("t1", user_id=None) == [] + async with session_factory() as session: + assert await session.get(RunEventThreadSeqRow, "t1") is None + assert (await session.get(RunEventThreadSeqRow, "t2")).seq == 1 + recreated = await s.put(thread_id="t1", run_id="r3", event_type="human_message", category="message") + assert recreated["seq"] == 1 + finally: + reset_current_user(owner_token) + await close_engine() + + @pytest.mark.anyio + @pytest.mark.parametrize("remaining_category", ["message", "trace"]) + async def test_owner_scoped_thread_delete_preserves_watermark_for_remaining_events(self, tmp_path, remaining_category): + from types import SimpleNamespace + + from deerflow.persistence.engine import close_engine, get_session_factory, init_engine + from deerflow.persistence.models.run_event import RunEventThreadSeqRow + from deerflow.runtime.events.store.db import DbRunEventStore + from deerflow.runtime.user_context import reset_current_user, set_current_user + + await init_engine("sqlite", url=f"sqlite+aiosqlite:///{tmp_path / 'test.db'}", sqlite_dir=str(tmp_path)) + owner_b_token = set_current_user(SimpleNamespace(id="owner-b")) + try: + session_factory = get_session_factory() + s = DbRunEventStore(session_factory) + await s.put(thread_id="t1", run_id="r1", event_type="human_message", category=remaining_category) + owner_a_token = set_current_user(SimpleNamespace(id="owner-a")) + try: + for _ in range(2): + last_deleted = await s.put(thread_id="t1", run_id="r2", event_type="human_message", category="message") + finally: + reset_current_user(owner_a_token) + + assert await s.delete_by_thread("t1", user_id="owner-a") == 2 + assert await s.delete_by_thread("t1", user_id="missing-owner") == 0 + remaining = await s.list_events("t1", "r1", user_id=None) + assert [(event["user_id"], event["seq"]) for event in remaining] == [("owner-b", 1)] + async with session_factory() as session: + assert (await session.get(RunEventThreadSeqRow, "t1")).seq == last_deleted["seq"] + + appended = await s.put(thread_id="t1", run_id="r3", event_type="human_message", category="message") + assert appended["seq"] == last_deleted["seq"] + 1 + page = await s.list_messages("t1", after_seq=last_deleted["seq"], user_id="owner-b") + assert [event["seq"] for event in page] == [appended["seq"]] + finally: + reset_current_user(owner_b_token) + await close_engine() + + @pytest.mark.anyio + @pytest.mark.parametrize("write_path", ["put", "put_batch", "put_if_absent"]) + async def test_delete_by_run_preserves_watermark_for_every_write_path(self, tmp_path, write_path): + from deerflow.persistence.engine import close_engine, get_session_factory, init_engine + from deerflow.runtime.events.store.db import DbRunEventStore + + url = f"sqlite+aiosqlite:///{tmp_path / 'test.db'}" + await init_engine("sqlite", url=url, sqlite_dir=str(tmp_path)) + s = DbRunEventStore(get_session_factory()) + + await s.put(thread_id="t1", run_id="r1", event_type="human_message", category="message") + r2 = await s.put(thread_id="t1", run_id="r2", event_type="human_message", category="message") + await s.delete_by_run("t1", "r2", user_id=None) + + if write_path == "put": + r3 = await s.put(thread_id="t1", run_id="r3", event_type="human_message", category="message", content="three") + elif write_path == "put_batch": + r3 = ( + await s.put_batch( + [ + { + "thread_id": "t1", + "run_id": "r3", + "event_type": "human_message", + "category": "message", + "content": "three", + } + ] + ) + )[0] + else: + r3, created = await s.put_if_absent( + thread_id="t1", + run_id="r3", + event_type="run.delivery", + category="message", + content="three", + ) + assert created is True + + assert r3["seq"] > r2["seq"] + messages = await s.list_messages("t1", limit=50, after_seq=r2["seq"], user_id=None) + assert [message["seq"] for message in messages] == [r3["seq"]] + + await close_engine() + @pytest.mark.anyio async def test_put_batch_seq_continuity(self, tmp_path): """Batch write produces continuous seq values with no gaps.""" @@ -1084,7 +1260,7 @@ class TestDbRunEventStoreWriteLock: # Deleting the thread must evict the now-orphaned lock so the registry # does not grow unbounded across the singleton store's lifetime. - await s.delete_by_thread("t1") + await s.delete_by_thread("t1", user_id=None) assert "t1" not in s._write_locks # A subsequent write recreates a fresh lock and seq restarts from 1.