Multi-Server Orchestration

"""
Prefix-based tool namespacing: db::search, ml::search
Each server gets an alias; the host rewrites tool names on registration.
"""

import asyncio
import json
from contextlib import AsyncExitStack
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from mcp.types import Tool


class MCPHost:
    """
    WHAT: Manages N server connections; rewrites tool names with server prefixes.
    WHY:  Prevents silent collisions when two servers expose the same tool name.
    """

    def __init__(self):
        self._sessions: dict[str, ClientSession] = {}
        self._namespaced_tools: list[dict] = []  # merged list sent to LLM
        self._tool_to_server: dict[str, str] = {}  # namespaced_name → alias
        self._exit_stack = AsyncExitStack()

    async def add_server(
        self,
        alias: str,                # e.g. "db", "ml", "storage"
        command: str,              # e.g. "python"
        args: list[str],           # e.g. ["data_server.py"]
        env: dict | None = None,
    ) -> None:
        """
        WHAT: Connect to a stdio MCP server and register its tools under alias::.
        GOTCHA: alias must be unique; duplicate aliases overwrite the previous entry.
        """
        params = StdioServerParameters(command=command, args=args, env=env)
        read, write = await self._exit_stack.enter_async_context(
            stdio_client(params)
        )
        session = await self._exit_stack.enter_async_context(
            ClientSession(read, write)
        )
        await session.initialize()
        self._sessions[alias] = session

        # Fetch tools and namespace them
        tools_result = await session.list_tools()
        for tool in tools_result.tools:
            namespaced_name = f"{alias}::{tool.name}"
            self._namespaced_tools.append({
                "name": namespaced_name,
                "description": f"[{alias}] {tool.description}",
                "input_schema": tool.inputSchema,
            })
            self._tool_to_server[namespaced_name] = alias

    def get_tools_for_llm(self) -> list[dict]:
        """Return the full namespaced tool list to pass to the Anthropic API."""
        return [
            {
                "name": t["name"],
                "description": t["description"],
                "input_schema": t["input_schema"],
            }
            for t in self._namespaced_tools
        ]

    async def call_tool(self, namespaced_name: str, tool_input: dict) -> str:
        """
        WHAT: Parse alias::tool_name, look up the right session, forward the call.
        WHY:  The LLM always uses the namespaced name; this is the single routing point.
        GOTCHA: If alias is unknown, raise immediately — never silently drop calls.
        """
        if "::" not in namespaced_name:
            raise ValueError(
                f"Tool name '{namespaced_name}' missing namespace prefix (expected alias::tool_name)"
            )
        alias, original_name = namespaced_name.split("::", 1)
        if alias not in self._sessions:
            raise KeyError(f"No server registered with alias '{alias}'")
        session = self._sessions[alias]
        result = await session.call_tool(original_name, tool_input)
        # Extract text content from MCP result
        if result.content:
            return result.content[0].text if hasattr(result.content[0], "text") else str(result.content[0])
        return ""

    async def close(self) -> None:
        await self._exit_stack.aclose()


# ── Usage ────────────────────────────────────────────────────────────────────
async def demo():
    host = MCPHost()
    try:
        await host.add_server("db", "python", ["data_server.py"])
        await host.add_server("ml", "python", ["compute_server.py"])
        await host.add_server("store", "python", ["storage_server.py"])

        tools = host.get_tools_for_llm()
        # Tools are now: db::search_filings, ml::fuzzy_match_score, store::save_result, …
        for t in tools:
            print(f"  {t['name']}: {t['description'][:60]}")

        # Route a call — alias extracted automatically
        result = await host.call_tool("db::search_filings", {"query": "Acme Corp"})
        print("Result:", result)
    finally:
        await host.close()

asyncio.run(demo())

"""
SmartRouter — three configurable routing strategies.
WHAT: Wraps MCPHost and selects server-specific tool subsets based on strategy.
WHY:  Reduces tool list noise, improves routing accuracy, lowers token cost.
"""

import anthropic
from enum import Enum
from typing import Literal


