aixtools 0.2.5__py3-none-any.whl → 0.2.7__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.5'
-__version_tuple__ = version_tuple = (0, 2, 5)
+__version__ = version = '0.2.7'
+__version_tuple__ = version_tuple = (0, 2, 7)
 
 __commit_id__ = commit_id = None

aixtools/a2a/google_sdk/utils.py

@@ -1,3 +1,5 @@
+"""Utilities for handling A2A SDK agent cards and connections."""
+
 import asyncio
 
 import httpx
@@ -8,13 +10,34 @@ 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
+
+logger = get_logger(__name__)
+
+DEFAULT_A2A_TIMEOUT = 60.0
 
 
 class AgentCardLoadFailedError(Exception):
     pass
 
 
+async def get_agent_card(client: httpx.AsyncClient, address: str) -> AgentCard:
+    """Retrieve the agent card from the given agent address."""
+    for card_path in [AGENT_CARD_WELL_KNOWN_PATH, PREV_AGENT_CARD_WELL_KNOWN_PATH]:
+        try:
+            card_resolver = A2ACardResolver(client, address, card_path)
+            card = await card_resolver.get_agent_card()
+            card.url = address
+            return card
+        except Exception as e:
+            logger.warning(f"Error retrieving agent card from {address} at path {card_path}: {e}")
+
+    raise AgentCardLoadFailedError(f"Failed to load agent card from {address}")
+
+
 class _AgentCardResolver:
+    """Helper class to resolve and manage agent cards and their connections."""
+
     def __init__(self, client: httpx.AsyncClient):
         self._httpx_client = client
         self._a2a_client_factory = ClientFactory(ClientConfig(httpx_client=self._httpx_client))
@@ -25,17 +48,13 @@ class _AgentCardResolver:
         self.clients[card.name] = remote_connection
 
     async def retrieve_card(self, address: str):
-        for card_path in [AGENT_CARD_WELL_KNOWN_PATH, PREV_AGENT_CARD_WELL_KNOWN_PATH]:
-            try:
-                card_resolver = A2ACardResolver(self._httpx_client, address, card_path)
-                card = await card_resolver.get_agent_card()
-                card.url = address
-                self.register_agent_card(card)
-                return
-            except Exception as e:
-                print(f"Error retrieving agent card from {address} at path {card_path}: {e}")
-
-        raise AgentCardLoadFailedError(f"Failed to load agent card from {address}")
+        try:
+            card = await get_agent_card(self._httpx_client, address)
+            self.register_agent_card(card)
+            return
+        except Exception as e:
+            logger.error(f"Error retrieving agent card from {address}: {e}")
+            return
 
     async def get_a2a_clients(self, agent_hosts: list[str]) -> dict[str, RemoteAgentConnection]:
         async with asyncio.TaskGroup() as task_group:
@@ -45,15 +64,19 @@ class _AgentCardResolver:
         return self.clients
 
 
-async def get_a2a_clients(ctx: SessionIdTuple, agent_hosts: list[str]) -> dict[str, RemoteAgentConnection]:
+async def get_a2a_clients(
+    ctx: SessionIdTuple, agent_hosts: list[str], *, 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],
     }
-    httpx_client = httpx.AsyncClient(headers=headers, timeout=60.0)
+    httpx_client = httpx.AsyncClient(headers=headers, timeout=timeout, follow_redirects=True)
     return await _AgentCardResolver(httpx_client).get_a2a_clients(agent_hosts)
 
 
 def get_session_id_tuple(context: RequestContext) -> SessionIdTuple:
+    """Get the user_id, session_id tuple from the request context."""
     headers = context.call_context.state.get("headers", {})
     return headers.get("user-id", DEFAULT_USER_ID), headers.get("session-id", DEFAULT_SESSION_ID)

aixtools/auth/auth.py

