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

solace_agent_mesh/agent/proxies/a2a/config.py

@@ -11,11 +11,24 @@ from ..base.config import BaseProxyAppConfig, ProxiedAgentConfig
 from ....common.utils.pydantic_utils import SamConfigBase
 
 
+class HttpHeaderConfig(SamConfigBase):
+    """Configuration for a single HTTP header."""
+
+    name: str = Field(
+        ...,
+        description="The HTTP header name (e.g., 'X-API-Key', 'Authorization').",
+    )
+    value: str = Field(
+        ...,
+        description="The HTTP header value.",
+    )
+
+
 class AuthenticationConfig(SamConfigBase):
     """Authentication configuration for downstream A2A agents."""
 
     type: Optional[
-        Literal["static_bearer", "static_apikey", "oauth2_client_credentials"]
+        Literal["static_bearer", "static_apikey", "oauth2_client_credentials", "oauth2_authorization_code"]
     ] = Field(
         default=None,
         description="Authentication type. If not specified, inferred from 'scheme' for backward compatibility.",
@@ -34,11 +47,11 @@ class AuthenticationConfig(SamConfigBase):
     )
     client_id: Optional[str] = Field(
         default=None,
-        description="OAuth 2.0 client identifier (required for oauth2_client_credentials type).",
+        description="OAuth 2.0 client identifier (required for oauth2_client_credentials and oauth2_authorization_code types).",
     )
     client_secret: Optional[str] = Field(
         default=None,
-        description="OAuth 2.0 client secret (required for oauth2_client_credentials type).",
+        description="OAuth 2.0 client secret (required for oauth2_client_credentials type, optional for oauth2_authorization_code).",
     )
     scope: Optional[str] = Field(
         default=None,
@@ -50,6 +63,20 @@ class AuthenticationConfig(SamConfigBase):
         description="How long to cache OAuth 2.0 tokens before refresh, in seconds (default: 3300 = 55 minutes).",
     )
 
+    # NEW fields for oauth2_authorization_code
+    authorization_url: Optional[str] = Field(
+        default=None,
+        description="OAuth 2.0 authorization endpoint URL (for oauth2_authorization_code type). Can override agent card URL.",
+    )
+    redirect_uri: Optional[str] = Field(
+        default=None,
+        description="OAuth 2.0 redirect URI (required for oauth2_authorization_code type).",
+    )
+    scopes: Optional[List[str]] = Field(
+        default=None,
+        description="OAuth 2.0 scopes as a list of strings (optional for oauth2_authorization_code type).",
+    )
+
     @model_validator(mode="after")
     def validate_auth_config(self) -> "AuthenticationConfig":
         """Validates authentication configuration based on type."""
@@ -108,10 +135,35 @@ class AuthenticationConfig(SamConfigBase):
                     "OAuth 2.0 client credentials flow requires 'client_secret'."
                 )
 
+        elif auth_type == "oauth2_authorization_code":
+            # Validate client_id
+            if not self.client_id:
+                raise ValueError(
+                    "OAuth 2.0 authorization code flow requires 'client_id'."
+                )
+
+            # Validate redirect_uri
+            if not self.redirect_uri:
+                raise ValueError(
+                    "OAuth 2.0 authorization code flow requires 'redirect_uri'."
+                )
+
+            # Optional: Validate authorization_url if provided (can also come from agent card)
+            if self.authorization_url:
+                try:
+                    parsed_url = urlparse(self.authorization_url)
+                    if parsed_url.scheme not in ["https", "http"]:
+                        raise ValueError(
+                            f"OAuth 2.0 'authorization_url' must use HTTP(S). "
+                            f"Got scheme: '{parsed_url.scheme}'"
+                        )
+                except Exception as e:
+                    raise ValueError(f"Failed to parse 'authorization_url': {e}")
+
         else:
             raise ValueError(
                 f"Unsupported authentication type '{auth_type}'. "
-                f"Supported types: static_bearer, static_apikey, oauth2_client_credentials."
+                f"Supported types: static_bearer, static_apikey, oauth2_client_credentials, oauth2_authorization_code."
             )
 
         return self
