solace-agent-mesh 1.0.1__py3-none-any.whl → 1.0.3__py3-none-any.whl

solace_agent_mesh/agent/protocol/event_handlers.py

@@ -32,6 +32,7 @@ from ...common.types import (
     TaskStatus,
     TaskState,
     DataPart,
+    TextPart,
     A2ARequest,
 )
 from ...common.a2a_protocol import (
@@ -44,6 +45,7 @@ from ...common.a2a_protocol import (
     _extract_text_from_parts,
 )
 from ...agent.utils.artifact_helpers import (
+    generate_artifact_metadata_summary,
     load_artifact_content_or_metadata,
 )
 from ...agent.adk.runner import run_adk_async_task_thread_wrapper
@@ -57,114 +59,6 @@ from google.adk.events import Event as ADKEvent
 from google.genai import types as adk_types
 
 
-async def _format_artifact_summary_from_manifest(
-    component: "SamAgentComponent",
-    produced_artifacts: List[Dict[str, Any]],
-    peer_agent_name: str,
-    correlation_data: Dict[str, Any],
-) -> str:
-    """
-    Loads metadata for a list of produced artifacts and formats it into a
-    human-readable YAML summary string.
-    """
-    if not produced_artifacts:
-        return ""
-
-    artifact_summary_lines = [
-        f"Peer agent `{peer_agent_name}` created {len(produced_artifacts)} artifact(s):"
-    ]
-
-    original_task_context = correlation_data.get("original_task_context", {})
-    user_id = original_task_context.get("user_id")
-    session_id = original_task_context.get("session_id")
-
-    if not (component.artifact_service and user_id and session_id):
-        log.warning(
-            "%s Cannot load artifact metadata: missing artifact_service or context.",
-            component.log_identifier,
-        )
-        for artifact_ref in produced_artifacts:
-            artifact_summary_lines.append(
-                f"- `{artifact_ref.get('filename')}` (v{artifact_ref.get('version')})"
-            )
-        return "\n".join(artifact_summary_lines)
-
-    peer_agent_name_for_artifact = peer_agent_name
-    if (
-        not peer_agent_name_for_artifact
-        or peer_agent_name_for_artifact == "A peer agent"
-    ):
-        log.warning(
-            "%s Peer agent name not in task metadata, using self agent name for artifact loading.",
-            component.log_identifier,
-        )
-        peer_agent_name_for_artifact = component.agent_name
-
-    for artifact_ref in produced_artifacts:
-        filename = artifact_ref.get("filename")
-        version = artifact_ref.get("version")
-        if not filename or version is None:
-            continue
-
-        try:
-            metadata_result = await load_artifact_content_or_metadata(
-                artifact_service=component.artifact_service,
-                app_name=peer_agent_name_for_artifact,
-                user_id=user_id,
-                session_id=session_id,
-                filename=filename,
-                version=version,
-                load_metadata_only=True,
-            )
-            if metadata_result.get("status") == "success":
-                metadata = metadata_result.get("metadata", {})
-                TRUNCATION_LIMIT_BYTES = 1024
-                TRUNCATION_MESSAGE = "\n... [truncated] ..."
-
-                try:
-                    formatted_metadata_str = yaml.safe_dump(
-                        metadata,
-                        default_flow_style=False,
-                        sort_keys=False,
-                        allow_unicode=True,
-                    )
-
-                    if (
-                        len(formatted_metadata_str.encode("utf-8"))
-                        > TRUNCATION_LIMIT_BYTES
-                    ):
-                        cutoff = TRUNCATION_LIMIT_BYTES - len(
-                            TRUNCATION_MESSAGE.encode("utf-8")
-                        )
-                        formatted_metadata_str = (
-                            formatted_metadata_str[:cutoff] + TRUNCATION_MESSAGE
-                        )
-
-                    summary_line = f"- `{filename}` (v{version}):\n  ```yaml\n{formatted_metadata_str}\n  ```"
-                    artifact_summary_lines.append(summary_line)
-                except Exception as e_format:
-                    log.error(
-                        "Error formatting metadata for %s v%s: %s",
-                        filename,
-                        version,
-                        e_format,
-                    )
-                    artifact_summary_lines.append(
-                        f"- `{filename}` (v{version}): Error formatting metadata."
-                    )
-            else:
-                artifact_summary_lines.append(
-                    f"- `{filename}` (v{version}): Could not load metadata."
-                )
-        except Exception as e_meta:
-            log.error(
-                "Error loading metadata for %s v%s: %s", filename, version, e_meta
-            )
-            artifact_summary_lines.append(
-                f"- `{filename}` (v{version}): Error loading metadata."
-            )
-
-    return "\n".join(artifact_summary_lines)
 
 
 def _register_peer_artifacts_in_parent_context(
@@ -248,7 +142,7 @@ async def process_event(component, event: Event):
             elif topic.startswith(agent_response_sub_prefix) or topic.startswith(
                 agent_status_sub_prefix
             ):
-                handle_a2a_response(component, message)
+                await handle_a2a_response(component, message)
             else:
                 log.warning(
                     "%s Received message on unhandled topic: %s",
@@ -264,31 +158,8 @@ async def process_event(component, event: Event):
             if timer_data.get("timer_id") == component._card_publish_timer_id:
                 publish_agent_card(component)
         elif event.event_type == EventType.CACHE_EXPIRY:
-            cache_data = event.data
-            log.info(
-                "%s Received cache expiry event: %s",
-                component.log_identifier,
-                cache_data,
-            )
-            sub_task_id = cache_data.get("key")
-            if sub_task_id and sub_task_id.startswith(
-                component.CORRELATION_DATA_PREFIX
-            ):
-                expired_data = cache_data.get("expired_data")
-                if expired_data:
-                    await component._handle_peer_timeout(sub_task_id, expired_data)
-                else:
-                    log.error(
-                        "%s Missing expired_data in cache expiry event for sub-task %s. Cannot process timeout.",
-                        component.log_identifier,
-                        sub_task_id,
-                    )
-            else:
-                log.debug(
-                    "%s Cache expiry for key '%s' is not a peer sub-task timeout.",
-                    component.log_identifier,
-                    sub_task_id,
-                )
+            # Delegate cache expiry handling to the component itself.
+            await component.handle_cache_expiry_event(event.data)
         else:
             log.warning(
                 "%s Received unknown event type: %s",
@@ -633,8 +504,51 @@ async def handle_a2a_request(component, message: SolaceMessage):
                 logical_task_id,
             )
 
+            a2a_message_for_adk = a2a_request.params.message
+            invoked_artifacts = (
+                a2a_message_for_adk.metadata.get("invoked_with_artifacts", [])
+                if a2a_message_for_adk.metadata
+                else []
+            )
+
+            if invoked_artifacts:
+                log.info(
+                    "%s Task %s invoked with %d artifact(s). Preparing context from metadata.",
+                    component.log_identifier,
+                    task_id,
+                    len(invoked_artifacts),
+                )
+                header_text = (
+                    "The user has provided the following artifacts as context for your task. "
+                    "Use the information contained within their metadata to complete your objective."
+                )
+                artifact_summary = await generate_artifact_metadata_summary(
+                    component=component,
+                    artifact_identifiers=invoked_artifacts,
+                    user_id=user_id,
+                    session_id=effective_session_id,
+                    app_name=agent_name,
+                    header_text=header_text,
+                )
+
+                task_description = _extract_text_from_parts(
+                    a2a_message_for_adk.parts
+                )
+                final_prompt = f"{task_description}\n\n{artifact_summary}"
+
+                a2a_message_for_adk = A2AMessage(
+                    role="user",
+                    parts=[TextPart(text=final_prompt)],
+                    metadata=a2a_message_for_adk.metadata,
+                )
+                log.debug(
+                    "%s Generated new prompt for task %s with artifact context.",
+                    component.log_identifier,
+                    task_id,
+                )
+
             adk_content = translate_a2a_to_adk_content(
-                a2a_request.params.message, component.log_identifier
+                a2a_message_for_adk, component.log_identifier
             )
 
             adk_session = await component.session_service.get_session(
@@ -854,10 +768,9 @@ def handle_agent_card_message(component, message: SolaceMessage):
         component.handle_error(e, Event(EventType.MESSAGE, message))
 
 
-def handle_a2a_response(component, message: SolaceMessage):
+async def handle_a2a_response(component, message: SolaceMessage):
     """Handles incoming responses/status updates from peer agents."""
     sub_task_id = None
-    agent_name = component.get_config("agent_name")
     payload_to_queue = None
     is_final_response = False
 
@@ -932,10 +845,8 @@ def handle_a2a_response(component, message: SolaceMessage):
                                             component.log_identifier,
                                             sub_task_id,
                                         )
-                                        correlation_data = (
-                                            component.cache_service.get_data(
-                                                sub_task_id
-                                            )
+                                        correlation_data = await component._get_correlation_data_for_sub_task(
+                                            sub_task_id
                                         )
                                         if not correlation_data:
                                             log.warning(
@@ -1227,7 +1138,12 @@ def handle_a2a_response(component, message: SolaceMessage):
                     "error": f"Failed to parse response from peer: {parse_error}",
                     "code": "PEER_PARSE_ERROR",
                 }
-                is_final_response = True
+                # Print out the stack trace for debugging
+                log.exception(
+                    "%s Exception stack trace: %s",
+                    component.log_identifier,
+                    parse_error,
+                )
 
         if not is_final_response:
             # This is an intermediate status update for monitoring.
@@ -1240,13 +1156,9 @@ def handle_a2a_response(component, message: SolaceMessage):
             message.call_acknowledgements()
             return
 
-        correlation_data = component.cache_service.get_data(sub_task_id)
+        correlation_data = await component._claim_peer_sub_task_completion(sub_task_id)
         if not correlation_data:
-            log.warning(
-                "%s No correlation data found for sub-task %s. Cannot process response. Ignoring.",
-                component.log_identifier,
-                sub_task_id,
-            )
+            # The helper method logs the reason (timeout, already claimed, etc.)
             message.call_acknowledgements()
             return
 
@@ -1309,14 +1221,31 @@ def handle_a2a_response(component, message: SolaceMessage):
                                 peer_agent_name = task_obj.metadata.get(
                                     "agent_name", "A peer agent"
                                 )
-                                artifact_summary = (
-                                    await _format_artifact_summary_from_manifest(
-                                        component,
-                                        produced_artifacts,
-                                        peer_agent_name,
-                                        correlation_data,
-                                    )
+                                original_task_context = correlation_data.get(
+                                    "original_task_context", {}
                                 )
+                                user_id = original_task_context.get("user_id")
+                                session_id = original_task_context.get("session_id")
+
+                                header_text = f"Peer agent `{peer_agent_name}` created {len(produced_artifacts)} artifact(s):"
+
+                                if user_id and session_id:
+                                    artifact_summary = (
+                                        await generate_artifact_metadata_summary(
+                                            component=component,
+                                            artifact_identifiers=produced_artifacts,
+                                            user_id=user_id,
+                                            session_id=session_id,
+                                            app_name=peer_agent_name,
+                                            header_text=header_text,
+                                        )
+                                    )
+                                else:
+                                    log.warning(
+                                        "%s Could not generate artifact summary: missing user_id or session_id in correlation data.",
+                                        log_retrigger,
+                                    )
+                                    artifact_summary = ""
                                 # Bubble up the peer's artifacts to the parent context
                                 _register_peer_artifacts_in_parent_context(
                                     task_context, task_obj, log_retrigger

solace_agent_mesh/agent/sac/app.py

@@ -246,8 +246,8 @@ class SamAgentApp(App):
                         "type": "string",
                         "required": False,
                         "default": "namespace",
-                        "enum": ["namespace", "app", "custom"],
-                        "description": "Scope for filesystem artifacts: 'namespace' (default, shared by namespace), 'app' (isolated by app name), 'custom' (use artifact_scope_value).",
+                        "enum": ["namespace", "app"],
+                        "description": "Process-wide scope for all artifact services. 'namespace' (default): shared by all components in the namespace. 'app': isolated by agent/gateway name. This setting must be consistent for all components in the same process.",
                     },
                     "artifact_scope_value": {
                         "type": "string",
@@ -340,6 +340,13 @@ class SamAgentApp(App):
                 "default": 0,
                 "description": "Minimum size in bytes for accumulated text from LLM stream before sending a status update. If 0 or less, batching is disabled and updates are sent per chunk. Final LLM chunks are always sent regardless of this threshold.",
             },
+            {
+                "name": "max_message_size_bytes",
+                "required": False,
+                "type": "integer",
+                "default": 10_000_000,
+                "description": "Maximum allowed message size in bytes before rejecting publication to prevent broker disconnections. Default: 10MB",
+            },
             {
                 "name": "enable_artifact_content_instruction",
                 "required": False,
@@ -551,6 +558,7 @@ class SamAgentApp(App):
         app_config = app_info.get("app_config", {})
         namespace = app_config.get("namespace")
         agent_name = app_config.get("agent_name")
+        broker_request_response = app_info.get("broker_request_response")
 
         if not namespace or not isinstance(namespace, str):
             raise ValueError(
@@ -619,6 +627,8 @@ class SamAgentApp(App):
             "component_config": {},
             "subscriptions": generated_subs,
         }
+        if broker_request_response:
+            component_definition["broker_request_response"] = broker_request_response
 
         app_info["components"] = [component_definition]
         log.debug("Replaced 'components' in app_info with programmatic definition.")

solace_agent_mesh/agent/sac/component.py

@@ -53,6 +53,8 @@ from ...common.types import (
     TaskStatusUpdateEvent,
     TaskArtifactUpdateEvent,
     SendTaskRequest,
+    CancelTaskRequest,
+    TaskIdParams,
 )
 from ...common.a2a_protocol import (
     get_a2a_base_topic,
@@ -65,7 +67,9 @@ from ...common.a2a_protocol import (
     get_gateway_status_topic,
 )
 from ...agent.utils.config_parser import resolve_instruction_provider
-from ...agent.utils.artifact_helpers import get_latest_artifact_version
+from ...agent.utils.artifact_helpers import (
+    get_latest_artifact_version,
+)
 from ...agent.adk.services import (
     initialize_session_service,
     initialize_artifact_service,
@@ -90,6 +94,8 @@ from ...agent.adk.invocation_monitor import InvocationMonitor
 from ...common.middleware.registry import MiddlewareRegistry
 from ...common.constants import DEFAULT_COMMUNICATION_TIMEOUT
 from ...agent.tools.registry import tool_registry
+from ...common.utils.message_utils import validate_message_size
+from ...common.exceptions import MessageSizeExceededError
 
 if TYPE_CHECKING:
     from .task_execution_context import TaskExecutionContext
@@ -230,6 +236,11 @@ class SamAgentComponent(ComponentBase):
                 raise ValueError(
                     "Internal Error: Inter-agent comms config missing after validation."
                 )
+
+            self.max_message_size_bytes = self.get_config(
+                "max_message_size_bytes", 10_000_000
+            )
+
             log.info("%s Configuration retrieved successfully.", self.log_identifier)
         except Exception as e:
             log.error(
@@ -392,6 +403,7 @@ class SamAgentComponent(ComponentBase):
                 self.session_service = initialize_session_service(self)
                 self.artifact_service = initialize_artifact_service(self)
                 self.memory_service = initialize_memory_service(self)
+
                 log.info(
                     "%s Synchronous ADK services initialized.", self.log_identifier
                 )
@@ -536,49 +548,123 @@ class SamAgentComponent(ComponentBase):
         if timer_data.get("timer_id") == self._card_publish_timer_id:
             publish_agent_card(self)
 
-    def handle_cache_expiry_event(self, cache_data: Dict[str, Any]):
-        """Handles cache expiry events, specifically for peer timeouts."""
+    async def handle_cache_expiry_event(self, cache_data: Dict[str, Any]):
+        """
+        Handles cache expiry events for peer timeouts by calling the atomic claim helper.
+        """
         log.debug("%s Received cache expiry event: %s", self.log_identifier, cache_data)
-        expired_key = cache_data.get("key")
-        expired_data = cache_data.get("expired_data")
+        sub_task_id = cache_data.get("key")
+        logical_task_id = cache_data.get("expired_data")
+
+        if not (
+            sub_task_id
+            and sub_task_id.startswith(CORRELATION_DATA_PREFIX)
+            and logical_task_id
+        ):
+            log.debug(
+                "%s Cache expiry for key '%s' is not a peer sub-task timeout or is missing data.",
+                self.log_identifier,
+                sub_task_id,
+            )
+            return
+
+        correlation_data = await self._claim_peer_sub_task_completion(
+            sub_task_id=sub_task_id, logical_task_id_from_event=logical_task_id
+        )
 
-        if expired_key and expired_key.startswith(CORRELATION_DATA_PREFIX):
-            sub_task_id = expired_key
+        if correlation_data:
             log.warning(
-                "%s Detected timeout for sub-task ID: %s",
+                "%s Detected timeout for sub-task %s (Main Task: %s). Claimed successfully.",
                 self.log_identifier,
                 sub_task_id,
+                logical_task_id,
             )
-            if expired_data:
-                try:
-                    original_task_context = expired_data.get("original_task_context")
-                    if original_task_context:
-                        self._handle_peer_timeout(sub_task_id, expired_data)
-                    else:
-                        log.error(
-                            "%s Missing 'original_task_context' in expired cache data for sub-task %s. Cannot process timeout.",
-                            self.log_identifier,
-                            sub_task_id,
-                        )
-                except Exception as e:
-                    log.exception(
-                        "%s Error handling peer timeout for sub-task %s: %s",
-                        self.log_identifier,
-                        sub_task_id,
-                        e,
-                    )
-            else:
-                log.error(
-                    "%s Missing expired_data in cache expiry event for sub-task %s. Cannot process timeout.",
-                    self.log_identifier,
-                    sub_task_id,
-                )
+            await self._handle_peer_timeout(sub_task_id, correlation_data)
         else:
-            log.debug(
-                "%s Cache expiry for key '%s' is not a peer sub-task timeout.",
+            log.info(
+                "%s Ignoring timeout event for sub-task %s as it was already completed.",
+                self.log_identifier,
+                sub_task_id,
+            )
+
+    async def _get_correlation_data_for_sub_task(
+        self, sub_task_id: str
+    ) -> Optional[Dict[str, Any]]:
+        """
+        Non-destructively retrieves correlation data for a sub-task.
+        Used for intermediate events where the sub-task should remain active.
+        """
+        logical_task_id = self.cache_service.get_data(sub_task_id)
+        if not logical_task_id:
+            log.warning(
+                "%s No cache entry for sub-task %s. Cannot get correlation data.",
                 self.log_identifier,
-                expired_key,
+                sub_task_id,
+            )
+            return None
+
+        with self.active_tasks_lock:
+            task_context = self.active_tasks.get(logical_task_id)
+
+        if not task_context:
+            log.error(
+                "%s TaskExecutionContext not found for task %s, but cache entry existed for sub-task %s. This may indicate a cleanup issue.",
+                self.log_identifier,
+                logical_task_id,
+                sub_task_id,
+            )
+            return None
+
+        with task_context.lock:
+            return task_context.active_peer_sub_tasks.get(sub_task_id)
+
+    async def _claim_peer_sub_task_completion(
+        self, sub_task_id: str, logical_task_id_from_event: Optional[str] = None
+    ) -> Optional[Dict[str, Any]]:
+        """
+        Atomically claims a sub-task as complete, preventing race conditions.
+        This is a destructive operation that removes state.
+
+        Args:
+            sub_task_id: The ID of the sub-task to claim.
+            logical_task_id_from_event: The parent task ID, if provided by the event (e.g., a timeout).
+                                        If not provided, it will be looked up from the cache.
+        """
+        log_id = f"{self.log_identifier}[ClaimSubTask:{sub_task_id}]"
+        logical_task_id = logical_task_id_from_event
+
+        if not logical_task_id:
+            logical_task_id = self.cache_service.get_data(sub_task_id)
+            if not logical_task_id:
+                log.warning(
+                    "%s No cache entry found. Task has likely timed out and been cleaned up. Cannot claim.",
+                    log_id,
+                )
+                return None
+
+        with self.active_tasks_lock:
+            task_context = self.active_tasks.get(logical_task_id)
+
+        if not task_context:
+            log.error(
+                "%s TaskExecutionContext not found for task %s. Cleaning up stale cache entry.",
+                log_id,
+                logical_task_id,
             )
+            self.cache_service.remove_data(sub_task_id)
+            return None
+
+        correlation_data = task_context.claim_sub_task_completion(sub_task_id)
+
+        if correlation_data:
+            # If we successfully claimed the task, remove the timeout tracker from the cache.
+            self.cache_service.remove_data(sub_task_id)
+            log.info("%s Successfully claimed completion.", log_id)
+            return correlation_data
+        else:
+            # This means the task was already claimed by a competing event (e.g., timeout vs. response).
+            log.warning("%s Failed to claim; it was already completed.", log_id)
+            return None
 
     async def _retrigger_agent_with_peer_responses(
         self,
@@ -686,8 +772,9 @@ class SamAgentComponent(ComponentBase):
         correlation_data: Dict[str, Any],
     ):
         """
-        Handles the timeout of a peer agent task by updating the completion counter
-        and potentially re-triggering the runner if all parallel tasks are now complete.
+        Handles the timeout of a peer agent task. It sends a cancellation request
+        to the peer, updates the local completion counter, and potentially
+        re-triggers the runner if all parallel tasks are now complete.
         """
         logical_task_id = correlation_data.get("logical_task_id")
         invocation_id = correlation_data.get("invocation_id")
@@ -700,6 +787,36 @@ class SamAgentComponent(ComponentBase):
             invocation_id,
         )
 
+        # Proactively send a cancellation request to the peer agent.
+        peer_agent_name = correlation_data.get("peer_agent_name")
+        if peer_agent_name:
+            try:
+                log.info(
+                    "%s Sending CancelTaskRequest to peer '%s' for timed-out sub-task %s.",
+                    log_retrigger,
+                    peer_agent_name,
+                    sub_task_id,
+                )
+                task_id_for_peer = sub_task_id.replace(CORRELATION_DATA_PREFIX, "", 1)
+                cancel_params = TaskIdParams(id=task_id_for_peer)
+                cancel_request = CancelTaskRequest(params=cancel_params)
+                user_props = {"clientId": self.agent_name}
+                peer_topic = self._get_agent_request_topic(peer_agent_name)
+                self._publish_a2a_message(
+                    payload=cancel_request.model_dump(exclude_none=True),
+                    topic=peer_topic,
+                    user_properties=user_props,
+                )
+            except Exception as e:
+                log.error(
+                    "%s Failed to send CancelTaskRequest to peer '%s' for sub-task %s: %s",
+                    log_retrigger,
+                    peer_agent_name,
+                    sub_task_id,
+                    e,
+                )
+
+        # Process the timeout locally.
         with self.active_tasks_lock:
             task_context = self.active_tasks.get(logical_task_id)
 
@@ -1841,7 +1958,8 @@ class SamAgentComponent(ComponentBase):
                     len(unprocessed_tail.encode("utf-8")),
                 )
         else:
-            resolved_text = resolved_text + unprocessed_tail
+            if unprocessed_tail is not None and unprocessed_tail != "":
+                resolved_text = resolved_text + unprocessed_tail
 
         if signals_found:
             log.info(
@@ -2860,6 +2978,29 @@ class SamAgentComponent(ComponentBase):
     ):
         """Helper to publish A2A messages via the SAC App."""
         try:
+            max_size_bytes = self.max_message_size_bytes
+
+            # Validate message size
+            is_valid, actual_size = validate_message_size(
+                payload, max_size_bytes, self.log_identifier
+            )
+
+            if not is_valid:
+                error_msg = (
+                    f"Message size validation failed: payload size ({actual_size} bytes) "
+                    f"exceeds maximum allowed size ({max_size_bytes} bytes)"
+                )
+                log.error("%s %s", self.log_identifier, error_msg)
+                raise MessageSizeExceededError(actual_size, max_size_bytes, error_msg)
+
+            # Debug logging to show message size when publishing
+            log.debug(
+                "%s Publishing message to topic %s (size: %d bytes)",
+                self.log_identifier,
+                topic,
+                actual_size,
+            )
+
             app = self.get_app()
             if app:
                 if self.invocation_monitor:
@@ -2877,6 +3018,9 @@ class SamAgentComponent(ComponentBase):
                     "%s Cannot publish message: Not running within a SAC App context.",
                     self.log_identifier,
                 )
+        except MessageSizeExceededError:
+            # Re-raise MessageSizeExceededError without wrapping
+            raise
         except Exception as e:
             log.exception(
                 "%s Failed to publish A2A message to topic %s: %s",