149 lines
5.3 KiB
Python
149 lines
5.3 KiB
Python
"""混合确认注册表 — 同进程 Event 唤醒 + 重启后 DB polling 兜底。
|
||
|
||
Agent 循环在工具需要确认时：
|
||
1. 写 PendingConfirmationRecord 到 DB（持久化挂起状态）
|
||
2. 调 ``ConfirmationRegistry.wait(run_id)`` — 同进程用 Event 0 延迟唤醒，
|
||
重启后 polling DB 兜底
|
||
|
||
approve/deny 端点统一调 ``registry.resolve(run_id, approved)``：
|
||
- 更新 PendingConfirmationRecord.status + resolved_at
|
||
- 同一进程直接 set Event
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import time
|
||
from datetime import UTC, datetime
|
||
|
||
from sqlalchemy import or_, select, update
|
||
|
||
from app.models.safety import PendingConfirmationRecord
|
||
|
||
|
||
class ConfirmationRegistry:
|
||
"""混合确认注册表。
|
||
|
||
- 同进程：``asyncio.Event`` 直接唤醒，0 延迟
|
||
- 重启后：Event 不存在 → polling ``PendingConfirmationRecord.status`` 兜底
|
||
- approve/deny 端点统一调用 ``resolve(run_id, approved)``
|
||
"""
|
||
|
||
def __init__(self) -> None:
|
||
self._events: dict[str, asyncio.Event] = {}
|
||
|
||
async def wait(
|
||
self,
|
||
run_id: str,
|
||
*,
|
||
poll_interval: float = 0.5,
|
||
timeout: float = 60,
|
||
confirmation_id: str | None = None,
|
||
) -> bool:
|
||
"""等待确认。
|
||
|
||
同一进程的事件直接唤醒；跨进程或重启后 polling 兜底。
|
||
|
||
``confirmation_id`` 可选，传入时会同时注册为该 Event 的别名，
|
||
使得外部以 ``confirmation_id`` 调 ``resolve`` 时也能立即唤醒同进程的 Event
|
||
（无需等 DB polling 兜底）。
|
||
|
||
返回：
|
||
True=批准, False=拒绝 ; 超时抛出 TimeoutError
|
||
"""
|
||
event = asyncio.Event()
|
||
self._events[run_id] = event
|
||
if confirmation_id:
|
||
self._events[confirmation_id] = event
|
||
|
||
try:
|
||
deadline = time.monotonic() + timeout
|
||
while time.monotonic() < deadline:
|
||
# 1) polling DB 获取真实状态（Event 只用于唤醒，不携带批准/拒绝值）
|
||
from app.db.session import AsyncSessionLocal
|
||
|
||
async with AsyncSessionLocal() as session:
|
||
if confirmation_id is None:
|
||
stmt = select(PendingConfirmationRecord).where(
|
||
or_(
|
||
PendingConfirmationRecord.id == run_id,
|
||
PendingConfirmationRecord.run_id == run_id,
|
||
)
|
||
)
|
||
else:
|
||
stmt = select(PendingConfirmationRecord).where(
|
||
PendingConfirmationRecord.id == confirmation_id,
|
||
PendingConfirmationRecord.run_id == run_id,
|
||
)
|
||
result = await session.execute(stmt)
|
||
record = result.scalar_one_or_none()
|
||
if record is not None and record.status != "pending":
|
||
return record.status == "approved"
|
||
|
||
# 2) 等一个间隔，同时监听 Event（Event set 只表示「有变化」，实际值靠 DB）
|
||
try:
|
||
await asyncio.wait_for(
|
||
asyncio.shield(event.wait()),
|
||
timeout=poll_interval,
|
||
)
|
||
# Event 被唤醒 → 回到循环顶部重新查 DB
|
||
continue
|
||
except TimeoutError:
|
||
continue
|
||
|
||
# 超时
|
||
raise TimeoutError(f"确认超时（等待 {timeout:.0f}s 未裁决）")
|
||
finally:
|
||
self._events.pop(run_id, None)
|
||
if confirmation_id:
|
||
self._events.pop(confirmation_id, None)
|
||
|
||
async def resolve(
|
||
self,
|
||
run_id: str,
|
||
approved: bool,
|
||
*,
|
||
confirmation_id: str | None = None,
|
||
) -> bool:
|
||
"""精确裁决一个待确认工具调用，并唤醒对应等待者。"""
|
||
from app.db.session import AsyncSessionLocal
|
||
|
||
target_id = confirmation_id or run_id
|
||
async with AsyncSessionLocal() as session:
|
||
status = "approved" if approved else "denied"
|
||
stmt = update(PendingConfirmationRecord).where(
|
||
PendingConfirmationRecord.status == "pending"
|
||
)
|
||
if confirmation_id is None:
|
||
stmt = stmt.where(
|
||
or_(
|
||
PendingConfirmationRecord.id == run_id,
|
||
PendingConfirmationRecord.run_id == run_id,
|
||
)
|
||
)
|
||
else:
|
||
stmt = stmt.where(
|
||
PendingConfirmationRecord.id == confirmation_id,
|
||
PendingConfirmationRecord.run_id == run_id,
|
||
)
|
||
result = await session.execute(
|
||
stmt.values(status=status, resolved_at=datetime.now(UTC))
|
||
)
|
||
await session.commit()
|
||
|
||
if result.rowcount == 0:
|
||
return False
|
||
|
||
event = self._events.get(target_id) or self._events.get(run_id)
|
||
if event is not None:
|
||
event.set()
|
||
return True
|
||
|
||
def has(self, run_id: str) -> bool:
|
||
"""是否仍在等待（同进程检查）。"""
|
||
event = self._events.get(run_id)
|
||
return event is not None and not event.is_set()
|
||
|
||
|
||
__all__ = ["ConfirmationRegistry"]
|