a2a_mcp_server.py

"""
Enhanced MCP A2A Bridge with improved task management and result retrieval.

This script implements a bridge between the MCP protocol and A2A protocol,
allowing MCP clients to interact with A2A agents.

Supports multiple transport types:
- stdio: Standard input/output transport
- streamable-http: Recommended HTTP transport
- sse: Server-Sent Events transport

Configure the transport type using the MCP_TRANSPORT environment variable.
"""

import asyncio
import os
import uuid
from typing import Any, AsyncGenerator, Dict, List, Optional, Union
import json
import logging
import atexit

import httpx
# Set required environment variable for FastMCP 2.8.1+
os.environ.setdefault('FASTMCP_LOG_LEVEL', 'INFO')
from fastmcp import Context, FastMCP
from pydantic import BaseModel, Field

from common.types import (
    AgentCard,
    AgentCapabilities,
    AgentSkill,
    Artifact,
    DataPart,
    Message,
    Part,
    TextPart,
    TaskState,
    TaskStatus,
    TaskStatusUpdateEvent,
    TaskArtifactUpdateEvent,
    JSONRPCResponse,
    SendTaskRequest,
    SendTaskResponse,
    SendTaskStreamingRequest,
    SendTaskStreamingResponse,
    GetTaskRequest,
    GetTaskResponse,
    TaskQueryParams,
    CancelTaskRequest, 
    CancelTaskResponse
)
from common.client.client import A2AClient
from common.server.task_manager import InMemoryTaskManager
from persistence_utils import save_to_json, load_from_json

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Create a FastMCP server
mcp = FastMCP("A2A Bridge Server")

# File paths for persistent storage
DATA_DIR = os.environ.get("A2A_MCP_DATA_DIR", ".")
REGISTERED_AGENTS_FILE = os.path.join(DATA_DIR, "registered_agents.json")
TASK_AGENT_MAPPING_FILE = os.path.join(DATA_DIR, "task_agent_mapping.json")

# Load stored data from disk if it exists
stored_agents = load_from_json(REGISTERED_AGENTS_FILE)
stored_tasks = load_from_json(TASK_AGENT_MAPPING_FILE)

# Initialize in-memory dictionaries with stored data
registered_agents = {}
task_agent_mapping = {}

# Register function to save data on exit
def save_data_on_exit():
    logger.info("Saving data before exit...")
    # Convert registered_agents to a serializable format
    agents_data = {url: agent.model_dump() for url, agent in registered_agents.items()}
    save_to_json(agents_data, REGISTERED_AGENTS_FILE)
    save_to_json(task_agent_mapping, TASK_AGENT_MAPPING_FILE)
    logger.info("Data saved successfully")

atexit.register(save_data_on_exit)

# Periodically save data (every 5 minutes)
async def periodic_save():
    while True:
        await asyncio.sleep(300)  # 5 minutes
        logger.info("Performing periodic data save...")
        agents_data = {url: agent.model_dump() for url, agent in registered_agents.items()}
        save_to_json(agents_data, REGISTERED_AGENTS_FILE)
        save_to_json(task_agent_mapping, TASK_AGENT_MAPPING_FILE)
        logger.info("Periodic save completed")

# Load transport configuration from environment variables
DEFAULT_TRANSPORT = "stdio"
TRANSPORT_TYPES = ["stdio", "streamable-http", "sse"]

# MCP server configuration
MCP_TRANSPORT = os.environ.get("MCP_TRANSPORT", DEFAULT_TRANSPORT).lower()
MCP_HOST = os.environ.get("MCP_HOST", "0.0.0.0")
MCP_PORT = int(os.environ.get("MCP_PORT", "8000"))
MCP_PATH = os.environ.get("MCP_PATH", "/mcp")  # For streamable-http
MCP_SSE_PATH = os.environ.get("MCP_SSE_PATH", "/sse")  # For sse

# A2A server configuration
A2A_HOST = os.environ.get("A2A_HOST", "0.0.0.0")
A2A_PORT = int(os.environ.get("A2A_PORT", "41241"))

# Validate transport type
if MCP_TRANSPORT not in TRANSPORT_TYPES:
    print(f"Warning: Invalid transport type '{MCP_TRANSPORT}'. Using default: {DEFAULT_TRANSPORT}")
    MCP_TRANSPORT = DEFAULT_TRANSPORT

