mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 13:16:24 +00:00
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 <willem.jiang@gmail.com>
This commit is contained in:
parent
1a88f29bf7
commit
85796f6492
@ -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
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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)
|
||||
@ -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",
|
||||
|
||||
@ -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"))
|
||||
|
||||
@ -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`.
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
@ -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"
|
||||
|
||||
|
||||
@ -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"}
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
125
backend/tests/test_migration_0034_run_event_seq_watermark.py
Normal file
125
backend/tests/test_migration_0034_run_event_seq_watermark.py
Normal file
@ -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()
|
||||
@ -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.
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user