class RoutingStrategy(str, Enum):
    STATIC = "static"          # fastest: pre-built name→server dict
    CAPABILITY = "capability"  # flexible: match tool name in per-server list
    INTENT = "intent"          # smartest: classify intent, show only relevant tools


class SmartRouter:
    def __init__(
        self,
        host: "MCPHost",                             # from Section 3
        strategy: RoutingStrategy = RoutingStrategy.CAPABILITY,
        anthropic_client: anthropic.Anthropic | None = None,
    ):
        self.host = host
        self.strategy = strategy
        self._client = anthropic_client or anthropic.Anthropic()
        # STATIC strategy: pre-built map of tool_name → server_alias
        self._static_map: dict[str, str] = {}
        # Per-server tool lists for CAPABILITY strategy
        self._server_tools: dict[str, list[str]] = {}

    def configure_static(self, tool_to_server: dict[str, str]) -> None:
        """Register explicit tool→server mappings for STATIC strategy."""
        self._static_map = tool_to_server

    def configure_server_tools(self, server_tools: dict[str, list[str]]) -> None:
        """
        Register per-server tool names for CAPABILITY strategy.
        e.g. {"db": ["search_filings", "get_filing_details"], "ml": ["fuzzy_match_score"]}
        """
        self._server_tools = server_tools

    def _route_static(self, tool_name: str) -> str:
        """O(1) dict lookup — fastest possible routing."""
        # Strip namespace prefix if present
        base = tool_name.split("::")[-1] if "::" in tool_name else tool_name
        server = self._static_map.get(base)
        if server is None:
            raise KeyError(f"No static route for tool '{tool_name}'")
        return server

    def _route_capability(self, tool_name: str) -> str:
        """Iterate per-server tool lists; return first match."""
        base = tool_name.split("::")[-1] if "::" in tool_name else tool_name
        for server_alias, tools in self._server_tools.items():
            if base in tools:
                return server_alias
        raise KeyError(f"No capability match for tool '{tool_name}'")

    async def _classify_intent(self, user_message: str) -> list[str]:
        """
        Use a cheap haiku model to classify which servers are relevant.
        Returns a list of server aliases, e.g. ["db", "ml"].
        WHY: Reduces the tool list shown to the primary model from N to a subset.
        """
        server_list = ", ".join(self._server_tools.keys())
        prompt = (
            f"Which of these MCP servers are needed to answer this request: {server_list}?\n"
            f"Request: {user_message}\n"
            "Respond with ONLY a JSON array of server names, e.g. [\"db\", \"ml\"]."
        )
        resp = self._client.messages.create(
            model="claude-haiku-4-5",
            max_tokens=64,
            messages=[{"role": "user", "content": prompt}],
        )
        import json
        text = resp.content[0].text.strip()
        # Safely parse; fall back to all servers on parse error
        try:
            servers = json.loads(text)
            return [s for s in servers if s in self._server_tools]
        except (json.JSONDecodeError, TypeError):
            return list(self._server_tools.keys())

    def get_tools_for_strategy(self, servers: list[str] | None = None) -> list[dict]:
        """
        Return the tool subset for the given servers (intent strategy).
        If servers is None, return full list (static/capability strategies).
        """
        if servers is None:
            return self.host.get_tools_for_llm()
        all_tools = self.host.get_tools_for_llm()
        return [
            t for t in all_tools
            if t["name"].split("::")[0] in servers
        ]

    async def run_agent(
        self,
        user_message: str,
        system_prompt: str = "You are a helpful research agent.",
        max_turns: int = 10,
    ) -> str:
        """
        WHAT: Full agent loop using the configured routing strategy.
        WHY:  Single entry point — caller doesn't need to know which strategy is active.
        """
        # Determine tool list based on strategy
        if self.strategy == RoutingStrategy.INTENT:
            relevant_servers = await self._classify_intent(user_message)
            tools = self.get_tools_for_strategy(relevant_servers)
        else:
            tools = self.get_tools_for_strategy()

        messages = [{"role": "user", "content": user_message}]

        for _ in range(max_turns):
            resp = self._client.messages.create(
                model="claude-sonnet-4-5",
                max_tokens=4096,
                system=system_prompt,
                tools=tools,
                messages=messages,
            )

            if resp.stop_reason == "end_turn":
                # Extract final text
                for block in resp.content:
                    if hasattr(block, "text"):
                        return block.text
                return ""

            if resp.stop_reason != "tool_use":
                break

            # Process all tool calls in this turn
            messages.append({"role": "assistant", "content": resp.content})
            tool_results = []

            for block in resp.content:
                if block.type != "tool_use":
                    continue
                tool_name = block.name  # already namespaced: db::search_filings

                # Route based on strategy
                if self.strategy == RoutingStrategy.STATIC:
                    alias = self._route_static(tool_name)
                    call_name = f"{alias}::{tool_name}" if "::" not in tool_name else tool_name
                elif self.strategy == RoutingStrategy.CAPABILITY:
                    alias = self._route_capability(tool_name)
                    call_name = f"{alias}::{tool_name}" if "::" not in tool_name else tool_name
                else:
                    # INTENT and prefix-based: tool_name is already namespaced
                    call_name = tool_name

                try:
                    result = await self.host.call_tool(call_name, block.input)
                    tool_results.append({
                        "type": "tool_result",
                        "tool_use_id": block.id,
                        "content": result,
                    })
                except Exception as e:
                    tool_results.append({
                        "type": "tool_result",
                        "tool_use_id": block.id,
                        "content": f"Error: {e}",
                        "is_error": True,
                    })

            messages.append({"role": "user", "content": tool_results})

        return "Max turns reached without final answer."