@@ -2,21 +2,60 @@
 Module that manages OAuth2 functions for authentication
 """
 
+import enum
 import logging
 
 import jwt
-from jwt import ExpiredSignatureError, InvalidAudienceError, InvalidIssuerError, PyJWKClient
+from fastapi import HTTPException
+from jwt import ExpiredSignatureError, InvalidAudienceError, InvalidIssuerError, InvalidSignatureError, PyJWKClient
 
 from aixtools.utils import config
 
 logger = logging.getLogger(__name__)
 
 
+class AuthTokenErrorCode(str, enum.Enum):
+    """Enum for error codes returned by the AuthTokenError exception."""
+
+    TOKEN_EXPIRED = "Token expired"
+    INVALID_AUDIENCE = "Token not for expected audience"
+    INVALID_ISSUER = "Token not for expected issuer"
+    INVALID_SIGNATURE = "Token signature error"
+    INVALID_TOKEN = "Invalid token"
+    JWT_ERROR = "Generic JWT error"
+    MISSING_GROUPS_ERROR = "Missing authorized groups"
+    INVALID_TOKEN_SCOPE = "Token scope does not match configured scope"
+
+
 class AuthTokenError(Exception):
     """Exception raised for authentication token errors."""
 
+    def __init__(self, error_code: AuthTokenErrorCode, msg: str = None):
+        self.error_code = error_code
+        error_msg = error_code.value if msg is None else msg
+        super().__init__(error_msg)
+
+    def to_http_exception(self, required_scope: str = None, realm: str = "MCP") -> HTTPException:
+        """
+        Returns an HTTPException with 401 status for all AuthTokenErrorCode,
+        including MCP JSON body and WWW-Authenticate header.
+        """
+        status_code = 401
+        www_error = (
+            "insufficient_scope" if self.error_code == AuthTokenErrorCode.INVALID_TOKEN_SCOPE else "invalid_token"
+        )
+
+        header_value = f'Bearer realm="{realm}", error="{www_error}", error_description="{self.error_code.value}"'
+        if self.error_code == AuthTokenErrorCode.INVALID_TOKEN_SCOPE and required_scope:
+            header_value += f', scope="{required_scope}"'
+
+        detail = {"error": {"code": self.error_code.name, "message": self.error_code.value}}
+        if self.error_code == AuthTokenErrorCode.INVALID_TOKEN_SCOPE and required_scope:
+            detail["error"]["required_scope"] = required_scope
+
+        return HTTPException(status_code=status_code, detail=detail, headers={"WWW-Authenticate": header_value})
+
 
-# pylint: disable=too-few-public-methods
 class AccessTokenVerifier:
     """
     Verifies Microsoft SSO JWT token against the configured Tenant ID, Audience, API ID and Issuer URL.
@@ -25,7 +64,12 @@ class AccessTokenVerifier:
     def __init__(self):
         tenant_id = config.APP_TENANT_ID
         self.api_id = config.APP_API_ID
