aixtools 0.2.18__py3-none-any.whl → 0.2.20__py3-none-any.whl

aixtools/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.2.18'
-__version_tuple__ = version_tuple = (0, 2, 18)
+__version__ = version = '0.2.20'
+__version_tuple__ = version_tuple = (0, 2, 20)
 
 __commit_id__ = commit_id = None

aixtools/a2a/google_sdk/utils.py

@@ -11,6 +11,7 @@ from a2a.utils import AGENT_CARD_WELL_KNOWN_PATH, PREV_AGENT_CARD_WELL_KNOWN_PAT
 from aixtools.a2a.google_sdk.remote_agent_connection import RemoteAgentConnection
 from aixtools.context import DEFAULT_SESSION_ID, DEFAULT_USER_ID, SessionIdTuple
 from aixtools.logging.logging_config import get_logger
+from aixtools.server.utils import create_session_headers
 
 logger = get_logger(__name__)
 
@@ -23,6 +24,7 @@ class AgentCardLoadFailedError(Exception):
 
 async def get_agent_card(client: httpx.AsyncClient, address: str) -> AgentCard:
     """Retrieve the agent card from the given agent address."""
+    warnings = []
     for card_path in [AGENT_CARD_WELL_KNOWN_PATH, PREV_AGENT_CARD_WELL_KNOWN_PATH]:
         try:
             card_resolver = A2ACardResolver(client, address, card_path)
@@ -30,8 +32,10 @@ async def get_agent_card(client: httpx.AsyncClient, address: str) -> AgentCard:
             card.url = address
             return card
         except Exception as e:
-            logger.warning(f"Error retrieving agent card from {address} at path {card_path}: {e}")
+            warnings.append(f"Error retrieving agent card from {address} at path {card_path}: {e}")
 
+    for warning in warnings:
+        logger.warning(warning)
     raise AgentCardLoadFailedError(f"Failed to load agent card from {address}")
 
 
@@ -65,15 +69,19 @@ class _AgentCardResolver:
 
 
 async def get_a2a_clients(
-    ctx: SessionIdTuple, agent_hosts: list[str], *, timeout: float = DEFAULT_A2A_TIMEOUT
+    agent_hosts: list[str],
+    session_id_tuple: SessionIdTuple,
+    auth_token: str = None,
+    *,
+    timeout: float = DEFAULT_A2A_TIMEOUT,
 ) -> dict[str, RemoteAgentConnection]:
     """Get A2A clients for all agents defined in the configuration."""
-    headers = {
-        "user-id": ctx[0],
-        "session-id": ctx[1],
-    }
+    headers = create_session_headers(session_id_tuple, auth_token)
     httpx_client = httpx.AsyncClient(headers=headers, timeout=timeout, follow_redirects=True)
-    return await _AgentCardResolver(httpx_client).get_a2a_clients(agent_hosts)
+    clients = await _AgentCardResolver(httpx_client).get_a2a_clients(agent_hosts)
+    for client in clients.values():
+        logger.info("Using A2A server at: %s", client.get_agent_card().url)
+    return clients
 
 
 def card2description(card: AgentCard) -> str:

aixtools/mcp/__init__.py

@@ -4,8 +4,12 @@ Model Context Protocol (MCP) implementation for AI agent communication.
 
 from aixtools.mcp.exceptions import AixToolError
 from aixtools.mcp.fast_mcp_log import FastMcpLog
+from aixtools.mcp.middleware import AixErrorHandlingMiddleware
+from aixtools.mcp.server import create_mcp_server
 
 __all__ = [
-    "FastMcpLog",
+    "AixErrorHandlingMiddleware",
     "AixToolError",
+    "FastMcpLog",
+    "create_mcp_server",
 ]

aixtools/mcp/client.py

@@ -3,6 +3,7 @@
 import asyncio
 import logging
 from contextlib import asynccontextmanager
+from dataclasses import dataclass, field
 from datetime import timedelta
 from typing import Any, AsyncGenerator
 