@@ -128,6 +180,30 @@ class A2AProxiedAgentConfig(ProxiedAgentConfig):
         default=None,
         description="Authentication details for the downstream agent.",
     )
+    use_auth_for_agent_card: bool = Field(
+        default=False,
+        description="If true, applies the configured authentication to agent card fetching. "
+        "If false, agent card requests are made without authentication.",
+    )
+    use_agent_card_url: bool = Field(
+        default=True,
+        description="If true, uses the URL from the agent card for task invocations. "
+        "If false, uses the configured URL directly for task invocations. "
+        "Note: The configured URL is always used to fetch the agent card itself.",
+    )
+    agent_card_headers: Optional[List[HttpHeaderConfig]] = Field(
+        default=None,
+        description="Custom HTTP headers to include when fetching the agent card. "
+        "These headers are added alongside authentication headers.",
+    )
+    task_headers: Optional[List[HttpHeaderConfig]] = Field(
+        default=None,
+        description="Custom HTTP headers to include when invoking A2A tasks. "
+        "These headers are added alongside authentication headers. Note: The A2A SDK's "
+        "AuthInterceptor applies authentication headers after these are set, so custom "
+        "headers cannot override authentication. For custom auth, omit the 'authentication' "
+        "config and use task_headers to set auth headers directly.",
+    )
 
 
 class A2AProxyAppConfig(BaseProxyAppConfig):

solace_agent_mesh/agent/proxies/base/component.py

@@ -7,8 +7,9 @@ from __future__ import annotations
 import asyncio
 import concurrent.futures
 import threading
+import time
 from abc import ABC, abstractmethod
-from typing import Any, Dict, Optional, TYPE_CHECKING
+from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING
 
 import httpx
 
