151 lines
6.3 KiB
Python
151 lines
6.3 KiB
Python
"""Run repository for the MAF roundtable experiment.
|
|
|
|
A run is one execution of a session. Phase 1 only persists the run row and its
|
|
status lifecycle — actual workflow execution (MAF SDK) is Phase 2+. The run
|
|
row carries the MAF checkpoint (for resume-after-restart) and a denormalized
|
|
``last_event_seq`` cache for fast polling; the authoritative seq lives in
|
|
``maf_roundtable_events``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from datetime import UTC, datetime
|
|
from typing import Any
|
|
from uuid import uuid4
|
|
|
|
from sqlalchemy import select, update
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
|
|
|
from deerflow.persistence.maf_roundtable.model import MafRoundtableRunRow
|
|
|
|
|
|
def _now() -> datetime:
|
|
return datetime.now(UTC)
|
|
|
|
|
|
def run_to_dict(row: MafRoundtableRunRow) -> dict[str, Any]:
|
|
"""Serialise a run ORM row to a plain dict (public — reused by the store facade)."""
|
|
return {
|
|
"id": row.id,
|
|
"session_id": row.session_id,
|
|
"status": row.status,
|
|
"workflow_version": row.workflow_version,
|
|
"maf_checkpoint": row.maf_checkpoint_json,
|
|
"last_event_seq": row.last_event_seq,
|
|
"started_at": row.started_at.isoformat() if isinstance(row.started_at, datetime) else row.started_at,
|
|
"finished_at": row.finished_at.isoformat() if isinstance(row.finished_at, datetime) else row.finished_at,
|
|
"error_code": row.error_code,
|
|
"error_message_safe": row.error_message_safe,
|
|
"token_usage": row.token_usage_json,
|
|
"latency_ms": row.latency_ms,
|
|
"created_at": row.created_at.isoformat() if isinstance(row.created_at, datetime) else row.created_at,
|
|
"updated_at": row.updated_at.isoformat() if isinstance(row.updated_at, datetime) else row.updated_at,
|
|
}
|
|
|
|
|
|
# Back-compat alias for internal callers that used the old private name.
|
|
_to_dict = run_to_dict
|
|
|
|
|
|
class MafRunRepository:
|
|
def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None:
|
|
self._sf = session_factory
|
|
|
|
async def create_run(self, session_id: str, *, workflow_version: str | None = None) -> dict[str, Any]:
|
|
now = _now()
|
|
row = MafRoundtableRunRow(
|
|
id=uuid4().hex,
|
|
session_id=session_id,
|
|
status="pending",
|
|
workflow_version=workflow_version,
|
|
started_at=None, # pending = queued, not yet handed to the executor
|
|
created_at=now,
|
|
updated_at=now,
|
|
)
|
|
async with self._sf() as session:
|
|
session.add(row)
|
|
await session.commit()
|
|
return await self.get_run(row.id) # type: ignore[arg-type]
|
|
|
|
async def get_run(self, run_id: str) -> dict[str, Any] | None:
|
|
async with self._sf() as session:
|
|
stmt = select(MafRoundtableRunRow).where(MafRoundtableRunRow.id == run_id).limit(1)
|
|
row = (await session.execute(stmt)).scalars().first()
|
|
return _to_dict(row) if row is not None else None
|
|
|
|
async def list_runs(self, session_id: str) -> list[dict[str, Any]]:
|
|
stmt = (
|
|
select(MafRoundtableRunRow)
|
|
.where(MafRoundtableRunRow.session_id == session_id)
|
|
.order_by(MafRoundtableRunRow.created_at.desc())
|
|
)
|
|
async with self._sf() as session:
|
|
rows = (await session.execute(stmt)).scalars().all()
|
|
return [_to_dict(r) for r in rows]
|
|
|
|
async def update_run(self, run_id: str, fields: dict[str, Any]) -> dict[str, Any] | None:
|
|
async with self._sf() as session:
|
|
stmt = select(MafRoundtableRunRow).where(MafRoundtableRunRow.id == run_id).limit(1)
|
|
row = (await session.execute(stmt)).scalars().first()
|
|
if row is None:
|
|
return None
|
|
col_map = {
|
|
"status": "status",
|
|
"workflow_version": "workflow_version",
|
|
"maf_checkpoint": "maf_checkpoint_json",
|
|
"last_event_seq": "last_event_seq",
|
|
"started_at": "started_at",
|
|
"finished_at": "finished_at",
|
|
"error_code": "error_code",
|
|
"error_message_safe": "error_message_safe",
|
|
"token_usage": "token_usage_json",
|
|
"latency_ms": "latency_ms",
|
|
}
|
|
for key, value in fields.items():
|
|
col = col_map.get(key)
|
|
if col is not None:
|
|
setattr(row, col, value)
|
|
row.updated_at = _now()
|
|
await session.commit()
|
|
return _to_dict(row)
|
|
|
|
async def set_status(self, run_id: str, status: str) -> None:
|
|
async with self._sf() as session:
|
|
await session.execute(
|
|
update(MafRoundtableRunRow)
|
|
.where(MafRoundtableRunRow.id == run_id)
|
|
.values(status=status, updated_at=_now())
|
|
)
|
|
await session.commit()
|
|
|
|
async def claim_for_resume(self, run_id: str) -> bool:
|
|
"""Run-level CAS gating post-approval execution: ``waiting_approval → running``.
|
|
|
|
After every approval is decided, exactly one concurrent ``resume()`` may
|
|
proceed to execute the approved actions and write the parent run's final
|
|
state. The winner (``rowcount == 1``) executes; every other concurrent
|
|
caller gets False and must return the current run WITHOUT writing a
|
|
terminal status (so it can never race a ``completed``/``failed`` write).
|
|
"""
|
|
async with self._sf() as session:
|
|
result = await session.execute(
|
|
update(MafRoundtableRunRow)
|
|
.where(
|
|
MafRoundtableRunRow.id == run_id,
|
|
MafRoundtableRunRow.status == "waiting_approval",
|
|
)
|
|
.values(status="running", started_at=_now(), updated_at=_now())
|
|
)
|
|
await session.commit()
|
|
return result.rowcount == 1 # type: ignore[union-attr]
|
|
|
|
async def bump_last_event_seq(self, run_id: str, seq: int) -> None:
|
|
"""Denormalized cache update — only bumps forward (max wins)."""
|
|
async with self._sf() as session:
|
|
await session.execute(
|
|
update(MafRoundtableRunRow)
|
|
.where(MafRoundtableRunRow.id == run_id, MafRoundtableRunRow.last_event_seq < seq)
|
|
.values(last_event_seq=seq, updated_at=_now())
|
|
)
|
|
await session.commit()
|