288 lines
12 KiB
Python
288 lines
12 KiB
Python
"""Action repository for the MAF roundtable experiment (external writes).
|
|
|
|
Records every external write / tool action with an idempotency key so the
|
|
Action Coordinator (Phase 5) can guarantee at-most-once execution across
|
|
retries / restarts. Phase 1 ships the storage + CRUD; the Coordinator itself
|
|
is Phase 5. ``arguments_redacted_json`` / ``result_redacted_json`` must
|
|
already have credentials / tokens / cookies stripped by the writer — the
|
|
store does not re-redact.
|
|
"""
|
|
|
|
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 MafRoundtableActionRow
|
|
|
|
|
|
def _now() -> datetime:
|
|
return datetime.now(UTC)
|
|
|
|
|
|
def action_to_dict(row: MafRoundtableActionRow) -> dict[str, Any]:
|
|
"""Serialise an action ORM row to a plain dict (public — reused by the store facade)."""
|
|
return {
|
|
"id": row.id,
|
|
"session_id": row.session_id,
|
|
"run_id": row.run_id,
|
|
"agent_id": row.agent_id,
|
|
"tool_name": row.tool_name,
|
|
"logical_action_id": row.logical_action_id,
|
|
"turn_index": row.turn_index,
|
|
"idempotency_key": row.idempotency_key,
|
|
"target_key": row.target_key,
|
|
"risk_level": row.risk_level,
|
|
"status": row.status,
|
|
"arguments": row.arguments_redacted_json,
|
|
"result": row.result_redacted_json,
|
|
"external_object_id": row.external_object_id,
|
|
"approval_id": row.approval_id,
|
|
"attempt_count": row.attempt_count,
|
|
"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,
|
|
"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 = action_to_dict
|
|
|
|
|
|
class MafActionStateMachineError(RuntimeError):
|
|
"""An action was claimed or transitioned from an unexpected state.
|
|
|
|
Raised when a claim loses the ``approved → running`` CAS and the row is in
|
|
a state that is neither ``running`` (in progress elsewhere) nor a recognized
|
|
terminal (``succeeded`` / ``failed`` / ``unknown`` / ``rejected``) — e.g.
|
|
still ``waiting_approval`` / ``pending`` — meaning the caller tried to
|
|
execute an action that was never properly approved.
|
|
"""
|
|
|
|
|
|
# Terminal statuses a failed claim may reflect back to the caller.
|
|
_TERMINAL_CLAIM_STATUSES = ("succeeded", "failed", "unknown", "rejected")
|
|
|
|
|
|
class MafActionRepository:
|
|
def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None:
|
|
self._sf = session_factory
|
|
|
|
async def create_action(self, data: dict[str, Any]) -> dict[str, Any]:
|
|
"""Create an action. Raises ``IntegrityError`` if the ``idempotency_key``
|
|
already exists (caller should treat that as "already submitted" and
|
|
fetch the existing row via ``get_by_idempotency_key``).
|
|
"""
|
|
now = _now()
|
|
row = MafRoundtableActionRow(
|
|
id=data.get("id") or uuid4().hex,
|
|
session_id=str(data["session_id"]),
|
|
run_id=data.get("run_id"),
|
|
agent_id=data.get("agent_id"),
|
|
tool_name=str(data["tool_name"]),
|
|
logical_action_id=data.get("logical_action_id"),
|
|
turn_index=data.get("turn_index"),
|
|
idempotency_key=str(data["idempotency_key"]),
|
|
target_key=data.get("target_key"),
|
|
risk_level=str(data.get("risk_level") or "R0"),
|
|
status=str(data.get("status") or "pending"),
|
|
arguments_redacted_json=data.get("arguments"),
|
|
result_redacted_json=data.get("result"),
|
|
external_object_id=data.get("external_object_id"),
|
|
approval_id=data.get("approval_id"),
|
|
attempt_count=int(data.get("attempt_count", 0)),
|
|
started_at=data.get("started_at"),
|
|
finished_at=data.get("finished_at"),
|
|
created_at=now,
|
|
updated_at=now,
|
|
)
|
|
async with self._sf() as session:
|
|
session.add(row)
|
|
await session.commit()
|
|
return await self.get_action(row.id) # type: ignore[arg-type]
|
|
|
|
async def get_action(self, action_id: str) -> dict[str, Any] | None:
|
|
async with self._sf() as session:
|
|
stmt = select(MafRoundtableActionRow).where(MafRoundtableActionRow.id == action_id).limit(1)
|
|
row = (await session.execute(stmt)).scalars().first()
|
|
return _to_dict(row) if row is not None else None
|
|
|
|
async def get_by_idempotency_key(self, key: str) -> dict[str, Any] | None:
|
|
async with self._sf() as session:
|
|
stmt = select(MafRoundtableActionRow).where(MafRoundtableActionRow.idempotency_key == key).limit(1)
|
|
row = (await session.execute(stmt)).scalars().first()
|
|
return _to_dict(row) if row is not None else None
|
|
|
|
async def list_actions(
|
|
self, *, session_id: str | None = None, run_id: str | None = None, status: str | None = None
|
|
) -> list[dict[str, Any]]:
|
|
stmt = select(MafRoundtableActionRow)
|
|
if session_id is not None:
|
|
stmt = stmt.where(MafRoundtableActionRow.session_id == session_id)
|
|
if run_id is not None:
|
|
stmt = stmt.where(MafRoundtableActionRow.run_id == run_id)
|
|
if status is not None:
|
|
stmt = stmt.where(MafRoundtableActionRow.status == status)
|
|
stmt = stmt.order_by(MafRoundtableActionRow.created_at.asc())
|
|
async with self._sf() as session:
|
|
rows = (await session.execute(stmt)).scalars().all()
|
|
return [_to_dict(r) for r in rows]
|
|
|
|
async def update_action(self, action_id: str, fields: dict[str, Any]) -> dict[str, Any] | None:
|
|
async with self._sf() as session:
|
|
stmt = select(MafRoundtableActionRow).where(MafRoundtableActionRow.id == action_id).limit(1)
|
|
row = (await session.execute(stmt)).scalars().first()
|
|
if row is None:
|
|
return None
|
|
col_map = {
|
|
"status": "status",
|
|
"risk_level": "risk_level",
|
|
"arguments": "arguments_redacted_json",
|
|
"result": "result_redacted_json",
|
|
"external_object_id": "external_object_id",
|
|
"approval_id": "approval_id",
|
|
"attempt_count": "attempt_count",
|
|
"started_at": "started_at",
|
|
"finished_at": "finished_at",
|
|
}
|
|
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 claim_for_execution(self, action_id: str) -> str:
|
|
"""Atomically transition an action from ``approved`` to ``running``.
|
|
|
|
Single conditional UPDATE — only the caller whose UPDATE affects
|
|
exactly one row wins the right to execute the external write. This
|
|
guarantees at-most-once execution even under concurrent ``resume()``
|
|
calls or retries after a restart.
|
|
|
|
Returns the effective claim status (callers must NOT treat a lost claim
|
|
as a blanket "skip"):
|
|
|
|
- ``"claimed"`` — won the CAS; the caller may execute.
|
|
- ``"in_progress"`` — another executor already moved it to ``running``.
|
|
- ``"succeeded"`` / ``"failed"`` / ``"unknown"`` / ``"rejected"`` — the
|
|
action is already in that terminal state.
|
|
|
|
Raises ``MafActionStateMachineError`` when the row is missing or in any
|
|
other state (e.g. still ``waiting_approval`` / ``pending``) — i.e. it was
|
|
never properly approved and must not be executed.
|
|
"""
|
|
now = _now()
|
|
async with self._sf() as session:
|
|
result = await session.execute(
|
|
update(MafRoundtableActionRow)
|
|
.where(
|
|
MafRoundtableActionRow.id == action_id,
|
|
MafRoundtableActionRow.status == "approved",
|
|
)
|
|
.values(status="running", started_at=now, updated_at=now)
|
|
)
|
|
await session.commit()
|
|
if result.rowcount == 1: # type: ignore[union-attr]
|
|
return "claimed"
|
|
|
|
# Lost the CAS (or never approved) — report the ACTUAL state.
|
|
current = await self.get_action(action_id)
|
|
if current is None:
|
|
raise MafActionStateMachineError(f"action {action_id} not found during claim")
|
|
status = current.get("status")
|
|
if status == "running":
|
|
return "in_progress"
|
|
if status in _TERMINAL_CLAIM_STATUSES:
|
|
return str(status)
|
|
raise MafActionStateMachineError(
|
|
f"action {action_id} cannot be claimed from state '{status}'"
|
|
)
|
|
|
|
async def mark_approved(self, action_id: str) -> bool:
|
|
"""Conditionally transition ``waiting_approval → approved``.
|
|
|
|
Returns True only if this call performed the transition (``rowcount==1``).
|
|
A concurrent decider or an already-terminal action yields False — the
|
|
caller must NOT overwrite a terminal state via a generic update.
|
|
"""
|
|
now = _now()
|
|
async with self._sf() as session:
|
|
result = await session.execute(
|
|
update(MafRoundtableActionRow)
|
|
.where(
|
|
MafRoundtableActionRow.id == action_id,
|
|
MafRoundtableActionRow.status == "waiting_approval",
|
|
)
|
|
.values(status="approved", updated_at=now)
|
|
)
|
|
await session.commit()
|
|
return result.rowcount == 1 # type: ignore[union-attr]
|
|
|
|
async def mark_rejected(self, action_id: str) -> bool:
|
|
"""Conditionally transition ``waiting_approval → rejected`` (+ finished_at).
|
|
|
|
Returns True only if this call performed the transition. Terminal states
|
|
are never overwritten.
|
|
"""
|
|
now = _now()
|
|
async with self._sf() as session:
|
|
result = await session.execute(
|
|
update(MafRoundtableActionRow)
|
|
.where(
|
|
MafRoundtableActionRow.id == action_id,
|
|
MafRoundtableActionRow.status == "waiting_approval",
|
|
)
|
|
.values(status="rejected", finished_at=now, updated_at=now)
|
|
)
|
|
await session.commit()
|
|
return result.rowcount == 1 # type: ignore[union-attr]
|
|
|
|
async def record_terminal_result(
|
|
self,
|
|
action_id: str,
|
|
*,
|
|
status: str,
|
|
result: dict[str, Any] | None = None,
|
|
external_object_id: str | None = None,
|
|
) -> dict[str, Any] | None:
|
|
"""Record the terminal outcome of an action execution.
|
|
|
|
State-machine constrained: only a ``running`` action may transition
|
|
to a terminal status (``succeeded`` / ``failed`` / ``unknown``).
|
|
A concurrent caller that lost the race (or a replay after restart)
|
|
finds the row already terminal and its UPDATE is a no-op
|
|
(``rowcount == 0``) — it must NOT overwrite the real result.
|
|
"""
|
|
if status not in ("succeeded", "failed", "unknown"):
|
|
raise ValueError(f"invalid action result status: {status}")
|
|
now = _now()
|
|
async with self._sf() as session:
|
|
result_q = await session.execute(
|
|
update(MafRoundtableActionRow)
|
|
.where(
|
|
MafRoundtableActionRow.id == action_id,
|
|
MafRoundtableActionRow.status == "running",
|
|
)
|
|
.values(
|
|
status=status,
|
|
result_redacted_json=result,
|
|
external_object_id=external_object_id,
|
|
finished_at=now,
|
|
updated_at=now,
|
|
)
|
|
)
|
|
await session.commit()
|
|
# Use rowcount instead of .returning() for cross-DB compatibility
|
|
# (SQLite, PostgreSQL, MySQL).
|
|
if result_q.rowcount != 1: # type: ignore[union-attr]
|
|
# Lost the race or not in running state — do not overwrite.
|
|
return None
|
|
return await self.get_action(action_id)
|