"""A2A-compatible wrapper for Strands Agent.

This module provides the A2AAgent class, which adapts a Strands Agent to the A2A protocol,
allowing it to be used in A2A-compatible systems.
"""

import logging
- from typing import Any, Literal
+ from typing import Any, Callable, Literal
from urllib.parse import urlparse


class A2AServer:
    """A2A-compatible wrapper for Strands Agent."""

    def __init__(
        self,
-        agent: SAAgent,
+       agent_generator: Callable[[], SAAgent],
        *,
        # AgentCard
        host: str = "127.0.0.1",
        port: int = 9000,
        http_url: str | None = None,
        serve_at_root: bool = False,
        version: str = "0.0.1",
        skills: list[AgentSkill] | None = None,
        # RequestHandler
        task_store: TaskStore | None = None,
        queue_manager: QueueManager | None = None,
        push_config_store: PushNotificationConfigStore | None = None,
        push_sender: PushNotificationSender | None = None,
    ):

-        self.strands_agent = agent
+       self.agent_generator = agent_generator
+       self.strands_agent = self.agent_generator()
        self.name = self.strands_agent.name
        self.description = self.strands_agent.description
        self.capabilities = AgentCapabilities(streaming=True)
        self.request_handler = DefaultRequestHandler(
-            agent_executor=StrandsA2AExecutor(self.strands_agent),
+            agent_executor=StrandsA2AExecutor(self.agent_generator),
            task_store=task_store or InMemoryTaskStore(),
            queue_manager=queue_manager,
            push_config_store=push_config_store,
            push_sender=push_sender,
        )
        self._agent_skills = skills
 

class StrandsA2AExecutor(AgentExecutor):

-     def __init__(self, agent: SAAgent):
+     def __init__(self, agent_generator: Callable[[], SAAgent]):

    async def _execute_streaming(self, context: RequestContext, updater: TaskUpdater) -> None:
        """Execute request in streaming mode.

        Streams the agent's response in real-time, sending incremental updates
        as they become available from the agent.

        Args:
            context: The A2A request context, containing the user's input and other metadata.
            updater: The task updater for managing task state and sending updates.
        """
        # Convert A2A message parts to Strands ContentBlocks
        if context.message and hasattr(context.message, "parts"):
            content_blocks = self._convert_a2a_parts_to_content_blocks(context.message.parts)
            if not content_blocks:
                raise ValueError("No content blocks available")
        else:
            raise ValueError("No content blocks available")

+         agent = self.agent_generator()

        try:
-            async for event in self.agent.stream_async(content_blocks):
+           async for event in agent.stream_async(content_blocks):
                await self._handle_streaming_event(event, updater)
        except Exception:
            logger.exception("Error in streaming execution")
            raise