# ── Usage ─────────────────────────────────────────────────────────────────
async def main():
    host = MCPHost()
    await host.add_server("db", "python", ["data_server.py"])
    await host.add_server("ml", "python", ["compute_server.py"])
    await host.add_server("store", "python", ["storage_server.py"])

    router = SmartRouter(host, strategy=RoutingStrategy.INTENT)
    router.configure_server_tools({
        "db": ["search_filings", "get_filing_details"],
        "ml": ["fuzzy_match_score", "calculate_risk_score"],
        "store": ["save_result", "recall_similar"],
    })

    answer = await router.run_agent(
        "Find all filings for Acme Corp and calculate their risk score."
    )
    print(answer)
    await host.close()

asyncio.run(main())

"""
LoadBalancedMCPPool — multiple instances of the same logical server.
WHAT: Round-robin or least-connections across N replicas; auto-failover.
WHY:  Zero-downtime deployments, high-throughput parallelism, HA.
"""

import asyncio
import time
from contextlib import AsyncExitStack
from dataclasses import dataclass, field
from enum import Enum
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client


class LBStrategy(str, Enum):
    ROUND_ROBIN = "round_robin"
    LEAST_CONNECTIONS = "least_connections"


@dataclass
class PoolEntry:
    session: ClientSession
    alias: str
    healthy: bool = True
    in_flight: int = 0
    last_check: float = field(default_factory=time.time)