-        self.issuer_url = config.APP_ISSUER_URL
+        self.issuer_url = f"https://sts.windows.net/{tenant_id}/"
+
+        self.authorized_groups = set(config.APP_AUTHORIZED_GROUPS.split(",")) if config.APP_AUTHORIZED_GROUPS else set()
+        if not self.authorized_groups:
+            logger.warning("No authorized groups configured")
+
         # Azure AD endpoints
         jwks_url = f"https://login.microsoftonline.com/{tenant_id}/discovery/v2.0/keys"
         self.jwks_client = PyJWKClient(
@@ -48,7 +92,7 @@ class AccessTokenVerifier:
         """
         try:
             signing_key = self.jwks_client.get_signing_key_from_jwt(token)
-
+            logger.info("Verifying JWT token")
             claims = jwt.decode(
                 token,
                 signing_key.key,
@@ -58,13 +102,48 @@ class AccessTokenVerifier:
                 # ensure audience verification is carried out
                 options={"verify_aud": True},
             )
+            logger.info("Verified JWT token")
             return claims
 
         except ExpiredSignatureError as e:
-            raise AuthTokenError("Token expired") from e
+            raise AuthTokenError(AuthTokenErrorCode.TOKEN_EXPIRED) from e
         except InvalidAudienceError as e:
-            raise AuthTokenError(f"Token not for expected audience: {e}") from e
+            raise AuthTokenError(AuthTokenErrorCode.INVALID_AUDIENCE) from e
         except InvalidIssuerError as e:
-            raise AuthTokenError(f"Token not for expected issuer: {e}") from e
+            raise AuthTokenError(AuthTokenErrorCode.INVALID_ISSUER) from e
+        except InvalidSignatureError as e:
+            raise AuthTokenError(AuthTokenErrorCode.INVALID_SIGNATURE) from e
         except jwt.exceptions.PyJWTError as e:
-            raise AuthTokenError(f"Invalid token: {e}") from e
+            raise AuthTokenError(AuthTokenErrorCode.JWT_ERROR) from e
+
+    def authorize_claims(self, claims: dict, expected_scope: str):
+        """
+        Authorize claims based on token scope, expected scope and authorized groups
+        claims: decoded JWT claims
+        expected_scope: expected scope for the token
+        Raises AuthTokenError if authorization fails.
+        """
+        logger.info("Checking JWT token claims")
+        if expected_scope:
+            token_scopes = claims.get("scp", "").split()
+            if expected_scope not in token_scopes:
+                logger.error("Expected token scope: %s, got: %s", expected_scope, token_scopes)
+                raise AuthTokenError(
+                    AuthTokenErrorCode.INVALID_TOKEN_SCOPE,
+                    f"Expected token scope: {expected_scope}, got: {token_scopes}",
+                )
+
+        if not self.authorized_groups:
+            logger.info("Authorized JWT token, no authorized groups configured")
+            return
+
+        groups = claims.get("groups", [])
+        if self.authorized_groups & set(groups):
+            logger.info("Authorized JWT token, against %s", groups)
+            return
+
+        logger.error("Could not find any group in JWT token, matching: %s", self.authorized_groups)
+        raise AuthTokenError(
+            AuthTokenErrorCode.MISSING_GROUPS_ERROR,
+            f"Could not find any group in JWT token, matching: {self.authorized_groups}",
+        )

aixtools/utils/config.py

@@ -135,5 +135,6 @@ APP_CLIENT_ID = get_variable_env("APP_CLIENT_ID")
 # used for token audience check
 APP_API_ID = get_variable_env("APP_API_ID")
 APP_TENANT_ID = get_variable_env("APP_TENANT_ID")
-# used for token issuer check
-APP_ISSUER_URL = get_variable_env("APP_ISSUER_URL")
+
+# used for token authorization check
+APP_AUTHORIZED_GROUPS = get_variable_env("APP_AUTHORIZED_GROUPS", allow_empty=True)

{aixtools-0.2.5.dist-info → aixtools-0.2.7.dist-info}/METADATA

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

{aixtools-0.2.5.dist-info → aixtools-0.2.7.dist-info}/RECORD

@@ -1,5 +1,5 @@
 aixtools/__init__.py,sha256=9NGHm7LjsQmsvjTZvw6QFJexSvAU4bCoN_KBk9SCa00,260
-aixtools/_version.py,sha256=9wrJ_4Dlc0arUzKiaIqvTY85rMJma3eb1nNlF3uHAxU,704
+aixtools/_version.py,sha256=yXzK2akXKIKUAfJk0WCQothqygqvndys6GBuXxo-wk0,704
 aixtools/app.py,sha256=JzQ0nrv_bjDQokllIlGHOV0HEb-V8N6k_nGQH-TEsVU,5227
 aixtools/chainlit.md,sha256=yC37Ly57vjKyiIvK4oUvf4DYxZCwH7iocTlx7bLeGLU,761
 aixtools/context.py,sha256=I_MD40ZnvRm5WPKAKqBUAdXIf8YaurkYUUHSVVy-QvU,598
@@ -22,7 +22,7 @@ aixtools/a2a/utils.py,sha256=EHr3IyyBJn23ni-JcfAf6i3VpQmPs0g1TSnAZazvY_8,4039
 aixtools/a2a/google_sdk/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 aixtools/a2a/google_sdk/card.py,sha256=P0L3bKbm28HaRkcIxIvjuSGUKOOc0ymyRAFHKm3a5GQ,996
 aixtools/a2a/google_sdk/remote_agent_connection.py,sha256=oDCRSN3gfONY1Ibp8BrtysIVqfQQ-lWe5N7lr1ymHxY,2819
-aixtools/a2a/google_sdk/utils.py,sha256=hjrNRZywJEUxHHaOttJFQU0FLzteg0Ggtm3qAeXMSVw,2430
+aixtools/a2a/google_sdk/utils.py,sha256=kag7KWNRatpw9bf9ThqBdDlyn0QuBx8-Pja1id1Gk30,3242
 aixtools/a2a/google_sdk/pydantic_ai_adapter/agent_executor.py,sha256=MMbhbEnUL6NwSYnisJrDdHW8zJoSyJ3Pzzkt8jqwNdI,7066
 aixtools/a2a/google_sdk/pydantic_ai_adapter/storage.py,sha256=nGoVL7MPoZJW7iVR71laqpUYP308yFKZIifJtvUgpiU,878
 aixtools/agents/__init__.py,sha256=MAW196S2_G7uGqv-VNjvlOETRfuV44WlU1leO7SiR0A,282
@@ -31,7 +31,7 @@ aixtools/agents/agent_batch.py,sha256=0Zu9yNCRPAQZPjXQ-dIUAmP1uGTVbxVt7xvnMpoJMj
 aixtools/agents/print_nodes.py,sha256=wVTngNfqM0As845WTRz6G3Rei_Gr3HuBlvu-G_eXuig,1665
 aixtools/agents/prompt.py,sha256=p9OYnyJ4-MyGXwHPrQeJBhZ2a3RV2HqhtdUUCrTMsAQ,3361
 aixtools/auth/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-aixtools/auth/auth.py,sha256=aKYCKJRjSNrVZmIWN2h2p1zYqkhMLLBXBfk_Qy5NKik,2365
+aixtools/auth/auth.py,sha256=4vNcljdpFzvIBXqjK7i3xAy-CC8Hc7ii11t-11z6ofY,6041
 aixtools/compliance/__init__.py,sha256=vnw0zEdySIJWvDAJ8DCRRaWmY_agEOz1qlpAdhmtiuo,191
 aixtools/compliance/private_data.py,sha256=OOM9mIp3_w0fNgj3VAEWBl7-jrPc19_Ls1pC5dfF5UY,5323
 aixtools/db/__init__.py,sha256=b8vRhme3egV-aUZbAntnOaDkSXB8UT0Xy5oqQhU_z0Q,399
@@ -79,7 +79,7 @@ aixtools/tools/doctor/mcp_tool_doctor.py,sha256=sX2q5GfNkmUYxnXrqMpeGIwGfeL1LpYJ
 aixtools/tools/doctor/tool_doctor.py,sha256=EY1pshjLGLD0j6cc1ZFtbc0G19I5IbOZwHFDqypE49Q,2661
 aixtools/tools/doctor/tool_recommendation.py,sha256=LYyVOSXdAorWiY4P-ucSA1vLlV5BTEfX4GzBXNE_X0M,1569
 aixtools/utils/__init__.py,sha256=xT6almZBQYMfj4h7Hq9QXDHyVXbOOTxqLsmJsxYYnSw,757
-aixtools/utils/config.py,sha256=t32731F53Cv1YYoX95wksoreE0Zn0B8UKyEiKWne4ec,5147
+aixtools/utils/config.py,sha256=fCATakEfeASo6eFMh5Me_idwMQgiITS3LMYc2Z1bGzU,5187
 aixtools/utils/config_util.py,sha256=3Ya4Qqhj1RJ1qtTTykQ6iayf5uxlpigPXgEJlTi1wn4,2229
 aixtools/utils/enum_with_description.py,sha256=zjSzWxG74eR4x7dpmb74pLTYCWNSMvauHd7_9LpDYIw,1088
 aixtools/utils/files.py,sha256=8JnxwHJRJcjWCdFpjzWmo0po2fRg8esj4H7sOxElYXU,517
@@ -89,8 +89,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.5.dist-info/METADATA,sha256=BHPUgnHXs7ET3BvwAkPxYRkkXnxLdptFwbYNDkoBMbw,27229
-aixtools-0.2.5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-aixtools-0.2.5.dist-info/entry_points.txt,sha256=q8412TG4T0S8K0SKeWp2vkVPIDYQs0jNoHqcQ7qxOiA,155
-aixtools-0.2.5.dist-info/top_level.txt,sha256=wBn-rw9bCtxrR4AYEYgjilNCUVmKY0LWby9Zan2PRJM,9
-aixtools-0.2.5.dist-info/RECORD,,
+aixtools-0.2.7.dist-info/METADATA,sha256=V3-mduY2Z8lURHq7Vhe4KVyieG1snJAHeSMmvVL2k5k,27229
+aixtools-0.2.7.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+aixtools-0.2.7.dist-info/entry_points.txt,sha256=q8412TG4T0S8K0SKeWp2vkVPIDYQs0jNoHqcQ7qxOiA,155
+aixtools-0.2.7.dist-info/top_level.txt,sha256=wBn-rw9bCtxrR4AYEYgjilNCUVmKY0LWby9Zan2PRJM,9
+aixtools-0.2.7.dist-info/RECORD,,