class AgentInfo(BaseModel):
    """Information about an A2A agent."""
    url: str = Field(description="URL of the A2A agent")
    name: str = Field(description="Name of the A2A agent")
    description: str = Field(description="Description of the A2A agent")

class A2ABridgeTaskManager(InMemoryTaskManager):
    """Task manager that forwards tasks to A2A agents."""
    
    def __init__(self):
        super().__init__()
        self.agent_clients = {}  # Maps agent URLs to A2AClient instances
        
    def get_or_create_client(self, agent_url: str) -> A2AClient:
        """Get an existing client or create a new one."""
        if agent_url not in self.agent_clients:
            self.agent_clients[agent_url] = A2AClient(url=agent_url)  # Use named parameter
        return self.agent_clients[agent_url]
    
    async def on_send_task(self, request: SendTaskRequest) -> SendTaskResponse:
        """Handle a task send request by forwarding to the appropriate A2A agent."""
        task_id = request.params.id
        # Extract the agent URL from metadata
        agent_url = request.params.metadata.get("agent_url") if request.params.metadata else None
        
        if not agent_url:
            # No agent URL provided, return error
            return SendTaskResponse(
                id=request.id,
                error={
                    "code": -32602,
                    "message": "Agent URL not provided in task metadata",
                }
            )
        
        client = self.get_or_create_client(agent_url)
        
        # Forward the message to the A2A agent
        try:
            # Create payload as a single dictionary
            payload = {
                "id": task_id,
                "message": request.params.message,
            }
            if request.params.sessionId:
                payload["sessionId"] = request.params.sessionId
            if request.params.metadata:
                payload["metadata"] = request.params.metadata
                
            result = await client.send_task(payload)
            
            # Store the task result
            self.tasks[task_id] = result
            
            # Return the response
            return SendTaskResponse(id=request.id, result=result)
        except Exception as e:
            return SendTaskResponse(
                id=request.id,
                error={
                    "code": -32603,
                    "message": f"Error forwarding task to A2A agent: {str(e)}",
                }
            )
    
    async def on_send_task_subscribe(
        self, request: SendTaskStreamingRequest
    ) -> AsyncGenerator[SendTaskStreamingResponse, None] | JSONRPCResponse:
        """Handle a task subscription request by forwarding to the appropriate A2A agent."""
        task_id = request.params.id
        # Extract the agent URL from metadata
        agent_url = request.params.metadata.get("agent_url") if request.params.metadata else None
        
        if not agent_url:
            # No agent URL provided, return error
            return JSONRPCResponse(
                id=request.id,
                error={
                    "code": -32602,
                    "message": "Agent URL not provided in task metadata",
                }
            )
        
        client = self.get_or_create_client(agent_url)
        
        # Set up SSE consumer
        sse_event_queue = await self.setup_sse_consumer(task_id=task_id)
        
        # Start forwarding the task in a background task
        asyncio.create_task(self._forward_task_stream(
            client=client,
            request=request,
            task_id=task_id,
        ))
        
        # Return the SSE consumer
        return self.dequeue_events_for_sse(
            request_id=request.id,
            task_id=task_id,
            sse_event_queue=sse_event_queue,
        )
    
    async def _forward_task_stream(
        self, client: A2AClient, request: SendTaskStreamingRequest, task_id: str
    ):
        """Forward a task stream to an A2A agent and relay the responses."""
        try:
            # Create payload as a single dictionary
            payload = {
                "id": task_id,
                "message": request.params.message,
            }
            if request.params.sessionId:
                payload["sessionId"] = request.params.sessionId
            if request.params.metadata:
                payload["metadata"] = request.params.metadata
                
            # Send the task and subscribe to updates
            stream = client.send_task_subscribe(payload)
            
            # Process the stream events
            async for event in stream:
                # Forward the event to our SSE queue
                await self.enqueue_events_for_sse(task_id, event)
                
                # If this is the final event, break
                if hasattr(event, "final") and event.final:
                    break
                    
        except Exception as e:
            # Create an error event and enqueue it
            error_event = TaskStatusUpdateEvent(
                id=task_id,
                status=TaskStatus(
                    state=TaskState.FAILED,
                    message=Message(
                        role="agent",
                        parts=[TextPart(text=f"Error forwarding task: {str(e)}")],
                    ),
                ),
                final=True,
            )
            await self.enqueue_events_for_sse(task_id, error_event)

