Agent/cyber_agent/core/agent_mode.py

"""Agent 委托模式及其统一运行策略。

Runner 仅在创建根 Trace 时读取环境模式，随后将策略持久化并传给所有本地子孙。
Runner、Sub-Agent 工具和 Goal 工具共同使用该策略，保证运行中的任务树不被环境变更改变。
"""

from __future__ import annotations

from dataclasses import dataclass
from enum import Enum
import os
from typing import Any, Mapping, MutableMapping


AGENT_MODE_ENV = "AGENT_MODE"
REMOVED_RECURSION_ENV = "AGENT_RECURSION_ENABLED"
AGENT_MODE_CONTEXT_KEY = "agent_mode"
AGENT_MODE_REVISION_CONTEXT_KEY = "agent_mode_revision"

# Runner 注入给本地 agent 工具的内部上下文。它们不来自 LLM 参数，也不持久化
# 到 Trace；Recursive 子级只能继承调用方在本次运行中实际拥有的能力与调度策略。
RECURSIVE_CAPABILITY_TOOLS_CONTEXT_KEY = "_recursive_capability_tool_names"
RECURSIVE_CHILD_EXECUTION_MODE_CONTEXT_KEY = "_recursive_child_execution_mode"
RECURSIVE_MAX_PARALLEL_CHILDREN_CONTEXT_KEY = "_recursive_max_parallel_children"
MAX_PARALLEL_RECURSIVE_CHILDREN = 2


class AgentMode(str, Enum):
    """Agent 对外暴露的两种运行模式。

    `legacy` 保留旧的一层委托，`recursive` 启用递归任务协议与资源边界。
    """

    LEGACY = "legacy"
    RECURSIVE = "recursive"


@dataclass(frozen=True)
class AgentPolicy:
    """某棵 Trace 树不可变的模式策略快照。

    Runner 将其写入根 Trace，Sub-Agent 创建子 Trace 时继承，各运行门禁据此判断能力。
    """

    mode: AgentMode
    revision: int
    max_depth: int
    max_children_per_parent: int | None
    accumulate_sub_trace_ids: bool
    show_sub_trace_shortcuts: bool

    @property
    def requires_task_protocol(self) -> bool:
        """Recursive revision 2 是否必须走结构化任务报告和审核流程。"""
        return self.mode is AgentMode.RECURSIVE and self.revision >= 2

    def to_context(self) -> dict[str, Any]:
        return {
            AGENT_MODE_CONTEXT_KEY: self.mode.value,
            AGENT_MODE_REVISION_CONTEXT_KEY: self.revision,
        }


def _policy(mode: AgentMode, revision: int = 1) -> AgentPolicy:
    if mode is AgentMode.LEGACY:
        return AgentPolicy(
            mode=mode,
            revision=1,
            max_depth=1,
            max_children_per_parent=None,
            accumulate_sub_trace_ids=False,
            show_sub_trace_shortcuts=False,
        )
    return AgentPolicy(
        mode=mode,
        revision=revision,
        max_depth=5,
        max_children_per_parent=6,
        accumulate_sub_trace_ids=True,
        show_sub_trace_shortcuts=True,
    )


def assert_removed_config_absent() -> None:
    """拒绝已删除的递归开关，避免运行时继续使用含糊配置。

    Runner 准备 Trace 和 ``agent`` 工具入口都会调用，新根 Trace 解析模式时也会间接调用。
    """
    if REMOVED_RECURSION_ENV in os.environ:
        raise ValueError(
            f"{REMOVED_RECURSION_ENV} has been removed; "
            f"set {AGENT_MODE_ENV}=legacy or {AGENT_MODE_ENV}=recursive instead"
        )


def policy_from_environment(*, recursive_revision: int = 1) -> AgentPolicy:
    """为新建根 Trace 解析一次环境模式。

    `AgentRunner._prepare_new_trace` 在任何 Trace 落库前调用，返回的策略随后被持久化。
    """
    assert_removed_config_absent()
    raw_mode = os.getenv(AGENT_MODE_ENV, AgentMode.LEGACY.value).strip().lower()
    try:
        mode = AgentMode(raw_mode)
    except ValueError as exc:
        allowed = ", ".join(mode.value for mode in AgentMode)
        raise ValueError(
            f"Invalid {AGENT_MODE_ENV}={raw_mode!r}; expected one of: {allowed}"
        ) from exc
    revision = recursive_revision if mode is AgentMode.RECURSIVE else 1
    return _policy(mode, revision)


def policy_from_context(context: Mapping[str, Any] | None) -> AgentPolicy:
    """从 Trace context 恢复已持久化的模式策略。

    Runner、Sub-Agent 和 Goal 工具在恢复或执行时调用；无模式字段的历史 Trace 按 Legacy 处理。
    """
    context = context or {}
    raw_mode = context.get(AGENT_MODE_CONTEXT_KEY, AgentMode.LEGACY.value)
    try:
        mode = AgentMode(str(raw_mode).strip().lower())
    except ValueError as exc:
        raise ValueError(f"Invalid persisted agent mode: {raw_mode!r}") from exc

    raw_revision = context.get(AGENT_MODE_REVISION_CONTEXT_KEY, 1)
    if not isinstance(raw_revision, int) or raw_revision < 1:
        raise ValueError(f"Invalid persisted agent mode revision: {raw_revision!r}")
    return _policy(mode, raw_revision)


def apply_policy_to_context(
    context: MutableMapping[str, Any] | None,
    policy: AgentPolicy,
) -> dict[str, Any]:
    """将模式和 revision 覆盖写入一份 Trace context 副本。

    Runner 创建根 Trace、Sub-Agent 创建或续跑子 Trace 时调用，保证子孙继承同一策略。
    """
    merged = dict(context or {})
    merged.update(policy.to_context())
    return merged


def validate_recursive_child_execution(
    mode: Any,
    max_parallel_children: Any,
) -> tuple[str, int]:
    """校验单批 Recursive 子 Agent 的串并行策略和硬上限。

    Runner 创建/恢复根 Trace 以及 `agent()` 调度子任务前调用，防止模型或 context 提高并发。
    """
    if mode not in {"sequential", "parallel"}:
        raise ValueError(
            "child_execution_mode must be 'sequential' or 'parallel'"
        )
    if (
        isinstance(max_parallel_children, bool)
        or not isinstance(max_parallel_children, int)
        or not 1 <= max_parallel_children <= MAX_PARALLEL_RECURSIVE_CHILDREN
    ):
        raise ValueError(
            "max_parallel_children must be an integer between 1 and "
            f"{MAX_PARALLEL_RECURSIVE_CHILDREN}"
        )
    return mode, max_parallel_children