class LoadBalancedMCPPool:
    """
    Pool of identical MCP server instances for HA and throughput.
    GOTCHA: All instances must expose identical tool schemas. If they diverge
    after a partial upgrade, the LLM may call a tool that doesn't exist on
    the routed instance. Use versioned deployments to prevent this.
    """

    HEALTH_CHECK_INTERVAL = 30  # seconds

    def __init__(self, strategy: LBStrategy = LBStrategy.ROUND_ROBIN):
        self.strategy = strategy
        self._entries: list[PoolEntry] = []
        self._rr_index: int = 0
        self._exit_stack = AsyncExitStack()
        self._health_task: asyncio.Task | None = None

    async def add_instance(
        self, command: str, args: list[str], env: dict | None = None, alias: str | None = None
    ) -> None:
        """Connect a new server instance to the pool."""
        params = StdioServerParameters(command=command, args=args, env=env)
        read, write = await self._exit_stack.enter_async_context(stdio_client(params))
        session = await self._exit_stack.enter_async_context(ClientSession(read, write))
        await session.initialize()
        idx = len(self._entries)
        self._entries.append(PoolEntry(
            session=session,
            alias=alias or f"instance-{idx}",
        ))

    async def start_health_checks(self) -> None:
        """Launch background health-check loop."""
        self._health_task = asyncio.create_task(self._health_loop())

    async def _health_loop(self) -> None:
        """Ping each instance every 30 s; mark unhealthy on failure."""
        while True:
            await asyncio.sleep(self.HEALTH_CHECK_INTERVAL)
            for entry in self._entries:
                try:
                    # list_tools() is a lightweight ping — no side effects
                    await asyncio.wait_for(entry.session.list_tools(), timeout=5.0)
                    entry.healthy = True
                    entry.last_check = time.time()
                except Exception as e:
                    entry.healthy = False
                    print(f"[HealthCheck] {entry.alias} UNHEALTHY: {e}")

    def _pick_round_robin(self) -> PoolEntry:
        """
        WHAT: Advance index, skip unhealthy, wrap around.
        GOTCHA: If all instances are unhealthy, raise immediately — never loop forever.
        """
        healthy = [e for e in self._entries if e.healthy]
        if not healthy:
            raise RuntimeError("All pool instances are unhealthy — no call possible")
        self._rr_index = (self._rr_index + 1) % len(healthy)
        return healthy[self._rr_index % len(healthy)]

    def _pick_least_connections(self) -> PoolEntry:
        """Pick the healthy instance with the fewest in-flight calls."""
        healthy = [e for e in self._entries if e.healthy]
        if not healthy:
            raise RuntimeError("All pool instances are unhealthy")
        return min(healthy, key=lambda e: e.in_flight)

    def _pick(self) -> PoolEntry:
        if self.strategy == LBStrategy.ROUND_ROBIN:
            return self._pick_round_robin()
        return self._pick_least_connections()

    async def call_tool(self, tool_name: str, tool_input: dict) -> str:
        """
        WHAT: Pick an instance, increment in-flight counter, call, decrement.
        WHY:  in_flight tracking enables least-connections strategy.
        GOTCHA: Always decrement in-flight even on exception (use try/finally).
        """
        entry = self._pick()
        entry.in_flight += 1
        try:
            result = await entry.session.call_tool(tool_name, tool_input)
            content = result.content
            if content:
                return content[0].text if hasattr(content[0], "text") else str(content[0])
            return ""
        except Exception:
            # Mark instance unhealthy on tool call failure
            entry.healthy = False
            raise
        finally:
            entry.in_flight = max(0, entry.in_flight - 1)

    async def close(self) -> None:
        if self._health_task:
            self._health_task.cancel()
        await self._exit_stack.aclose()


# ── Usage ─────────────────────────────────────────────────────────────────
async def main():
    pool = LoadBalancedMCPPool(strategy=LBStrategy.LEAST_CONNECTIONS)

    # Start 3 replicas of the same compute server
    for port_suffix in range(3):
        await pool.add_instance(
            "python", ["compute_server.py"],
            alias=f"compute-{port_suffix}",
        )

    await pool.start_health_checks()

    # Calls are automatically distributed across healthy instances
    results = await asyncio.gather(
        pool.call_tool("fuzzy_match_score", {"a": "Acme Inc", "b": "Acme Corp"}),
        pool.call_tool("calculate_risk_score", {"entity_id": "E001"}),
        pool.call_tool("fuzzy_match_score", {"a": "GlobalTech", "b": "Global Technologies"}),
    )
    for r in results:
        print(r)

    await pool.close()

asyncio.run(main())

"""data_server.py — UCC filing search and retrieval."""
import json
from mcp.server.fastmcp import FastMCP

mcp = FastMCP("DataServer")

# Demo data — replace with real DB connection
FILINGS = {
    "F001": {"id": "F001", "debtor": "Acme Corp", "secured_party": "First National Bank",
             "amount": 500_000, "state": "CA", "status": "active"},
    "F002": {"id": "F002", "debtor": "Acme Inc",  "secured_party": "Regional Credit Union",
             "amount": 125_000, "state": "CA", "status": "active"},
    "F003": {"id": "F003", "debtor": "GlobalTech LLC", "secured_party": "Capital Partners",
             "amount": 2_000_000, "state": "NY", "status": "terminated"},
}