async def fetch_agent_card(url: str) -> AgentCard:
    """
    Fetch the agent card from the agent's URL.
    First try the main URL, then the well-known location.
    """
    async with httpx.AsyncClient() as client:
        # First try the main endpoint
        try:
            response = await client.get(url)
            if response.status_code == 200:
                try:
                    data = response.json()
                    if isinstance(data, dict) and "name" in data and "url" in data:
                        return AgentCard(**data)
                except json.JSONDecodeError:
                    pass  # Not a valid JSON response, try the well-known URL
        except Exception:
            pass  # Connection error, try the well-known URL
        
        # Try the well-known location
        well_known_url = f"{url.rstrip('/')}/.well-known/agent.json"
        try:
            response = await client.get(well_known_url)
            if response.status_code == 200:
                try:
                    data = response.json()
                    return AgentCard(**data)
                except json.JSONDecodeError:
                    raise ValueError(f"Invalid JSON in agent card from {well_known_url}")
        except httpx.RequestError as e:
            raise ValueError(f"Failed to fetch agent card from {well_known_url}: {str(e)}")
    
    # If we can't get the agent card, create a minimal one with default values
    return AgentCard(
        name="Unknown Agent",
        url=url,
        version="0.1.0",
        capabilities=AgentCapabilities(streaming=False),
        skills=[
            AgentSkill(
                id="unknown",
                name="Unknown Skill",
                description="Unknown agent capabilities",
            )
        ],
    )


@mcp.tool()
async def register_agent(url: str, ctx: Context) -> Dict[str, Any]:
    """
    Register an A2A agent with the bridge server.
    
    Args:
        url: URL of the A2A agent
        
    Returns:
        Dictionary with registration status
    """
    try:
        # Fetch the agent card directly
        agent_card = await fetch_agent_card(url)
        
        # Store the agent information
        agent_info = AgentInfo(
            url=url,
            name=agent_card.name,
            description=agent_card.description or "No description provided",
        )
        
        registered_agents[url] = agent_info
        
        # Save to disk immediately
        agents_data = {url: agent.model_dump() for url, agent in registered_agents.items()}
        save_to_json(agents_data, REGISTERED_AGENTS_FILE)
        
        await ctx.info(f"Successfully registered agent: {agent_card.name}")
        return {
            "status": "success",
            "agent": agent_info.model_dump(),
        }
    except Exception as e:
        return {
            "status": "error",
            "message": f"Failed to register agent: {str(e)}",
        }


@mcp.tool()
async def list_agents() -> List[Dict[str, Any]]:
    """
    List all registered A2A agents.
    
    Returns:
        List of registered agents
    """
    return [agent.model_dump() for agent in registered_agents.values()]


@mcp.tool()
async def unregister_agent(url: str, ctx: Context = None) -> Dict[str, Any]:
    """
    Unregister an A2A agent from the bridge server.
    
    Args:
        url: URL of the A2A agent to unregister
        
    Returns:
        Dictionary with unregistration status
    """
    if url not in registered_agents:
        return {
            "status": "error",
            "message": f"Agent not registered: {url}",
        }
    
    try:
        # Get agent name before removing it
        agent_name = registered_agents[url].name
        
        # Remove from registered agents
        del registered_agents[url]
        
        # Clean up any task mappings related to this agent
        # Create a list of task_ids to remove to avoid modifying the dictionary during iteration
        tasks_to_remove = []
        for task_id, agent_url in task_agent_mapping.items():
            if agent_url == url:
                tasks_to_remove.append(task_id)
        
        # Now remove the task mappings
        for task_id in tasks_to_remove:
            del task_agent_mapping[task_id]
        
        # Save changes to disk immediately
        agents_data = {url: agent.model_dump() for url, agent in registered_agents.items()}
        save_to_json(agents_data, REGISTERED_AGENTS_FILE)
        save_to_json(task_agent_mapping, TASK_AGENT_MAPPING_FILE)
        
        if ctx:
            await ctx.info(f"Successfully unregistered agent: {agent_name}")
        
        return {
            "status": "success",
            "message": f"Successfully unregistered agent: {agent_name}",
            "removed_tasks": len(tasks_to_remove),
        }
    except Exception as e:
        return {
            "status": "error",
            "message": f"Error unregistering agent: {str(e)}",
        }


