executor.py

"""Strands Agent executor for the A2A protocol.

This module provides the StrandsA2AExecutor class, which adapts a Strands Agent
to be used as an executor in the A2A protocol. It handles the execution of agent
requests and the conversion of Strands Agent streamed responses to A2A events.

The A2A AgentExecutor ensures clients receive responses for synchronous and
streamed requests to the A2AServer.
"""

import asyncio
import base64
import json
import logging
import mimetypes
import uuid
import warnings
from typing import Any, Literal

from a2a.server.agent_execution import AgentExecutor, RequestContext
from a2a.server.events import EventQueue
from a2a.server.tasks import TaskUpdater
from a2a.types import DataPart, FilePart, InternalError, Part, TaskState, TextPart, UnsupportedOperationError
from a2a.utils import new_agent_text_message, new_task
from a2a.utils.errors import ServerError

from ...agent.agent import Agent as SAAgent
from ...agent.agent import AgentResult as SAAgentResult
from ...types.content import ContentBlock
from ...types.media import (
    DocumentContent,
    DocumentSource,
    ImageContent,
    ImageSource,
    VideoContent,
    VideoSource,
)

logger = logging.getLogger(__name__)


class StrandsA2AExecutor(AgentExecutor):
    """Executor that adapts a Strands Agent to the A2A protocol.

    This executor uses streaming mode to handle the execution of agent requests
    and converts Strands Agent responses to A2A protocol events. It supports the
    full A2A task lifecycle including error handling (failed state), cancellation,
    and interrupt-based input_required flows.
    """

    # Default formats for each file type when MIME type is unavailable or unrecognized
    DEFAULT_FORMATS = {"document": "txt", "image": "png", "video": "mp4", "unknown": "txt"}

    # Handle special cases where format differs from extension
    FORMAT_MAPPINGS = {"jpg": "jpeg", "htm": "html", "3gp": "three_gp", "3gpp": "three_gp", "3g2": "three_gp"}

    # A2A-compliant streaming mode
    _current_artifact_id: str | None
    _is_first_chunk: bool

    def __init__(self, agent: SAAgent, *, enable_a2a_compliant_streaming: bool = False):
        """Initialize a StrandsA2AExecutor.

        Args:
            agent: The Strands Agent instance to adapt to the A2A protocol.
            enable_a2a_compliant_streaming: If True, uses A2A-compliant streaming with
                artifact updates. If False, uses legacy status updates streaming behavior
                for backwards compatibility. Defaults to False.
        """
        self.agent = agent
        self.enable_a2a_compliant_streaming = enable_a2a_compliant_streaming

    async def execute(
        self,
        context: RequestContext,
        event_queue: EventQueue,
    ) -> None:
        """Execute a request using the Strands Agent and send the response as A2A events.

        This method executes the user's input using the Strands Agent in streaming mode
        and converts the agent's response to A2A events. If the agent raises an exception,
        the task transitions to the `failed` state. If the agent returns with interrupts,
        the task transitions to the `input_required` state.

        Args:
            context: The A2A request context, containing the user's input and task metadata.
            event_queue: The A2A event queue used to send response events back to the client.

        Raises:
            ServerError: If an unrecoverable error occurs during agent execution setup
                (e.g., missing input). Agent execution errors are handled gracefully
                by transitioning the task to the failed state.
        """
        task = context.current_task
        if not task:
            task = new_task(context.message)  # type: ignore
            await event_queue.enqueue_event(task)

        updater = TaskUpdater(event_queue, task.id, task.context_id)

        try:
            await self._execute_streaming(context, updater)
        except ServerError:
            # Re-raise ServerErrors (setup failures like missing input)
            raise
        except asyncio.CancelledError:
            # asyncio.CancelledError is a BaseException (not Exception) — raised when
            # the asyncio task is cancelled (e.g., HTTP client disconnect, server shutdown).
            # We transition to canceled state so the task doesn't remain a zombie in "working".
            logger.warning("task_id=<%s> | asyncio task cancelled, transitioning to canceled state", task.id)
            try:
                await updater.cancel(
                    message=updater.new_agent_message(
                        parts=[Part(root=TextPart(text="Task cancelled due to connection termination"))]
                    )
                )
            except RuntimeError:
                # Task already in terminal state
                logger.debug("task_id=<%s> | task already in terminal state, cannot transition to canceled", task.id)
            raise
        except Exception:
            # Agent execution failures transition to failed state
            logger.exception("task_id=<%s> | agent execution failed, transitioning to failed state", task.id)
            try:
                await updater.failed(
                    message=updater.new_agent_message(parts=[Part(root=TextPart(text="Agent execution failed"))])
                )
            except RuntimeError:
                # Task already in terminal state (e.g., completed before error in cleanup)
                logger.debug("task_id=<%s> | task already in terminal state, cannot transition to failed", task.id)

    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.

        Raises:
            ServerError: If input conversion fails (missing or empty content).
        """
        # 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 ServerError(
                    error=InternalError(message="No valid content found in request message parts")
                ) from None
        else:
            raise ServerError(error=InternalError(message="Request message is missing or has no parts")) from None

        if not self.enable_a2a_compliant_streaming:
            warnings.warn(
                "The default A2A response stream implemented in the strands sdk does not conform to "
                "what is expected in the A2A spec. Please set the `enable_a2a_compliant_streaming` "
                "boolean to `True` on your `A2AServer` class to properly conform to the spec. "
                "In the next major version release, this will be the default behavior.",
                UserWarning,
                stacklevel=3,
            )

        if self.enable_a2a_compliant_streaming:
            self._current_artifact_id = str(uuid.uuid4())
            self._is_first_chunk = True

        # Pass the A2A RequestContext through invocation state so downstream
        # tools and hooks can access request metadata, task info, configuration, etc.
        invocation_state: dict[str, Any] = {"a2a_request_context": context}

        try:
            result: SAAgentResult | None = None
            async for event in self.agent.stream_async(content_blocks, invocation_state=invocation_state):
                if "result" in event:
                    result = event["result"]
                else:
                    await self._handle_streaming_event(event, updater)

            # Check if agent returned with interrupts (input_required)
            # Note: stop_reason="interrupt" is the authoritative signal. Even if interrupts
            # list is empty (edge case), the agent still indicated it needs input.
            if result is not None and result.stop_reason == "interrupt":
                await self._handle_interrupt_result(result, updater)
            else:
                await self._handle_agent_result(result, updater)
        except Exception:
            logger.exception("Error in streaming execution")
            raise
        finally:
            if self.enable_a2a_compliant_streaming:
                self._current_artifact_id = None
                self._is_first_chunk = True

    async def _handle_interrupt_result(self, result: SAAgentResult, updater: TaskUpdater) -> None:
        """Handle an agent result that contains interrupts.

        When the Strands Agent returns with stop_reason="interrupt", this maps to
        the A2A `input_required` state. The interrupt details are communicated to
        the client via the status message.

        Args:
            result: The agent result containing interrupts.
            updater: The task updater for managing task state.
        """
        # Build a descriptive message about what input is needed
        interrupt_descriptions = []
        for interrupt in result.interrupts or []:
            desc = f"- {interrupt.name}"
            if interrupt.reason:
                desc += f": {interrupt.reason}"
            interrupt_descriptions.append(desc)

        if interrupt_descriptions:
            input_message = "Agent requires input:\n" + "\n".join(interrupt_descriptions)
        else:
            # Edge case: stop_reason="interrupt" but no interrupt details provided.
            # Still transition to input_required — the agent signaled it needs input.
            input_message = "Agent requires additional input to continue"

        await updater.requires_input(message=updater.new_agent_message(parts=[Part(root=TextPart(text=input_message))]))

    async def _handle_streaming_event(self, event: dict[str, Any], updater: TaskUpdater) -> None:
        """Handle a single streaming event from the Strands Agent.

        Processes streaming events from the agent, converting data chunks to A2A
        task updates and handling the final result when streaming is complete.

        Args:
            event: The streaming event from the agent, containing either 'data' for
                incremental content or 'result' for the final response.
            updater: The task updater for managing task state and sending updates.
        """
        logger.debug("Streaming event: %s", event)
        if "data" in event:
            if text_content := event["data"]:
                if self.enable_a2a_compliant_streaming:
                    await updater.add_artifact(
                        [Part(root=TextPart(text=text_content))],
                        artifact_id=self._current_artifact_id,
                        name="agent_response",
                        append=not self._is_first_chunk,
                    )
                    self._is_first_chunk = False
                else:
                    # Legacy use update_status with agent message
                    await updater.update_status(
                        TaskState.working,
                        new_agent_text_message(
                            text_content,
                            updater.context_id,
                            updater.task_id,
                        ),
                    )

    async def _handle_agent_result(self, result: SAAgentResult | None, updater: TaskUpdater) -> None:
        """Handle the final result from the Strands Agent.

        For A2A-compliant streaming: sends the final artifact chunk marker and marks
        the task as complete. If no data chunks were previously sent, includes the
        result content.

        For legacy streaming: adds the final result as a simple artifact without
        artifact_id tracking.

        Args:
            result: The agent result object containing the final response, or None if no result.
            updater: The task updater for managing task state and adding the final artifact.
        """
        if self.enable_a2a_compliant_streaming:
            if self._is_first_chunk:
                final_content = str(result) if result else ""
                await updater.add_artifact(
                    [Part(root=TextPart(text=final_content))],
                    artifact_id=self._current_artifact_id,
                    name="agent_response",
                    last_chunk=True,
                )
            else:
                await updater.add_artifact(
                    [Part(root=TextPart(text=""))],
                    artifact_id=self._current_artifact_id,
                    name="agent_response",
                    append=True,
                    last_chunk=True,
                )
        elif final_content := str(result):
            await updater.add_artifact(
                [Part(root=TextPart(text=final_content))],
                name="agent_response",
            )
        await updater.complete()

    async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None:
        """Cancel an ongoing execution.

        Transitions the task to the canceled state and attempts to stop the agent.
        The agent's cancel() method is called if available to signal cooperative
        cancellation of in-flight execution.

        Note: This transitions the A2A task state. The underlying agent execution
        may still complete its current model call before stopping.

        Args:
            context: The A2A request context.
            event_queue: The A2A event queue.

        Raises:
            ServerError: If no current task exists or the task is already in a terminal state.
        """
        task = context.current_task
        if not task:
            logger.warning("context_id=<%s> | cancel requested but no current task found", context.context_id)
            raise ServerError(error=UnsupportedOperationError()) from None

        # Attempt to cooperatively cancel the agent's execution (best-effort).
        # Agent.cancel() may not exist on all implementations, so we guard with hasattr.
        try:
            self.agent.cancel()
        except (AttributeError, NotImplementedError):
            # Agent doesn't support cancel — proceed with state transition only
            pass
        except Exception:
            logger.debug("task_id=<%s> | agent cancel signal failed (non-critical)", task.id)

        updater = TaskUpdater(event_queue, task.id, task.context_id)

        try:
            await updater.cancel(
                message=updater.new_agent_message(parts=[Part(root=TextPart(text="Task cancelled by client request"))])
            )
        except RuntimeError:
            # TaskUpdater raises RuntimeError when task is already in a terminal state
            logger.warning("task_id=<%s> | cannot cancel, already in terminal state", task.id)
            raise ServerError(error=UnsupportedOperationError()) from None

    def _get_file_type_from_mime_type(self, mime_type: str | None) -> Literal["document", "image", "video", "unknown"]:
        """Classify file type based on MIME type.

        Args:
            mime_type: The MIME type of the file

        Returns:
            The classified file type
        """
        if not mime_type:
            return "unknown"

        mime_type = mime_type.lower()

        if mime_type.startswith("image/"):
            return "image"
        elif mime_type.startswith("video/"):
            return "video"
        elif (
            mime_type.startswith("text/")
            or mime_type.startswith("application/")
            or mime_type in ["application/pdf", "application/json", "application/xml"]
        ):
            return "document"
        else:
            return "unknown"

    def _get_file_format_from_mime_type(self, mime_type: str | None, file_type: str) -> str:
        """Extract file format from MIME type using Python's mimetypes library.

        Args:
            mime_type: The MIME type of the file
            file_type: The classified file type (image, video, document, txt)

        Returns:
            The file format string
        """
        if not mime_type:
            return self.DEFAULT_FORMATS.get(file_type, "txt")

        mime_type = mime_type.lower()

        # Extract subtype from MIME type and check existing format mappings
        if "/" in mime_type:
            subtype = mime_type.split("/")[-1]
            if subtype in self.FORMAT_MAPPINGS:
                return self.FORMAT_MAPPINGS[subtype]

        # Use mimetypes library to find extensions for the MIME type
        extensions = mimetypes.guess_all_extensions(mime_type)

        if extensions:
            extension = extensions[0][1:]  # Remove the leading dot
            return self.FORMAT_MAPPINGS.get(extension, extension)

        # Fallback to defaults for unknown MIME types
        return self.DEFAULT_FORMATS.get(file_type, "txt")

    def _strip_file_extension(self, file_name: str) -> str:
        """Strip the file extension from a file name.

        Args:
            file_name: The original file name with extension

        Returns:
            The file name without extension
        """
        if "." in file_name:
            return file_name.rsplit(".", 1)[0]
        return file_name

    def _convert_a2a_parts_to_content_blocks(self, parts: list[Part]) -> list[ContentBlock]:
        """Convert A2A message parts to Strands ContentBlocks.

        Args:
            parts: List of A2A Part objects

        Returns:
            List of Strands ContentBlock objects
        """
        content_blocks: list[ContentBlock] = []

        for part in parts:
            try:
                part_root = part.root

                if isinstance(part_root, TextPart):
                    # Handle TextPart
                    content_blocks.append(ContentBlock(text=part_root.text))

                elif isinstance(part_root, FilePart):
                    # Handle FilePart
                    file_obj = part_root.file
                    mime_type = getattr(file_obj, "mime_type", None)
                    raw_file_name = getattr(file_obj, "name", "FileNameNotProvided")
                    file_name = self._strip_file_extension(raw_file_name)
                    file_type = self._get_file_type_from_mime_type(mime_type)
                    file_format = self._get_file_format_from_mime_type(mime_type, file_type)

                    # Handle FileWithBytes vs FileWithUri
                    bytes_data = getattr(file_obj, "bytes", None)
                    uri_data = getattr(file_obj, "uri", None)

                    if bytes_data:
                        try:
                            # A2A bytes are always base64-encoded strings
                            decoded_bytes = base64.b64decode(bytes_data)
                        except Exception as e:
                            raise ValueError(f"Failed to decode base64 data for file '{raw_file_name}': {e}") from e

                        if file_type == "image":
                            content_blocks.append(
                                ContentBlock(
                                    image=ImageContent(
                                        format=file_format,  # type: ignore
                                        source=ImageSource(bytes=decoded_bytes),
                                    )
                                )
                            )
                        elif file_type == "video":
                            content_blocks.append(
                                ContentBlock(
                                    video=VideoContent(
                                        format=file_format,  # type: ignore
                                        source=VideoSource(bytes=decoded_bytes),
                                    )
                                )
                            )
                        else:  # document or unknown
                            content_blocks.append(
                                ContentBlock(
                                    document=DocumentContent(
                                        format=file_format,  # type: ignore
                                        name=file_name,
                                        source=DocumentSource(bytes=decoded_bytes),
                                    )
                                )
                            )
                    # Handle FileWithUri
                    elif uri_data:
                        # For URI files, create a text representation since Strands ContentBlocks expect bytes
                        content_blocks.append(
                            ContentBlock(text=f"[File: {file_name} ({mime_type})] - Referenced file at: {uri_data}")
                        )
                elif isinstance(part_root, DataPart):
                    # Handle DataPart - convert structured data to JSON text
                    try:
                        data_text = json.dumps(part_root.data, indent=2)
                        content_blocks.append(ContentBlock(text=f"[Structured Data]\n{data_text}"))
                    except Exception:
                        logger.exception("Failed to serialize data part")
            except Exception:
                logger.exception("Error processing part")

        return content_blocks