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

solace_agent_mesh/agent/adk/models/oauth2_token_manager.py

@@ -6,12 +6,9 @@ It handles token acquisition, caching, and automatic refresh with proper error h
 
 import asyncio
 import logging
-import random
-import time
-from typing import Any, Dict, Optional
-
-import httpx
+from typing import Optional
 
+from solace_agent_mesh.common.oauth import OAuth2RetryClient, is_token_expired
 from solace_agent_mesh.common.utils.in_memory_cache import InMemoryCache
 
 logger = logging.getLogger(__name__)
@@ -72,17 +69,19 @@ class OAuth2ClientCredentialsTokenManager:
         self.scope = scope
         self.ca_cert_path = ca_cert_path
         self.refresh_buffer_seconds = refresh_buffer_seconds
-        self.max_retries = max_retries
-        
+
         # Thread-safe token access
         self._lock = asyncio.Lock()
-        
+
         # Token cache using existing InMemoryCache singleton
         self._cache = InMemoryCache()
-        
+
         # Cache key for this token manager instance
         self._cache_key = f"oauth_token_{hash((token_url, client_id))}"
-        
+
+        # OAuth client with retry logic
+        self._oauth_client = OAuth2RetryClient(max_retries=max_retries)
+
         logger.info(
             "OAuth2ClientCredentialsTokenManager initialized for endpoint: %s",
             token_url
@@ -104,14 +103,24 @@ class OAuth2ClientCredentialsTokenManager:
         async with self._lock:
             # Check if we have a cached token
             cached_token_data = self._cache.get(self._cache_key)
-            
-            if cached_token_data and not self._is_token_expired(cached_token_data):
+
+            if cached_token_data and not is_token_expired(
+                cached_token_data["expires_at"], buffer_seconds=self.refresh_buffer_seconds
+            ):
                 logger.debug("Using cached OAuth token")
                 return cached_token_data["access_token"]
-            
-            # Fetch new token
+
+            # Fetch new token using common OAuth client
             logger.info("Fetching new OAuth token from %s", self.token_url)
-            token_data = await self._fetch_token()
+            verify = self.ca_cert_path if self.ca_cert_path else True
+            token_data = await self._oauth_client.fetch_client_credentials_token(
+                token_url=self.token_url,
+                client_id=self.client_id,
+                client_secret=self.client_secret,
+                scope=self.scope,
+                verify=verify,
+                timeout=30.0,
+            )
             
             # Cache the token with TTL
             expires_in = token_data.get("expires_in", 3600)  # Default 1 hour
@@ -121,125 +130,3 @@ class OAuth2ClientCredentialsTokenManager:
             
             logger.info("OAuth token cached with TTL: %d seconds", cache_ttl)
             return token_data["access_token"]
-
-    def _is_token_expired(self, token_data: Dict[str, Any]) -> bool:
-        """Check if a token is expired or near expiry.
-        
-        Args:
-            token_data: Token data dictionary with 'expires_at' timestamp
-            
-        Returns:
-            True if token is expired or near expiry, False otherwise
-        """
-        if "expires_at" not in token_data:
-            return True
-            
-        current_time = time.time()
-        expires_at = token_data["expires_at"]
-        
-        # Consider token expired if it expires within the buffer time
-        return current_time >= (expires_at - self.refresh_buffer_seconds)
-
-    async def _fetch_token(self) -> Dict[str, Any]:
-        """Fetch a new OAuth 2.0 access token from the token endpoint.
-
-        Implements retry logic with exponential backoff for transient failures.
-
-        Returns:
-            Token data dictionary containing access_token, expires_in, etc.
-
-        Raises:
-            httpx.HTTPError: If HTTP request fails after all retries
-            ValueError: If response is invalid or missing required fields
-        """
-        # Prepare request payload
-        payload = {
-            "grant_type": "client_credentials",
-            "client_id": self.client_id,
-            "client_secret": self.client_secret,
-        }
-        
-        if self.scope:
-            payload["scope"] = self.scope
-        
-        # Configure HTTP client with SSL settings
-        verify = True
-        if self.ca_cert_path:
-            verify = self.ca_cert_path
-            
-        headers = {
-            "Content-Type": "application/x-www-form-urlencoded",
-            "Accept": "application/json",
-        }
-        
-        last_exception = None
-
-        for attempt in range(self.max_retries + 1):
-            try:
-                async with httpx.AsyncClient(verify=verify) as client:
-                    response = await client.post(
-                        self.token_url,
-                        data=payload,
-                        headers=headers,
-                        timeout=30.0,
-                    )
-                    response.raise_for_status()
-
-                    token_data = response.json()
-
-                    # Validate response
-                    if "access_token" not in token_data:
-                        raise ValueError("Token response missing 'access_token' field")
-
-                    # Add expiration timestamp for cache management
-                    expires_in = token_data.get("expires_in", 3600)
-                    token_data["expires_at"] = time.time() + expires_in
-
-                    logger.info("Successfully fetched OAuth token, expires in %d seconds", expires_in)
-                    return token_data
-
-            except httpx.HTTPStatusError as e:
-                last_exception = e
-                # Don't retry on 4xx errors (client errors)
-                if 400 <= e.response.status_code < 500:
-                    logger.error(
-                        "OAuth token request failed with client error %d: %s",
-                        e.response.status_code,
-                        e.response.text
-                    )
-                    raise
-
-                logger.warning(
-                    "OAuth token request failed with status %d (attempt %d/%d): %s",
-                    e.response.status_code,
-                    attempt + 1,
-                    self.max_retries + 1,
-                    e.response.text
-                )
-
-            except httpx.RequestError as e:
-                last_exception = e
-                logger.warning(
-                    "OAuth token request failed (attempt %d/%d): %s",
-                    attempt + 1,
-                    self.max_retries + 1,
-                    str(e)
-                )
-
-            except Exception as e:
-                last_exception = e
-                logger.error("Unexpected error during OAuth token fetch: %s", str(e))
-                raise
-
-            # Exponential backoff with jitter for retries
-            if attempt < self.max_retries:
-                delay = (2 ** attempt) + random.uniform(0, 1)
-                logger.info("Retrying OAuth token request in %.2f seconds", delay)
-                await asyncio.sleep(delay)
-
-        # All retries exhausted
-        logger.error("OAuth token request failed after %d attempts", self.max_retries + 1)
-        if last_exception:
-            raise last_exception
-        else:
-            raise RuntimeError("OAuth token request failed after all retries")

solace_agent_mesh/agent/adk/runner.py

@@ -38,6 +38,7 @@ async def run_adk_async_task_thread_wrapper(
     adk_content: adk_types.Content,
     run_config: RunConfig,
     a2a_context: dict[str, Any],
+    append_context_event: bool = True,
 ):
     """
     Wrapper to run the async ADK task.
@@ -73,9 +74,8 @@ async def run_adk_async_task_thread_wrapper(
             logical_task_id,
         )
 
-        if adk_session and component.session_service:
+        if adk_session and component.session_service and append_context_event:
             context_setting_invocation_id = logical_task_id
-            original_message = a2a_context.pop("original_solace_message", None)
             try:
                 context_setting_event = ADKEvent(
                     invocation_id=context_setting_invocation_id,
@@ -107,15 +107,13 @@ async def run_adk_async_task_thread_wrapper(
                     e_append,
                     exc_info=True,
                 )
-            finally:
-                if original_message:
-                    a2a_context["original_solace_message"] = original_message
         else:
-            log.warning(
-                "%s Could not inject a2a_context into ADK session state via event for task %s (session or session_service invalid). Tool scope filtering might not work.",
-                component.log_identifier,
-                logical_task_id,
-            )
+            if append_context_event:
+                log.warning(
+                    "%s Could not inject a2a_context into ADK session state via event for task %s (session or session_service invalid). Tool scope filtering might not work.",
+                    component.log_identifier,
+                    logical_task_id,
+                )
 
         is_paused = await run_adk_async_task(
             component,
@@ -126,6 +124,15 @@ async def run_adk_async_task_thread_wrapper(
             a2a_context,
         )
 
+        # Mark task as paused if it's waiting for peer response or user input
+        if task_context and is_paused:
+            task_context.set_paused(True)
+            log.debug(
+                "%s Task %s marked as paused, waiting for peer response or user input.",
+                component.log_identifier,
+                logical_task_id,
+            )
+
         log.debug(
             "%s ADK task %s awaited and completed (Paused: %s).",
             component.log_identifier,
@@ -194,13 +201,13 @@ async def run_adk_async_task_thread_wrapper(
             llm_limit_e,
         )
     except BadRequestError as e:
+        exception_to_finalize_with = e
         log.error(
-            "%s Bad Request for task %s: %s.",
+            "%s Bad Request for task %s: %s. Scheduling finalization.",
             component.log_identifier,
             logical_task_id,
-            e.message
+            e.message,
         )
-        raise
     except Exception as e:
         exception_to_finalize_with = e
         log.exception(
@@ -254,7 +261,11 @@ async def run_adk_async_task(
     logical_task_id = a2a_context.get("logical_task_id", "unknown_task")
     event_loop_stored = False
     current_loop = asyncio.get_running_loop()
-    is_paused = False
+    # Track pending long-running tools by their function_call IDs
+    # This replaces the simple boolean is_paused to properly handle sync returns
+    pending_long_running_tools: set[str] = set()
+    # Collect synchronous responses from long-running tools for potential re-run
+    sync_tool_responses: list[adk_types.Part] = []
 
     adk_event_generator = component.runner.run_async(
         user_id=adk_session.user_id,
@@ -309,7 +320,10 @@ async def run_adk_async_task(
                 break
 
             if event.long_running_tool_ids:
-                is_paused = True
+                # Track which long-running tool calls are pending (waiting for async response)
+                pending_long_running_tools = pending_long_running_tools.union(
+                    event.long_running_tool_ids
+                )
 
             if not event_loop_stored and event.invocation_id:
                 task_context.set_event_loop(current_loop)
@@ -335,8 +349,18 @@ async def run_adk_async_task(
             if event.content and event.content.parts:
                 for part in event.content.parts:
                     if part.function_response:
-                        if part.function_response.name.startswith("peer_"):
-                            pass
+                        # Check if this is a sync response from a long-running tool
+                        # (i.e., the tool returned immediately instead of async)
+                        response_id = part.function_response.id
+                        if response_id and response_id in pending_long_running_tools:
+                            pending_long_running_tools.discard(response_id)
+                            sync_tool_responses.append(part)
+                            log.info(
+                                "%s Long-running tool %s (id=%s) returned synchronously.",
+                                component.log_identifier,
+                                part.function_response.name,
+                                response_id,
+                            )
 
     except TaskCancelledError:
         raise
@@ -345,7 +369,7 @@ async def run_adk_async_task(
             "%s Bad Request for task %s: %s.",
             component.log_identifier,
             logical_task_id,
-            e.message
+            e.message,
         )
         raise
     except Exception as e:
@@ -367,14 +391,55 @@ async def run_adk_async_task(
             f"Task {logical_task_id} was cancelled before finalization."
         )
 
-    if is_paused:
+    invocation_id = a2a_context.get("invocation_id")
+
+    # Check if we still have pending long-running tools (waiting for async responses)
+    if pending_long_running_tools:
+        # Store any sync responses using the SAME format as event_handlers.py
+        # This ensures they're combined with async responses when _retrigger is called
+        if sync_tool_responses:
+            for part in sync_tool_responses:
+                result = {
+                    "adk_function_call_id": part.function_response.id,
+                    "peer_tool_name": part.function_response.name,
+                    "payload": part.function_response.response,  # Already a dict from ADK
+                }
+                task_context.record_parallel_result(result, invocation_id)
+            log.info(
+                "%s Stored %d sync tool response(s) for later combination. Waiting for: %s",
+                component.log_identifier,
+                len(sync_tool_responses),
+                pending_long_running_tools,
+            )
+
         log.info(
-            "%s ADK run completed by invoking a long-running tool. Task %s will remain open, awaiting peer response.",
+            "%s Task %s paused, waiting for %d async tool response(s).",
             component.log_identifier,
             logical_task_id,
+            len(pending_long_running_tools),
         )
         return True
 
+    # All tools returned synchronously - re-run ADK with their responses
+    # The ADK already created the Part objects, so we use them directly (no duplication)
+    if sync_tool_responses:
+        log.info(
+            "%s All %d long-running tool(s) returned synchronously for task %s. Re-running ADK.",
+            component.log_identifier,
+            len(sync_tool_responses),
+            logical_task_id,
+        )
+        # Use the Part objects directly from the ADK (already properly formatted)
+        tool_response_content = adk_types.Content(role="tool", parts=sync_tool_responses)
+        return await run_adk_async_task(
+            component=component,
+            task_context=task_context,
+            adk_session=adk_session,
+            adk_content=tool_response_content,
+            run_config=run_config,
+            a2a_context=a2a_context,
+        )
+
     log.debug(
         "%s ADK run_async completed for task %s. Returning to wrapper for finalization.",
         component.log_identifier,

solace_agent_mesh/agent/adk/schema_migration.py

@@ -0,0 +1,88 @@
+"""
+ADK Database Schema Migrations
+
+Automatically runs Alembic migrations on agent startup to ensure
+database schema compatibility with the installed ADK version.
+
+This uses Google ADK's official migration approach via Alembic's
+autogenerate feature to detect and apply schema changes.
+"""
+
+import logging
+import re
+from pathlib import Path
+from alembic.config import Config
+from alembic import command
+
+log = logging.getLogger(__name__)
+
+
+def run_migrations(db_service, component):
+    """
+    Run Alembic database migrations programmatically.
+
+    Executes any pending migrations to ensure the database schema
+    matches ADK's model definitions. This is equivalent to running:
+        alembic upgrade head
+
+    Args:
+        db_service: DatabaseSessionService instance
+        component: Component that owns this service (for logging)
+
+    Raises:
+        RuntimeError: If migration fails
+    """
+
+    try:
+        # Get paths to alembic directory and config
+        module_dir = Path(__file__).parent
+        alembic_ini = module_dir / "alembic.ini"
+        alembic_dir = module_dir / "alembic"
+
+        # Verify files exist
+        if not alembic_ini.exists():
+            log.warning(
+                "%s alembic.ini not found at %s, skipping migration",
+                component.log_identifier,
+                alembic_ini
+            )
+            return
+
+        if not alembic_dir.exists():
+            log.warning(
+                "%s alembic/ directory not found at %s, skipping migration",
+                component.log_identifier,
+                alembic_dir
+            )
+            return
+
+        # Create Alembic config
+        alembic_cfg = Config(str(alembic_ini))
+        alembic_cfg.set_main_option("script_location", str(alembic_dir))
+        
+        # IMPORTANT: Store the engine in config attributes so env.py can access it
+        # This avoids URL encoding issues entirely
+        alembic_cfg.attributes['connection'] = db_service.db_engine
+
+        log.info(
+            "%s Running Alembic migrations for ADK schema compatibility...",
+            component.log_identifier
+        )
+
+        # Run migrations (equivalent to: alembic upgrade head)
+        # Run migrations (env.py will use the engine from attributes)
+        command.upgrade(alembic_cfg, "head")
+
+        log.info(
+            "%s Database schema migration complete",
+            component.log_identifier
+        )
+
+    except Exception as e:
+        log.error(
+            "%s Database migration failed: %s",
+            component.log_identifier,
+            e
+        )
+        raise RuntimeError(f"ADK database migration failed: {e}") from e
+    

solace_agent_mesh/agent/adk/services.py

@@ -5,30 +5,36 @@ Initializes ADK Services based on configuration.
 import logging
 import os
 import re
-from typing import Dict, Optional, List, Any
-from typing_extensions import override
-
-from google.genai import types as adk_types
+from typing import Any, Dict, List, Optional
 
-
-from google.adk.sessions import (
-    BaseSessionService,
-    InMemorySessionService,
-    DatabaseSessionService,
-    VertexAiSessionService,
-)
 from google.adk.artifacts import (
     BaseArtifactService,
-    InMemoryArtifactService,
     GcsArtifactService,
+    InMemoryArtifactService,
+)
+from google.adk.artifacts.base_artifact_service import ArtifactVersion
+from google.adk.auth.credential_service.base_credential_service import (
+    BaseCredentialService,
+)
+from google.adk.auth.credential_service.in_memory_credential_service import (
+    InMemoryCredentialService,
 )
 from google.adk.memory import (
     BaseMemoryService,
     InMemoryMemoryService,
     VertexAiRagMemoryService,
 )
+from google.adk.sessions import (
+    BaseSessionService,
+    DatabaseSessionService,
+    InMemorySessionService,
+    VertexAiSessionService,
+)
+from google.genai import types as adk_types
+from typing_extensions import override
 
 from .artifacts.filesystem_artifact_service import FilesystemArtifactService
+from .schema_migration import run_migrations
 
 log = logging.getLogger(__name__)
 
@@ -152,6 +158,42 @@ class ScopedArtifactServiceWrapper(BaseArtifactService):
             filename=filename,
         )
 
+    @override
+    async def list_artifact_versions(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        filename: str,
+        session_id: str,
+    ) -> List[ArtifactVersion]:
+        scoped_app_name = self._get_scoped_app_name(app_name)
+        return await self.wrapped_service.list_artifact_versions(
+            app_name=scoped_app_name,
+            user_id=user_id,
+            filename=filename,
+            session_id=session_id,
+        )
+
+    @override
+    async def get_artifact_version(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        filename: str,
+        session_id: str,
+        version: Optional[int] = None,
+    ) -> Optional[ArtifactVersion]:
+        scoped_app_name = self._get_scoped_app_name(app_name)
+        return await self.wrapped_service.get_artifact_version(
+            app_name=scoped_app_name,
+            user_id=user_id,
+            filename=filename,
+            session_id=session_id,
+            version=version,
+        )
+
 
 def _sanitize_for_path(identifier: str) -> str:
     """Sanitizes a string to be safe for use as a directory name."""
@@ -170,9 +212,9 @@ def initialize_session_service(component) -> BaseSessionService:
     config = component.get_config("session_service", {})
 
     # Handle both dict and SessionServiceConfig object
-    if hasattr(config, 'type'):
+    if hasattr(config, "type"):
         service_type = config.type.lower()
-        db_url = getattr(config, 'database_url', None)
+        db_url = getattr(config, "database_url", None)
     else:
         service_type = config.get("type", "memory").lower()
         db_url = config.get("database_url")
@@ -191,7 +233,9 @@ def initialize_session_service(component) -> BaseSessionService:
                 f"{component.log_identifier} 'database_url' is required for sql session service."
             )
         try:
-            return DatabaseSessionService(db_url=db_url)
+            db_service = DatabaseSessionService(db_url=db_url)
+            run_migrations(db_service, component)
+            return db_service
         except ImportError:
             log.error(
                 "%s SQLAlchemy not installed. Please install 'google-adk[database]' or 'sqlalchemy'.",
@@ -280,7 +324,6 @@ def initialize_artifact_service(component) -> BaseArtifactService:
                 "aws_secret_access_key",
                 "aws_session_token",
                 "region_name",
-                "endpoint_url",
                 "config",
             ]
 
@@ -296,8 +339,9 @@ def initialize_artifact_service(component) -> BaseArtifactService:
                     s3_config[key] = config[key]
 
             # Set credentials from environment variables as a fallback.
-            if "endpoint_url" not in s3_config:
-                s3_config["endpoint_url"] = os.environ.get("S3_ENDPOINT_URL")
+            endpoint_url = config.get("endpoint_url") or os.environ.get("S3_ENDPOINT_URL") or "https://s3.amazonaws.com"
+            s3_config["endpoint_url"] = endpoint_url
+
             if "aws_access_key_id" not in s3_config:
                 env_access_key = os.environ.get("AWS_ACCESS_KEY_ID")
                 if env_access_key is not None:
@@ -390,3 +434,35 @@ def initialize_memory_service(component) -> BaseMemoryService:
         raise ValueError(
             f"{component.log_identifier} Unsupported memory service type: {service_type}"
         )
+
+
+def initialize_credential_service(component) -> BaseCredentialService | None:
+    """Initializes the ADK Credential Service based on configuration."""
+    config = component.get_config("credential_service", None)
+
+    # If no credential service is configured, return None
+    if config is None:
+        log.info(
+            "%s No credential service configured, skipping initialization",
+            component.log_identifier,
+        )
+        return None
+
+    # Handle both dict and CredentialServiceConfig object
+    if hasattr(config, "type"):
+        service_type = config.type.lower()
+    else:
+        service_type = config.get("type", "memory").lower()
+
+    log.info(
+        "%s Initializing Credential Service of type: %s",
+        component.log_identifier,
+        service_type,
+    )
+
+    if service_type == "memory":
+        return InMemoryCredentialService()
+    else:
+        raise ValueError(
+            f"{component.log_identifier} Unsupported credential service type: {service_type}"
+        )