MCP Integration

MCP Integration

Overview

The Model Context Protocol (MCP) integration enables agents to access external tools and services during prototype generation. MCP handlers connect to external servers that expose tools via the MCP protocol, and those tools become available to agents during their execution.

The integration follows a handler-based plugin pattern: each handler owns its transport, authentication, and protocol logic. The extension provides lifecycle management, tool routing, scoping, and a circuit breaker for reliability.

Architecture

The MCP system has three layers:

MCPHandler (Abstract Base Class)

Defined in mcp/base.py. Each handler implements:

• connect() -- establish connection to the MCP server
• disconnect() -- clean shutdown
• list_tools() -- return available tool definitions
• call_tool(name, arguments) -- invoke a tool and return results

Handlers are fully responsible for transport (HTTP, stdio, WebSocket), authentication, protocol handshake, error handling, retries, and timeouts.

MCPRegistry

Defined in mcp/registry.py. Follows the same builtin/custom resolution pattern as the AgentRegistry:

1. Custom handlers (loaded from .prototype/mcp/ directory)
2. Built-in handlers (shipped with the extension)

Custom handlers override built-in handlers of the same name.

MCPManager

Defined in mcp/manager.py. The primary interface for the rest of the extension. Provides:

• Lazy connection: handlers connect on first tool access, not at startup
• Tool routing: maps tool names to handlers and dispatches calls
• Circuit breaker: marks handlers as failed after consecutive errors
• OpenAI schema conversion: formats tools for AI provider consumption
• Context manager: clean shutdown when the session ends
• Thread safety: safe for concurrent tool calls from parallel agent execution

Configuration

MCP servers are configured in prototype.yaml under mcp.servers. Sensitive fields (API keys, tokens) route to prototype.secrets.yaml automatically via the mcp.servers prefix in SECRET_KEY_PREFIXES.

mcp:
  servers:
    - name: lightpanda
      enabled: true
      stages: ["build", "deploy"]   # null = all stages
      agents: ["terraform-agent"]   # null = all agents
      timeout: 30
      max_retries: 2
      max_result_bytes: 8192
      settings:
        base_url: "http://localhost:8080"
        api_key: "..."   # routed to secrets file

Handler Configuration Options

The MCPHandlerConfig dataclass defines the configuration fields:

Field	Type	Default	Description
name	string	(required)	Unique handler name
stages	list of strings or null	null (all)	Stages where this handler's tools are available
agents	list of strings or null	null (all)	Agents that can use this handler's tools
enabled	bool	true	Whether the handler is active
timeout	int	30	Seconds per tool call
max_retries	int	2	Maximum retry attempts on failure
max_result_bytes	int	8192	Truncate results exceeding this size
settings	dict	{}	Handler-specific settings (URLs, credentials, etc.)

Custom Handlers

Create custom MCP handlers by placing Python files in the .prototype/mcp/ directory within your project.

Naming Convention

The filename determines the handler name: lightpanda_handler.py registers as lightpanda (the _handler suffix is stripped automatically). Alternatively, define MCP_HANDLER_CLASS in the module to specify the handler class explicitly; otherwise the loader auto-discovers the first MCPHandler subclass.

Implementation

A custom handler extends MCPHandler and implements four abstract methods: connect(), disconnect(), list_tools(), and call_tool(). The handler is fully responsible for transport, authentication, error handling, and retries.

Minimal example -- a handler that wraps a REST API:

# .prototype/mcp/cost_lookup_handler.py
import requests

from azext_prototype.mcp.base import (
    MCPHandler,
    MCPHandlerConfig,
    MCPToolDefinition,
    MCPToolResult,
)


class CostLookupHandler(MCPHandler):
    """Handler that looks up Azure retail prices."""

    name = "cost-lookup"
    description = "Azure Retail Prices API for real-time cost data"

    def __init__(self, config: MCPHandlerConfig, **kwargs):
        super().__init__(config, **kwargs)
        self._session: requests.Session | None = None

    def connect(self) -> None:
        self._session = requests.Session()
        self._session.headers["Accept"] = "application/json"
        # No auth required for Azure Retail Prices API
        self._connected = True
        self.logger.info("Connected to Azure Retail Prices API")

    def disconnect(self) -> None:
        if self._session:
            self._session.close()
            self._session = None
        self._connected = False

    def list_tools(self) -> list[MCPToolDefinition]:
        return [
            MCPToolDefinition(
                name="lookup_price",
                description="Look up Azure retail price for a service and SKU",
                input_schema={
                    "type": "object",
                    "properties": {
                        "service_name": {
                            "type": "string",
                            "description": "Azure service name (e.g., 'Virtual Machines')",
                        },
                        "sku_name": {
                            "type": "string",
                            "description": "SKU name (e.g., 'Standard_B2s')",
                        },
                        "region": {
                            "type": "string",
                            "description": "Azure region (e.g., 'eastus')",
                        },
                    },
                    "required": ["service_name"],
                },
                handler_name=self.name,
            )
        ]

    def call_tool(self, name: str, arguments: dict) -> MCPToolResult:
        if name != "lookup_price":
            return MCPToolResult(
                content="", is_error=True,
                error_message=f"Unknown tool: {name}",
            )

        try:
            service = arguments["service_name"]
            sku = arguments.get("sku_name", "")
            region = arguments.get("region", "eastus")

            filter_parts = [
                f"serviceName eq '{service}'",
                f"armRegionName eq '{region}'",
            ]
            if sku:
                filter_parts.append(f"skuName eq '{sku}'")

            resp = self._session.get(
                "https://prices.azure.com/api/retail/prices",
                params={"$filter": " and ".join(filter_parts)},
                timeout=self.config.timeout,
            )
            resp.raise_for_status()
            data = resp.json()

            items = data.get("Items", [])[:5]  # Top 5 results
            if not items:
                return MCPToolResult(content="No pricing data found.")

            lines = []
            for item in items:
                lines.append(
                    f"- {item['productName']} / {item['skuName']}: "
                    f"${item['retailPrice']:.4f} {item['unitOfMeasure']} "
                    f"({item['type']})"
                )

            content = "\n".join(lines)
            if len(content) > self.config.max_result_bytes:
                content = content[: self.config.max_result_bytes] + "...(truncated)"

            return MCPToolResult(content=content)

        except requests.Timeout:
            return MCPToolResult(
                content="", is_error=True,
                error_message=f"Timeout after {self.config.timeout}s",
            )
        except Exception as exc:
            return MCPToolResult(
                content="", is_error=True,
                error_message=f"Failed: {exc}",
            )