@@ -32,6 +33,8 @@ from ....common.a2a.protocol import (
 from a2a.types import (
     A2ARequest,
     AgentCard,
+    AgentCapabilities,
+    AgentExtension,
     CancelTaskRequest,
     FilePart,
     InternalError,
@@ -77,16 +80,22 @@ class BaseProxyComponent(ComponentBase, ABC):
             "artifact_service", {"type": "memory"}
         )
         self.discovery_interval_sec = self.get_config("discovery_interval_seconds", 60)
+        self.task_state_ttl_minutes = self.get_config("task_state_ttl_minutes", 60)
+        self.task_cleanup_interval_minutes = self.get_config(
+            "task_cleanup_interval_minutes", 10
+        )
 
         self.agent_registry = AgentRegistry()
         self.artifact_service: Optional[BaseArtifactService] = None
-        self.active_tasks: Dict[str, ProxyTaskContext] = {}
+        # Store (timestamp, ProxyTaskContext) tuples for TTL-based cleanup
+        self.active_tasks: Dict[str, Tuple[float, ProxyTaskContext]] = {}
         self.active_tasks_lock = threading.Lock()
 
         self._async_loop: Optional[asyncio.AbstractEventLoop] = None
         self._async_thread: Optional[threading.Thread] = None
         self._async_init_future: Optional[concurrent.futures.Future] = None
         self._discovery_timer_id = f"proxy_discovery_{self.name}"
+        self._task_cleanup_timer_id = f"proxy_task_cleanup_{self.name}"
 
         try:
             # Initialize synchronous services first
@@ -158,8 +167,11 @@ class BaseProxyComponent(ComponentBase, ABC):
                 if not hasattr(event, "_proxy_message_handled"):
                     event._proxy_message_handled = message_handled
         elif event.event_type == EventType.TIMER:
-            if event.data.get("timer_id") == self._discovery_timer_id:
+            timer_id = event.data.get("timer_id")
+            if timer_id == self._discovery_timer_id:
                 await self._discover_and_publish_agents()
+            elif timer_id == self._task_cleanup_timer_id:
+                await self._cleanup_stale_tasks()
         else:
             log.debug(
                 "%s Ignoring unhandled event type: %s",
@@ -209,12 +221,13 @@ class BaseProxyComponent(ComponentBase, ABC):
                     "status_topic": message.get_user_properties().get("a2aStatusTopic"),
                     "reply_to_topic": message.get_user_properties().get("replyTo"),
                     "is_streaming": isinstance(a2a_request.root, SendStreamingMessageRequest),
+                    "user_properties": message.get_user_properties(),
                 }
                 task_context = ProxyTaskContext(
                     task_id=logical_task_id, a2a_context=a2a_context
                 )
                 with self.active_tasks_lock:
-                    self.active_tasks[logical_task_id] = task_context
+                    self.active_tasks[logical_task_id] = (time.time(), task_context)
 
                 log.info(
                     "%s Forwarding request for task %s to agent %s.",
@@ -228,12 +241,13 @@ class BaseProxyComponent(ComponentBase, ABC):
 
             elif isinstance(a2a_request.root, CancelTaskRequest):
                 logical_task_id = get_task_id_from_cancel_request(a2a_request)
-                
+
                 # Get the agent name from the topic (same as for send_message)
                 target_agent_name = topic.split("/")[-1]
-                
+
                 with self.active_tasks_lock:
-                    task_context = self.active_tasks.get(logical_task_id)
+                    task_entry = self.active_tasks.get(logical_task_id)
+                    task_context = task_entry[1] if task_entry else None
                 
                 if task_context:
                     log.info(
@@ -319,6 +333,73 @@ class BaseProxyComponent(ComponentBase, ABC):
 
         return update_message_parts(original_message, resolved_parts)
 
+    def _update_agent_card_for_proxy(self, agent_card: AgentCard, agent_alias: str) -> AgentCard:
+        """
+        Updates an agent card for proxying by:
+        1. Setting the name to the proxy alias
+        2. Adding/updating the display-name extension to preserve the original name
+
+        Args:
+            agent_card: The original agent card fetched from the remote agent
+            agent_alias: The alias/name to use for this agent in SAM
+
+        Returns:
+            A modified copy of the agent card with updated name and display-name extension
+        """
+        # Create a deep copy to avoid modifying the original
+        card_copy = agent_card.model_copy(deep=True)
+
+        # Store the original name as the display name (if not already set)
+        # This preserves the agent's identity while allowing it to be proxied under an alias
+        original_display_name = agent_card.name
+
+        # Check if there's already a display-name extension to preserve
+        display_name_uri = "https://solace.com/a2a/extensions/display-name"
+        if card_copy.capabilities and card_copy.capabilities.extensions:
+            for ext in card_copy.capabilities.extensions:
+                if ext.uri == display_name_uri and ext.params and ext.params.get("display_name"):
+                    # Use the existing display name from the extension
+                    original_display_name = ext.params["display_name"]
+                    break
+
+        # Update the card's name to the proxy alias
+        card_copy.name = agent_alias
+
+        # Ensure capabilities and extensions exist
+        if not card_copy.capabilities:
+            card_copy.capabilities = AgentCapabilities(extensions=[])
+        if not card_copy.capabilities.extensions:
+            card_copy.capabilities.extensions = []
+
+        # Find or create the display-name extension
+        display_name_ext = None
+        for ext in card_copy.capabilities.extensions:
+            if ext.uri == display_name_uri:
+                display_name_ext = ext
+                break
+
+        if display_name_ext:
+            # Update existing extension
+            if not display_name_ext.params:
+                display_name_ext.params = {}
+            display_name_ext.params["display_name"] = original_display_name
+        else:
+            # Create new extension
+            new_ext = AgentExtension(
+                uri=display_name_uri,
+                params={"display_name": original_display_name}
+            )
+            card_copy.capabilities.extensions.append(new_ext)
+
+        log.debug(
+            "%s Updated agent card: name='%s', display_name='%s'",
+            self.log_identifier,
+            agent_alias,
+            original_display_name
+        )
+
+        return card_copy
+
     def _initial_discovery_sync(self):
         """
         Synchronously fetches agent cards to populate the registry at startup.
@@ -340,12 +421,12 @@ class BaseProxyComponent(ComponentBase, ABC):
                     continue
                 try:
                     # Use a synchronous client for this initial blocking call
-                    response = client.get(f"{agent_url}/.well-known/agent.json")
+                    response = client.get(f"{agent_url}/.well-known/agent-card.json")
                     response.raise_for_status()
                     agent_card = AgentCard.model_validate(response.json())
 
-                    card_for_proxy = agent_card.model_copy(deep=True)
-                    card_for_proxy.name = agent_alias
+                    # Update the card for proxying (preserves display name)
+                    card_for_proxy = self._update_agent_card_for_proxy(agent_card, agent_alias)
                     self.agent_registry.add_or_update_agent(card_for_proxy)
                     log.info(
                         "%s Initial discovery successful for alias '%s' (actual name: '%s').",
@@ -376,12 +457,12 @@ class BaseProxyComponent(ComponentBase, ABC):
                     continue
 
                 agent_alias = agent_config["name"]
-                card_for_proxy = modern_card.model_copy(deep=True)
-                card_for_proxy.name = agent_alias
-                self.agent_registry.add_or_update_agent(card_for_proxy)
+                # Update the card for proxying (preserves display name)
+                card_for_registry = self._update_agent_card_for_proxy(modern_card, agent_alias)
+                self.agent_registry.add_or_update_agent(card_for_registry)
 
-                # Publish the modern card directly after updating its URL
-                card_to_publish = card_for_proxy.model_copy(deep=True)
+                # Create a separate copy for publishing
+                card_to_publish = card_for_registry.model_copy(deep=True)
                 card_to_publish.url = (
                     f"solace:{a2a.get_agent_request_topic(self.namespace, agent_alias)}"
                 )
@@ -418,10 +499,14 @@ class BaseProxyComponent(ComponentBase, ABC):
         response = create_success_response(
             result=event, request_id=a2a_context.get("jsonrpc_request_id")
         )
-        self._publish_a2a_message(response.model_dump(exclude_none=True), target_topic)
+        self._publish_a2a_message(
+            response.model_dump(exclude_none=True),
+            target_topic,
+            user_properties=a2a_context.get("user_properties")
+        )
 
-    async def _publish_final_response(self, task: Task, a2a_context: Dict):
-        """Publishes the final Task object to the appropriate Solace topic."""
+    async def _publish_task_response(self, task: Task, a2a_context: Dict):
+        """Publishes a Task object to the reply topic."""
         target_topic = a2a_context.get("reply_to_topic")
         if not target_topic:
             log.warning(
@@ -434,7 +519,11 @@ class BaseProxyComponent(ComponentBase, ABC):
         response = create_success_response(
             result=task, request_id=a2a_context.get("jsonrpc_request_id")
         )
-        self._publish_a2a_message(response.model_dump(exclude_none=True), target_topic)
+        self._publish_a2a_message(
+            response.model_dump(exclude_none=True),
+            target_topic,
+            user_properties=a2a_context.get("user_properties")
+        )
 
     async def _publish_artifact_update(
         self, event: TaskArtifactUpdateEvent, a2a_context: Dict
@@ -452,7 +541,11 @@ class BaseProxyComponent(ComponentBase, ABC):
         response = create_success_response(
             result=event, request_id=a2a_context.get("jsonrpc_request_id")
         )
-        self._publish_a2a_message(response.model_dump(exclude_none=True), target_topic)
+        self._publish_a2a_message(
+            response.model_dump(exclude_none=True),
+            target_topic,
+            user_properties=a2a_context.get("user_properties")
+        )
 
     async def _publish_error_response(
         self,
@@ -551,16 +644,73 @@ class BaseProxyComponent(ComponentBase, ABC):
                             nack_e,
                         )
 
+    async def _cleanup_stale_tasks(self):
+        """
+        Removes task state older than configured TTL.
+        This prevents memory leaks from tasks that complete without sending terminal events
+        (e.g., due to agent crashes or network failures).
+        """
+        ttl_seconds = self.task_state_ttl_minutes * 60
+        cutoff_time = time.time() - ttl_seconds
+
+        with self.active_tasks_lock:
+            stale_task_ids = [
+                task_id
+                for task_id, (timestamp, _) in self.active_tasks.items()
+                if timestamp < cutoff_time
+            ]
+            for task_id in stale_task_ids:
+                del self.active_tasks[task_id]
+                log.warning(
+                    "%s Cleaned up stale task %s (exceeded TTL of %d minutes)",
+                    self.log_identifier,
+                    task_id,
+                    self.task_state_ttl_minutes,
+                )
+
+        if stale_task_ids:
+            log.info(
+                "%s Stale task cleanup removed %d tasks",
+                self.log_identifier,
+                len(stale_task_ids),
+            )
+
+    def _cleanup_task_state(self, task_id: str) -> None:
+        """
+        Cleans up state for a completed task.
+        Called when a terminal event is detected (Task with terminal state,
+        or TaskStatusUpdateEvent with final=true).
+
+        Args:
+            task_id: The ID of the task to clean up
+        """
+        with self.active_tasks_lock:
+            entry = self.active_tasks.pop(task_id, None)
+            if entry:
+                log.info(
+                    "%s Removed task %s from active_tasks (terminal event detected)",
+                    self.log_identifier,
+                    task_id,
+                )
+            else:
+                log.debug(
+                    "%s Task %s not found in active_tasks during cleanup (already removed)",
+                    self.log_identifier,
+                    task_id,
+                )
+
     def _publish_discovered_cards(self):
         """Publishes all agent cards currently in the registry."""
         log.info(
             "%s Publishing initially discovered agent cards...", self.log_identifier
         )
         for agent_alias in self.agent_registry.get_agent_names():
-            card_to_publish = self.agent_registry.get_agent(agent_alias)
-            if not card_to_publish:
+            original_card = self.agent_registry.get_agent(agent_alias)
+            if not original_card:
                 continue
 
+            # Create a copy for publishing to avoid modifying the card in the registry
+            card_to_publish = original_card.model_copy(deep=True)
             card_to_publish.url = (
                 f"solace:{a2a.get_agent_request_topic(self.namespace, agent_alias)}"
             )
@@ -599,6 +749,21 @@ class BaseProxyComponent(ComponentBase, ABC):
                 self.discovery_interval_sec,
             )
 
+        # Schedule the recurring task cleanup timer
+        if self.task_cleanup_interval_minutes > 0:
+            cleanup_interval_ms = self.task_cleanup_interval_minutes * 60 * 1000
+            self.add_timer(
+                delay_ms=cleanup_interval_ms,
+                timer_id=self._task_cleanup_timer_id,
+                interval_ms=cleanup_interval_ms,
+            )
+            log.info(
+                "%s Scheduled recurring stale task cleanup every %d minutes (TTL: %d minutes).",
+                self.log_identifier,
+                self.task_cleanup_interval_minutes,
+                self.task_state_ttl_minutes,
+            )
+
         super().run()
 
     def clear_client_cache(self):
@@ -614,6 +779,7 @@ class BaseProxyComponent(ComponentBase, ABC):
         """Cleans up resources on component shutdown."""
         log.info("%s Cleaning up proxy component.", self.log_identifier)
         self.cancel_timer(self._discovery_timer_id)
+        self.cancel_timer(self._task_cleanup_timer_id)
 
         # Clear active tasks (no need to signal cancellation - downstream agents own their tasks)
         with self.active_tasks_lock:

solace_agent_mesh/agent/proxies/base/proxy_task_context.py

@@ -13,5 +13,7 @@ class ProxyTaskContext:
     This object is created when a task is initiated and destroyed when it completes.
     """
 
-    task_id: str
+    task_id: str  # SAM's task ID (used for upstream communication)
     a2a_context: Dict[str, Any]
+    downstream_task_id: str | None = None  # Downstream agent's task ID (used for cancellation)
+    original_request: Any = None  # Original A2A request (for task pause/resume in OAuth2 flows)

solace_agent_mesh/agent/sac/app.py

@@ -230,33 +230,33 @@ class ArtifactServiceConfig(SamConfigBase):
             )
         return self
 
+
 class AgentIdentityConfig(SamConfigBase):
     """Configuration for agent identity and key management."""
+
     key_mode: Literal["auto", "manual"] = Field(
         default="auto",
-        description="Key mode for agent identity: 'auto' for automatic generation, 'manual' for user-provided."
+        description="Key mode for agent identity: 'auto' for automatic generation, 'manual' for user-provided.",
     )
     key_identity: Optional[str] = Field(
-        default=None,
-        description="Actual key value when key_mode is 'manual'."
+        default=None, description="Actual key value when key_mode is 'manual'."
     )
     key_persistence: Optional[str] = Field(
         default=None,
-        description="Path to the key file, e.g. '/path/to/keys/agent_{name}.key'."
+        description="Path to the key file, e.g. '/path/to/keys/agent_{name}.key'.",
     )
 
     @model_validator(mode="after")
     def check_key_mode_and_identity(self) -> "AgentIdentityConfig":
         if self.key_mode == "manual" and not self.key_identity:
-            raise ValueError(
-                "'key_identity' is required when 'key_mode' is 'manual'."
-            )
+            raise ValueError("'key_identity' is required when 'key_mode' is 'manual'.")
         if self.key_mode == "auto" and self.key_identity:
             log.warning(
                 "Configuration Warning: 'key_identity' is ignored when 'key_mode' is 'auto'."
             )
         return self
 
+
 class SessionServiceConfig(SamConfigBase):
     """Configuration for the ADK Session Service."""
 
@@ -298,7 +298,7 @@ class SamAgentAppConfig(SamConfigBase):
     )
     agent_identity: Optional[AgentIdentityConfig] = Field(
         default_factory=lambda: AgentIdentityConfig(key_mode="auto"),
-        description="Configuration for agent identity and key management."
+        description="Configuration for agent identity and key management.",
     )
     trust_manager: Optional[Union[TrustManagerConfig, Dict[str, Any]]] = Field(
         default=None,
@@ -375,6 +375,12 @@ class SamAgentAppConfig(SamConfigBase):
         ge=0,
         description="Maximum number of dictionary keys to inspect during schema inference.",
     )
+    schema_inference_depth: int = Field(
+        default=4,
+        ge=1,
+        le=10,
+        description="Maximum depth for schema inference on nested structures. Higher values reveal more nested field names but increase metadata size.",
+    )
     enable_embed_resolution: bool = Field(
         default=True,
         description="Enable early-stage processing of dynamic embeds and inject related instructions.",
@@ -384,7 +390,7 @@ class SamAgentAppConfig(SamConfigBase):
         description="If true, automatically attempts to continue LLM generation if it is interrupted by a token limit.",
     )
     stream_batching_threshold_bytes: int = Field(
-        default=0,
+        default=100,
         description="Minimum size in bytes for accumulated text from LLM stream before sending a status update.",
     )
     max_message_size_bytes: int = Field(
@@ -476,7 +482,9 @@ class SamAgentApp(App):
             # Overwrite the raw dict with the validated object for downstream use
             app_info["app_config"] = app_config
         except ValidationError as e:
-            message = SamAgentAppConfig.format_validation_error_message(e, app_info['name'], app_config_dict.get('agent_name'))
+            message = SamAgentAppConfig.format_validation_error_message(
+                e, app_info["name"], app_config_dict.get("agent_name")
+            )
             log.error("Invalid Agent configuration:\n%s", message)
             raise
 
@@ -552,13 +560,30 @@ class SamAgentApp(App):
         super().__init__(app_info, **kwargs)
         log.debug("%s Agent initialization complete.", agent_name)
 
+    def run(self):
+        """
+        Override run to ensure component initialization failures cause application failure.
+
+        This is critical for containerized deployments where the process must exit with
+        a non-zero code if initialization fails. By re-raising the exception, we allow
+        the SAC framework's main() to handle cleanup (component.cleanup(), broker
+        disconnection, etc.) before exiting with code 1.
+        """
+        try:
+            super().run()
+        except Exception as e:
+            log.critical(
+                "Failed to start agent application '%s': %s", self.name, e, exc_info=e
+            )
+            raise
+
     def get_component(self, component_name: str = None) -> "SamAgentComponent":
         """
         Retrieves the running SamAgentComponent instance from the app's flow.
-        
+
         Args:
             component_name: Optional component name (for compatibility, but ignored since there's only one component)
-            
+
         Returns:
             The SamAgentComponent instance or None if not found
         """