初始化提交

backend/app/application/agent/confirmation.py

@@ -0,0 +1,148 @@
+"""混合确认注册表 — 同进程 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"]