@@ -21,6 +22,7 @@ from pydantic_ai.toolsets.abstract import ToolsetTool
 
 from aixtools.context import SessionIdTuple
 from aixtools.logging.logging_config import get_logger
+from aixtools.server.utils import create_session_headers
 
 MCP_TOOL_CACHE_TTL = 300  # 5 minutes
 DEFAULT_MCP_CONNECTION_TIMEOUT = 30
@@ -49,6 +51,14 @@ async def default_mcp_log_handler(message: LogMessage):
     logger.log(level, msg, extra=extra)
 
 
+@dataclass
+class MCPConfig:
+    """Configuration for an MCP server retrieved from config.yaml"""
+
+    url: str
+    read_timeout: float = field(default=DEFAULT_MCP_READ_TIMEOUT)
+
+
 def get_mcp_client(
     url: str | None = None,
     command: str | None = None,
@@ -73,57 +83,50 @@ def get_mcp_client(
     raise ValueError("Either url or command must be provided to create MCP client.")
 
 
-def get_mcp_headers(session_id_tuple: SessionIdTuple) -> dict[str, str] | None:
+def get_mcp_servers(
+    mcp_configs: list[MCPConfig],
+    session_id_tuple: SessionIdTuple,
+    auth_token: str = None,
+    *,
+    timeout: float = DEFAULT_MCP_CONNECTION_TIMEOUT,
+):
     """
-    Generate headers for MCP server requests.
+    Create cached MCP server instances with robust error handling and isolation.
 
-    This function creates a dictionary of headers to be used in requests to
-    the MCP servers. If a `user_id` or `session_id` is provided, they are
-    included in the headers.
+    This function creates and returns a list of `CachedMCPServerStreamableHTTP` instances
+    based on the provided URLs. Each server instance includes:
+    - TTL-based caching for tool lists (5 minutes default)
+    - Complete task isolation to prevent cancellation propagation
+    - Comprehensive error handling and fallback mechanisms
+    - Optional user/session headers for request authentication
 
     Args:
-        session_id_tuple (SessionIdTuple): user_id and session_id tuple
+        mcp_configs (list[MCPConfig]): A list of MCP server configurations to use.
+        session_id_tuple (SessionIdTuple): A tuple containing (user_id, session_id).
+        auth_token (str, optional): The authentication token for the user. Defaults to None.
+        timeout (float, optional): Timeout in seconds for MCP server connections.
     Returns:
-        dict[str, str] | None: A dictionary of headers for MCP server requests,
-                               or None if neither user_id nor session_id is
-                               provided. When None is returned, default headers
-                               from the client or transport will be used.
+        list[CachedMCPServerStreamableHTTP]: List of cached MCP server instances with
+                                            isolation and error handling. Each server
+                                            operates independently - failures in one
+                                            server won't affect others.
     """
-    headers = None
-    user_id, session_id = session_id_tuple
-    if session_id or user_id:
-        headers = {}
-        if session_id:
-            headers["session-id"] = session_id
-        if user_id:
-            headers["user-id"] = user_id
-    return headers
+    headers = create_session_headers(session_id_tuple, auth_token)
+    servers = []
+    for config in mcp_configs:
+        server = CachedMCPServerStreamableHTTP(
+            url=config.url, headers=headers, timeout=timeout, read_timeout=config.read_timeout
+        )
+        logger.info("Using MCP server at %s", config.url)
+        servers.append(server)
+    return servers
 
 
 def get_configured_mcp_servers(
     session_id_tuple: SessionIdTuple, mcp_urls: list[str], timeout: int = DEFAULT_MCP_CONNECTION_TIMEOUT
 ):
-    """
-    Retrieve the configured MCP server instances with optional caching.
-
-    Context values `user_id` and `session_id` are included in the headers for each server request.
-
-    Each server is wrapped in a try-except block to isolate them from each other.
-    If one server fails, it won't affect the others.
-
-    Args:
-        session_id_tuple (SessionIdTuple): A tuple containing (user_id, session_id).
-        mcp_urls: (list[str], optional): A list of MCP server URLs to use.
-        timeout (int, optional): Timeout in seconds for MCP server connections. Defaults to 30 seconds.
-    Returns:
-        list[MCPServerStreamableHTTP]: A list of configured MCP server instances. If
-                              neither user_id nor session_id is provided, the
-                              server instances will use default headers defined
-                              by the underlying HTTP implementation.
-    """
-    headers = get_mcp_headers(session_id_tuple)
-
-    return [CachedMCPServerStreamableHTTP(url=url, headers=headers, timeout=timeout) for url in mcp_urls]
+    """Create MCP server instances from a list of URLs."""
+    return get_mcp_servers([MCPConfig(url=url) for url in mcp_urls], session_id_tuple, timeout=timeout)
 
 
 class CachedMCPServerStreamableHTTP(MCPServerStreamableHTTP):
@@ -259,7 +262,7 @@ class CachedMCPServerStreamableHTTP(MCPServerStreamableHTTP):
 
         # First, check if we have a valid cached result
         if CACHE_KEY in self._tools_cache:
-            logger.info("Using cached tools for %s", self.url)
+            logger.debug("Using cached tools for %s", self.url)
             return self._tools_cache[CACHE_KEY]
 
         # Create isolated task to prevent cancellation propagation

aixtools/{logging/mcp_middleware.py → mcp/middleware.py}

@@ -7,7 +7,7 @@ import traceback
 from fastmcp.server.middleware.error_handling import ErrorHandlingMiddleware
 from fastmcp.server.middleware.middleware import MiddlewareContext
 
-from aixtools.mcp import AixToolError
+from aixtools.mcp.exceptions import AixToolError
 
 
 class AixErrorHandlingMiddleware(ErrorHandlingMiddleware):

aixtools/mcp/server.py

@@ -0,0 +1,72 @@
+"""Utilities for FastMCP servers."""
+
+from typing import Any
+
+from fastmcp import FastMCP
+from fastmcp.server.middleware.logging import LoggingMiddleware
+from fastmcp.server.middleware.timing import TimingMiddleware
+from fastmcp.utilities.types import NotSet
+
+from aixtools.auth.auth import AccessTokenAuthProvider
+from aixtools.logging.logging_config import get_logger
+from aixtools.mcp.middleware import AixErrorHandlingMiddleware
+
+
+def create_mcp_server(
+    *,
+    name: str,
+    instructions: str | None = None,
+    **kwargs: Any,
+) -> FastMCP:
+    """
+    MCP server instance with preconfigured auth and middleware.
+
+    All FastMCP constructor parameters are supported via **kwargs.
+
+    Args:
+        name: Server name
+        instructions: Optional server instructions
+        **kwargs: All other FastMCP constructor parameters:
+            - version: str | None
+            - auth: AuthProvider | None (AccessTokenAuthProvider if not set, pass None to disable)
+            - middleware: list[Middleware] | None (custom middleware if not set, pass None to disable)
+            - lifespan: Callable | None
+            - tool_serializer: Callable[[Any], str] | None
+            - cache_expiration_seconds: float | None
+            - on_duplicate_tools: DuplicateBehavior | None
+            - on_duplicate_resources: DuplicateBehavior | None
+            - on_duplicate_prompts: DuplicateBehavior | None
+            - resource_prefix_format: Literal["protocol", "path"] | None
+            - mask_error_details: bool | None
+            - tools: list[Tool | Callable] | None
+            - tool_transformations: dict[str, ToolTransformConfig] | None
+            - dependencies: list[str] | None
+            - include_tags: set[str] | None
+            - exclude_tags: set[str] | None
+            - include_fastmcp_meta: bool | None
+
+    Returns:
+        Configured FastMCP server instance
+    """
+    middleware = kwargs.pop("middleware", NotSet)
+    auth = kwargs.pop("auth", NotSet)
+
+    if middleware is NotSet:
+        middleware = [
+            LoggingMiddleware(include_payloads=True, logger=get_logger("middleware.log")),
+            AixErrorHandlingMiddleware(include_traceback=True, logger=get_logger("middleware.err")),
+            TimingMiddleware(logger=get_logger("middleware.timing")),
+        ]
+
+    if auth is NotSet:
+        auth = AccessTokenAuthProvider()
+
+    mcp_args = {
+        "name": name,
+        "instructions": instructions,
+        "auth": auth,
+        "middleware": middleware,
+        **kwargs,
+    }
+
+    return FastMCP(**mcp_args)

aixtools/server/utils.py

@@ -8,10 +8,13 @@ from functools import wraps
 from fastmcp import Context
 from fastmcp.server import dependencies
 
-from ..context import DEFAULT_SESSION_ID, DEFAULT_USER_ID, session_id_var, user_id_var
+from ..context import DEFAULT_SESSION_ID, DEFAULT_USER_ID, SessionIdTuple, session_id_var, user_id_var
+from ..logging.logging_config import get_logger
 
+logger = get_logger(__name__)
 
-def get_session_id_tuple(ctx: Context | None = None) -> tuple[str, str]:
+
+def get_session_id_tuple(ctx: Context | None = None) -> SessionIdTuple:
     """
     Get the user and session IDs from the user session.
     If `ctx` is None, the current FastMCP request HTTP headers are used.
@@ -60,6 +63,37 @@ def get_session_id_str(ctx: Context | None = None) -> str:
     return f"{user_id}:{session_id}"
 
 
+def create_session_headers(session_id_tuple: SessionIdTuple, auth_token: str | None = None) -> dict[str, str]:
+    """
+    Generate headers for MCP or A2A server requests.
+
+    This function creates a dictionary of headers to be used in requests to
+    the MCP servers. If a `user_id` or `session_id` is provided, they are
+    included in the headers.
+
+    Args:
+        session_id_tuple (SessionIdTuple): user_id and session_id tuple
+        auth_token (str | None): Optional authorization token to include in headers
+
+    Returns:
+        dict[str, str]: A dictionary of headers for MCP server requests.
+                       May be empty if no user_id, session_id, or auth_token is provided.
+    """
+    headers = {}
+    user_id, session_id = session_id_tuple
+    if auth_token:
+        logger.debug("Using auth token for MCP server authentication for user:%s, session_id:%s", user_id, session_id)
+        headers["Authorization"] = f"Bearer {auth_token}"
+    else:
+        logger.warning("No auth token found to forward to MCP/A2A servers.")
+
+    if session_id:
+        headers["session-id"] = session_id
+    if user_id:
+        headers["user-id"] = user_id
+    return headers
+
+
 def run_in_thread(func):
     """decorator to run blocking function with `asyncio.to_thread`"""
 

{aixtools-0.2.18.dist-info → aixtools-0.2.20.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aixtools
-Version: 0.2.18
+Version: 0.2.20
 Summary: Tools for AI exploration and debugging
 Requires-Python: >=3.11.2
 Description-Content-Type: text/markdown

{aixtools-0.2.18.dist-info → aixtools-0.2.20.dist-info}/RECORD

@@ -1,5 +1,5 @@
 aixtools/__init__.py,sha256=9NGHm7LjsQmsvjTZvw6QFJexSvAU4bCoN_KBk9SCa00,260
-aixtools/_version.py,sha256=dCdYCOiddErhHj0N8E7eGD9CFxW2AjdU7N-JosKZW0o,706
+aixtools/_version.py,sha256=4N3ayuoZZJYPEGMvrxu7tnGigRTxbAdCyp5a8y7c6aw,706
 aixtools/app.py,sha256=JzQ0nrv_bjDQokllIlGHOV0HEb-V8N6k_nGQH-TEsVU,5227
 aixtools/chainlit.md,sha256=yC37Ly57vjKyiIvK4oUvf4DYxZCwH7iocTlx7bLeGLU,761
 aixtools/context.py,sha256=I_MD40ZnvRm5WPKAKqBUAdXIf8YaurkYUUHSVVy-QvU,598
@@ -21,7 +21,7 @@ aixtools/a2a/app.py,sha256=p18G7fAInl9dcNYq6RStBjv1C3aD6oilQq3WXtBuk30,5069
 aixtools/a2a/utils.py,sha256=EHr3IyyBJn23ni-JcfAf6i3VpQmPs0g1TSnAZazvY_8,4039
 aixtools/a2a/google_sdk/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 aixtools/a2a/google_sdk/remote_agent_connection.py,sha256=gFsIja_nIQsAzxElv3A7_hG9GpDI5pkV-B_KDlcYlnc,3329
-aixtools/a2a/google_sdk/utils.py,sha256=4VIPV2GtG5IRY-KSZN6iRMS3ntxG2uKgd_d5tQvnZHw,3515
+aixtools/a2a/google_sdk/utils.py,sha256=rMezLwv7P-Drn2TPkuwTBkbW23yPhwOIywGb5CrFKE0,3828
 aixtools/a2a/google_sdk/pydantic_ai_adapter/agent_executor.py,sha256=8VuU2WXeSHUK3_rRm_mjX6elqdC9NA2uz1aELzeC8BU,9784
 aixtools/a2a/google_sdk/pydantic_ai_adapter/storage.py,sha256=nGoVL7MPoZJW7iVR71laqpUYP308yFKZIifJtvUgpiU,878
 aixtools/agents/__init__.py,sha256=MAW196S2_G7uGqv-VNjvlOETRfuV44WlU1leO7SiR0A,282
@@ -59,21 +59,22 @@ aixtools/logging/log_objects.py,sha256=gohsgcfyr8vsY7G_hfmj973-Ek1_PN-bMMLEUA-4u
 aixtools/logging/logging_config.py,sha256=LvxV3C75-I0096PpcCIbgM-Cp998LzWXeMM14HYbU20,4985
 aixtools/logging/mcp_log_models.py,sha256=7-H2GJXiiyLhpImuyLLftAGG4skxJal8Swax0ob04MY,3463
 aixtools/logging/mcp_logger.py,sha256=d2I5l4t0d6rQH17w23FpE1IUD8Ax-mSaKfByCH86q4I,6257
-aixtools/logging/mcp_middleware.py,sha256=0kpTAwvz9Fd_mFDP3J9ldH4dPP-bhcUjLJKELGOx0IQ,2257
 aixtools/logging/model_patch_logging.py,sha256=CW5-kKI-zNEgZhNV4vx3EQu6fbrEtX7VjA6fE5loRLQ,2916
 aixtools/logging/open_telemetry.py,sha256=fJjF1ou_8GyfNfbyWDQPGK6JAUrUaPwURYPHhXEtDBE,1121
-aixtools/mcp/__init__.py,sha256=Qp4uD1RtypCYgzWrt_ThBVxa5-CFBgcwMWkHbHFUQ54,232
-aixtools/mcp/client.py,sha256=0UR3KiPCg7WAEf8u4-ewp3H_hiaMwD9KKmx2XRzdA_8,18271
+aixtools/mcp/__init__.py,sha256=y6Akm5xeEFJ0K0TNM6aiTsUSe5k5kft8LSTaSEaY7vE,404
+aixtools/mcp/client.py,sha256=1flZ49irNzM6kdoKKocwFslFR3exDVWccoR0fp9SPvk,18396
 aixtools/mcp/example_client.py,sha256=QCFGP3NCNJMOKWjUOnFwjnbJhUSb879IA1ZYmwjRnmc,889
 aixtools/mcp/example_server.py,sha256=1SWCyrLWsAnOa81HC4QbPJo_lBVu0b3SZBWI-qDh1vQ,458
 aixtools/mcp/exceptions.py,sha256=9m3hy9fRINbXTSgkE71XW_luklgH6xBRp7VnSRfpyzQ,173
 aixtools/mcp/fast_mcp_log.py,sha256=XYOS406dVjn5YTHyGRsRvVNQ0SKlRObfrKj6EeLFjHg,1057
 aixtools/mcp/faulty_mcp.py,sha256=uU9vlNGCS_i2k20wocVMaDHTlYjMQxuzjILad9O1cjA,12807
+aixtools/mcp/middleware.py,sha256=V-bWhSL6F-tv74EtqaoYjbvxIUCIfG3XKz1khU7ebU4,2268
+aixtools/mcp/server.py,sha256=wpN7tJmUZqGsb_rOVfXbyTFJ9VPnQl7nG5wa1JAz-No,2566
 aixtools/model_patch/model_patch.py,sha256=JT-oHubIn2LeoSwWbzEQ5vLH7crJmFUecHyQfaAFHa0,1813
 aixtools/server/__init__.py,sha256=rwPx020YpOzCnrxA80Lc4yLLcIp-Mpe9hNqVO9wDPv0,448
 aixtools/server/app_mounter.py,sha256=0tJ0tC140ezAjnYdlhpLJQjY-TO8NVw7D8LseYCCVY8,3336
 aixtools/server/path.py,sha256=nI4yRQcE6gjKx5GG3PmHf7iT1FelT6Q8Xhw4ol9O1e0,5219
-aixtools/server/utils.py,sha256=tZWITIx6M-luV9yve4j3rPtYGSSA6zWS0JWEAySne_M,2276
+aixtools/server/utils.py,sha256=jIRgFT_tdgwuyI8vjc9iQJyDvC0CAoi09Gc5of6I1dg,3554
 aixtools/testing/__init__.py,sha256=mlmaAR2gmS4SbsYNCxnIprmFpFp-syjgVUkpUszo3mE,166
 aixtools/testing/agent_mock.py,sha256=D4DrRGaeKuWEhq8hdB5s9PdfLAXHMKiC3JA7M8yN13o,4244
 aixtools/testing/aix_test_model.py,sha256=i0xBdmpKoEfJHle6JDmcoJLUENN8Eqt181_WZ7XtDdU,6240
@@ -94,8 +95,8 @@ aixtools/utils/chainlit/cl_agent_show.py,sha256=vaRuowp4BRvhxEr5hw0zHEJ7iaSF_5bo
 aixtools/utils/chainlit/cl_utils.py,sha256=fxaxdkcZg6uHdM8uztxdPowg3a2f7VR7B26VPY4t-3c,5738
 aixtools/vault/__init__.py,sha256=fsr_NuX3GZ9WZ7dGfe0gp_5-z3URxAfwVRXw7Xyc0dU,141
 aixtools/vault/vault.py,sha256=9dZLWdZQk9qN_Q9Djkofw9LUKnJqnrX5H0fGusVLBhA,6037
-aixtools-0.2.18.dist-info/METADATA,sha256=h94-SI_IsK2rFnfBJs4R1k6qESnQVGd9T5UKMZV-sUw,27958
-aixtools-0.2.18.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-aixtools-0.2.18.dist-info/entry_points.txt,sha256=q8412TG4T0S8K0SKeWp2vkVPIDYQs0jNoHqcQ7qxOiA,155
-aixtools-0.2.18.dist-info/top_level.txt,sha256=wBn-rw9bCtxrR4AYEYgjilNCUVmKY0LWby9Zan2PRJM,9
-aixtools-0.2.18.dist-info/RECORD,,
+aixtools-0.2.20.dist-info/METADATA,sha256=Z5i3qRQdNfrFrKutnd6j0UMFab38xzXcKl7Hr2GKnos,27958
+aixtools-0.2.20.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+aixtools-0.2.20.dist-info/entry_points.txt,sha256=q8412TG4T0S8K0SKeWp2vkVPIDYQs0jNoHqcQ7qxOiA,155
+aixtools-0.2.20.dist-info/top_level.txt,sha256=wBn-rw9bCtxrR4AYEYgjilNCUVmKY0LWby9Zan2PRJM,9
+aixtools-0.2.20.dist-info/RECORD,,