solace-agent-mesh 1.6.1__py3-none-any.whl → 1.13.2__py3-none-any.whl

solace_agent_mesh/agent/proxies/a2a/component.py

@@ -27,6 +27,8 @@ from a2a.types import (
     AgentCard,
     Artifact,
     CancelTaskRequest,
+    DataPart,
+    InternalError,
     Message,
     SendMessageRequest,
     SendStreamingMessageRequest,
@@ -44,6 +46,8 @@ from solace_ai_connector.common.log import log
 from datetime import datetime, timezone
 
 from ....common import a2a
+from ....common.oauth import OAuth2Client, validate_https_url
+from ....common.data_parts import AgentProgressUpdateData
 from ....agent.utils.artifact_helpers import format_artifact_uri
 from ..base.component import BaseProxyComponent
 
@@ -80,11 +84,29 @@ class A2AProxyComponent(BaseProxyComponent):
         # when multiple concurrent requests target the same agent
         self._oauth_token_cache: OAuth2TokenCache = OAuth2TokenCache()
 
+        # OAuth 2.0 client for protocol operations (no retry for A2A)
+        self._oauth_client = OAuth2Client()
+
         # Index agent configs by name for O(1) lookup (performance optimization)
         self._agent_config_by_name: Dict[str, Dict[str, Any]] = {
             agent["name"]: agent for agent in self.proxied_agents_config
         }
 
+        # NEW: OAuth 2.0 authorization code support (enterprise feature)
+        # Stores paused tasks waiting for user authorization
+        self._paused_a2a_oauth2_tasks: Dict[str, Dict[str, Any]] = {}
+        # Caches CredentialManagerWithDiscovery instances per agent
+        self._a2a_oauth2_credential_managers: Dict[str, Any] = {}
+
+        # NEW: Initialize enterprise features for OAuth2 support
+        try:
+            from solace_agent_mesh_enterprise.init_enterprise_component import (
+                init_enterprise_proxy_features
+            )
+            init_enterprise_proxy_features(self)
+        except ImportError:
+            pass  # Enterprise not installed
+
         # OAuth 2.0 configuration is now validated by Pydantic models at app initialization
         # No need for separate _validate_oauth_config() method
 
@@ -100,11 +122,82 @@ class A2AProxyComponent(BaseProxyComponent):
         """
         return self._agent_config_by_name.get(agent_name)
 
+    async def _build_headers(
+        self,
+        agent_name: str,
+        agent_config: Dict[str, Any],
+        custom_headers_key: str,
+        use_auth: bool = True,
+    ) -> Dict[str, str]:
+        """
+        Builds HTTP headers for requests, applying authentication and custom headers.
+
+        Args:
+            agent_name: The name of the agent.
+            agent_config: The agent configuration dictionary.
+            custom_headers_key: Key to look up custom headers in config ('agent_card_headers' or 'task_headers').
+            use_auth: Whether to apply authentication headers.
+
+        Returns:
+            Dictionary of HTTP headers. Custom headers are applied after auth headers.
+            Note: For task invocations, the A2A SDK's AuthInterceptor may further
+            modify authentication headers after these are set.
+        """
+        headers: Dict[str, str] = {}
+
+        # Step 1: Add authentication headers if requested
+        if use_auth:
+            auth_config = agent_config.get("authentication")
+            if auth_config:
+                auth_type = auth_config.get("type")
+
+                # Determine auth type (with backward compatibility)
+                if not auth_type:
+                    scheme = auth_config.get("scheme", "bearer")
+                    auth_type = "static_bearer" if scheme == "bearer" else "static_apikey"
+
+                # Apply authentication based on type
+                if auth_type == "static_bearer":
+                    token = auth_config.get("token")
+                    if token:
+                        headers["Authorization"] = f"Bearer {token}"
+                elif auth_type == "static_apikey":
+                    token = auth_config.get("token")
+                    if token:
+                        headers["X-API-Key"] = token
+                elif auth_type == "oauth2_client_credentials":
+                    # Fetch OAuth token
+                    try:
+                        access_token = await self._fetch_oauth2_token(agent_name, auth_config)
+                        headers["Authorization"] = f"Bearer {access_token}"
+                    except Exception as e:
+                        log.error(
+                            "%s Failed to obtain OAuth 2.0 token for headers: %s",
+                            self.log_identifier,
+                            e,
+                        )
+                        # Continue without auth header - let the request fail downstream
+
+        # Step 2: Add custom headers (these override auth headers)
+        custom_headers_list = agent_config.get(custom_headers_key)
+        if custom_headers_list:
+            for header_config in custom_headers_list:
+                header_name = header_config.get("name")
+                header_value = header_config.get("value")
+                if header_name and header_value:
+                    headers[header_name] = header_value
+
+        return headers
+
     async def _fetch_agent_card(
         self, agent_config: Dict[str, Any]
     ) -> Optional[AgentCard]:
         """
         Fetches the AgentCard from a downstream A2A agent via HTTPS.
+
+        Applies authentication and custom headers based on configuration:
+        - If use_auth_for_agent_card=true, applies the configured authentication
+        - Custom agent_card_headers override authentication headers
         """
         agent_name = agent_config.get("name")
         agent_url = agent_config.get("url")
@@ -116,8 +209,27 @@ class A2AProxyComponent(BaseProxyComponent):
             return None
 
         try:
+            # Build headers based on configuration
+            use_auth = agent_config.get("use_auth_for_agent_card", False)
+            headers = await self._build_headers(
+                agent_name=agent_name,
+                agent_config=agent_config,
+                custom_headers_key="agent_card_headers",
+                use_auth=use_auth,
+            )
+
+            if headers:
+                log.debug(
+                    "%s Fetching agent card with %d custom header(s) (auth=%s)",
+                    log_identifier,
+                    len(headers),
+                    use_auth,
+                )
+            else:
+                log.debug("%s Fetching agent card without authentication", log_identifier)
+
             log.info("%s Fetching agent card from %s", log_identifier, agent_url)
-            async with httpx.AsyncClient() as client:
+            async with httpx.AsyncClient(headers=headers) as client:
                 resolver = A2ACardResolver(httpx_client=client, base_url=agent_url)
                 agent_card = await resolver.get_agent_card()
                 return agent_card
@@ -152,6 +264,9 @@ class A2AProxyComponent(BaseProxyComponent):
             f"{self.log_identifier}[ForwardRequest:{task_context.task_id}:{agent_name}]"
         )
 
+        # Store original request for potential resumption (OAuth2 authorization code flow)
+        task_context.original_request = request
+
         # Step 1: Initialize retry counter
         # Why only retry once: Prevents infinite loops on persistent auth failures.
         # First 401 may be due to token expiration between cache check and request;
@@ -159,7 +274,57 @@ class A2AProxyComponent(BaseProxyComponent):
         max_auth_retries: int = 1
         auth_retry_count: int = 0
 
-        # Step 2: Create while loop for retry logic
+        # Step 2: Check for OAuth2 authorization code flow
+        # This auth type requires user interaction and can pause the task,
+        # so we check it before attempting normal request flow
+        agent_config = self._get_agent_config(agent_name)
+        auth_config = agent_config.get("authentication") if agent_config else None
+        auth_type = auth_config.get("type") if auth_config else None
+
+        if auth_type == "oauth2_authorization_code":
+            try:
+                from solace_agent_mesh_enterprise.auth.a2a import (
+                    check_authorization_required,
+                    request_authorization,
+                )
+
+                # Check if user authorization is needed
+                needs_auth = await check_authorization_required(
+                    component=self,
+                    agent_name=agent_name,
+                    task_context=task_context,
+                )
+
+                if needs_auth:
+                    # Pause task and request authorization
+                    log.info(
+                        "%s User authorization required for agent '%s'. Pausing task.",
+                        log_identifier,
+                        agent_name,
+                    )
+                    await request_authorization(
+                        component=self,
+                        agent_name=agent_name,
+                        task_context=task_context,
+                    )
+                    return  # Exit - task paused, will resume after OAuth callback
+
+            except ImportError:
+                log.error(
+                    "%s Agent '%s' requires OAuth2 authorization code flow, "
+                    "but solace-agent-mesh-enterprise is not installed.",
+                    log_identifier,
+                    agent_name,
+                )
+                raise ValueError(
+                    f"Agent '{agent_name}' requires OAuth2 authorization code flow, "
+                    "but solace-agent-mesh-enterprise is not installed."
+                )
+
+        # Step 3: Normal request flow for all other auth types
+        # (static_bearer, static_apikey, oauth2_client_credentials, or authorized oauth2_authorization_code)
+        received_final_task = False
+
         while auth_retry_count <= max_auth_retries:
             try:
                 # Get or create A2AClient
@@ -184,6 +349,24 @@ class A2AProxyComponent(BaseProxyComponent):
                     # Extract the Message from the request params
                     message_to_send = request.params.message
 
+                    # Check if this is a RUN_BASED request by inspecting message metadata
+                    # For RUN_BASED requests, omit context_id to indicate independent tasks
+                    if message_to_send.metadata:
+                        session_behavior = message_to_send.metadata.get("sessionBehavior")
+                        if session_behavior:
+                            session_behavior = str(session_behavior).upper()
+                            if session_behavior == "RUN_BASED" and message_to_send.context_id:
+                                # For RUN_BASED requests, omit context_id entirely
+                                # Each request is independent with no logical grouping
+                                log.debug(
+                                    "%s RUN_BASED request detected. Omitting context_id "
+                                    "(independent task)",
+                                    log_identifier,
+                                )
+                                message_to_send = message_to_send.model_copy(
+                                    update={"context_id": None}
+                                )
+
                     # WORKAROUND: The A2A SDK has a bug in ClientTaskManager that breaks streaming.
                     # For streaming requests, we bypass the Client.send_message() method and call
                     # the transport directly to avoid the buggy ClientTaskManager.
@@ -202,6 +385,10 @@ class A2AProxyComponent(BaseProxyComponent):
                             await self._process_downstream_response(
                                 raw_event, task_context, client, agent_name
                             )
+                            # Check if this is a final task
+                            if isinstance(raw_event, Task) and raw_event.status:
+                                if raw_event.status.state in [TaskState.completed, TaskState.failed, TaskState.canceled]:
+                                    received_final_task = True
                     else:
                         # Non-streaming: use normal client method (works fine)
                         log.debug(
@@ -214,21 +401,51 @@ class A2AProxyComponent(BaseProxyComponent):
                             await self._process_downstream_response(
                                 event, task_context, client, agent_name
                             )
+                            # Check if this is a final task (event is tuple of (Task, Optional[UpdateEvent]))
+                            if isinstance(event, tuple) and len(event) > 0:
+                                task = event[0]
+                                if isinstance(task, Task) and task.status:
+                                    if task.status.state in [TaskState.completed, TaskState.failed, TaskState.canceled]:
+                                        received_final_task = True
                 elif isinstance(request, CancelTaskRequest):
-                    # Forward cancel request to downstream agent
+                    # Forward cancel request to downstream agent using the downstream task ID
+                    # The request.params.id contains SAM's task ID, but we need to send
+                    # the downstream agent's task ID for the cancel to work
+
+                    if not task_context.downstream_task_id:
+                        log.error(
+                            "%s Cannot forward cancel request: downstream task ID not yet captured for SAM task %s",
+                            log_identifier,
+                            task_context.task_id,
+                        )
+                        # Create an error response
+                        from a2a.types import InvalidRequestError
+                        error = InvalidRequestError(
+                            message=f"Cannot cancel task {task_context.task_id}: downstream task ID not available",
+                            data={"taskId": task_context.task_id}
+                        )
+                        # Publish error response
+                        await self._publish_error_response(error, task_context.a2a_context)
+                        break
+
                     log.info(
-                        "%s Forwarding cancel request for task %s to downstream agent.",
+                        "%s Forwarding cancel request for task %s (SAM ID: %s, downstream ID: %s) to downstream agent.",
                         log_identifier,
-                        request.params.id,
+                        task_context.downstream_task_id,
+                        task_context.task_id,
+                        task_context.downstream_task_id,
                     )
-                    # Use the modern client's cancel_task method
-                    # Note: Pass the entire params object (TaskIdParams) instead of just the id string
-                    # to work around an SDK bug where it doesn't properly handle string inputs
+
+                    # Create new params with the downstream task ID
+                    from a2a.types import TaskIdParams
+                    downstream_params = TaskIdParams(id=task_context.downstream_task_id)
+
+                    # Use the modern client's cancel_task method with the downstream task ID
                     result = await client.cancel_task(
-                        request.params, context=call_context
+                        downstream_params, context=call_context
                     )
                     # Publish the canceled task response
-                    await self._publish_final_response(result, task_context.a2a_context)
+                    await self._publish_task_response(result, task_context.a2a_context)
                 else:
                     log.warning(
                         "%s Unhandled request type for forwarding: %s",
@@ -262,6 +479,22 @@ class A2AProxyComponent(BaseProxyComponent):
 
             except A2AClientJSONRPCError as e:
                 # Handle JSON-RPC protocol errors
+
+                # Special case: Task already in terminal state (canceled/completed/failed)
+                # This is not a fatal error - the cancellation is effectively a no-op
+                if (e.error.code == -32002 and
+                    "cannot be canceled" in e.error.message.lower() and
+                    isinstance(request, CancelTaskRequest)):
+                    log.warning(
+                        "%s Task %s is already in terminal state: %s. Treating as successful cancellation.",
+                        log_identifier,
+                        task_context.downstream_task_id,
+                        e.error.message,
+                    )
+                    # Task is already done - return success (cancellation is effectively complete)
+                    # We don't need to publish a response because the task already sent its final response
+                    break
+
                 log.error(
                     "%s JSON-RPC error from agent '%s': %s",
                     log_identifier,
@@ -335,6 +568,20 @@ class A2AProxyComponent(BaseProxyComponent):
                 # and publish an error response.
                 raise
 
+        # After retry loop completes - check if we received a final task
+        # This detects cases where the stream closes without error but also without final response
+        if not received_final_task and isinstance(request, (SendStreamingMessageRequest, SendMessageRequest)):
+            from ....common.a2a import create_error_response
+
+            error_msg = f"Remote agent '{agent_name}' disconnected without completing the task. Agent may have crashed."
+            log.error("%s %s", log_identifier, error_msg)
+
+            error = InternalError(message=error_msg, data={"agent_name": agent_name})
+            reply_topic = task_context.a2a_context.get("reply_to_topic")
+            if reply_topic:
+                response = create_error_response(error=error, request_id=task_context.a2a_context.get("jsonrpc_request_id"))
+                self._publish_a2a_message(response.model_dump(exclude_none=True), reply_topic)
+
     async def _handle_auth_error(
         self, agent_name: str, task_context: ProxyTaskContext
     ) -> bool:
@@ -486,18 +733,8 @@ class A2AProxyComponent(BaseProxyComponent):
                 "'token_url', 'client_id', and 'client_secret'."
             )
 
-        # SECURITY: Enforce HTTPS for token URL
-        parsed_url = urlparse(token_url)
-        if parsed_url.scheme != "https":
-            log.error(
-                "%s OAuth 2.0 token_url must use HTTPS for security. Got scheme: '%s'",
-                log_identifier,
-                parsed_url.scheme,
-            )
-            raise ValueError(
-                f"{log_identifier} OAuth 2.0 token_url must use HTTPS for security. "
-                f"Got: {parsed_url.scheme}://"
-            )
+        # SECURITY: Enforce HTTPS for token URL using common utility
+        validate_https_url(token_url)
 
         # Step 3: Extract optional parameters
         scope = auth_config.get("scope", "")
@@ -515,49 +752,32 @@ class A2AProxyComponent(BaseProxyComponent):
         )
 
         try:
-            # Step 5: Create temporary httpx client with 30-second timeout
-            async with httpx.AsyncClient(timeout=30.0) as client:
-                # Step 6: Execute POST request
-                # SECURITY: client_secret is sent in POST body (not logged or in URL)
-                response = await client.post(
-                    token_url,
-                    data={
-                        "grant_type": "client_credentials",
-                        "client_id": client_id,
-                        "client_secret": client_secret,
-                        "scope": scope,
-                    },
-                    headers={
-                        "Content-Type": "application/x-www-form-urlencoded",
-                        "Accept": "application/json",
-                    },
-                )
-                response.raise_for_status()
-
-                # Step 7: Parse response
-                token_response = response.json()
-                access_token = token_response.get("access_token")
+            # Step 5: Fetch token using common OAuth client (no retry for A2A)
+            token_data = await self._oauth_client.fetch_client_credentials_token(
+                token_url=token_url,
+                client_id=client_id,
+                client_secret=client_secret,
+                scope=scope,
+                verify=True,
+                timeout=30.0,
+            )
 
-                if not access_token:
-                    raise ValueError(
-                        f"{log_identifier} Token response missing 'access_token' field. "
-                        f"Response keys: {list(token_response.keys())}"
-                    )
+            access_token = token_data["access_token"]
 
-                # Step 8: Cache the token
-                await self._oauth_token_cache.set(
-                    agent_name, access_token, cache_duration
-                )
+            # Step 6: Cache the token
+            await self._oauth_token_cache.set(
+                agent_name, access_token, cache_duration
+            )
 
-                # Step 9: Log success
-                log.info(
-                    "%s Successfully obtained OAuth 2.0 token (cached for %ds)",
-                    log_identifier,
-                    cache_duration,
-                )
+            # Step 7: Log success
+            log.info(
+                "%s Successfully obtained OAuth 2.0 token (cached for %ds)",
+                log_identifier,
+                cache_duration,
+            )
 
-                # Step 10: Return access token
-                return access_token
+            # Step 8: Return access token
+            return access_token
 
         except httpx.HTTPStatusError as e:
             log.error(
@@ -597,6 +817,7 @@ class A2AProxyComponent(BaseProxyComponent):
         - static_bearer: Static bearer token authentication
         - static_apikey: Static API key authentication
         - oauth2_client_credentials: OAuth 2.0 Client Credentials flow with automatic token refresh
+        - oauth2_authorization_code: OAuth 2.0 Authorization Code flow
 
         For backward compatibility, legacy configurations without a 'type' field
         will have their type inferred from the 'scheme' field.
@@ -622,6 +843,20 @@ class A2AProxyComponent(BaseProxyComponent):
             log.error(f"Agent card not found for '{agent_name}' in registry.")
             return None
 
+        # Check if we should use the configured URL or the agent card URL
+        use_agent_card_url = agent_config.get("use_agent_card_url", True)
+        if not use_agent_card_url:
+            # Override the agent card URL with the configured URL
+            configured_url = agent_config.get("url")
+            log.info(
+                "%s Overriding agent card URL with configured URL for agent '%s': %s",
+                self.log_identifier,
+                agent_name,
+                configured_url,
+            )
+            # Create a modified copy of the agent card with the configured URL
+            agent_card = agent_card.model_copy(update={"url": configured_url})
+
         # Resolve timeout - ensure we always have a valid timeout value
         default_timeout = self.get_config("default_request_timeout_seconds", 300)
         agent_timeout = agent_config.get("request_timeout_seconds")
@@ -629,7 +864,18 @@ class A2AProxyComponent(BaseProxyComponent):
             agent_timeout = default_timeout
         log.info("Using timeout of %ss for agent '%s'.", agent_timeout, agent_name)
 
-        # Create a new httpx client with the specific timeout for this agent
+        # Build custom headers for task invocation
+        # Note: We build headers here but apply them via the httpx client
+        # The A2A SDK's AuthInterceptor will add auth headers via the credential store,
+        # but we want custom headers to override those, so we apply them at the client level
+        task_headers = await self._build_headers(
+            agent_name=agent_name,
+            agent_config=agent_config,
+            custom_headers_key="task_headers",
+            use_auth=False,  # Auth will be handled by AuthInterceptor below
+        )
+
+        # Create a new httpx client with the specific timeout and custom headers for this agent
         # httpx.Timeout requires explicit values for connect, read, write, and pool
         httpx_client_for_agent = httpx.AsyncClient(
             timeout=httpx.Timeout(
@@ -637,9 +883,18 @@ class A2AProxyComponent(BaseProxyComponent):
                 read=agent_timeout,
                 write=agent_timeout,
                 pool=agent_timeout,
-            )
+            ),
+            headers=task_headers if task_headers else None,
         )
 
+        if task_headers:
+            log.info(
+                "%s Applied %d custom task header(s) for agent '%s'",
+                self.log_identifier,
+                len(task_headers),
+                agent_name,
+            )
+
         # Setup authentication if configured
         auth_config = agent_config.get("authentication")
         if auth_config:
@@ -712,10 +967,63 @@ class A2AProxyComponent(BaseProxyComponent):
                     )
                     raise
 
+            elif auth_type == "oauth2_authorization_code":
+                # NEW: OAuth 2.0 Authorization Code Flow (enterprise feature)
+                # At this point, user has already authorized (checked in _forward_request)
+                # We just need to get the access token from enterprise helpers
+                try:
+                    from solace_agent_mesh_enterprise.auth.a2a import get_access_token
+
+                    # Get access token (enterprise handles refresh if needed)
+                    access_token = await get_access_token(
+                        component=self,
+                        agent_name=agent_name,
+                        task_context=task_context,
+                    )
+
+                    if not access_token:
+                        raise ValueError(
+                            f"No OAuth2 credential found for agent '{agent_name}'. "
+                            "User authorization should have completed in _forward_request()."
+                        )
+
+                    # Find the OAuth2 authorization code scheme name from agent card
+                    oauth_scheme_name = None
+                    if agent_card and agent_card.security_schemes:
+                        for scheme_name, scheme_wrapper in agent_card.security_schemes.items():
+                            scheme = scheme_wrapper.root
+                            if (hasattr(scheme, 'type') and scheme.type.lower() == 'oauth2' and
+                                    hasattr(scheme, 'flows') and scheme.flows and
+                                    scheme.flows.authorization_code):
+                                oauth_scheme_name = scheme_name
+                                break
+
+                    # Fallback if not found
+                    if not oauth_scheme_name:
+                        oauth_scheme_name = "oauth2_authorization_code"
+                        log.warning(
+                            "%s No OAuth2 authorization code scheme found in agent card, using default name",
+                            self.log_identifier
+                        )
+
+                    # Store in credential store for AuthInterceptor
+                    await self._credential_store.set_credentials(
+                        session_id, oauth_scheme_name, access_token
+                    )
+
+                except ImportError:
+                    log.error(
+                        "%s OAuth2 authorization code requires solace-agent-mesh-enterprise package",
+                        self.log_identifier,
+                    )
+                    raise ValueError(
+                        "OAuth2 authorization code requires solace-agent-mesh-enterprise package"
+                    )
+
             else:
                 raise ValueError(
                     f"Unsupported authentication type '{auth_type}' for agent '{agent_name}'. "
-                    f"Supported types: static_bearer, static_apikey, oauth2_client_credentials."
+                    f"Supported types: static_bearer, static_apikey, oauth2_client_credentials, oauth2_authorization_code."
                 )
 
         # Create ClientConfig for the modern client
@@ -981,6 +1289,79 @@ class A2AProxyComponent(BaseProxyComponent):
                 type(event_payload).__name__,
             )
 
+        # Convert TextParts to AgentProgressUpdateData for intermediate status updates if configured
+        # Only convert non-final status updates; final status updates are used to construct the final Task
+        if isinstance(event_payload, TaskStatusUpdateEvent) and not event_payload.final:
+            agent_config = self._get_agent_config(agent_name)
+            convert_progress = agent_config.get("convert_progress_updates", True) if agent_config else True
+
+            # DEBUG: Log config lookup results
+            log.info(
+                "%s DEBUG convert_progress_updates: agent_name='%s', agent_config_name='%s', agent_config_keys=%s, convert_progress_value=%s, convert_progress=%s",
+                log_identifier,
+                agent_name,
+                agent_config.get('name') if agent_config else None,
+                list(agent_config.keys()) if agent_config else None,
+                agent_config.get("convert_progress_updates") if agent_config else None,
+                convert_progress,
+            )
+
+            if convert_progress and event_payload.status and event_payload.status.message:
+                message = event_payload.status.message
+                original_parts = a2a.get_parts_from_message(message)
+
+                if original_parts:
+                    converted_parts = []
+                    text_parts_converted = 0
+
+                    for part in original_parts:
+                        if isinstance(part, TextPart) and part.text:
+                            # Convert TextPart to DataPart with AgentProgressUpdateData
+                            progress_data = AgentProgressUpdateData(
+                                type="agent_progress_update",
+                                status_text=part.text
+                            )
+                            data_part = DataPart(
+                                kind="data",
+                                data=progress_data.model_dump(),
+                                metadata=part.metadata
+                            )
+                            converted_parts.append(data_part)
+                            text_parts_converted += 1
+                        else:
+                            # Keep non-text parts as-is
+                            converted_parts.append(part)
+
+                    if text_parts_converted > 0:
+                        # Update the message with converted parts
+                        event_payload.status.message = a2a.update_message_parts(
+                            message, converted_parts
+                        )
+                        log.debug(
+                            "%s Converted %d TextPart(s) to AgentProgressUpdateData in status update",
+                            log_identifier,
+                            text_parts_converted,
+                        )
+
+        # Capture the downstream task ID before we replace it
+        # This is needed for forwarding cancellation requests to the downstream agent
+        downstream_id = None
+        if hasattr(event_payload, "task_id") and event_payload.task_id:
+            downstream_id = event_payload.task_id
+        elif hasattr(event_payload, "id") and event_payload.id:
+            downstream_id = event_payload.id
+
+        # Store the downstream task ID in the context if we haven't already
+        if downstream_id and not task_context.downstream_task_id:
+            task_context.downstream_task_id = downstream_id
+            log.debug(
+                "%s Captured downstream task ID: %s (SAM task ID: %s)",
+                log_identifier,
+                downstream_id,
+                task_context.task_id,
+            )
+
+        # Replace the downstream task ID with SAM's task ID for upstream responses
         original_task_id = task_context.task_id
         if hasattr(event_payload, "task_id") and event_payload.task_id:
             event_payload.task_id = original_task_id
@@ -1031,20 +1412,123 @@ class A2AProxyComponent(BaseProxyComponent):
                         Part(root=summary_message_part)
                     )
 
-        if isinstance(event_payload, (Task, TaskStatusUpdateEvent)):
-            if isinstance(event_payload, Task):
-                await self._publish_final_response(
-                    event_payload, task_context.a2a_context
+        # Convert text-only TaskArtifactUpdateEvents to TaskStatusUpdateEvents
+        # Some A2A agents send text content as artifacts, which SAM expects as status updates
+        if isinstance(event_payload, TaskArtifactUpdateEvent):
+            artifact = event_payload.artifact
+            if a2a.is_text_only_artifact(artifact):
+                log.info(
+                    "%s Converting text-only artifact to status update",
+                    log_identifier,
+                )
+                # Extract text from text-only artifact
+                text_content = "\n".join(a2a.get_text_content_from_artifact(artifact))
+
+                # Convert to status update
+                text_message = a2a.create_agent_text_message(
+                    text=text_content,
+                    task_id=event_payload.task_id,
+                    context_id=event_payload.context_id,
+                )
+
+                status_event = TaskStatusUpdateEvent(
+                    task_id=event_payload.task_id,
+                    context_id=event_payload.context_id,
+                    kind="status-update",
+                    status=TaskStatus(state=TaskState.working, message=text_message),
+                    final=False,
+                    metadata=event_payload.metadata,
                 )
-            else:
-                await self._publish_status_update(
-                    event_payload, task_context.a2a_context
+
+                # Replace event_payload with the converted status update
+                event_payload = status_event
+                log.info(
+                    "%s Converted text-only artifact (length: %d bytes) to status update",
+                    log_identifier,
+                    len(text_content.encode("utf-8")),
+                )
+
+        # Determine if this is a terminal event requiring cleanup
+        should_cleanup_task = False
+
+        # Route based on event type
+        if isinstance(event_payload, Task):
+            # Discard initial Task events (non-completed states)
+            # The final Task will be constructed from the final status update
+            if event_payload.status.state != TaskState.completed:
+                log.debug(
+                    "%s Discarding Task event with state=%s (not completed). Final Task will be constructed from final status update.",
+                    log_identifier,
+                    event_payload.status.state,
+                )
+                # Don't publish, don't cleanup - wait for final status update
+                return
+
+            # Forward completed Task to reply topic
+            await self._publish_task_response(event_payload, task_context.a2a_context)
+
+            # Completed Task is terminal - cleanup
+            should_cleanup_task = True
+            log.debug(
+                "%s Task in terminal state: %s",
+                log_identifier,
+                event_payload.status.state,
+            )
+
+        elif isinstance(event_payload, TaskStatusUpdateEvent):
+            # Forward status update to status topic
+            await self._publish_status_update(event_payload, task_context.a2a_context)
+
+            # Check if final event - construct and send Task
+            if event_payload.final:
+                log.info(
+                    "%s Received final status update (final=true). Constructing completed Task.",
+                    log_identifier,
                 )
+
+                # Construct Task from final status update
+                # Copy the status but ensure state is "completed"
+                final_task_status = TaskStatus(
+                    state=TaskState.completed,
+                    message=event_payload.status.message if event_payload.status else None,
+                )
+
+                final_task = Task(
+                    id=event_payload.task_id,
+                    context_id=event_payload.context_id,
+                    status=final_task_status,
+                    artifacts=None,  # Artifacts come via separate events
+                    metadata=event_payload.metadata,
+                )
+
+                # Add produced_artifacts metadata if any artifacts were processed
+                if produced_artifacts:
+                    if not final_task.metadata:
+                        final_task.metadata = {}
+                    final_task.metadata["produced_artifacts"] = produced_artifacts
+                    log.info(
+                        "%s Added manifest of %d produced artifacts to final Task metadata.",
+                        log_identifier,
+                        len(produced_artifacts),
+                    )
+
+                # Publish the constructed Task
+                await self._publish_task_response(final_task, task_context.a2a_context)
+
+                should_cleanup_task = True
+                log.debug(
+                    "%s Published final Task constructed from status update",
+                    log_identifier,
+                )
+
         elif isinstance(event_payload, TaskArtifactUpdateEvent):
+            # Forward artifact update to status topic
             await self._publish_artifact_update(event_payload, task_context.a2a_context)
+
         elif isinstance(event_payload, Message):
+            # Wrap Message in Task for gateway compatibility
             log.info(
-                "%s Received a direct Message response. Wrapping in a completed Task.",
+                "%s Received direct Message response. Wrapping in completed Task.",
                 log_identifier,
             )
             final_task = Task(
@@ -1053,7 +1537,7 @@ class A2AProxyComponent(BaseProxyComponent):
                 status=TaskStatus(state=TaskState.completed, message=event_payload),
             )
 
-            # Add produced_artifacts metadata to the wrapped Task if any artifacts were processed
+            # Add produced_artifacts metadata if any artifacts were processed
             if produced_artifacts:
                 final_task.metadata = {"produced_artifacts": produced_artifacts}
                 log.info(
@@ -1062,11 +1546,24 @@ class A2AProxyComponent(BaseProxyComponent):
                     len(produced_artifacts),
                 )
 
-            await self._publish_final_response(final_task, task_context.a2a_context)
+            await self._publish_task_response(final_task, task_context.a2a_context)
+            should_cleanup_task = True
+
         else:
             log.warning(
-                f"Received unhandled response payload type: {type(event_payload)}"
+                "%s Received unhandled response payload type: %s",
+                log_identifier,
+                type(event_payload).__name__,
+            )
+
+        # Cleanup task state if terminal event detected
+        if should_cleanup_task:
+            log.info(
+                "%s Terminal event detected for task %s. Cleaning up state.",
+                log_identifier,
+                task_context.task_id,
             )
+            self._cleanup_task_state(task_context.task_id)
 
     def clear_client_cache(self):
         """