orchestrator/agent-team/agent_team/topology.py

"""Pipeline topology for the status dashboard, derived from the real graph.

The map the WebUI draws is **introspected from the compiled LangGraph** rather
than hand-laid: :func:`build_topology` assembles the maximal graph wiring (all
optional P2/P3 nodes injected as lightweight stubs — we want the *shape*, not
the behaviour), calls ``compiled.get_graph()`` for its nodes + edges, and merges
that with :data:`NODE_META`, a display-only sidecar (label / owning agent /
process tree / kind / gated).

Why derive instead of hardcode: adding a new agent node in
:mod:`agent_team.graph` makes it appear on the map automatically — the
"easy to add nodes" goal. ``NODE_META`` supplies only presentation; a graph node
missing from it still renders (raw id, ``unassigned`` tree) so a new agent is
never silently dropped, and :func:`missing_meta` lets a test fail loudly until
its metadata is filled in.

``tree`` groups nodes into processes branching off the ``intake`` coordinator
(today: the ``core`` root + the ``sdlc`` pipeline); new processes are new
``tree`` values across their nodes' meta entries.

This module is import-light and has no live model / DB dependency: the stub
nodes/routes are never executed (``get_graph()`` reads the wiring statically), so
building the topology is a pure, cheap, deterministic operation.
"""

from __future__ import annotations

from collections import deque
from typing import Any

from agent_team import graph as graph_mod
from agent_team.task_model import Phase

__all__ = [
    "NODE_META",
    "PHASE_TO_NODE",
    "TREES",
    "build_topology",
    "missing_meta",
    "node_for_phase",
]

# LangGraph's synthetic terminal vertices — excluded from the drawn map.
_START = "__start__"
_END = "__end__"

# Display sidecar: graph node id -> presentation. ``kind`` is "phase" for work
# stages and "gate" for the clarifier (which holds the human interrupt). ``tree``
# groups nodes into processes branching off intake. ``gated`` marks nodes that
# are drawn but inert in the default production deploy (the P3 build->verify->
# dispatch path), so the UI can dim them. Keep ids in sync with agent_team.graph.
NODE_META: dict[str, dict[str, Any]] = {
    graph_mod.INTAKE: {
        "label": "Intake",
        "agent": "coordinator",
        "tree": "core",
        "kind": "phase",
        "gated": False,
    },
    graph_mod.CLARIFY: {
        "label": "Clarify",
        "agent": "Claude (sub)",
        "tree": "sdlc",
        "kind": "gate",  # the human gate (LangGraph interrupt) lives here
        "gated": False,
    },
    graph_mod.PLAN: {
        "label": "Plan",
        "agent": "Claude (sub)",
        "tree": "sdlc",
        "kind": "phase",
        "gated": False,
    },
    graph_mod.REVIEW: {
        "label": "Review",
        "agent": "GPT-4.1 (cross)",
        "tree": "sdlc",
        "kind": "phase",
        "gated": False,
    },
    graph_mod.BUILD_NODE: {
        "label": "Build",
        "agent": "DeepSeek (fast)",
        "tree": "sdlc",
        "kind": "phase",
        "gated": True,  # P3, opt-in/inert in the default deploy
    },
    graph_mod.VERIFY_NODE: {
        "label": "Verify",
        "agent": "Claude (sub)",
        "tree": "sdlc",
        "kind": "phase",
        "gated": True,
    },
    graph_mod.DISPATCH_NODE: {
        "label": "Dispatch",
        "agent": "GitHub PR",
        "tree": "sdlc",
        "kind": "phase",
        "gated": True,
    },
}

# Process trees, in display order. ``root`` flags the tree that owns intake (the
# branch point); other trees hang off it. New processes append here.
TREES: tuple[dict[str, Any], ...] = (
    {"id": "core", "label": "Core", "root": True},
    {"id": "sdlc", "label": "SDLC Pipeline", "root": False},
)

# Default tree/meta for a graph node with no NODE_META entry (a newly added
# agent whose metadata has not been filled in yet) — rendered, never dropped.
_UNASSIGNED_META: dict[str, Any] = {
    "label": "",  # filled with the raw id at build time
    "agent": "",
    "tree": "unassigned",
    "kind": "phase",
    "gated": False,
}

# Phase value (TaskStatus current_phase) -> graph node id, so /api/state can
# bucket a live task onto its node. BUILD/VERIFY phases map to the P3 vertex ids
# (build_node/verify_node). DONE/PARKED are terminal/exception states with no
# vertex — a task in them is shown in the list, not on a node.
PHASE_TO_NODE: dict[str, str] = {
    Phase.INTAKE.value: graph_mod.INTAKE,
    Phase.CLARIFY.value: graph_mod.CLARIFY,
    Phase.PLAN.value: graph_mod.PLAN,
    Phase.REVIEW.value: graph_mod.REVIEW,
    Phase.BUILD.value: graph_mod.BUILD_NODE,
    Phase.VERIFY.value: graph_mod.VERIFY_NODE,
}