@mcp.tool()
def search_filings(query: str, state: str = "", max_results: int = 10) -> str:
    """
    Search UCC filings by debtor name (case-insensitive substring match).
    Returns a JSON list of matching filings with id, debtor, amount, and status.
    """
    q = query.lower()
    results = [
        {"id": f["id"], "debtor": f["debtor"], "amount": f["amount"], "status": f["status"]}
        for f in FILINGS.values()
        if q in f["debtor"].lower() and (not state or f["state"] == state)
    ][:max_results]
    if not results:
        return json.dumps({"results": [], "message": f"No filings found for '{query}'"})
    return json.dumps({"results": results, "total": len(results)})


@mcp.tool()
def get_filing_details(filing_id: str) -> str:
    """Retrieve full details for a single UCC filing by its ID."""
    filing = FILINGS.get(filing_id)
    if not filing:
        return json.dumps({"error": f"Filing '{filing_id}' not found"})
    return json.dumps(filing)


if __name__ == "__main__":
    mcp.run()

"""compute_server.py — fuzzy matching and risk scoring."""
import json
from mcp.server.fastmcp import FastMCP

mcp = FastMCP("ComputeServer")


def _levenshtein(a: str, b: str) -> int:
    """Classic edit distance — no external deps."""
    if len(a) < len(b):
        a, b = b, a
    if not b:
        return len(a)
    prev = list(range(len(b) + 1))
    for i, ca in enumerate(a, 1):
        curr = [i]
        for j, cb in enumerate(b, 1):
            curr.append(min(prev[j] + 1, curr[j - 1] + 1, prev[j - 1] + (ca != cb)))
        prev = curr
    return prev[-1]


@mcp.tool()
def fuzzy_match_score(entity_a: str, entity_b: str) -> str:
    """
    Compute a 0-100 similarity score between two entity name strings.
    Uses normalized Levenshtein distance. Score 100 = identical, 0 = completely different.
    """
    if not entity_a or not entity_b:
        return json.dumps({"error": "Both entity_a and entity_b are required", "score": 0})
    a, b = entity_a.lower().strip(), entity_b.lower().strip()
    max_len = max(len(a), len(b))
    if max_len == 0:
        return json.dumps({"score": 100, "entity_a": entity_a, "entity_b": entity_b})
    dist = _levenshtein(a, b)
    score = round((1 - dist / max_len) * 100, 1)
    return json.dumps({"score": score, "entity_a": entity_a, "entity_b": entity_b, "edit_distance": dist})


@mcp.tool()
def calculate_risk_score(
    amount: float,
    status: str,
    match_score: float = 100.0,
) -> str:
    """
    Calculate a 0-100 lien risk score.
    Factors: filing amount, status (active=high risk, terminated=low), entity match confidence.
    """
    if amount < 0:
        return json.dumps({"error": "amount must be non-negative", "risk_score": None})
    # Normalize amount: 0 = $0, 100 = $5M+
    amount_score = min(amount / 50_000, 100.0)
    status_score = 80.0 if status == "active" else 10.0
    # Lower match_score increases uncertainty, reducing risk confidence
    confidence = match_score / 100.0
    raw = (amount_score * 0.5 + status_score * 0.5) * confidence
    risk = round(min(raw, 100.0), 1)
    return json.dumps({
        "risk_score": risk,
        "factors": {"amount_score": amount_score, "status_score": status_score, "confidence": confidence},
    })


if __name__ == "__main__":
    mcp.run()

"""
storage_server.py — save and recall analysis results.
In-memory for demo. Replace save/recall with ChromaDB for production vector search.
"""
import json
import time
from mcp.server.fastmcp import FastMCP

mcp = FastMCP("StorageServer")

# In-memory store: {result_id: {entity, risk_score, timestamp, ...}}
RESULTS: dict[str, dict] = {}