# Tell the loader which class to instantiate
MCP_HANDLER_CLASS = CostLookupHandler

Configure it in prototype.yaml:

mcp:
  servers:
    - name: cost-lookup
      enabled: true
      stages: ["build", "design"]
      agents: ["cost-analyst", "cloud-architect"]
      timeout: 15
      max_retries: 1
      max_result_bytes: 4096
      settings: {}

Full JSON-RPC Example (Lightpanda)

For handlers that communicate with MCP servers using the JSON-RPC protocol, see the reference implementation at mcp/examples/lightpanda_handler.py. It demonstrates:

• MCP initialize handshake with protocolVersion and clientInfo
• Tool discovery via tools/list JSON-RPC call
• Tool invocation via tools/call with retry logic
• Session URL management for stateful MCP servers
• Proper error handling and result truncation

Configuration for JSON-RPC handlers:

# prototype.yaml
mcp:
  servers:
    - name: lightpanda
      stages: ["build", "deploy"]
      agents: ["qa-engineer", "app-developer"]
      timeout: 30
      max_retries: 2
      settings:
        url: "https://mcp.pipedream.net/v2"

# prototype.secrets.yaml (sensitive settings auto-routed here)
mcp:
  servers:
    - name: lightpanda
      settings:
        api_key: "lpd_xxxxxxxxxxxx"

Handler Base Class

The MCPHandler base class provides these built-in attributes (available via self):

Attribute	Description
self.config	MCPHandlerConfig with all settings from prototype.yaml
self.client_info	MCPClientInfo for MCP initialize handshake
self.console	Console object for user-facing messages
self.logger	Python logger for debug/info/warning output
self.project_config	Full project config dict (read-only)
self._connected	Connection state (set to True in connect())

Override health_check() -> bool for custom health validation beyond connection state.

Scoping

Tools are scoped along two dimensions:

Per-Stage Scoping

Set stages: ["build", "deploy"] to restrict a handler's tools to specific stages. Set stages: null (or omit the field) to make tools available in all stages.

Per-Agent Scoping

Set agents: ["terraform-agent", "bicep-agent"] to restrict a handler's tools to specific agents. Set agents: null (or omit the field) to make tools available to all agents.

Both dimensions are AND-combined: a tool is available only when both the current stage and the requesting agent match the handler's scope configuration.

Circuit Breaker

The MCPManager tracks consecutive errors per handler. After 3 consecutive errors, the handler is marked as failed and its tools are no longer offered to agents for the remainder of the session. A warning is displayed when the circuit breaker trips.

This prevents a misbehaving MCP server from degrading the entire session with repeated timeouts or errors.

Two Invocation Modes

AI-Driven (Tool Call Loop)

The default mode. When an agent's AI response includes tool calls, BaseAgent.execute() enters a loop:

1. Get scoped tools from MCPManager, formatted as OpenAI function-calling schema
2. Pass tools to the AI provider alongside the conversation
3. Detect tool_calls in the AI response
4. Invoke each tool via MCPManager.call_tool()
5. Feed tool results back to the AI as role="tool" messages
6. Repeat until the AI produces a final text response (up to _max_tool_iterations = 10)

Agents opt in via _enable_mcp_tools = True (the default for all agents).

Code-Driven (Proactive)

Stages or other code can invoke tools directly via manager.call_tool(tool_name, arguments) without going through the AI loop. This is useful for predetermined tool calls where the stage knows exactly what tool to invoke and with what arguments.

Data Model

Class	Purpose
MCPClientInfo	Client identity for the MCP initialize handshake
MCPHandlerConfig	Configuration from prototype.yaml
MCPToolDefinition	Tool name, description, and JSON Schema for arguments
MCPToolResult	Tool invocation result with content, error state, and metadata
MCPToolCall	A tool call request with id, name, and arguments

Related

• Knowledge System -- another source of external information for agents
• Escalation -- MCP failures can trigger escalation