def node_for_phase(phase: str | None) -> str | None:
    """Return the graph node id a live ``current_phase`` belongs to, or ``None``."""
    if not phase:
        return None
    return PHASE_TO_NODE.get(phase)


def _stub_node(state: Any) -> dict[str, Any]:
    """A no-op node used only to assemble the maximal graph shape (never run)."""
    return {}


def _stub_route(state: Any) -> str:
    """A no-op router; never executed — get_graph() reads the edge map statically."""
    return graph_mod.APPROVED_ROUTE


def _maximal_compiled() -> Any:
    """Compile the full P3+ wiring with stub nodes, for shape introspection."""
    return graph_mod.build_graph(
        None,
        live_clarify_node=_stub_node,
        live_plan_node=_stub_node,
        review_node=_stub_node,
        route_review=_stub_route,
        build_verify=(_stub_node, _stub_node, _stub_route),
        dispatch_node=_stub_node,
    )


def _bfs_order(node_ids: list[str], edges: list[tuple[str, str]]) -> dict[str, int]:
    """Assign each node a forward rank via BFS from ``__start__``.

    Used to classify edges: an edge whose target ranks at or before its source
    goes backward (a retry loop). Nodes unreachable in the BFS get a large rank
    so they never spuriously read as loop targets.
    """
    adj: dict[str, list[str]] = {}
    for src, dst in edges:
        adj.setdefault(src, []).append(dst)
    order: dict[str, int] = {}
    queue: deque[str] = deque([_START])
    order[_START] = 0
    while queue:
        node = queue.popleft()
        for nxt in adj.get(node, []):
            if nxt not in order:
                order[nxt] = order[node] + 1
                queue.append(nxt)
    big = len(node_ids) + len(edges) + 1
    for nid in node_ids:
        order.setdefault(nid, big)
    return order


def _edge_kind(src: str, dst: str, conditional: bool, order: dict[str, int]) -> str:
    """Classify an edge as spine | branch | loopback for styling."""
    if conditional and order.get(dst, 0) <= order.get(src, 0):
        return "loopback"
    if conditional:
        return "branch"
    return "spine"


def missing_meta() -> list[str]:
    """Return graph node ids (excluding start/end) that lack a NODE_META entry.

    A test asserts this is empty so a newly added agent fails loudly until its
    display metadata is filled in (the node still renders meanwhile).
    """
    compiled = _maximal_compiled()
    drawable = compiled.get_graph()
    ids = [n for n in drawable.nodes if n not in (_START, _END)]
    return [nid for nid in ids if nid not in NODE_META]


def build_topology() -> dict[str, Any]:
    """Return the dashboard topology: ``{trees, nodes, edges}``.

    Nodes/edges are derived from the compiled LangGraph (so new graph nodes
    appear automatically) and enriched with :data:`NODE_META`. Edges are
    de-duplicated and classified spine/branch/loopback. The synthetic
    ``__start__``/``__end__`` vertices are dropped; an edge touching them is
    dropped too (the map shows agent nodes, not the framework terminals).
    """
    compiled = _maximal_compiled()
    drawable = compiled.get_graph()

    raw_edges = [(e.source, e.target, bool(e.conditional)) for e in drawable.edges]
    order = _bfs_order(list(drawable.nodes), [(s, t) for s, t, _ in raw_edges])

    node_ids = [n for n in drawable.nodes if n not in (_START, _END)]
    nodes: list[dict[str, Any]] = []
    for nid in node_ids:
        meta = NODE_META.get(nid)
        if meta is None:
            meta = {**_UNASSIGNED_META, "label": nid}
        nodes.append(
            {
                "id": nid,
                "label": meta["label"] or nid,
                "agent": meta["agent"],
                "tree": meta["tree"],
                "kind": meta["kind"],
                "gated": bool(meta["gated"]),
            }
        )

    seen: set[tuple[str, str]] = set()
    edges: list[dict[str, Any]] = []
    for src, dst, conditional in raw_edges:
        if src in (_START, _END) or dst in (_START, _END):
            continue
        key = (src, dst)
        if key in seen:
            continue
        seen.add(key)
        edges.append(
            {"from": src, "to": dst, "kind": _edge_kind(src, dst, conditional, order)}
        )

    return {
        "trees": [dict(t) for t in TREES],
        "nodes": nodes,
        "edges": edges,
    }