Agent/agent/execution/protocols.py

"""
Trace Storage Protocol - Trace 存储接口定义

使用 Protocol 定义接口，允许不同的存储实现（内存、PostgreSQL、Neo4j 等）
"""

from typing import Protocol, List, Optional, Dict, Any, runtime_checkable

from agent.execution.models import Trace, Step


@runtime_checkable
class TraceStore(Protocol):
    """Trace + Step 存储接口"""

    # ===== Trace 操作 =====

    async def create_trace(self, trace: Trace) -> str:
        """
        创建新的 Trace

        Args:
            trace: Trace 对象

        Returns:
            trace_id
        """
        ...

    async def get_trace(self, trace_id: str) -> Optional[Trace]:
        """获取 Trace"""
        ...

    async def update_trace(self, trace_id: str, **updates) -> None:
        """
        更新 Trace

        Args:
            trace_id: Trace ID
            **updates: 要更新的字段
        """
        ...

    async def list_traces(
        self,
        mode: Optional[str] = None,
        agent_type: Optional[str] = None,
        uid: Optional[str] = None,
        status: Optional[str] = None,
        limit: int = 50
    ) -> List[Trace]:
        """列出 Traces"""
        ...

    # ===== Step 操作 =====

    async def add_step(self, step: Step) -> str:
        """
        添加 Step

        Args:
            step: Step 对象

        Returns:
            step_id
        """
        ...

    async def get_step(self, step_id: str) -> Optional[Step]:
        """获取 Step"""
        ...

    async def get_trace_steps(self, trace_id: str) -> List[Step]:
        """获取 Trace 的所有 Steps（按 sequence 排序）"""
        ...

    async def get_step_children(self, step_id: str) -> List[Step]:
        """获取 Step 的子节点"""
        ...

    async def update_step(self, step_id: str, **updates) -> None:
        """
        更新 Step 字段（用于状态变更、错误记录等）

        Args:
            step_id: Step ID
            **updates: 要更新的字段
        """
        ...

    # ===== 事件流操作（用于 WebSocket 断线续传）=====

    async def get_events(
        self,
        trace_id: str,
        since_event_id: int = 0
    ) -> List[Dict[str, Any]]:
        """
        获取事件流（用于 WS 断线续传）

        Args:
            trace_id: Trace ID
            since_event_id: 从哪个事件 ID 开始（0 表示全部）

        Returns:
            事件列表（按 event_id 排序）
        """
        ...

    async def append_event(
        self,
        trace_id: str,
        event_type: str,
        payload: Dict[str, Any]
    ) -> int:
        """
        追加事件，返回 event_id

        Args:
            trace_id: Trace ID
            event_type: 事件类型（step_added/step_updated/trace_completed）
            payload: 事件数据

        Returns:
            event_id: 新事件的 ID
        """
        ...