200 lines
7.9 KiB
Python
200 lines
7.9 KiB
Python
"""Event repository for the MAF roundtable experiment — monotonic per-run seq.
|
|
|
|
``(run_id, seq)`` is unique. The store allocates seq via ``SELECT MAX(seq) +
|
|
1`` and retries on ``IntegrityError`` so concurrent appends from parallel
|
|
seats cannot collide or skip seq values. This is the standard portable
|
|
approach (works on sqlite / postgres / mysql — no ``RETURNING`` / ``UPSERT``
|
|
dialect differences).
|
|
|
|
Pull API: ``list_events(run_id, after_seq=N)`` returns events with ``seq > N``
|
|
in ascending seq order — the断线补拉 contract. The router exposes this both
|
|
as a one-shot ``GET /runs/{run_id}/events?after_seq=N`` and as an SSE
|
|
``GET /runs/{run_id}/stream?after_seq=N`` (Phase 6 wires the SSE reducer;
|
|
Phase 1 ships the storage + one-shot pull).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from datetime import UTC, datetime
|
|
from typing import Any
|
|
from uuid import uuid4
|
|
|
|
from sqlalchemy import func, select
|
|
from sqlalchemy.exc import IntegrityError
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
|
|
|
from deerflow.persistence.maf_roundtable.model import MafRoundtableEventRow
|
|
from deerflow.persistence.maf_roundtable.runs import MafRunRepository
|
|
|
|
# Retry budget for the MAX+1 race: a parallel seat that wins the INSERT between
|
|
# our SELECT and our INSERT forces one retry. 8 is far more than the realistic
|
|
# concurrency (one run has at most ~8 seats), so a true deadlock here would
|
|
# indicate a bug, not contention.
|
|
_SEQ_INSERT_RETRIES = 8
|
|
|
|
|
|
def _now() -> datetime:
|
|
return datetime.now(UTC)
|
|
|
|
|
|
def _to_dict(row: MafRoundtableEventRow) -> dict[str, Any]:
|
|
return {
|
|
"id": row.id,
|
|
"session_id": row.session_id,
|
|
"run_id": row.run_id,
|
|
"seq": row.seq,
|
|
"event_type": row.event_type,
|
|
"agent_id": row.agent_id,
|
|
"turn_index": row.turn_index,
|
|
"source_event_id": row.source_event_id,
|
|
"payload": row.payload_json,
|
|
"created_at": row.created_at.isoformat() if isinstance(row.created_at, datetime) else row.created_at,
|
|
}
|
|
|
|
|
|
class MafEventRepository:
|
|
def __init__(
|
|
self,
|
|
session_factory: async_sessionmaker[AsyncSession],
|
|
run_repo: MafRunRepository | None = None,
|
|
) -> None:
|
|
self._sf = session_factory
|
|
self._run_repo = run_repo
|
|
|
|
async def append_event(
|
|
self,
|
|
*,
|
|
session_id: str,
|
|
run_id: str,
|
|
event_type: str,
|
|
agent_id: str | None = None,
|
|
turn_index: int | None = None,
|
|
payload: dict[str, Any] | None = None,
|
|
source_event_id: str | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Append one event with a monotonic per-run seq.
|
|
|
|
Concurrent appends are safe: each attempts ``MAX(seq)+1`` and retries
|
|
on the ``(run_id, seq)`` unique violation. The run's
|
|
``last_event_seq`` cache is bumped forward (max wins) so pollers can
|
|
cheaply check ``runs.last_event_seq`` without scanning the events
|
|
table.
|
|
|
|
Durable dedup: when ``source_event_id`` is given, the
|
|
``(run_id, source_event_id)`` unique constraint guarantees at-most-once
|
|
persistence. A duplicate insert returns the already-stored row instead
|
|
of raising — this survives process restarts (unlike an in-memory seen
|
|
set), so a replayed runtime stream never double-writes an event.
|
|
"""
|
|
# Fast-path dedup: if this source_event_id already landed, return it.
|
|
if source_event_id:
|
|
existing = await self._get_by_source(run_id, source_event_id)
|
|
if existing is not None:
|
|
return existing
|
|
|
|
last_exc: Exception | None = None
|
|
for _ in range(_SEQ_INSERT_RETRIES):
|
|
async with self._sf() as session:
|
|
max_seq = (
|
|
await session.execute(
|
|
select(func.max(MafRoundtableEventRow.seq)).where(
|
|
MafRoundtableEventRow.run_id == run_id
|
|
)
|
|
)
|
|
).scalar_one()
|
|
next_seq = int(max_seq or 0) + 1
|
|
row = MafRoundtableEventRow(
|
|
id=uuid4().hex,
|
|
session_id=session_id,
|
|
run_id=run_id,
|
|
seq=next_seq,
|
|
event_type=event_type,
|
|
agent_id=agent_id,
|
|
turn_index=turn_index,
|
|
source_event_id=source_event_id,
|
|
payload_json=payload,
|
|
created_at=_now(),
|
|
)
|
|
session.add(row)
|
|
try:
|
|
await session.commit()
|
|
except IntegrityError as exc: # noqa: PERF203 - retry path
|
|
await session.rollback()
|
|
# A source_event_id collision means a concurrent writer
|
|
# already persisted this exact event — return its row.
|
|
if source_event_id:
|
|
dup = await self._get_by_source(run_id, source_event_id)
|
|
if dup is not None:
|
|
return dup
|
|
last_exc = exc
|
|
continue
|
|
if self._run_repo is not None:
|
|
# best-effort cache bump; never let a cache failure abort
|
|
# an event the user already sees.
|
|
try:
|
|
await self._run_repo.bump_last_event_seq(run_id, next_seq)
|
|
except Exception: # noqa: BLE001
|
|
pass
|
|
return _to_dict(row)
|
|
raise IntegrityError(
|
|
"Could not append MAF event after retries (run_id seq race)",
|
|
params=None,
|
|
orig=last_exc,
|
|
)
|
|
|
|
async def _get_by_source(self, run_id: str, source_event_id: str) -> dict[str, Any] | None:
|
|
stmt = select(MafRoundtableEventRow).where(
|
|
MafRoundtableEventRow.run_id == run_id,
|
|
MafRoundtableEventRow.source_event_id == source_event_id,
|
|
)
|
|
async with self._sf() as session:
|
|
row = (await session.execute(stmt)).scalars().first()
|
|
return _to_dict(row) if row is not None else None
|
|
|
|
async def append_events(
|
|
self, *, session_id: str, run_id: str, events: list[dict[str, Any]]
|
|
) -> list[dict[str, Any]]:
|
|
"""Append a batch in order. Each event is appended sequentially so the
|
|
per-run seq stays monotonic; a parallel writer can interleave but never
|
|
reorder within this batch.
|
|
"""
|
|
out: list[dict[str, Any]] = []
|
|
for ev in events:
|
|
out.append(
|
|
await self.append_event(
|
|
session_id=session_id,
|
|
run_id=run_id,
|
|
event_type=str(ev["event_type"]),
|
|
agent_id=ev.get("agent_id"),
|
|
turn_index=ev.get("turn_index"),
|
|
payload=ev.get("payload"),
|
|
source_event_id=ev.get("source_event_id"),
|
|
)
|
|
)
|
|
return out
|
|
|
|
async def list_events(
|
|
self, run_id: str, *, after_seq: int = 0, limit: int = 1000
|
|
) -> list[dict[str, Any]]:
|
|
"""断线补拉: events with ``seq > after_seq`` in ascending seq order."""
|
|
stmt = (
|
|
select(MafRoundtableEventRow)
|
|
.where(MafRoundtableEventRow.run_id == run_id, MafRoundtableEventRow.seq > after_seq)
|
|
.order_by(MafRoundtableEventRow.seq.asc())
|
|
.limit(limit)
|
|
)
|
|
async with self._sf() as session:
|
|
rows = (await session.execute(stmt)).scalars().all()
|
|
return [_to_dict(r) for r in rows]
|
|
|
|
async def latest_seq(self, run_id: str) -> int:
|
|
async with self._sf() as session:
|
|
max_seq = (
|
|
await session.execute(
|
|
select(func.max(MafRoundtableEventRow.seq)).where(
|
|
MafRoundtableEventRow.run_id == run_id
|
|
)
|
|
)
|
|
).scalar_one()
|
|
return int(max_seq or 0)
|