Agent/cyber_agent/trace/protocols.py

"""
Trace 存储协议。

使用 Protocol 约束 Runner、Sub-Agent、资源预算和 API 依赖的存储能力，
使文件系统或未来的数据库实现可替换而不改上层调用。
"""

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

from .models import Trace, Message
from .goal_models import GoalTree, Goal


@runtime_checkable
class TraceStore(Protocol):
    """Trace、Message、GoalTree 与 Recursive 树级用量的统一存储接口。

    Runner 负责主路径读写，Sub-Agent 用父子过滤查询直接孩子，预算控制器读写根树用量。"""

    # ===== 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,
        parent_trace_id: Optional[str] = None,
        created_by_tool: Optional[str] = None,
    ) -> List[Trace]:
        """按条件列出 Trace，过滤必须在 limit 之前完成。

        Recursive Spawn Guard 用 parent_trace_id + created_by_tool 精确统计直接业务孩子。"""
        ...

    # ===== Recursive 树级资源用量 =====

    async def get_resource_usage(
        self,
        root_trace_id: str,
    ) -> Optional[Dict[str, Any]]:
        """读取 Recursive 根 Trace 的树级累计用量。

        ResourceBudgetController 在预留 Agent/LLM 或登记 Token/成本前后调用。"""
        ...

    async def replace_resource_usage(
        self,
        root_trace_id: str,
        usage: Dict[str, Any],
    ) -> None:
        """原子替换 Recursive 根 Trace 的树级累计用量。

        ResourceBudgetController 在单进程根树锁内调用，存储实现需避免留下半写入快照。"""
        ...

    async def get_candidate_ledger(
        self,
        root_trace_id: str,
    ) -> Optional[Dict[str, Any]]:
        """Read the root Trace's framework-owned candidate ledger."""
        ...

    async def replace_candidate_ledger(
        self,
        root_trace_id: str,
        ledger: Dict[str, Any],
    ) -> None:
        """Atomically replace the root Trace's candidate ledger."""
        ...

    # ===== GoalTree 操作 =====

    async def get_goal_tree(self, trace_id: str) -> Optional[GoalTree]:
        """
        获取 GoalTree

        Args:
            trace_id: Trace ID

        Returns:
            GoalTree 对象，如果不存在返回 None
        """
        ...

    async def update_goal_tree(self, trace_id: str, tree: GoalTree) -> None:
        """
        更新完整 GoalTree

        Args:
            trace_id: Trace ID
            tree: GoalTree 对象
        """
        ...

    async def add_goal(self, trace_id: str, goal: Goal) -> None:
        """
        添加 Goal 到 GoalTree

        Args:
            trace_id: Trace ID
            goal: Goal 对象
        """
        ...

    async def update_goal(self, trace_id: str, goal_id: str, **updates) -> None:
        """
        更新 Goal 字段

        Args:
            trace_id: Trace ID
            goal_id: Goal ID
            **updates: 要更新的字段（如 status, summary, self_stats, cumulative_stats）
        """
        ...

    # ===== Message 操作 =====

    async def add_message(self, message: Message) -> str:
        """
        添加 Message

        自动更新关联 Goal 的 stats（self_stats 和祖先的 cumulative_stats）

        Args:
            message: Message 对象

        Returns:
            message_id
        """
        ...

    async def get_message(self, message_id: str) -> Optional[Message]:
        """获取 Message"""
        ...

    async def get_trace_messages(
        self,
        trace_id: str,
    ) -> List[Message]:
        """
        获取 Trace 的所有 Messages（按 sequence 排序）

        返回该 Trace 下所有消息（包含所有分支）。
        如需获取特定主路径的消息，使用 get_main_path_messages()。

        Args:
            trace_id: Trace ID

        Returns:
            Message 列表
        """
        ...

    async def get_main_path_messages(
        self,
        trace_id: str,
        head_sequence: int
    ) -> List[Message]:
        """
        获取主路径上的消息（从 head_sequence 沿 parent_sequence 链回溯到 root）

        Args:
            trace_id: Trace ID
            head_sequence: 主路径头节点的 sequence

        Returns:
            按 sequence 正序排列的主路径 Message 列表
        """
        ...

    async def get_messages_by_goal(
        self,
        trace_id: str,
        goal_id: str
    ) -> List[Message]:
        """
        获取指定 Goal 关联的所有 Messages

        Args:
            trace_id: Trace ID
            goal_id: Goal ID

        Returns:
            Message 列表
        """
        ...

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

        Args:
            message_id: Message ID
            **updates: 要更新的字段
        """
        ...

    async def abandon_messages_after(self, trace_id: str, cutoff_sequence: int) -> List[str]:
        """
        将 cutoff_sequence 之后的所有 active 消息标记为 abandoned（回溯专用）

        Args:
            trace_id: Trace ID
            cutoff_sequence: 截断点（该 sequence 及之前的消息保留）

        Returns:
            被标记为 abandoned 的 message_id 列表
        """
        ...

    # ===== 事件流操作（用于 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: 事件类型
            payload: 事件数据

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

    async def get_tool_approval_batch(self, trace_id: str):
        """Load the persisted approval batch for a Trace, if any."""
        ...

    async def replace_tool_approval_batch(self, trace_id: str, batch) -> None:
        """Atomically replace a Trace's tool approval batch."""
        ...

    async def mark_reflections_consumed(
        self,
        trace_id: str,
        reflections: List[Dict[str, Any]],
        consumed_at: str,
    ) -> int:
        """Atomically mark selected reflection cognition events consumed."""
        ...

    async def write_message_attachment(
        self,
        trace_id: str,
        message_id: str,
        *,
        suffix: str,
        data: bytes,
    ):
        """Atomically persist a Message-owned binary attachment."""
        ...