feat: Addressed PR review comments and fixed spelling and linting issues

Author: Rajesh Ramamoorthy

src/a2a/server/events/distributed_event_queue.py

@@ -35,17 +35,17 @@
 }
 
 
-def _serialise_event(
+def _serialize_event(
     event: Event,
     task_id: str,
     instance_id: str,
 ) -> str:
-    """Serialises an event into the SNS wire-format JSON string.
+    """Serializes an event into the SNS wire-format JSON string.
 
     Args:
-        event: The event to serialise.
+        event: The event to serialize.
         task_id: The task ID this event belongs to.
-        instance_id: The originating instance ID (for dedup).
+        instance_id: The originating instance ID (for deduplication).
 
     Returns:
         A JSON string suitable for use as an SNS ``Message`` payload.
@@ -55,13 +55,13 @@ def _serialise_event(
         'task_id': task_id,
         'type': _EVENT_TYPE,
         'event_kind': event.kind,
-        'event_data': json.loads(event.model_dump_json()),
+        'event_data': event.model_dump(mode='json'),
     }
     return json.dumps(payload)
 
 
-def _serialise_close(task_id: str, instance_id: str) -> str:
-    """Serialises a close signal into the SNS wire-format JSON string.
+def _serialize_close(task_id: str, instance_id: str) -> str:
+    """Serializes a close signal into the SNS wire-format JSON string.
 
     Args:
         task_id: The task ID whose queue is being closed.
@@ -78,7 +78,7 @@ def _serialise_close(task_id: str, instance_id: str) -> str:
     return json.dumps(payload)
 
 
-def deserialise_wire_message(
+def deserialize_wire_message(
     raw: str,
 ) -> dict[str, Any]:
     """Parses a raw SNS/SQS wire-format JSON string.
@@ -110,7 +110,7 @@ def decode_event(msg: dict[str, Any]) -> Event | None:
             ``event_data`` fields.
 
     Returns:
-        The decoded Event, or ``None`` if the ``kind`` is unrecognised.
+        The decoded Event, or ``None`` if the ``kind`` is unrecognized.
     """
     kind = msg.get('event_kind')
     event_data = msg.get('event_data')
@@ -138,10 +138,11 @@ class DistributedEventQueue(EventQueue):
 
     Args:
         publish_fn: Async callable ``(message: str) -> None`` that publishes
-            the serialised wire message to SNS. Provided by
+            the serialized wire message to SNS. Provided by
             :class:`SnsQueueManager` and injected at construction time.
         task_id: The task ID this queue serves.
-        instance_id: The unique ID of the local instance (used for dedup).
+        instance_id: The unique ID of the local instance (used for
+            deduplication of self-published messages).
         max_queue_size: Maximum number of events to buffer locally.
             Defaults to ``DEFAULT_MAX_QUEUE_SIZE``.
     """
@@ -154,13 +155,13 @@ def __init__(
         *,
         max_queue_size: int = DEFAULT_MAX_QUEUE_SIZE,
     ) -> None:
-        """Initialises the DistributedEventQueue."""
+        """Initializes the DistributedEventQueue."""
         super().__init__(max_queue_size=max_queue_size)
         self._publish_fn = publish_fn
         self._task_id = task_id
         self._instance_id = instance_id
         logger.debug(
-            'DistributedEventQueue initialised (task_id=%s, instance=%s).',
+            'DistributedEventQueue initialized (task_id=%s, instance=%s).',
             task_id,
             instance_id,
         )
@@ -204,13 +205,13 @@ async def close(self, immediate: bool = False) -> None:
         await super().close(immediate=immediate)
 
     async def _publish_event(self, event: Event) -> None:
-        """Fire-and-forget coroutine: serialises and publishes one event.
+        """Fire-and-forget coroutine: serializes and publishes one event.
 
         Args:
             event: The event to publish.
         """
         try:
-            message = _serialise_event(event, self._task_id, self._instance_id)
+            message = _serialize_event(event, self._task_id, self._instance_id)
             await self._publish_fn(message)
             logger.debug(
                 'Event published to SNS (task_id=%s, kind=%s).',
@@ -225,7 +226,7 @@ async def _publish_event(self, event: Event) -> None:
     async def _publish_close(self) -> None:
         """Fire-and-forget coroutine: publishes the close signal to SNS."""
         try:
-            message = _serialise_close(self._task_id, self._instance_id)
+            message = _serialize_close(self._task_id, self._instance_id)
             await self._publish_fn(message)
             logger.debug(
                 'Close signal published to SNS (task_id=%s).', self._task_id

src/a2a/server/events/queue_lifecycle_manager.py

@@ -1,31 +1,20 @@
 """SQS queue lifecycle manager for per-instance ECS auto-scaling support."""
 
+from __future__ import annotations
+
 import json
 import logging
 import uuid
 
 from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
 
 
-logger = logging.getLogger(__name__)
+if TYPE_CHECKING:
+    import aioboto3
 
-# SQS policy template granting SNS permission to send messages to the queue.
-_SNS_POLICY_TEMPLATE = json.dumps(
-    {
-        'Version': '2012-10-17',
-        'Statement': [
-            {
-                'Effect': 'Allow',
-                'Principal': {'Service': 'sns.amazonaws.com'},
-                'Action': 'SQS:SendMessage',
-                'Resource': '{queue_arn}',
-                'Condition': {
-                    'ArnEquals': {'aws:SourceArn': '{topic_arn}'},
-                },
-            }
-        ],
-    }
-)
+
+logger = logging.getLogger(__name__)
 
 
 @dataclass
@@ -77,7 +66,7 @@ class QueueLifecycleManager:
     queue_name_prefix: str = 'a2a-instance'
     instance_id: str = field(default_factory=lambda: str(uuid.uuid4()))
     region_name: str = 'us-east-1'
-    session: object | None = field(default=None, repr=False)
+    session: aioboto3.Session | None = field(default=None, repr=False)
 
     # Set after provision(); read via properties.
     _provision_result: QueueProvisionResult | None = field(
@@ -163,7 +152,7 @@ async def provision(self) -> QueueProvisionResult:
             self.region_name,
         )
 
-        async with self.session.client(  # type: ignore[union-attr]
+        async with self.session.client(
             'sqs', region_name=self.region_name
         ) as sqs:
             # Step 1: Create the SQS queue.
@@ -179,10 +168,27 @@ async def provision(self) -> QueueProvisionResult:
             queue_arn: str = attr_resp['Attributes']['QueueArn']
             logger.debug('SQS queue ARN: %s.', queue_arn)
 
-            # Step 3: Attach SNS → SQS access policy.
-            policy = _SNS_POLICY_TEMPLATE.replace(
-                '{queue_arn}', queue_arn
-            ).replace('{topic_arn}', self.topic_arn)
+            # Step 3: Attach SNS → SQS access policy.  ARN values are
+            # set as dict entries so json.dumps() escapes them safely,
+            # preventing any injection via crafted ARN strings.
+            policy = json.dumps(
+                {
+                    'Version': '2012-10-17',
+                    'Statement': [
+                        {
+                            'Effect': 'Allow',
+                            'Principal': {'Service': 'sns.amazonaws.com'},
+                            'Action': 'SQS:SendMessage',
+                            'Resource': queue_arn,
+                            'Condition': {
+                                'ArnEquals': {
+                                    'aws:SourceArn': self.topic_arn,
+                                },
+                            },
+                        }
+                    ],
+                }
+            )
             await sqs.set_queue_attributes(
                 QueueUrl=queue_url,
                 Attributes={'Policy': policy},
@@ -192,7 +198,7 @@ async def provision(self) -> QueueProvisionResult:
         # Step 4: Subscribe the SQS queue to the SNS topic.
         subscription_arn = ''
         try:
-            async with self.session.client(  # type: ignore[union-attr]
+            async with self.session.client(
                 'sns', region_name=self.region_name
             ) as sns:
                 sub_resp = await sns.subscribe(
@@ -210,7 +216,7 @@ async def provision(self) -> QueueProvisionResult:
                 queue_url,
             )
             try:
-                async with self.session.client(  # type: ignore[union-attr]
+                async with self.session.client(
                     'sqs', region_name=self.region_name
                 ) as sqs:
                     await sqs.delete_queue(QueueUrl=queue_url)
@@ -251,7 +257,7 @@ async def teardown(self) -> None:
 
         # Step 1: Unsubscribe from SNS (best-effort).
         try:
-            async with self.session.client(  # type: ignore[union-attr]
+            async with self.session.client(
                 'sns', region_name=self.region_name
             ) as sns:
                 await sns.unsubscribe(SubscriptionArn=result.subscription_arn)
@@ -266,7 +272,7 @@ async def teardown(self) -> None:
 
         # Step 2: Delete the SQS queue.
         try:
-            async with self.session.client(  # type: ignore[union-attr]
+            async with self.session.client(
                 'sqs', region_name=self.region_name
             ) as sqs:
                 await sqs.delete_queue(QueueUrl=result.queue_url)
@@ -278,7 +284,7 @@ async def teardown(self) -> None:
     # Async context manager
     # ------------------------------------------------------------------
 
-    async def __aenter__(self) -> 'QueueLifecycleManager':  # noqa: PYI034
+    async def __aenter__(self) -> QueueLifecycleManager:  # noqa: PYI034
         """Provisions resources and returns ``self``."""
         await self.provision()
         return self

src/a2a/server/events/sns_queue_manager.py

@@ -1,18 +1,25 @@
 """SnsQueueManager — distributed QueueManager using SNS/SQS fan-out."""
 
+from __future__ import annotations
+
 import asyncio
 import json
 import logging
 import uuid
 
-from typing import Any
+from typing import TYPE_CHECKING, Any
+
+
+if TYPE_CHECKING:
+    import aioboto3
+
+    from a2a.server.events.event_queue import EventQueue
 
 from a2a.server.events.distributed_event_queue import (
     DistributedEventQueue,
     decode_event,
-    deserialise_wire_message,
+    deserialize_wire_message,
 )
-from a2a.server.events.event_queue import EventQueue
 from a2a.server.events.queue_manager import (
     NoTaskQueue,
     QueueManager,
@@ -78,12 +85,12 @@ def __init__(  # noqa: PLR0913
         *,
         instance_id: str | None = None,
         region_name: str = 'us-east-1',
-        session: object | None = None,
+        session: aioboto3.Session | None = None,
         poll_interval_seconds: float = 1.0,
         max_messages: int = 10,
         visibility_timeout_seconds: int = 30,
     ) -> None:
-        """Initialises the SnsQueueManager."""
+        """Initializes the SnsQueueManager."""
         try:
             import aioboto3  # noqa: PLC0415
         except ImportError as exc:
@@ -157,7 +164,7 @@ async def stop(self) -> None:
     # ------------------------------------------------------------------
 
     async def add(self, task_id: str, queue: EventQueue) -> None:
-        """Adds a pre-existing EventQueue for *task_id*.
+        """Adds an already-created EventQueue for *task_id*.
 
         Raises:
             TaskQueueExists: If a queue for *task_id* already exists.
@@ -233,12 +240,12 @@ async def create_or_tap(self, task_id: str) -> EventQueue:
     # ------------------------------------------------------------------
 
     async def _sns_publish(self, message: str) -> None:
-        """Publishes a serialised wire message to the SNS topic.
+        """Publishes a serialized wire message to the SNS topic.
 
         Args:
             message: JSON string in the distributed wire format.
         """
-        async with self._session.client(  # type: ignore[union-attr]
+        async with self._session.client(
             'sns', region_name=self._region_name
         ) as sns:
             await sns.publish(TopicArn=self._topic_arn, Message=message)
@@ -261,7 +268,7 @@ async def _poll_loop(self) -> None:
             self._instance_id,
             self._sqs_queue_url,
         )
-        async with self._session.client(  # type: ignore[union-attr]
+        async with self._session.client(
             'sqs', region_name=self._region_name
         ) as sqs:
             while not self._stop_event.is_set():
@@ -332,27 +339,37 @@ async def _handle_sqs_message(self, sqs_msg: dict[str, Any]) -> None:  # noqa: P
         Args:
             sqs_msg: A single SQS message dictionary from ReceiveMessage.
         """
+        # Use message ID in warnings to avoid logging potentially sensitive
+        # message body content (PII in Task/Message payloads).
+        msg_id = sqs_msg.get('MessageId', '<no-id>')
         body_str = sqs_msg.get('Body', '{}')
+
         try:
             body: dict[str, Any] = json.loads(body_str)
         except json.JSONDecodeError:
-            logger.warning('SQS message body is not valid JSON: %s', body_str)
+            logger.warning(
+                'SQS message body is not valid JSON (msg_id=%s).', msg_id
+            )
             return
 
         # Unwrap SNS notification envelope if present.
         if body.get('Type') == 'Notification':
             inner_str = body.get('Message', '{}')
             try:
-                wire_msg = deserialise_wire_message(inner_str)
+                wire_msg = deserialize_wire_message(inner_str)
             except ValueError:
-                logger.warning('Malformed inner SNS message: %s', inner_str)
+                logger.warning(
+                    'Malformed inner SNS message (msg_id=%s).', msg_id
+                )
                 return
         else:
             # Raw delivery — body itself is the wire message.
             try:
-                wire_msg = deserialise_wire_message(body_str)
+                wire_msg = deserialize_wire_message(body_str)
             except ValueError:
-                logger.warning('Malformed raw SQS message body: %s', body_str)
+                logger.warning(
+                    'Malformed raw SQS message body (msg_id=%s).', msg_id
+                )
                 return
 
         # Deduplicate: ignore messages we published ourselves.