solace-agent-mesh 1.0.9__py3-none-any.whl → 1.3.0__py3-none-any.whl

solace_agent_mesh/gateway/base/component.py

@@ -4,19 +4,16 @@ Base Component class for Gateway implementations in the Solace AI Connector.
 
 import asyncio
 import queue
-import re
-import threading
 import base64
 import uuid
 from datetime import datetime, timezone
 from typing import Any, Dict, Optional, List, Tuple, Union
-from urllib.parse import urlparse, parse_qs
 
-from solace_ai_connector.components.component_base import ComponentBase
 from solace_ai_connector.common.log import log
 from google.adk.artifacts import BaseArtifactService
 
 from ...common.agent_registry import AgentRegistry
+from ...common.sac.sam_component_base import SamComponentBase
 from ...core_a2a.service import CoreA2AService
 from ...agent.adk.services import initialize_artifact_service
 from ...common.services.identity_service import (
@@ -24,8 +21,8 @@ from ...common.services.identity_service import (
     create_identity_service,
 )
 from .task_context import TaskContextManager
-from ...common.types import (
-    Part as A2APart,
+from ...common.a2a.types import ContentPart
+from a2a.types import (
     Message as A2AMessage,
     AgentCard,
     JSONRPCResponse,
@@ -34,21 +31,11 @@ from ...common.types import (
     TaskArtifactUpdateEvent,
     JSONRPCError,
     TextPart,
-    TaskStatus,
-    TaskState,
     FilePart,
-    DataPart,
+    FileWithBytes,
     Artifact as A2AArtifact,
 )
-from ...common.a2a_protocol import (
-    get_gateway_response_topic,
-    get_gateway_response_subscription_topic,
-    get_gateway_status_topic,
-    get_gateway_status_subscription_topic,
-    get_discovery_topic,
-    _topic_matches_subscription,
-    _subscription_to_regex,
-)
+from ...common import a2a
 from ...common.utils import is_text_based_mime_type
 from ...common.utils.embeds import (
     resolve_embeds_in_string,
@@ -65,9 +52,6 @@ from solace_ai_connector.common.event import Event, EventType
 from abc import abstractmethod
 
 from ...common.middleware.registry import MiddlewareRegistry
-from ...agent.utils.artifact_helpers import (
-    load_artifact_content_or_metadata,
-)
 
 info = {
     "class_name": "BaseGatewayComponent",
@@ -88,7 +72,7 @@ info = {
 }
 
 
-class BaseGatewayComponent(ComponentBase):
+class BaseGatewayComponent(SamComponentBase):
     """
     Abstract base class for Gateway components.
 
@@ -111,17 +95,16 @@ class BaseGatewayComponent(ComponentBase):
 
         return super().get_config(key, default)
 
-    def __init__(self, **kwargs: Any):
+    def __init__(self, resolve_artifact_uris_in_gateway: bool = True, **kwargs: Any):
         super().__init__(info, **kwargs)
+        self.resolve_artifact_uris_in_gateway = resolve_artifact_uris_in_gateway
         log.info("%s Initializing Base Gateway Component...", self.log_identifier)
 
         try:
-            self.namespace: str = self.get_config("namespace")
+            # Note: self.namespace and self.max_message_size_bytes are initialized in SamComponentBase
             self.gateway_id: str = self.get_config("gateway_id")
-            if not self.namespace or not self.gateway_id:
-                raise ValueError(
-                    "Namespace and Gateway ID must be configured in the app_config."
-                )
+            if not self.gateway_id:
+                raise ValueError("Gateway ID must be configured in the app_config.")
 
             self.enable_embed_resolution: bool = self.get_config(
                 "enable_embed_resolution", True
@@ -132,8 +115,8 @@ class BaseGatewayComponent(ComponentBase):
             self.gateway_recursive_embed_depth: int = self.get_config(
                 "gateway_recursive_embed_depth"
             )
-            self.gateway_artifact_content_limit_bytes: int = self.get_config(
-                "gateway_artifact_content_limit_bytes"
+            self.artifact_handling_mode: str = self.get_config(
+                "artifact_handling_mode", "embed"
             )
             _ = self.get_config("artifact_service")
 
@@ -162,9 +145,6 @@ class BaseGatewayComponent(ComponentBase):
 
         self.task_context_manager: TaskContextManager = TaskContextManager()
         self.internal_event_queue: queue.Queue = queue.Queue()
-        self.message_processor_thread: Optional[threading.Thread] = None
-        self.async_loop: Optional[asyncio.AbstractEventLoop] = None
-        self.async_thread: Optional[threading.Thread] = None
 
         identity_service_config = self.get_config("identity_service")
         self.identity_service: Optional[BaseIdentityService] = create_identity_service(
@@ -181,40 +161,6 @@ class BaseGatewayComponent(ComponentBase):
             "%s Base Gateway Component initialized successfully.", self.log_identifier
         )
 
-    def publish_a2a_message(
-        self, topic: str, payload: Dict, user_properties: Optional[Dict] = None
-    ) -> None:
-        log.debug(
-            "%s Publishing A2A message to topic: %s via App", self.log_identifier, topic
-        )
-        try:
-            app = self.get_app()
-            if app:
-                app.send_message(
-                    payload=payload, topic=topic, user_properties=user_properties
-                )
-                log.debug(
-                    "%s Successfully published message to %s via App",
-                    self.log_identifier,
-                    topic,
-                )
-            else:
-                log.error(
-                    "%s Cannot publish message: Not running within a SAC App context.",
-                    self.log_identifier,
-                )
-                raise RuntimeError(
-                    "Cannot publish message: Not running within a SAC App context."
-                )
-        except Exception as e:
-            log.exception(
-                "%s Failed to publish A2A message to topic %s via App: %s",
-                self.log_identifier,
-                topic,
-                e,
-            )
-            raise
-
     async def authenticate_and_enrich_user(
         self, external_event_data: Any
     ) -> Optional[Dict[str, Any]]:
@@ -255,7 +201,7 @@ class BaseGatewayComponent(ComponentBase):
     async def submit_a2a_task(
         self,
         target_agent_name: str,
-        a2a_parts: List[A2APart],
+        a2a_parts: List[ContentPart],
         external_request_context: Dict[str, Any],
         user_identity: Any,
         is_streaming: bool = True,
@@ -343,7 +289,7 @@ class BaseGatewayComponent(ComponentBase):
             )
             external_request_context["a2a_session_id"] = a2a_session_id
 
-        a2a_metadata = {}
+        a2a_metadata = {"agent_name": target_agent_name}
         invoked_artifacts = external_request_context.get("invoked_with_artifacts")
         if invoked_artifacts:
             a2a_metadata["invoked_with_artifacts"] = invoked_artifacts
@@ -353,67 +299,46 @@ class BaseGatewayComponent(ComponentBase):
                 len(invoked_artifacts),
             )
 
-        a2a_message = A2AMessage(role="user", parts=a2a_parts, metadata=a2a_metadata)
-        reply_topic_pattern = get_gateway_response_topic(
-            self.namespace, self.gateway_id, "{task_id}"
-        )
-        status_topic_pattern = get_gateway_status_topic(
-            self.namespace, self.gateway_id, "{task_id}"
-        )
+        # This correlation ID is used by the gateway to track the task
+        task_id = f"gdk-task-{uuid.uuid4().hex}"
 
-        task_metadata_override: Dict[str, Any] = {}
-        system_purpose = self.get_config("system_purpose", "")
-        response_format = self.get_config("response_format", "")
+        prepared_a2a_parts = await self._prepare_parts_for_publishing(
+            parts=a2a_parts,
+            user_id=user_id_for_a2a,
+            session_id=a2a_session_id,
+            target_agent_name=target_agent_name,
+        )
 
-        if system_purpose:
-            task_metadata_override["system_purpose"] = system_purpose
-            log.debug("%s Adding system_purpose to task metadata.", log_id_prefix)
-        if response_format:
-            task_metadata_override["response_format"] = response_format
-            log.debug("%s Adding response_format to task metadata.", log_id_prefix)
+        a2a_message = a2a.create_user_message(
+            parts=prepared_a2a_parts,
+            metadata=a2a_metadata,
+            context_id=a2a_session_id,
+        )
 
         if is_streaming:
-            target_topic, payload, user_properties = (
-                self.core_a2a_service.submit_streaming_task(
-                    agent_name=target_agent_name,
-                    a2a_message=a2a_message,
-                    session_id=a2a_session_id,
-                    client_id=self.gateway_id,
-                    reply_to_topic=reply_topic_pattern,
-                    status_to_topic=status_topic_pattern,
-                    user_id=user_id_for_a2a,
-                    a2a_user_config=user_config,
-                    metadata_override=task_metadata_override,
-                )
+            a2a_request = a2a.create_send_streaming_message_request(
+                message=a2a_message, task_id=task_id
             )
         else:
-            target_topic, payload, user_properties = self.core_a2a_service.submit_task(
-                agent_name=target_agent_name,
-                a2a_message=a2a_message,
-                session_id=a2a_session_id,
-                client_id=self.gateway_id,
-                reply_to_topic=reply_topic_pattern,
-                user_id=user_id_for_a2a,
-                a2a_user_config=user_config,
-                metadata_override=task_metadata_override,
+            a2a_request = a2a.create_send_message_request(
+                message=a2a_message, task_id=task_id
             )
 
-        task_id = payload.get("params", {}).get("id")
-        if not task_id:
-            log.error(
-                "%s CoreA2AService did not return a task ID in the payload.",
-                log_id_prefix,
-            )
-            raise ValueError("CoreA2AService did not return a task ID in the payload.")
+        payload = a2a_request.model_dump(by_alias=True, exclude_none=True)
+        target_topic = a2a.get_agent_request_topic(self.namespace, target_agent_name)
 
-        if user_properties is None:
-            user_properties = {}
+        user_properties = {
+            "clientId": self.gateway_id,
+            "userId": user_id_for_a2a,
+        }
+        if user_config:
+            user_properties["a2aUserConfig"] = user_config
 
-        user_properties["replyTo"] = get_gateway_response_topic(
+        user_properties["replyTo"] = a2a.get_gateway_response_topic(
             self.namespace, self.gateway_id, task_id
         )
         if is_streaming:
-            user_properties["a2aStatusTopic"] = get_gateway_status_topic(
+            user_properties["a2aStatusTopic"] = a2a.get_gateway_status_topic(
                 self.namespace, self.gateway_id, task_id
             )
 
@@ -421,7 +346,7 @@ class BaseGatewayComponent(ComponentBase):
         log.info("%s Stored external context for task_id: %s", log_id_prefix, task_id)
 
         self.publish_a2a_message(
-            topic=target_topic, payload=payload, user_properties=user_properties
+            payload=payload, topic=target_topic, user_properties=user_properties
         )
         log.info(
             "%s Submitted A2A task %s to agent %s. Streaming: %s",
@@ -511,15 +436,12 @@ class BaseGatewayComponent(ComponentBase):
                         )
                         continue
                     try:
-                        signal_data_part = DataPart(
-                            data={"type": "agent_status", "text": status_text},
-                            metadata={"source": "agent_progress_update"},
-                        )
-                        signal_a2a_message = A2AMessage(
-                            role="agent", parts=[signal_data_part]
-                        )
-                        signal_task_status = TaskStatus(
-                            state=TaskState.WORKING, message=signal_a2a_message
+                        signal_a2a_message = a2a.create_agent_data_message(
+                            data={
+                                "type": "agent_progress_update",
+                                "status_text": status_text,
+                            },
+                            part_metadata={"source": "gateway_signal"},
                         )
                         a2a_task_id_for_signal = external_request_context.get(
                             "a2a_task_id_for_event", original_rpc_id
@@ -531,10 +453,11 @@ class BaseGatewayComponent(ComponentBase):
                             )
                             continue
 
-                        signal_event = TaskStatusUpdateEvent(
-                            id=a2a_task_id_for_signal,
-                            status=signal_task_status,
-                            final=False,
+                        signal_event = a2a.create_status_update(
+                            task_id=a2a_task_id_for_signal,
+                            context_id=external_request_context.get("a2a_session_id"),
+                            message=signal_a2a_message,
+                            is_final=False,
                         )
                         await self._send_update_to_external(
                             external_request_context=external_request_context,
@@ -556,104 +479,53 @@ class BaseGatewayComponent(ComponentBase):
                         signal_type,
                     )
 
-    async def _resolve_uri_in_file_part(self, part: A2APart):
+    async def _resolve_uri_in_file_part(self, file_part: FilePart):
         """
-        Checks if a part is a FilePart with a resolvable URI and, if so,
-        resolves it and mutates the part in-place.
+        Checks if a FilePart has a resolvable URI and, if so,
+        resolves it and mutates the part in-place by calling the common utility.
         """
-        if not (
-            isinstance(part, FilePart)
-            and part.file
-            and part.file.uri
-            and part.file.uri.startswith("artifact://")
-        ):
-            return
-
-        if not self.shared_artifact_service:
-            log.warning(
-                "%s Cannot resolve artifact URI, shared_artifact_service is not configured.",
-                self.log_identifier,
-            )
-            return
-
-        uri = part.file.uri
-        log_id_prefix = f"{self.log_identifier}[ResolveURI]"
-        try:
-            log.info("%s Found artifact URI to resolve: %s", log_id_prefix, uri)
-            parsed_uri = urlparse(uri)
-            app_name = parsed_uri.netloc
-            path_parts = parsed_uri.path.strip("/").split("/")
-
-            if not app_name or len(path_parts) != 3:
-                raise ValueError(
-                    "Invalid URI structure. Expected artifact://app_name/user_id/session_id/filename"
-                )
-
-            user_id, session_id, filename = path_parts
-            version = int(parse_qs(parsed_uri.query).get("version", [None])[0])
-
-            loaded_artifact = await load_artifact_content_or_metadata(
-                artifact_service=self.shared_artifact_service,
-                app_name=app_name,
-                user_id=user_id,
-                session_id=session_id,
-                filename=filename,
-                version=version,
-                return_raw_bytes=True,
-            )
-
-            if loaded_artifact.get("status") == "success":
-                content_bytes = loaded_artifact.get("raw_bytes")
-                part.file.bytes = base64.b64encode(content_bytes).decode("utf-8")
-                part.file.uri = None
-                log.info(
-                    "%s Successfully resolved and embedded artifact: %s",
-                    log_id_prefix,
-                    uri,
-                )
-            else:
-                log.error(
-                    "%s Failed to resolve artifact URI '%s': %s",
-                    log_id_prefix,
-                    uri,
-                    loaded_artifact.get("message"),
-                )
-        except Exception as e:
-            log.exception(
-                "%s Error resolving artifact URI '%s': %s", log_id_prefix, uri, e
-            )
+        await a2a.resolve_file_part_uri(
+            part=file_part,
+            artifact_service=self.shared_artifact_service,
+            log_identifier=self.log_identifier,
+        )
 
-    async def _resolve_uris_in_parts_list(self, parts: List[A2APart]):
-        """Iterates over a list of A2APart objects and resolves any FilePart URIs."""
+    async def _resolve_uris_in_parts_list(self, parts: List[ContentPart]):
+        """Iterates over a list of part objects and resolves any FilePart URIs."""
         if not parts:
             return
         for part in parts:
-            await self._resolve_uri_in_file_part(part)
+            if isinstance(part, FilePart):
+                await self._resolve_uri_in_file_part(part)
 
     async def _resolve_uris_in_payload(self, parsed_event: Any):
         """
         Dispatcher that calls the appropriate targeted URI resolver based on the
         Pydantic model type of the event.
         """
+        parts_to_resolve: List[ContentPart] = []
         if isinstance(parsed_event, TaskStatusUpdateEvent):
-            if parsed_event.status and parsed_event.status.message:
-                await self._resolve_uris_in_parts_list(
-                    parsed_event.status.message.parts
-                )
+            message = a2a.get_message_from_status_update(parsed_event)
+            if message:
+                parts_to_resolve.extend(a2a.get_parts_from_message(message))
         elif isinstance(parsed_event, TaskArtifactUpdateEvent):
-            if parsed_event.artifact:
-                await self._resolve_uris_in_parts_list(parsed_event.artifact.parts)
+            artifact = a2a.get_artifact_from_artifact_update(parsed_event)
+            if artifact:
+                parts_to_resolve.extend(a2a.get_parts_from_artifact(artifact))
         elif isinstance(parsed_event, Task):
             if parsed_event.status and parsed_event.status.message:
-                await self._resolve_uris_in_parts_list(
-                    parsed_event.status.message.parts
+                parts_to_resolve.extend(
+                    a2a.get_parts_from_message(parsed_event.status.message)
                 )
             if parsed_event.artifacts:
                 for artifact in parsed_event.artifacts:
-                    await self._resolve_uris_in_parts_list(artifact.parts)
+                    parts_to_resolve.extend(a2a.get_parts_from_artifact(artifact))
+
+        if parts_to_resolve:
+            await self._resolve_uris_in_parts_list(parts_to_resolve)
         else:
             log.debug(
-                "%s Payload type '%s' does not support targeted URI resolution. Skipping.",
+                "%s Payload type '%s' did not yield any parts for URI resolution. Skipping.",
                 self.log_identifier,
                 type(parsed_event).__name__,
             )
@@ -673,78 +545,34 @@ class BaseGatewayComponent(ComponentBase):
             )
             return False
 
-    def _extract_task_id_from_topic(
-        self, topic: str, subscription_pattern: str
-    ) -> Optional[str]:
-        """Extracts the task ID from the end of a topic string based on the subscription."""
-        base_regex_str = _subscription_to_regex(subscription_pattern).replace(r".*", "")
-        match = re.match(base_regex_str, topic)
-        if match:
-            task_id_part = topic[match.end() :]
-            task_id = task_id_part.lstrip("/")
-            if task_id:
-                log.debug(
-                    "%s Extracted Task ID '%s' from topic '%s'",
-                    self.log_identifier,
-                    task_id,
-                    topic,
-                )
-                return task_id
-        log.warning(
-            "%s Could not extract Task ID from topic '%s' using pattern '%s'",
-            self.log_identifier,
-            topic,
-            subscription_pattern,
-        )
-        return None
-
-    def _parse_a2a_event_from_rpc_result(
-        self, rpc_result: Dict, expected_task_id: Optional[str]
-    ) -> Optional[Union[Task, TaskStatusUpdateEvent, TaskArtifactUpdateEvent]]:
+    async def _prepare_parts_for_publishing(
+        self,
+        parts: List[ContentPart],
+        user_id: str,
+        session_id: str,
+        target_agent_name: str,
+    ) -> List[ContentPart]:
         """
-        Parses the result field of a JSONRPCResponse into a specific A2A Pydantic model.
-        Verifies task ID if expected_task_id is provided.
+        Prepares message parts for publishing according to the configured artifact_handling_mode
+        by calling the common utility function.
         """
-        if not isinstance(rpc_result, dict):
-            log.error(
-                "%s RPC result is not a dictionary. Cannot parse.", self.log_identifier
-            )
-            return None
-
-        actual_task_id = rpc_result.get("id")
-        if expected_task_id and actual_task_id != expected_task_id:
-            log.error(
-                "%s Task ID mismatch! Expected: %s, Got from payload: %s.",
-                self.log_identifier,
-                expected_task_id,
-                actual_task_id,
-            )
-            return None
-
-        try:
-            if "status" in rpc_result and "final" in rpc_result:
-                return TaskStatusUpdateEvent(**rpc_result)
-            elif "artifact" in rpc_result:
-                return TaskArtifactUpdateEvent(**rpc_result)
-            elif "status" in rpc_result and "sessionId" in rpc_result:
-                return Task(**rpc_result)
-            else:
-                log.warning(
-                    "%s Unknown result structure in RPC response for task %s: %s",
-                    self.log_identifier,
-                    actual_task_id or "unknown",
-                    rpc_result,
+        processed_parts: List[ContentPart] = []
+        for part in parts:
+            if isinstance(part, FilePart):
+                processed_part = await a2a.prepare_file_part_for_publishing(
+                    part=part,
+                    mode=self.artifact_handling_mode,
+                    artifact_service=self.shared_artifact_service,
+                    user_id=user_id,
+                    session_id=session_id,
+                    target_agent_name=target_agent_name,
+                    log_identifier=self.log_identifier,
                 )
-                return None
-        except Exception as e:
-            log.error(
-                "%s Failed to parse RPC result into A2A Pydantic model for task %s: %s. Result: %s",
-                self.log_identifier,
-                actual_task_id or "unknown",
-                e,
-                rpc_result,
-            )
-            return None
+                if processed_part:
+                    processed_parts.append(processed_part)
+            else:
+                processed_parts.append(part)
+        return processed_parts
 
     async def _resolve_embeds_and_handle_signals(
         self,
@@ -792,7 +620,7 @@ class BaseGatewayComponent(ComponentBase):
                 parts_owner = event_with_parts.artifact
 
         if parts_owner and parts_owner.parts:
-            new_parts_for_owner: List[A2APart] = []
+            new_parts: List[ContentPart] = []
             stream_buffer_key = f"{a2a_task_id}_stream_buffer"
             current_buffer = ""
 
@@ -801,13 +629,19 @@ class BaseGatewayComponent(ComponentBase):
                     self.task_context_manager.get_context(stream_buffer_key) or ""
                 )
 
-            for part_obj in parts_owner.parts:
-                if isinstance(part_obj, TextPart) and part_obj.text is not None:
-                    text_to_resolve = part_obj.text
-                    original_part_text = part_obj.text
+            parts: List[ContentPart] = []
+            if isinstance(parts_owner, A2AMessage):
+                parts = a2a.get_parts_from_message(parts_owner)
+            elif isinstance(parts_owner, A2AArtifact):
+                parts = a2a.get_parts_from_artifact(parts_owner)
+
+            for part in parts:
+                if isinstance(part, TextPart) and part.text is not None:
+                    text_to_resolve = part.text
+                    original_part_text = part.text
 
                     if is_streaming_status_update:
-                        current_buffer += part_obj.text
+                        current_buffer += part.text
                         text_to_resolve = current_buffer
 
                     resolved_text, processed_idx, signals = (
@@ -831,7 +665,7 @@ class BaseGatewayComponent(ComponentBase):
                         content_modified_or_signal_handled = True
 
                     if resolved_text is not None:
-                        new_parts_for_owner.append(TextPart(text=resolved_text))
+                        new_parts.append(a2a.create_text_part(text=resolved_text))
                         if is_streaming_status_update:
                             if resolved_text != text_to_resolve[:processed_idx]:
                                 content_modified_or_signal_handled = True
@@ -851,61 +685,73 @@ class BaseGatewayComponent(ComponentBase):
                         )
                         content_modified_or_signal_handled = True
 
-                elif (
-                    isinstance(part_obj, FilePart)
-                    and part_obj.file
-                    and part_obj.file.bytes
-                ):
-                    mime_type = part_obj.file.mimeType or ""
-                    is_container = is_text_based_mime_type(mime_type)
-                    try:
-                        decoded_content_for_check = base64.b64decode(
-                            part_obj.file.bytes
-                        ).decode("utf-8", errors="ignore")
-                        if (
-                            is_container
-                            and EMBED_DELIMITER_OPEN in decoded_content_for_check
-                        ):
-                            original_content = decoded_content_for_check
-                            resolved_content = (
-                                await resolve_embeds_recursively_in_string(
-                                    text=original_content,
-                                    context=embed_eval_context,
-                                    resolver_func=evaluate_embed,
-                                    types_to_resolve=LATE_EMBED_TYPES,
-                                    log_identifier=log_id_prefix,
-                                    config=embed_eval_config,
-                                    max_depth=self.gateway_recursive_embed_depth,
-                                )
-                            )
-                            if resolved_content != original_content:
-                                new_file_content = part_obj.file.model_copy()
-                                new_file_content.bytes = base64.b64encode(
-                                    resolved_content.encode("utf-8")
-                                ).decode("utf-8")
-                                new_parts_for_owner.append(
-                                    FilePart(
-                                        file=new_file_content,
-                                        metadata=part_obj.metadata,
+                elif isinstance(part, FilePart) and part.file:
+                    if isinstance(part.file, FileWithBytes) and part.file.bytes:
+                        mime_type = part.file.mime_type or ""
+                        is_container = is_text_based_mime_type(mime_type)
+                        try:
+                            decoded_content_for_check = base64.b64decode(
+                                part.file.bytes
+                            ).decode("utf-8", errors="ignore")
+                            if (
+                                is_container
+                                and EMBED_DELIMITER_OPEN in decoded_content_for_check
+                            ):
+                                original_content = decoded_content_for_check
+                                resolved_content = (
+                                    await resolve_embeds_recursively_in_string(
+                                        text=original_content,
+                                        context=embed_eval_context,
+                                        resolver_func=evaluate_embed,
+                                        types_to_resolve=LATE_EMBED_TYPES,
+                                        log_identifier=log_id_prefix,
+                                        config=embed_eval_config,
+                                        max_depth=self.gateway_recursive_embed_depth,
                                     )
                                 )
-                                content_modified_or_signal_handled = True
+                                if resolved_content != original_content:
+                                    new_file_content = part.file.model_copy()
+                                    new_file_content.bytes = base64.b64encode(
+                                        resolved_content.encode("utf-8")
+                                    ).decode("utf-8")
+                                    new_parts.append(
+                                        FilePart(
+                                            file=new_file_content,
+                                            metadata=part.metadata,
+                                        )
+                                    )
+                                    content_modified_or_signal_handled = True
+                                else:
+                                    new_parts.append(part)
                             else:
-                                new_parts_for_owner.append(part_obj)
-                        else:
-                            new_parts_for_owner.append(part_obj)
-                    except Exception as e:
-                        log.warning(
-                            "%s Error during recursive FilePart resolution for %s: %s. Using original.",
-                            log_id_prefix,
-                            part_obj.file.name,
-                            e,
-                        )
-                        new_parts_for_owner.append(part_obj)
+                                new_parts.append(part)
+                        except Exception as e:
+                            log.warning(
+                                "%s Error during recursive FilePart resolution for %s: %s. Using original.",
+                                log_id_prefix,
+                                part.file.name,
+                                e,
+                            )
+                            new_parts.append(part)
+                    else:
+                        # This is a FileWithUri or empty FileWithBytes, which we don't process for embeds here.
+                        new_parts.append(part)
                 else:
-                    new_parts_for_owner.append(part_obj)
+                    new_parts.append(part)
 
-            parts_owner.parts = new_parts_for_owner
+            if isinstance(parts_owner, A2AMessage):
+                if isinstance(event_with_parts, TaskStatusUpdateEvent):
+                    event_with_parts.status.message = a2a.update_message_parts(
+                        message=parts_owner, new_parts=new_parts
+                    )
+                elif isinstance(event_with_parts, Task):
+                    event_with_parts.status.message = a2a.update_message_parts(
+                        message=parts_owner, new_parts=new_parts
+                    )
+            elif isinstance(parts_owner, A2AArtifact):
+                event_with_parts.artifact = a2a.update_artifact_parts(
+                    artifact=parts_owner, new_parts=new_parts
+                )
 
             if is_streaming_status_update:
                 self.task_context_manager.store_context(
@@ -945,7 +791,7 @@ class BaseGatewayComponent(ComponentBase):
             elif isinstance(parsed_event, Task):
                 is_finalizing_context_for_embeds = True
 
-            if self.get_config("resolve_artifact_uris_in_gateway", False):
+            if self.resolve_artifact_uris_in_gateway:
                 log.debug(
                     "%s Resolving artifact URIs before sending to external...",
                     log_id_prefix,
@@ -1002,13 +848,11 @@ class BaseGatewayComponent(ComponentBase):
                     log.debug(
                         "%s Resolving embeds in final task response...", log_id_prefix
                     )
-                    combined_text = ""
-                    non_text_parts = []
-                    for part in parsed_event.status.message.parts:
-                        if isinstance(part, TextPart) and part.text:
-                            combined_text += part.text
-                        else:
-                            non_text_parts.append(part)
+                    message = parsed_event.status.message
+                    combined_text = a2a.get_text_from_message(message)
+                    data_parts = a2a.get_data_parts_from_message(message)
+                    file_parts = a2a.get_file_parts_from_message(message)
+                    non_text_parts = data_parts + file_parts
 
                     if combined_text:
                         embed_eval_context = {
@@ -1052,10 +896,15 @@ class BaseGatewayComponent(ComponentBase):
                             )
 
                         new_parts = (
-                            [TextPart(text=resolved_text)] if resolved_text else []
+                            [a2a.create_text_part(text=resolved_text)]
+                            if resolved_text
+                            else []
                         )
                         new_parts.extend(non_text_parts)
-                        parsed_event.status.message.parts = new_parts
+                        parsed_event.status.message = a2a.update_message_parts(
+                            message=parsed_event.status.message,
+                            new_parts=new_parts,
+                        )
                         log.info(
                             "%s Final response text updated with resolved embeds.",
                             log_id_prefix,
@@ -1086,7 +935,7 @@ class BaseGatewayComponent(ComponentBase):
                         },
                     }
                     embed_eval_config = {
-                        "gateway_artifact_content_limit_bytes": self.gateway_artifact_content_limit_bytes,
+                        "gateway_max_artifact_resolve_size_bytes": self.gateway_max_artifact_resolve_size_bytes,
                         "gateway_recursive_embed_depth": self.gateway_recursive_embed_depth,
                     }
                     resolved_remaining_text, _, signals = (
@@ -1106,17 +955,14 @@ class BaseGatewayComponent(ComponentBase):
                         is_finalizing_context=True,
                     )
                     if resolved_remaining_text:
-                        flush_status = TaskStatus(
-                            state=TaskState.WORKING,
-                            message=A2AMessage(
-                                role="agent",
-                                parts=[TextPart(text=resolved_remaining_text)],
-                            ),
+                        flush_message = a2a.create_agent_text_message(
+                            text=resolved_remaining_text
                         )
-                        flush_event = TaskStatusUpdateEvent(
-                            id=a2a_task_id,
-                            status=flush_status,
-                            final=False,
+                        flush_event = a2a.create_status_update(
+                            task_id=a2a_task_id,
+                            context_id=external_request_context.get("a2a_session_id"),
+                            message=flush_message,
+                            is_final=False,
                         )
                         await self._send_update_to_external(
                             external_request_context, flush_event, True
@@ -1157,7 +1003,7 @@ class BaseGatewayComponent(ComponentBase):
         Parses the payload, retrieves context using task_id_from_topic, and dispatches for processing.
         """
         try:
-            rpc_response = JSONRPCResponse(**payload)
+            rpc_response = JSONRPCResponse.model_validate(payload)
         except Exception as e:
             log.error(
                 "%s Failed to parse payload as JSONRPCResponse for topic %s (Task ID from topic: %s): %s. Payload: %s",
@@ -1169,7 +1015,7 @@ class BaseGatewayComponent(ComponentBase):
             )
             return False
 
-        original_rpc_id = str(rpc_response.id)
+        original_rpc_id = str(a2a.get_response_id(rpc_response))
 
         external_request_context = self.task_context_manager.get_context(
             task_id_from_topic
@@ -1190,19 +1036,43 @@ class BaseGatewayComponent(ComponentBase):
         parsed_event_obj: Union[
             Task, TaskStatusUpdateEvent, TaskArtifactUpdateEvent, JSONRPCError, None
         ] = None
-        if rpc_response.error:
-            parsed_event_obj = rpc_response.error
-        elif rpc_response.result:
-            parsed_event_obj = self._parse_a2a_event_from_rpc_result(
-                rpc_response.result, task_id_from_topic
-            )
+        error = a2a.get_response_error(rpc_response)
+        if error:
+            parsed_event_obj = error
+        else:
+            result = a2a.get_response_result(rpc_response)
+            if result:
+                # The result is already a parsed Pydantic model.
+                parsed_event_obj = result
+
+            # Validate task ID match
+            actual_task_id = None
+            if isinstance(parsed_event_obj, Task):
+                actual_task_id = parsed_event_obj.id
+            elif isinstance(
+                parsed_event_obj, (TaskStatusUpdateEvent, TaskArtifactUpdateEvent)
+            ):
+                actual_task_id = parsed_event_obj.task_id
+
+            if (
+                task_id_from_topic
+                and actual_task_id
+                and actual_task_id != task_id_from_topic
+            ):
+                log.error(
+                    "%s Task ID mismatch! Expected: %s, Got from payload: %s.",
+                    self.log_identifier,
+                    task_id_from_topic,
+                    actual_task_id,
+                )
+                parsed_event_obj = None
 
         if not parsed_event_obj:
             log.error(
                 "%s Failed to parse or validate A2A event from RPC result for task %s. Result: %s",
                 self.log_identifier,
                 task_id_from_topic,
-                rpc_response.result,
+                a2a.get_response_result(rpc_response) or "N/A",
             )
             generic_error = JSONRPCError(
                 code=-32000, message="Invalid event structure received from agent."
@@ -1239,9 +1109,36 @@ class BaseGatewayComponent(ComponentBase):
             )
             return False
 
+    async def _async_setup_and_run(self) -> None:
+        """Main async logic for the gateway component."""
+        log.info(
+            "%s Starting _start_listener() to initiate external platform connection.",
+            self.log_identifier,
+        )
+        self._start_listener()
+
+        log.info(
+            "%s Starting _message_processor_loop as an asyncio task.",
+            self.log_identifier,
+        )
+        await self._message_processor_loop()
+
+    def _pre_async_cleanup(self) -> None:
+        """Pre-cleanup actions for the gateway component."""
+        log.info("%s Calling _stop_listener()...", self.log_identifier)
+        self._stop_listener()
+
+        if self.internal_event_queue:
+            log.info(
+                "%s Signaling _message_processor_loop to stop by putting sentinel on queue...",
+                self.log_identifier,
+            )
+            # This unblocks the `self.internal_event_queue.get()` call in the loop
+            self.internal_event_queue.put(None)
+
     async def _message_processor_loop(self):
         log.info("%s Starting message processor loop...", self.log_identifier)
-        loop = asyncio.get_running_loop()
+        loop = self.get_async_loop()
 
         while not self.stop_signal.is_set():
             original_broker_message: Optional[SolaceMessage] = None
@@ -1272,38 +1169,38 @@ class BaseGatewayComponent(ComponentBase):
                     processed_successfully = False
                     continue
 
-                if _topic_matches_subscription(
-                    topic, get_discovery_topic(self.namespace)
+                if a2a.topic_matches_subscription(
+                    topic, a2a.get_discovery_topic(self.namespace)
                 ):
                     processed_successfully = await self._handle_discovery_message(
                         payload
                     )
-                elif _topic_matches_subscription(
+                elif a2a.topic_matches_subscription(
                     topic,
-                    get_gateway_response_subscription_topic(
+                    a2a.get_gateway_response_subscription_topic(
                         self.namespace, self.gateway_id
                     ),
-                ) or _topic_matches_subscription(
+                ) or a2a.topic_matches_subscription(
                     topic,
-                    get_gateway_status_subscription_topic(
+                    a2a.get_gateway_status_subscription_topic(
                         self.namespace, self.gateway_id
                     ),
                 ):
                     task_id_from_topic: Optional[str] = None
-                    response_sub = get_gateway_response_subscription_topic(
+                    response_sub = a2a.get_gateway_response_subscription_topic(
                         self.namespace, self.gateway_id
                     )
-                    status_sub = get_gateway_status_subscription_topic(
+                    status_sub = a2a.get_gateway_status_subscription_topic(
                         self.namespace, self.gateway_id
                     )
 
-                    if _topic_matches_subscription(topic, response_sub):
-                        task_id_from_topic = self._extract_task_id_from_topic(
-                            topic, response_sub
+                    if a2a.topic_matches_subscription(topic, response_sub):
+                        task_id_from_topic = a2a.extract_task_id_from_topic(
+                            topic, response_sub, self.log_identifier
                         )
-                    elif _topic_matches_subscription(topic, status_sub):
-                        task_id_from_topic = self._extract_task_id_from_topic(
-                            topic, status_sub
+                    elif a2a.topic_matches_subscription(topic, status_sub):
+                        task_id_from_topic = a2a.extract_task_id_from_topic(
+                            topic, status_sub, self.log_identifier
                         )
 
                     if task_id_from_topic:
@@ -1355,145 +1252,6 @@ class BaseGatewayComponent(ComponentBase):
 
         log.info("%s Message processor loop finished.", self.log_identifier)
 
-    def _run_async_operations(self):
-        log.info(
-            "%s Initializing asyncio event loop in dedicated thread...",
-            self.log_identifier,
-        )
-        self.async_loop = asyncio.new_event_loop()
-        asyncio.set_event_loop(self.async_loop)
-
-        processor_task = None
-        try:
-            log.info(
-                "%s Starting _message_processor_loop as an asyncio task.",
-                self.log_identifier,
-            )
-            processor_task = self.async_loop.create_task(self._message_processor_loop())
-
-            log.info(
-                "%s Calling _start_listener() to initiate external platform connection.",
-                self.log_identifier,
-            )
-            self._start_listener()
-
-            log.info(
-                "%s Running asyncio event loop forever (or until stop_signal).",
-                self.log_identifier,
-            )
-            self.async_loop.run_forever()
-
-        except Exception as e:
-            log.exception(
-                "%s Unhandled exception in _run_async_operations: %s",
-                self.log_identifier,
-                e,
-            )
-            self.stop_signal.set()
-        finally:
-            if processor_task and not processor_task.done():
-                log.info(
-                    "%s Cancelling _message_processor_loop task.", self.log_identifier
-                )
-                processor_task.cancel()
-                try:
-                    self.async_loop.run_until_complete(
-                        asyncio.gather(processor_task, return_exceptions=True)
-                    )
-                except RuntimeError as loop_err:
-                    log.warning(
-                        "%s Error awaiting processor task during cleanup (loop closed?): %s",
-                        self.log_identifier,
-                        loop_err,
-                    )
-
-            if self.async_loop.is_running():
-                log.info(
-                    "%s Stopping asyncio event loop from _run_async_operations finally block.",
-                    self.log_identifier,
-                )
-                self.async_loop.stop()
-            log.info(
-                "%s Async operations loop finished in dedicated thread.",
-                self.log_identifier,
-            )
-
-    def run(self):
-        log.info("%s Starting BaseGatewayComponent run method.", self.log_identifier)
-        if not self.async_thread or not self.async_thread.is_alive():
-            self.async_thread = threading.Thread(
-                target=self._run_async_operations,
-                name=f"{self.name}_AsyncOpsThread",
-                daemon=True,
-            )
-            self.async_thread.start()
-            log.info("%s Async operations thread started.", self.log_identifier)
-        else:
-            log.warning(
-                "%s Async operations thread already running.", self.log_identifier
-            )
-
-        super().run()
-        log.info("%s BaseGatewayComponent run method finished.", self.log_identifier)
-
-    def cleanup(self):
-        log.info("%s Starting cleanup for BaseGatewayComponent...", self.log_identifier)
-
-        log.info("%s Calling _stop_listener()...", self.log_identifier)
-        try:
-            if (
-                self.async_loop
-                and not self.async_loop.is_running()
-                and self.async_thread
-                and self.async_thread.is_alive()
-            ):
-                log.warning(
-                    "%s Async loop not running during cleanup, _stop_listener might face issues if it needs the loop.",
-                    self.log_identifier,
-                )
-            self._stop_listener()
-        except Exception as e:
-            log.exception(
-                "%s Error during _stop_listener(): %s", self.log_identifier, e
-            )
-
-        if self.internal_event_queue:
-            log.info(
-                "%s Signaling _message_processor_loop to stop...", self.log_identifier
-            )
-            self.internal_event_queue.put(None)
-
-        if self.async_loop and self.async_loop.is_running():
-            log.info("%s Requesting asyncio loop to stop...", self.log_identifier)
-            self.async_loop.call_soon_threadsafe(self.async_loop.stop)
-
-        if self.async_thread and self.async_thread.is_alive():
-            log.info(
-                "%s Joining async operations thread (timeout 10s)...",
-                self.log_identifier,
-            )
-            self.async_thread.join(timeout=10)
-            if self.async_thread.is_alive():
-                log.warning(
-                    "%s Async operations thread did not join cleanly.",
-                    self.log_identifier,
-                )
-
-        if self.async_loop and not self.async_loop.is_closed():
-            if self.async_loop.is_running():
-                self.async_loop.call_soon_threadsafe(self.async_loop.stop)
-            log.info(
-                "%s Closing asyncio event loop (if not already closed by its thread).",
-                self.log_identifier,
-            )
-            if not self.async_loop.is_running():
-                self.async_loop.close()
-            else:
-                self.async_loop.call_soon_threadsafe(self.async_loop.close)
-
-        super().cleanup()
-        log.info("%s BaseGatewayComponent cleanup finished.", self.log_identifier)
-
     @abstractmethod
     async def _extract_initial_claims(
         self, external_event_data: Any
@@ -1514,17 +1272,30 @@ class BaseGatewayComponent(ComponentBase):
         pass
 
     @abstractmethod
-    def _start_listener(self) -> None:
+    async def _translate_external_input(
+        self, external_event: Any
+    ) -> Tuple[str, List[ContentPart], Dict[str, Any]]:
+        """
+        Translates raw platform-specific event data into A2A task parameters.
+
+        Args:
+            external_event: Raw event data from the external platform
+                            (e.g., FastAPIRequest, Slack event dictionary).
+
+        Returns:
+            A tuple containing:
+            - target_agent_name (str): The name of the A2A agent to target.
+            - a2a_parts (List[ContentPart]): A list of A2A Part objects.
+            - external_request_context (Dict[str, Any]): Context for TaskContextManager.
+        """
         pass
 
     @abstractmethod
-    def _stop_listener(self) -> None:
+    def _start_listener(self) -> None:
         pass
 
     @abstractmethod
-    def _translate_external_input(
-        self, external_event: Any
-    ) -> Tuple[str, List[A2APart], Dict[str, Any]]:
+    def _stop_listener(self) -> None:
         pass
 
     @abstractmethod