@mcp.tool()
async def send_message(
    agent_url: str,
    message: str,
    session_id: Optional[str] = None,
    ctx: Context = None,
) -> Dict[str, Any]:
    """
    Send a message to an A2A agent.
    
    Args:
        agent_url: URL of the A2A agent
        message: Message to send
        session_id: Optional session ID for multi-turn conversations
        
    Returns:
        Agent's response with task_id for future reference
    """
    if agent_url not in registered_agents:
        return {
            "status": "error",
            "message": f"Agent not registered: {agent_url}",
        }
    
    # Create a client for the agent
    client = A2AClient(url=agent_url)
    
    try:
        # Generate a task ID
        task_id = str(uuid.uuid4())
        
        # Store the mapping of task_id to agent_url for later reference
        task_agent_mapping[task_id] = agent_url
        
        # Create the message
        a2a_message = Message(
            role="user",
            parts=[TextPart(text=message)],
        )
        
        if ctx:
            await ctx.info(f"Sending message to agent: {message}")
        
        # Create payload as a single dictionary
        payload = {
            "id": task_id,
            "message": a2a_message,
        }
        if session_id:
            payload["sessionId"] = session_id
        
        # Send the task with the payload
        result = await client.send_task(payload)
        
        # Save task mapping to disk
        save_to_json(task_agent_mapping, TASK_AGENT_MAPPING_FILE)
        
        # Debug: Print the raw response for analysis
        if ctx:
            await ctx.info(f"Raw response: {result}")
            
        # Create a response dictionary with as much info as we can extract
        response = {
            "status": "success",
            "task_id": task_id,
        }
        
        # Add any available fields from the result
        if hasattr(result, "sessionId"):
            response["session_id"] = result.sessionId
        else:
            response["session_id"] = None
            
        # Try to get the state
        try:
            if hasattr(result, "status") and hasattr(result.status, "state"):
                response["state"] = result.status.state
            else:
                response["state"] = "unknown"
        except Exception as e:
            response["state"] = f"error_getting_state: {str(e)}"
            
        # Try to extract response message
        try:
            if hasattr(result, "status") and hasattr(result.status, "message") and result.status.message:
                response_text = ""
                for part in result.status.message.parts:
                    if part.type == "text":
                        response_text += part.text
                if response_text:
                    response["message"] = response_text
        except Exception as e:
            response["message_error"] = f"Error extracting message: {str(e)}"
        
        # Try to get artifacts
        try:
            if hasattr(result, "artifacts") and result.artifacts:
                artifacts_data = []
                for artifact in result.artifacts:
                    artifact_data = {
                        "name": artifact.name if hasattr(artifact, "name") else "unnamed_artifact",
                        "contents": [],
                    }
                    
                    for part in artifact.parts:
                        if part.type == "text":
                            artifact_data["contents"].append({
                                "type": "text",
                                "text": part.text,
                            })
                        elif part.type == "data":
                            artifact_data["contents"].append({
                                "type": "data",
                                "data": part.data,
                            })
                    
                    artifacts_data.append(artifact_data)
                
                response["artifacts"] = artifacts_data
        except Exception as e:
            response["artifacts_error"] = f"Error extracting artifacts: {str(e)}"
            
        return response
    except Exception as e:
        return {
            "status": "error",
            "message": f"Error sending message: {str(e)}",
        }


