deer-flow/backend/packages/harness/deerflow/subagents/registry.py

"""Subagent registry for managing available subagents."""

import logging
import threading
import time
from collections.abc import Hashable
from dataclasses import replace
from typing import Any

from deerflow.persistence.managed_subagents import ManagedSubagentDefinition, get_managed_subagent_store
from deerflow.sandbox.security import is_host_bash_allowed
from deerflow.subagents.builtins import BUILTIN_SUBAGENTS
from deerflow.subagents.config import SubagentConfig

logger = logging.getLogger(__name__)
_MANAGED_SIGNATURE_TTL_SECONDS = 1.0
_managed_definitions_cache_lock = threading.RLock()
_managed_definitions_cache: dict[Hashable, tuple[float, Hashable, tuple[ManagedSubagentDefinition, ...]]] = {}


def _resolve_subagents_app_config(app_config: Any | None = None):
    if app_config is None:
        from deerflow.config.subagents_config import get_subagents_app_config

        return get_subagents_app_config()
    return getattr(app_config, "subagents", app_config)


def _build_custom_subagent_config(name: str, *, app_config: Any | None = None) -> SubagentConfig | None:
    """Build a SubagentConfig from config.yaml custom_agents section.

    Args:
        name: The name of the custom subagent.
        app_config: Optional AppConfig or SubagentsAppConfig to resolve from.

    Returns:
        SubagentConfig if found in custom_agents, None otherwise.
    """
    subagents_config = _resolve_subagents_app_config(app_config)
    custom = subagents_config.custom_agents.get(name)
    if custom is None:
        return None

    return SubagentConfig(
        name=name,
        description=custom.description,
        system_prompt=custom.system_prompt,
        tools=custom.tools,
        disallowed_tools=custom.disallowed_tools,
        skills=custom.skills,
        model=custom.model,
        max_turns=custom.max_turns,
        timeout_seconds=custom.timeout_seconds,
    )


def _clear_managed_definitions_cache() -> None:
    """Clear process-local registry snapshots (primarily for tests)."""
    with _managed_definitions_cache_lock:
        _managed_definitions_cache.clear()


def _managed_definitions(*, app_config: Any | None = None) -> tuple[ManagedSubagentDefinition, ...]:
    """Load and cache deployment-managed definitions until their signature changes."""
    store_config = app_config if hasattr(app_config, "agent_storage") else None
    store = get_managed_subagent_store(store_config)
    cache_key = store.cache_identity()

    with _managed_definitions_cache_lock:
        checked_at = time.monotonic()
        cached = _managed_definitions_cache.get(cache_key)
        # A prompt/catalog pass can resolve every managed name separately.
        # Avoid repeating the file stat sweep or SQL signature query for each
        # lookup while keeping cross-process changes visible within one second.
        if cached is not None and checked_at - cached[0] < _MANAGED_SIGNATURE_TTL_SECONDS:
            return cached[2]

        signature = store.signature()
        if cached is not None and cached[1] == signature:
            _managed_definitions_cache[cache_key] = (checked_at, signature, cached[2])
            return cached[2]

        definitions = tuple(store.list())
        _managed_definitions_cache[cache_key] = (checked_at, signature, definitions)
        return definitions


def _build_managed_subagent_config(name: str, *, app_config: Any | None = None) -> SubagentConfig | None:
    for definition in _managed_definitions(app_config=app_config):
        if definition.name != name or not definition.enabled:
            continue
        return SubagentConfig(
            name=definition.name,
            description=definition.description,
            system_prompt=definition.system_prompt,
            tools=definition.tools,
            disallowed_tools=definition.disallowed_tools,
            skills=definition.skills,
            model=definition.model,
            max_turns=definition.max_turns,
            timeout_seconds=definition.timeout_seconds,
        )
    return None


