179 lines
6.6 KiB
Python
179 lines
6.6 KiB
Python
"""待确认动作的存储（高风险工具调用的确认中转）。
|
||
|
||
司命决策为 ask 时不直接执行，而是登记一个 PendingConfirmation 并返回 confirmation_id；
|
||
前端展示确认卡片，用户 approve/deny 后由安全 API 消费。
|
||
|
||
提供两种实现：
|
||
- ``InMemoryConfirmationStore``：进程内 dict，默认开发模式，不依赖 DB
|
||
- ``DbConfirmationStore``：数据库持久化，生产多 worker 模式，确认项跨进程可见
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import uuid
|
||
from collections.abc import Sequence
|
||
from dataclasses import dataclass, field
|
||
from datetime import UTC, datetime
|
||
from typing import Any, Protocol
|
||
|
||
from sqlalchemy import select
|
||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
||
|
||
from app.models.safety import PendingConfirmationRecord
|
||
|
||
|
||
def _now_iso() -> str:
|
||
return datetime.now(UTC).isoformat()
|
||
|
||
|
||
def new_confirmation_id() -> str:
|
||
return f"confirm_{uuid.uuid4().hex}"
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class PendingConfirmation:
|
||
"""一次待确认的工具调用。"""
|
||
|
||
confirmation_id: str
|
||
tool: str
|
||
input: dict[str, Any]
|
||
risk_level: str
|
||
reason: str
|
||
run_id: str | None = None
|
||
character_id: str | None = None
|
||
trace_id: str | None = None
|
||
created_at: str = field(default_factory=_now_iso)
|
||
|
||
|
||
class ConfirmationStore(Protocol):
|
||
"""确认存储协议（全异步接口）。"""
|
||
|
||
async def create(self, pending: PendingConfirmation) -> PendingConfirmation: ...
|
||
|
||
async def get(self, confirmation_id: str) -> PendingConfirmation | None: ...
|
||
|
||
async def pop(self, confirmation_id: str) -> PendingConfirmation | None: ...
|
||
|
||
async def list_pending(self) -> list[PendingConfirmation]: ...
|
||
|
||
async def mark_timed_out(self, confirmation_id: str) -> None: ...
|
||
|
||
|
||
class InMemoryConfirmationStore:
|
||
"""内存确认存储，供本地与测试使用。"""
|
||
|
||
def __init__(self) -> None:
|
||
self._items: dict[str, PendingConfirmation] = {}
|
||
|
||
async def create(self, pending: PendingConfirmation) -> PendingConfirmation:
|
||
self._items[pending.confirmation_id] = pending
|
||
return pending
|
||
|
||
async def get(self, confirmation_id: str) -> PendingConfirmation | None:
|
||
return self._items.get(confirmation_id)
|
||
|
||
async def pop(self, confirmation_id: str) -> PendingConfirmation | None:
|
||
return self._items.pop(confirmation_id, None)
|
||
|
||
async def list_pending(self) -> list[PendingConfirmation]:
|
||
return sorted(self._items.values(), key=lambda item: item.created_at)
|
||
|
||
async def mark_timed_out(self, confirmation_id: str) -> None:
|
||
self._items.pop(confirmation_id, None)
|
||
|
||
|
||
class DbConfirmationStore:
|
||
"""数据库确认存储，供多 worker 生产环境使用。
|
||
|
||
待确认项持久化到 ``pending_confirmations`` 表，跨进程可见。
|
||
需要从 ``app.state.session_factory`` 或 ``async_sessionmaker`` 注入。
|
||
"""
|
||
|
||
def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None:
|
||
self._session_factory = session_factory
|
||
|
||
async def create(self, pending: PendingConfirmation) -> PendingConfirmation:
|
||
async with self._session_factory() as session:
|
||
record = PendingConfirmationRecord(
|
||
id=pending.confirmation_id,
|
||
tool=pending.tool,
|
||
input_json=json.dumps(pending.input, ensure_ascii=False),
|
||
risk_level=pending.risk_level,
|
||
reason=pending.reason,
|
||
run_id=pending.run_id,
|
||
character_id=pending.character_id,
|
||
trace_id=pending.trace_id,
|
||
)
|
||
session.add(record)
|
||
await session.commit()
|
||
return pending
|
||
|
||
async def get(self, confirmation_id: str) -> PendingConfirmation | None:
|
||
async with self._session_factory() as session:
|
||
result = await session.execute(
|
||
select(PendingConfirmationRecord).where(
|
||
PendingConfirmationRecord.id == confirmation_id,
|
||
PendingConfirmationRecord.status == "pending",
|
||
)
|
||
)
|
||
record = result.scalar_one_or_none()
|
||
if record is None:
|
||
return None
|
||
return self._record_to_pending(record)
|
||
|
||
async def pop(self, confirmation_id: str) -> PendingConfirmation | None:
|
||
async with self._session_factory() as session:
|
||
result = await session.execute(
|
||
select(PendingConfirmationRecord).where(
|
||
PendingConfirmationRecord.id == confirmation_id,
|
||
PendingConfirmationRecord.status == "pending",
|
||
)
|
||
)
|
||
record = result.scalar_one_or_none()
|
||
if record is None:
|
||
return None
|
||
record.status = "consumed"
|
||
await session.commit()
|
||
return self._record_to_pending(record)
|
||
|
||
async def list_pending(self) -> list[PendingConfirmation]:
|
||
async with self._session_factory() as session:
|
||
result = await session.execute(
|
||
select(PendingConfirmationRecord)
|
||
.where(PendingConfirmationRecord.status == "pending")
|
||
.order_by(PendingConfirmationRecord.created_at.asc())
|
||
)
|
||
records: Sequence[PendingConfirmationRecord] = result.scalars().all()
|
||
return [self._record_to_pending(r) for r in records]
|
||
|
||
async def mark_timed_out(self, confirmation_id: str) -> None:
|
||
"""将待确认记录标记为超时（Agent 循环确认等待超时）。"""
|
||
async with self._session_factory() as session:
|
||
result = await session.execute(
|
||
select(PendingConfirmationRecord).where(
|
||
PendingConfirmationRecord.id == confirmation_id,
|
||
PendingConfirmationRecord.status == "pending",
|
||
)
|
||
)
|
||
record = result.scalar_one_or_none()
|
||
if record is None:
|
||
return
|
||
record.status = "timed_out"
|
||
record.resolved_at = datetime.now(UTC)
|
||
await session.commit()
|
||
|
||
@staticmethod
|
||
def _record_to_pending(record: PendingConfirmationRecord) -> PendingConfirmation:
|
||
return PendingConfirmation(
|
||
confirmation_id=record.id,
|
||
tool=record.tool,
|
||
input=json.loads(record.input_json) if record.input_json else {},
|
||
risk_level=record.risk_level,
|
||
reason=record.reason,
|
||
run_id=record.run_id,
|
||
character_id=record.character_id,
|
||
trace_id=record.trace_id,
|
||
created_at=record.created_at.isoformat() if record.created_at else _now_iso(),
|
||
)
|