Agent/cyber_agent/core/agent_mode.py

"""Agent delegation mode policy.

The environment selects a mode only when a root Trace is created.  The chosen
mode is then persisted on the Trace and inherited by every local descendant so
that a running tree cannot change behaviour after a process configuration
change.
"""

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):
    LEGACY = "legacy"
    RECURSIVE = "recursive"


@dataclass(frozen=True)
class AgentPolicy:
    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:
        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:
    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:
    """Resolve the policy for a newly-created root 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:
    """Resolve a persisted Trace policy; missing metadata is 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]:
    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 的执行策略。"""
    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