def get_subagent_config(name: str, *, app_config: Any | None = None) -> SubagentConfig | None:
    """Get a subagent configuration by name, with config.yaml overrides applied.

    Resolution order (mirrors Codex's config layering):
    1. Built-in subagents (general-purpose, bash)
    2. Custom subagents from config.yaml custom_agents section
    3. Enabled administrator-managed subagents
    4. Per-agent overrides from config.yaml agents section (timeout, max_turns, model, skills)

    Args:
        name: The name of the subagent.
        app_config: Optional AppConfig or SubagentsAppConfig to resolve overrides from.

    Returns:
        SubagentConfig if found (with any config.yaml overrides applied), None otherwise.
    """
    # Step 1: Look up built-in, then fall back to custom_agents
    config = BUILTIN_SUBAGENTS.get(name)
    if config is None:
        config = _build_custom_subagent_config(name, app_config=app_config)
    if config is None:
        config = _build_managed_subagent_config(name, app_config=app_config)
    if config is None:
        return None

    # Step 2: Apply per-agent overrides from config.yaml agents section.
    # Only explicit per-agent overrides are applied here. Global defaults
    # (timeout_seconds, max_turns at the top level) apply to built-in agents
    # but must NOT override custom agents' own values — custom agents define
    # their own defaults in the custom_agents section.
    subagents_config = _resolve_subagents_app_config(app_config)
    is_builtin = name in BUILTIN_SUBAGENTS
    agent_override = subagents_config.agents.get(name)

    overrides = {}

    # Timeout: per-agent override > global default (builtins only) > config's own value
    if agent_override is not None and agent_override.timeout_seconds is not None:
        if agent_override.timeout_seconds != config.timeout_seconds:
            logger.debug("Subagent '%s': timeout overridden (%ss -> %ss)", name, config.timeout_seconds, agent_override.timeout_seconds)
            overrides["timeout_seconds"] = agent_override.timeout_seconds
    elif is_builtin and subagents_config.timeout_seconds != config.timeout_seconds:
        logger.debug("Subagent '%s': timeout from global default (%ss -> %ss)", name, config.timeout_seconds, subagents_config.timeout_seconds)
        overrides["timeout_seconds"] = subagents_config.timeout_seconds

    # Max turns: per-agent override > global default (builtins only) > config's own value
    if agent_override is not None and agent_override.max_turns is not None:
        if agent_override.max_turns != config.max_turns:
            logger.debug("Subagent '%s': max_turns overridden (%s -> %s)", name, config.max_turns, agent_override.max_turns)
            overrides["max_turns"] = agent_override.max_turns
    elif is_builtin and subagents_config.max_turns is not None and subagents_config.max_turns != config.max_turns:
        logger.debug("Subagent '%s': max_turns from global default (%s -> %s)", name, config.max_turns, subagents_config.max_turns)
        overrides["max_turns"] = subagents_config.max_turns

    # Model: per-agent override only (no global default for model)
    effective_model = subagents_config.get_model_for(name)
    if effective_model is not None and effective_model != config.model:
        logger.debug("Subagent '%s': model overridden (%s -> %s)", name, config.model, effective_model)
        overrides["model"] = effective_model

    # Skills: per-agent override only (no global default for skills)
    effective_skills = subagents_config.get_skills_for(name)
    if effective_skills is not None and effective_skills != config.skills:
        logger.debug("Subagent '%s': skills overridden (%s -> %s)", name, config.skills, effective_skills)
        overrides["skills"] = effective_skills

    if overrides:
        config = replace(config, **overrides)

    return config


def list_subagents(*, app_config: Any | None = None, allowed_subagents: list[str] | None = None) -> list[SubagentConfig]:
    """List all available subagent configurations (with config.yaml overrides applied).

    Returns:
        List of all registered SubagentConfig instances (built-in + custom).
    """
    configs = []
    for name in get_subagent_names(app_config=app_config, allowed_subagents=allowed_subagents):
        config = get_subagent_config(name, app_config=app_config)
        if config is not None:
            configs.append(config)
    return configs


def get_subagent_names(*, app_config: Any | None = None, allowed_subagents: list[str] | None = None) -> list[str]:
    """Get registered subagent names, optionally restricted by the caller policy.

    Returns:
        List of subagent names.
    """
    names = list(BUILTIN_SUBAGENTS.keys())

    # Merge custom_agents from config.yaml
    subagents_config = _resolve_subagents_app_config(app_config)
    for custom_name in subagents_config.custom_agents:
        if custom_name not in names:
            names.append(custom_name)

    # Built-in and config.yaml definitions have operator-controlled precedence.
    # A managed definition that later conflicts remains persisted for the
    # Settings UI, but is excluded from runtime discovery.
    for definition in _managed_definitions(app_config=app_config):
        if not definition.enabled:
            continue
        if definition.name in names:
            logger.debug("Managed subagent '%s' conflicts with a built-in or config.yaml definition and is excluded from runtime", definition.name)
            continue
        names.append(definition.name)

    if allowed_subagents is not None:
        allowed = set(allowed_subagents)
        names = [name for name in names if name in allowed]

    return names


def get_available_subagent_names(*, app_config: Any | None = None, allowed_subagents: list[str] | None = None) -> list[str]:
    """Get subagent names that should be exposed to the active runtime.

    Returns:
        List of subagent names visible to the current sandbox configuration.
    """
    names = get_subagent_names(app_config=app_config, allowed_subagents=allowed_subagents)
    try:
        host_bash_allowed = is_host_bash_allowed(app_config) if hasattr(app_config, "sandbox") else is_host_bash_allowed()
    except Exception:
        logger.debug("Could not determine host bash availability; exposing all subagents")
        return names

    if not host_bash_allowed:
        names = [name for name in names if name != "bash"]
    return names