@mcp.tool()
def save_result(
    entity_name: str,
    risk_score: float,
    filing_ids: list[str],
    notes: str = "",
) -> str:
    """
    Persist an entity risk analysis result.
    Returns the result_id that can be used for later retrieval.
    """
    if not entity_name:
        return json.dumps({"error": "entity_name is required"})
    result_id = f"R{len(RESULTS) + 1:04d}"
    RESULTS[result_id] = {
        "result_id": result_id,
        "entity_name": entity_name,
        "risk_score": risk_score,
        "filing_ids": filing_ids,
        "notes": notes,
        "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
    }
    return json.dumps({"result_id": result_id, "saved": True})


@mcp.tool()
def recall_similar(entity_name: str, min_risk_score: float = 0.0) -> str:
    """
    Retrieve previously saved results for entities similar to entity_name.
    Filters by optional minimum risk score threshold.
    """
    q = entity_name.lower()
    matches = [
        r for r in RESULTS.values()
        if q in r["entity_name"].lower() and r["risk_score"] >= min_risk_score
    ]
    return json.dumps({"matches": matches, "total": len(matches)})


if __name__ == "__main__":
    mcp.run()

"""
orchestrated_agent.py — three-server agent for UCC entity resolution.
WHAT: Connects DataServer + ComputeServer + StorageServer via MCPHost.
      Uses SmartRouter with intent routing to minimize tool list per turn.
WHY:  The LLM sees only the tools it needs; results are persisted automatically.
"""
import asyncio

# Import MCPHost and SmartRouter from Sections 3 and 4
from mcp_host import MCPHost
from smart_router import SmartRouter, RoutingStrategy


async def main():
    host = MCPHost()

    # Connect all three servers
    await host.add_server("db",    "python", ["data_server.py"])
    await host.add_server("ml",    "python", ["compute_server.py"])
    await host.add_server("store", "python", ["storage_server.py"])

    # Configure intent router so the classifier knows which tools live where
    router = SmartRouter(host, strategy=RoutingStrategy.INTENT)
    router.configure_server_tools({
        "db":    ["search_filings", "get_filing_details"],
        "ml":    ["fuzzy_match_score", "calculate_risk_score"],
        "store": ["save_result", "recall_similar"],
    })

    system = """You are an expert UCC lien risk analyst.
When given an entity name:
1. Search for filings using db::search_filings
2. For each match, score similarity with ml::fuzzy_match_score
3. Calculate risk with ml::calculate_risk_score
4. Save all high-risk results (score > 60) with store::save_result
5. Check for prior analyses with store::recall_similar
Provide a structured risk summary at the end."""

    answer = await router.run_agent(
        user_message="Analyze lien risk for 'Acme Corporation' — are there any high-risk filings?",
        system_prompt=system,
        max_turns=15,
    )
    print("=== Agent Answer ===")
    print(answer)

    await host.close()


asyncio.run(main())

"""
TracedMCPHost — distributed tracing for multi-server MCP agents.
WHAT: Injects trace_id into every tool call; collects per-server latency.
WHY:  Identify bottlenecks, reconstruct call graphs, isolate errors by server.
GOTCHA: trace_id must be a string, not a UUID object, for JSON serialization.
"""
import asyncio
import time
import uuid
from dataclasses import dataclass, field
from typing import Optional

# TracedMCPHost extends MCPHost from Section 3
from mcp_host import MCPHost


@dataclass
class SpanRecord:
    trace_id: str
    span_id: str
    server_alias: str
    tool_name: str
    start_time: float
    end_time: Optional[float] = None
    success: bool = True
    error: Optional[str] = None

    @property
    def duration_ms(self) -> float:
        if self.end_time is None:
            return 0.0
        return (self.end_time - self.start_time) * 1000