@mcp.tool()
async def get_task_result(
    task_id: str,
    history_length: Optional[int] = None,
    ctx: Context = None,
) -> Dict[str, Any]:
    """
    Retrieve the result of a task from an A2A agent.
    
    Args:
        task_id: ID of the task to retrieve
        history_length: Optional number of history items to include (null for all)
        
    Returns:
        Task result including status, message, and artifacts if available
    """
    if task_id not in task_agent_mapping:
        return {
            "status": "error",
            "message": f"Task ID not found: {task_id}",
        }
    
    agent_url = task_agent_mapping[task_id]
    
    # Create a client for the agent
    client = A2AClient(url=agent_url)
    
    try:
        # Create the request payload
        payload = {
            "id": task_id,
            "historyLength": history_length
        }
        
        if ctx:
            await ctx.info(f"Retrieving task result for task_id: {task_id}")
        
        # Send the get task request
        result = await client.get_task(payload)
        
        # Debug: Print the raw response for analysis
        if ctx:
            await ctx.info(f"Raw task result: {result}")
            
        # Create a response dictionary with as much info as we can extract
        response = {
            "status": "success",
            "task_id": task_id,
        }
        
        # Try to extract task data
        try:
            if hasattr(result, "result"):
                task = result.result
                
                # Add basic task info
                if hasattr(task, "sessionId"):
                    response["session_id"] = task.sessionId
                else:
                    response["session_id"] = None
                
                # Add task status
                if hasattr(task, "status"):
                    status = task.status
                    if hasattr(status, "state"):
                        response["state"] = status.state
                    
                    # Extract message from status
                    if hasattr(status, "message") and status.message:
                        response_text = ""
                        for part in status.message.parts:
                            if part.type == "text":
                                response_text += part.text
                        if response_text:
                            response["message"] = response_text
                
                # Extract artifacts
                if hasattr(task, "artifacts") and task.artifacts:
                    artifacts_data = []
                    for artifact in task.artifacts:
                        artifact_data = {
                            "name": artifact.name if hasattr(artifact, "name") else "unnamed_artifact",
                            "contents": [],
                        }
                        
                        for part in artifact.parts:
                            if part.type == "text":
                                artifact_data["contents"].append({
                                    "type": "text",
                                    "text": part.text,
                                })
                            elif part.type == "data":
                                artifact_data["contents"].append({
                                    "type": "data",
                                    "data": part.data,
                                })
                        
                        artifacts_data.append(artifact_data)
                    
                    response["artifacts"] = artifacts_data
                
                # Extract message history if available
                if hasattr(task, "history") and task.history:
                    history_data = []
                    for message in task.history:
                        message_data = {
                            "role": message.role,
                            "parts": [],
                        }
                        
                        for part in message.parts:
                            if part.type == "text":
                                message_data["parts"].append({
                                    "type": "text",
                                    "text": part.text,
                                })
                            elif hasattr(part, "data"):
                                message_data["parts"].append({
                                    "type": "data",
                                    "data": part.data,
                                })
                        
                        history_data.append(message_data)
                    
                    response["history"] = history_data
            else:
                response["error"] = "No result in response"
                
        except Exception as e:
            response["parsing_error"] = f"Error parsing task result: {str(e)}"
            
        return response
    except Exception as e:
        return {
            "status": "error",
            "message": f"Error retrieving task result: {str(e)}",
        }


@mcp.tool()
async def cancel_task(
    task_id: str,
    ctx: Context = None,
) -> Dict[str, Any]:
    """
    Cancel a running task on an A2A agent.
    
    Args:
        task_id: ID of the task to cancel
        
    Returns:
        Cancellation result
    """
    if task_id not in task_agent_mapping:
        return {
            "status": "error",
            "message": f"Task ID not found: {task_id}",
        }
    
    agent_url = task_agent_mapping[task_id]
    
    # Create a client for the agent
    client = A2AClient(url=agent_url)
    
    try:
        # Create the request payload
        payload = {
            "id": task_id
        }
        
        if ctx:
            await ctx.info(f"Cancelling task: {task_id}")
        
        # Send the cancel task request
        result = await client.cancel_task(payload)
        
        # Debug: Print the raw response for analysis
        if ctx:
            await ctx.info(f"Raw cancellation result: {result}")
            
        # Create a response dictionary
        if hasattr(result, "error"):
            return {
                "status": "error",
                "task_id": task_id,
                "message": result.error.message,
                "code": result.error.code
            }
        elif hasattr(result, "result"):
            return {
                "status": "success",
                "task_id": task_id,
                "message": "Task cancelled successfully"
            }
        else:
            return {
                "status": "unknown",
                "task_id": task_id,
                "message": "Unexpected response format"
            }
            
    except Exception as e:
        return {
            "status": "error",
            "message": f"Error cancelling task: {str(e)}",
        }


@mcp.tool()
async def send_message_stream(
    agent_url: str,
    message: str,
    session_id: Optional[str] = None,
    ctx: Context = None,
) -> Dict[str, Any]:
    """
    Send a message to an A2A agent and stream the response.
    
    Args:
        agent_url: URL of the A2A agent
        message: Message to send
        session_id: Optional session ID for multi-turn conversations
        
    Returns:
        Stream of agent's responses
    """
    if agent_url not in registered_agents:
        return {
            "status": "error",
            "message": f"Agent not registered: {agent_url}",
        }
    
    # Create a client for the agent
    client = A2AClient(url=agent_url)
    
    try:
        # Generate a task ID
        task_id = str(uuid.uuid4())
        
        # Store the mapping of task_id to agent_url for later reference
        task_agent_mapping[task_id] = agent_url
        
        # Save the task mapping to disk
        save_to_json(task_agent_mapping, TASK_AGENT_MAPPING_FILE)
        
        # Create the message
        a2a_message = Message(
            role="user",
            parts=[TextPart(text=message)],
        )
        
        if ctx:
            await ctx.info(f"Sending message to agent (streaming): {message}")
            
        # Start progress indication
        if ctx:
            await ctx.info("Processing...")
        
        # Dictionary to accumulate streaming responses
        complete_response = {
            "status": "success",
            "task_id": task_id,
            "session_id": session_id,
            "state": "working",
            "messages": [],
            "artifacts": [],
        }
        
        # Create payload as a single dictionary
        payload = {
            "id": task_id,
            "message": a2a_message,
        }
        if session_id:
            payload["sessionId"] = session_id
        
        # Send the task and subscribe to updates
        stream = client.send_task_streaming(payload)
        
        # Process and report stream events
        try:
            all_events = []
            
            async for event in stream:
                # Save all events for debugging
                all_events.append({
                    "type": str(type(event)),
                    "dir": str(dir(event)),
                })
                
                if hasattr(event, "result"):
                    if hasattr(event.result, "status"):
                        # It's a TaskStatusUpdateEvent
                        status_event = event.result
                        
                        # Update the state
                        if hasattr(status_event, "status") and hasattr(status_event.status, "state"):
                            complete_response["state"] = status_event.status.state
                        
                        # Extract any message
                        if hasattr(status_event, "status") and hasattr(status_event.status, "message") and status_event.status.message:
                            message_text = ""
                            for part in status_event.status.message.parts:
                                if part.type == "text":
                                    message_text += part.text
                            
                            if message_text:
                                complete_response["messages"].append(message_text)
                                if ctx:
                                    await ctx.info(f"Agent: {message_text}")
                        
                        # If this is the final event, set session ID
                        if hasattr(status_event, "final") and status_event.final:
                            complete_response["session_id"] = getattr(status_event, "sessionId", session_id)
                        
                    elif hasattr(event.result, "artifact"):
                        # It's a TaskArtifactUpdateEvent
                        artifact_event = event.result
                        
                        # Extract artifact content
                        artifact_data = {
                            "name": artifact_event.artifact.name if hasattr(artifact_event.artifact, "name") else "unnamed",
                            "contents": [],
                        }
                        
                        for part in artifact_event.artifact.parts:
                            if part.type == "text":
                                artifact_data["contents"].append({
                                    "type": "text",
                                    "text": part.text,
                                })
                            elif part.type == "data":
                                artifact_data["contents"].append({
                                    "type": "data",
                                    "data": part.data,
                                })
                        
                        complete_response["artifacts"].append(artifact_data)
                        
                        if ctx:
                            await ctx.info(f"Received artifact: {artifact_data['name']}")
                else:
                    # Unknown event type, try to extract what we can
                    complete_response["unknown_events"] = complete_response.get("unknown_events", []) + [
                        {
                            "type": str(type(event)),
                            "dir": str(dir(event))
                        }
                    ]
            
            # Include debug info            
            complete_response["_debug_info"] = {
                "all_events": all_events
            }
            
            return complete_response
        except Exception as e:
            return {
                "status": "error",
                "message": f"Error processing stream events: {str(e)}",
                "_debug_info": {
                    "all_events": all_events
                }
            }
            
    except Exception as e:
        return {
            "status": "error",
            "message": f"Error sending message (stream): {str(e)}",
        }


class CustomA2AServer:
    """
    A minimal A2A server implementation that uses the task manager.
    """
    
    def __init__(
        self,
        agent_card: AgentCard,
        task_manager: A2ABridgeTaskManager,
        host: str = "0.0.0.0",
        port: int = 41241,
    ):
        self.agent_card = agent_card
        self.task_manager = task_manager
        self.host = host
        self.port = port
        
    async def start_async(self):
        """Start the A2A server asynchronously."""
        # In a real implementation, this would start a FastAPI server
        # For now, just log that it's "started"
        print(f"A2A server 'started' at {self.host}:{self.port}")
        # Keep the server "running"
        while True:
            await asyncio.sleep(3600)  # Sleep for an hour
    
    def start(self):
        """Start the A2A server."""
        asyncio.create_task(self.start_async())