class TracedMCPHost(MCPHost):
    """
    Subclass of MCPHost with distributed tracing.
    Each agent run gets a unique trace_id; each tool call gets a span_id.
    """

    def __init__(self):
        super().__init__()
        self._spans: list[SpanRecord] = []
        self._current_trace_id: Optional[str] = None

    def start_trace(self, trace_id: Optional[str] = None) -> str:
        """Begin a new agent run trace. Returns the trace_id."""
        self._current_trace_id = trace_id or str(uuid.uuid4())[:8]
        self._spans.clear()
        print(f"[Trace] Starting run trace_id={self._current_trace_id}")
        return self._current_trace_id

    async def call_tool(self, namespaced_name: str, tool_input: dict) -> str:
        """
        Override call_tool to inject trace metadata and record latency.
        WHY: Every tool call is instrumented without changing individual servers.
        """
        if "::" not in namespaced_name:
            raise ValueError(f"Tool name '{namespaced_name}' missing namespace prefix")
        alias = namespaced_name.split("::")[0]
        span_id = str(uuid.uuid4())[:8]
        span = SpanRecord(
            trace_id=self._current_trace_id or "untraced",
            span_id=span_id,
            server_alias=alias,
            tool_name=namespaced_name,
            start_time=time.time(),
        )
        self._spans.append(span)

        # Inject trace metadata — servers that support it can log it
        traced_input = {**tool_input, "_trace_id": span.trace_id, "_span_id": span_id}

        try:
            result = await super().call_tool(namespaced_name, traced_input)
            span.end_time = time.time()
            print(f"  [Span] {namespaced_name} → {span.duration_ms:.1f}ms OK  trace={span.trace_id}")
            return result
        except Exception as e:
            span.end_time = time.time()
            span.success = False
            span.error = str(e)
            print(f"  [Span] {namespaced_name} → {span.duration_ms:.1f}ms ERR {e}  trace={span.trace_id}")
            raise

    def get_trace_summary(self) -> dict:
        """
        Return a full trace summary with per-server latency breakdown.
        WHY: One dict to paste into your log aggregator (Datadog, CloudWatch, etc.)
        """
        if not self._spans:
            return {"trace_id": self._current_trace_id, "spans": [], "total_ms": 0}

        total_ms = sum(s.duration_ms for s in self._spans)
        by_server: dict[str, dict] = {}
        for span in self._spans:
            if span.server_alias not in by_server:
                by_server[span.server_alias] = {"calls": 0, "total_ms": 0.0, "errors": 0}
            by_server[span.server_alias]["calls"] += 1
            by_server[span.server_alias]["total_ms"] += span.duration_ms
            if not span.success:
                by_server[span.server_alias]["errors"] += 1

        return {
            "trace_id": self._current_trace_id,
            "total_calls": len(self._spans),
            "total_ms": round(total_ms, 1),
            "by_server": {
                alias: {
                    "calls": data["calls"],
                    "avg_ms": round(data["total_ms"] / data["calls"], 1),
                    "errors": data["errors"],
                }
                for alias, data in by_server.items()
            },
            "spans": [
                {
                    "span_id": s.span_id,
                    "tool": s.tool_name,
                    "duration_ms": round(s.duration_ms, 1),
                    "success": s.success,
                    "error": s.error,
                }
                for s in self._spans
            ],
        }


# ── Usage ─────────────────────────────────────────────────────────────────
async def main():
    host = TracedMCPHost()
    await host.add_server("db",    "python", ["data_server.py"])
    await host.add_server("ml",    "python", ["compute_server.py"])
    await host.add_server("store", "python", ["storage_server.py"])

    trace_id = host.start_trace()
    print(f"Agent run: trace_id={trace_id}")

    # Run some tool calls (or use SmartRouter.run_agent with host=traced_host)
    await host.call_tool("db::search_filings", {"query": "Acme"})
    await host.call_tool("ml::fuzzy_match_score", {"entity_a": "Acme Corp", "entity_b": "Acme Inc"})
    await host.call_tool("store::save_result", {"entity_name": "Acme Corp", "risk_score": 73.0, "filing_ids": ["F001"]})

    import json
    summary = host.get_trace_summary()
    print("\n=== Trace Summary ===")
    print(json.dumps(summary, indent=2))

    await host.close()


asyncio.run(main())