def setup_a2a_server():
    """Set up the A2A server with our task manager."""
    # Create a sample agent card
    agent_card = AgentCard(
        name="MCP Bridge Agent",
        description="A bridge between MCP and A2A protocols",
        url=f"http://{A2A_HOST}:{A2A_PORT}",
        version="0.1.0",
        # Add the required capabilities field
        capabilities=AgentCapabilities(
            streaming=True,
            pushNotifications=False,
            stateTransitionHistory=False,
        ),
        # Add the required skills field with at least one skill
        skills=[
            AgentSkill(
                id="mcp-bridge",
                name="MCP Bridge",
                description="Allows MCP clients to communicate with A2A agents",
                tags=["bridge", "proxy", "mcp", "a2a"],
                examples=[
                    "Send a message to an A2A agent",
                    "Register an A2A agent with the bridge",
                    "List all registered A2A agents",
                ],
                inputModes=["text/plain"],
                outputModes=["text/plain", "application/json"],
            )
        ],
    )
    
    # Create our custom task manager
    task_manager = A2ABridgeTaskManager()
    
    # Create and return the A2A server
    return CustomA2AServer(
        agent_card=agent_card,
        task_manager=task_manager,
        host=A2A_HOST,
        port=A2A_PORT,
    )


async def main_async():
    """
    Main async function to start both the MCP and A2A servers.
    """
    # Load stored data into memory
    load_registered_agents()
    
    # Start periodic save task
    asyncio.create_task(periodic_save())
    
    # Set up and start the A2A server
    a2a_server = setup_a2a_server()
    a2a_task = asyncio.create_task(a2a_server.start_async())
    
    # Start the MCP server with the configured transport
    print(f"Starting MCP server with {MCP_TRANSPORT} transport...")
    
    if MCP_TRANSPORT == "stdio":
        # Use stdio transport (default)
        mcp_task = asyncio.create_task(
            mcp.run_async(transport="stdio")
        )
    elif MCP_TRANSPORT == "streamable-http":
        # Use streamable-http transport
        mcp_task = asyncio.create_task(
            mcp.run_async(
                transport="streamable-http",
                host=MCP_HOST,
                port=MCP_PORT,
                path=MCP_PATH,
            )
        )
    elif MCP_TRANSPORT == "sse":
        # Use sse transport (deprecated but still supported)
        mcp_task = asyncio.create_task(
            mcp.run_async(
                transport="sse",
                host=MCP_HOST,
                port=MCP_PORT,
                path=MCP_SSE_PATH,
            )
        )
    
    # Run both servers
    await asyncio.gather(a2a_task, mcp_task)


def load_registered_agents():
    """Load registered agents from stored data on startup."""
    global registered_agents, task_agent_mapping
    
    logger.info("Loading saved data...")
    
    # Load agents data
    agents_data = load_from_json(REGISTERED_AGENTS_FILE)
    for url, agent_data in agents_data.items():
        registered_agents[url] = AgentInfo(**agent_data)
    
    # Load task mappings
    task_agent_mapping = load_from_json(TASK_AGENT_MAPPING_FILE)
    
    logger.info(f"Loaded {len(registered_agents)} agents and {len(task_agent_mapping)} task mappings")


def main():
    """
    Main entry point.
    """
    # Print the configuration for debugging
    print(f"MCP Bridge Configuration:")
    print(f"- Transport: {MCP_TRANSPORT}")
    print(f"- Host: {MCP_HOST}")
    print(f"- Port: {MCP_PORT}")
    if MCP_TRANSPORT == "streamable-http":
        print(f"- Path: {MCP_PATH}")
    elif MCP_TRANSPORT == "sse":
        print(f"- SSE Path: {MCP_SSE_PATH}")
    print(f"A2A Server Configuration:")
    print(f"- Host: {A2A_HOST}")
    print(f"- Port: {A2A_PORT}")
    print(f"Data Storage:")
    print(f"- Registered Agents: {REGISTERED_AGENTS_FILE}")
    print(f"- Task Agent Mapping: {TASK_AGENT_MAPPING_FILE}")
    
    asyncio.run(main_async())


if __name__ == "__main__":
    main()