hmdl 0.0.1__py3-none-any.whl → 0.0.2__py3-none-any.whl

hmdl/__init__.py

@@ -6,7 +6,7 @@ OpenTelemetry-based observability tracking.
 """
 
 from hmdl.client import HeimdallClient
-from hmdl.decorators import trace_mcp_tool
+from hmdl.decorators import trace_mcp_tool, UserExtractor, SessionExtractor
 from hmdl.config import HeimdallConfig
 from hmdl.types import SpanKind, SpanStatus
 
@@ -17,6 +17,8 @@ __all__ = [
     "HeimdallClient",
     # Decorators
     "trace_mcp_tool",
+    "UserExtractor",
+    "SessionExtractor",
     # Configuration
     "HeimdallConfig",
     # Types

hmdl/client.py

@@ -20,27 +20,31 @@ logger = logging.getLogger(__name__)
 
 class HeimdallClient:
     """Client for sending observability data to Heimdall platform.
-    
+
     This client sets up OpenTelemetry tracing and provides methods for
     creating spans and recording MCP operations.
-    
+
     Example:
         >>> from hmdl import HeimdallClient
         >>> client = HeimdallClient(api_key="your-api-key")
         >>> with client.start_span("my-operation") as span:
         ...     # Your code here
         ...     span.set_attribute("custom.attribute", "value")
+
+        # Track user sessions
+        >>> client.set_session_id("session-123")
+        >>> client.set_user_id("user-456")
     """
-    
+
     _instance: Optional["HeimdallClient"] = None
     _initialized: bool = False
-    
+
     def __new__(cls, *args: Any, **kwargs: Any) -> "HeimdallClient":
         """Singleton pattern to ensure only one client instance."""
         if cls._instance is None:
             cls._instance = super().__new__(cls)
         return cls._instance
-    
+
     def __init__(
         self,
         config: Optional[HeimdallConfig] = None,
@@ -50,6 +54,8 @@ class HeimdallClient:
         environment: Optional[str] = None,
         org_id: Optional[str] = None,
         project_id: Optional[str] = None,
+        session_id: Optional[str] = None,
+        user_id: Optional[str] = None,
     ) -> None:
         """Initialize the Heimdall client.
 
@@ -61,6 +67,8 @@ class HeimdallClient:
             environment: Deployment environment.
             org_id: Organization ID from Heimdall dashboard.
             project_id: Project ID from Heimdall dashboard.
+            session_id: Session ID for tracking MCP client sessions.
+            user_id: User ID for tracking users.
         """
         if self._initialized:
             return
@@ -76,16 +84,22 @@ class HeimdallClient:
                 environment=environment or HeimdallConfig().environment,
                 org_id=org_id or HeimdallConfig().org_id,
                 project_id=project_id or HeimdallConfig().project_id,
+                session_id=session_id or HeimdallConfig().session_id,
+                user_id=user_id or HeimdallConfig().user_id,
             )
-        
+
         self._tracer: Optional[trace.Tracer] = None
         self._provider: Optional[TracerProvider] = None
-        
+
+        # Runtime session and user tracking (can be updated dynamically)
+        self._session_id: Optional[str] = self.config.session_id
+        self._user_id: Optional[str] = self.config.user_id
+
         if self.config.enabled:
             self._setup_tracing()
-        
+
         self._initialized = True
-        
+
         # Register cleanup on exit
         atexit.register(self.shutdown)
     
@@ -181,6 +195,45 @@ class HeimdallClient:
             self._provider.shutdown()
             logger.debug("Heimdall client shutdown complete")
 
+    def get_session_id(self) -> Optional[str]:
+        """Get the current session ID."""
+        return self._session_id
+
+    def set_session_id(self, session_id: Optional[str]) -> None:
+        """Set the session ID for all subsequent spans.
+
+        Call this when an MCP client connects to associate all operations
+        with that session.
+
+        Args:
+            session_id: The session identifier (e.g., from MCP client connection)
+
+        Example:
+            >>> # When MCP client connects
+            >>> client.set_session_id(ctx.session_id or ctx.client_info.name)
+        """
+        self._session_id = session_id
+        logger.debug(f"Session ID set to: {session_id}")
+
+    def get_user_id(self) -> Optional[str]:
+        """Get the current user ID."""
+        return self._user_id
+
+    def set_user_id(self, user_id: Optional[str]) -> None:
+        """Set the user ID for all subsequent spans.
+
+        Can be overridden per-span using user_extractor option in decorators.
+
+        Args:
+            user_id: The user identifier
+
+        Example:
+            >>> # When user is identified
+            >>> client.set_user_id("user-123")
+        """
+        self._user_id = user_id
+        logger.debug(f"User ID set to: {user_id}")
+
     @classmethod
     def get_instance(cls) -> Optional["HeimdallClient"]:
         """Get the singleton client instance."""

hmdl/config.py

@@ -18,6 +18,10 @@ class HeimdallConfig:
         environment: Deployment environment (e.g., 'production', 'staging').
         org_id: Organization ID from Heimdall dashboard.
         project_id: Project ID to associate traces with in Heimdall.
+        session_id: Session ID to associate with all spans. Useful for tracking
+            requests from the same MCP client session.
+        user_id: User ID to associate with all spans. Can be overridden per-span
+            using user_extractor option in decorators.
         enabled: Whether tracing is enabled.
         debug: Enable debug logging.
         batch_size: Number of spans to batch before sending.
@@ -31,7 +35,7 @@ class HeimdallConfig:
     )
     endpoint: str = field(
         default_factory=lambda: os.environ.get(
-            "HEIMDALL_ENDPOINT", "https://api.heimdall.dev"
+            "HEIMDALL_ENDPOINT", "http://localhost:4318"
         )
     )
     service_name: str = field(
@@ -46,6 +50,12 @@ class HeimdallConfig:
     project_id: str = field(
         default_factory=lambda: os.environ.get("HEIMDALL_PROJECT_ID", "default")
     )
+    session_id: Optional[str] = field(
+        default_factory=lambda: os.environ.get("HEIMDALL_SESSION_ID")
+    )
+    user_id: Optional[str] = field(
+        default_factory=lambda: os.environ.get("HEIMDALL_USER_ID")
+    )
     enabled: bool = field(
         default_factory=lambda: os.environ.get("HEIMDALL_ENABLED", "true").lower() == "true"
     )

hmdl/decorators.py

@@ -2,11 +2,12 @@
 
 from __future__ import annotations
 
+import base64
 import functools
 import inspect
 import json
 import time
-from typing import Any, Callable, Optional, TypeVar, Union, overload
+from typing import Any, Callable, Dict, Optional, TypeVar, Union, overload
 
 from opentelemetry import trace
 from opentelemetry.trace import Status, StatusCode
@@ -15,6 +16,64 @@ from hmdl.types import HeimdallAttributes, SpanKind, SpanStatus
 
 F = TypeVar("F", bound=Callable[..., Any])
 
+# Type alias for user extractor function
+# Takes (args, kwargs) and returns user ID string or None
+UserExtractor = Callable[[tuple, dict], Optional[str]]
+
+# Type alias for session extractor function
+# Takes (args, kwargs) and returns session ID string or None
+SessionExtractor = Callable[[tuple, dict], Optional[str]]
+
+# MCP header names
+MCP_SESSION_ID_HEADER = "Mcp-Session-Id"
+AUTHORIZATION_HEADER = "Authorization"
+
+
+def _parse_jwt_claims(token: str) -> Dict[str, Any]:
+    """Parse JWT claims from a token string (without verification)."""
+    try:
+        token_str = token
+        if token_str.lower().startswith("bearer "):
+            token_str = token_str[7:]
+        parts = token_str.split(".")
+        if len(parts) != 3:
+            return {}
+        payload = parts[1]
+        # Add padding if needed
+        padding = 4 - len(payload) % 4
+        if padding != 4:
+            payload += "=" * padding
+        decoded = base64.urlsafe_b64decode(payload).decode("utf-8")
+        return json.loads(decoded)
+    except Exception:
+        return {}
+
+
+def _extract_user_id_from_token(token: str) -> Optional[str]:
+    """Extract user ID from a JWT token."""
+    claims = _parse_jwt_claims(token)
+    # Check common user ID claims in order of preference
+    for claim in ["sub", "user_id", "userId", "uid"]:
+        if claim in claims and isinstance(claims[claim], str):
+            return claims[claim]
+    return None
+
+
+def _extract_from_headers(headers: Dict[str, str]) -> tuple[Optional[str], Optional[str]]:
+    """Extract session ID and user ID from HTTP headers.
+
+    Returns:
+        Tuple of (session_id, user_id)
+    """
+    # Normalize headers to lowercase for case-insensitive lookup
+    normalized = {k.lower(): v for k, v in headers.items()}
+
+    session_id = normalized.get(MCP_SESSION_ID_HEADER.lower())
+    auth_header = normalized.get(AUTHORIZATION_HEADER.lower())
+    user_id = _extract_user_id_from_token(auth_header) if auth_header else None
+
+    return session_id, user_id
+
 
 def _serialize_value(value: Any) -> str:
     """Safely serialize a value to string for span attributes."""
@@ -30,52 +89,128 @@ def _get_client() -> Any:
     return HeimdallClient.get_instance()
 
 
+def _extract_session_id(
+    args: tuple,
+    kwargs: dict,
+    session_extractor: Optional[SessionExtractor],
+    header_session_id: Optional[str],
+) -> Optional[str]:
+    """Extract session ID using the extractor callback or headers.
+
+    Priority: session_extractor callback > headers > None
+    """
+    # Try extractor callback first
+    if session_extractor:
+        try:
+            result = session_extractor(args, kwargs)
+            if result:
+                return result
+        except Exception:
+            # Ignore extraction errors
+            pass
+
+    # Fall back to headers
+    if header_session_id:
+        return header_session_id
+
+    return None
+
+
+def _extract_user_id(
+    args: tuple,
+    kwargs: dict,
+    user_extractor: Optional[UserExtractor],
+    header_user_id: Optional[str],
+) -> Optional[str]:
+    """Extract user ID using the extractor callback or headers.
+
+    Priority: user_extractor callback > headers > None
+    """
+    # Try extractor callback first
+    if user_extractor:
+        try:
+            result = user_extractor(args, kwargs)
+            if result:
+                return result
+        except Exception:
+            # Ignore extraction errors
+            pass
+
+    # Fall back to headers
+    if header_user_id:
+        return header_user_id
+
+    return None
+
+
 def _create_span_decorator(
     span_kind: SpanKind,
     name_attr: str,
     args_attr: str,
     result_attr: str,
-) -> Callable[[Optional[str]], Callable[[F], F]]:
+) -> Callable[..., Callable[[F], F]]:
     """Factory for creating MCP-specific decorators."""
-    
-    def decorator(name: Optional[str] = None) -> Callable[[F], F]:
+
+    def decorator(
+        name: Optional[str] = None,
+        *,
+        headers: Optional[Dict[str, str]] = None,
+        user_extractor: Optional[UserExtractor] = None,
+        session_extractor: Optional[SessionExtractor] = None,
+    ) -> Callable[[F], F]:
+        # Pre-extract from headers if provided
+        header_session_id, header_user_id = _extract_from_headers(headers) if headers else (None, None)
+
         def wrapper(func: F) -> F:
             span_name = name or func.__name__
             is_async = inspect.iscoroutinefunction(func)
-            
+
             if is_async:
                 @functools.wraps(func)
                 async def async_wrapped(*args: Any, **kwargs: Any) -> Any:
                     client = _get_client()
                     if client is None:
                         return await func(*args, **kwargs)
-                    
+
                     tracer = client.tracer
                     with tracer.start_as_current_span(
                         name=span_name,
                         kind=trace.SpanKind.SERVER,
                     ) as span:
                         start_time = time.perf_counter()
-                        
+
                         # Set input attributes
                         span.set_attribute(name_attr, span_name)
                         span.set_attribute("heimdall.span_kind", span_kind.value)
-                        
+
+                        # Extract session ID - priority: extractor > headers > client
+                        session_id = _extract_session_id(args, kwargs, session_extractor, header_session_id)
+                        if not session_id:
+                            session_id = client.get_session_id()
+                        if session_id:
+                            span.set_attribute(HeimdallAttributes.HEIMDALL_SESSION_ID, session_id)
+
+                        # Extract user ID - priority: extractor > headers > client > "anonymous"
+                        user_id = _extract_user_id(args, kwargs, user_extractor, header_user_id)
+                        if not user_id:
+                            user_id = client.get_user_id()
+                        span.set_attribute(HeimdallAttributes.HEIMDALL_USER_ID, user_id or "anonymous")
+
                         # Capture arguments
                         try:
                             all_args = _capture_arguments(func, args, kwargs)
                             span.set_attribute(args_attr, _serialize_value(all_args))
                         except Exception:
                             pass
-                        
+
                         try:
                             result = await func(*args, **kwargs)
-                            
+
                             # Set output attributes
                             span.set_attribute(result_attr, _serialize_value(result))
                             span.set_attribute(HeimdallAttributes.STATUS, SpanStatus.OK.value)
                             span.set_status(Status(StatusCode.OK))
-                            
+
                             return result
                         except Exception as e:
                             _record_error(span, e)
@@ -83,7 +218,7 @@ def _create_span_decorator(
                         finally:
                             duration_ms = (time.perf_counter() - start_time) * 1000
                             span.set_attribute(HeimdallAttributes.DURATION_MS, duration_ms)
-                
+
                 return async_wrapped  # type: ignore
             else:
                 @functools.wraps(func)
@@ -91,33 +226,46 @@ def _create_span_decorator(
                     client = _get_client()
                     if client is None:
                         return func(*args, **kwargs)
-                    
+
                     tracer = client.tracer
                     with tracer.start_as_current_span(
                         name=span_name,
                         kind=trace.SpanKind.SERVER,
                     ) as span:
                         start_time = time.perf_counter()
-                        
+
                         # Set input attributes
                         span.set_attribute(name_attr, span_name)
                         span.set_attribute("heimdall.span_kind", span_kind.value)
-                        
+
+                        # Extract session ID - priority: extractor > headers > client
+                        session_id = _extract_session_id(args, kwargs, session_extractor, header_session_id)
+                        if not session_id:
+                            session_id = client.get_session_id()
+                        if session_id:
+                            span.set_attribute(HeimdallAttributes.HEIMDALL_SESSION_ID, session_id)
+
+                        # Extract user ID - priority: extractor > headers > client > "anonymous"
+                        user_id = _extract_user_id(args, kwargs, user_extractor, header_user_id)
+                        if not user_id:
+                            user_id = client.get_user_id()
+                        span.set_attribute(HeimdallAttributes.HEIMDALL_USER_ID, user_id or "anonymous")
+
                         # Capture arguments
                         try:
                             all_args = _capture_arguments(func, args, kwargs)
                             span.set_attribute(args_attr, _serialize_value(all_args))
                         except Exception:
                             pass
-                        
+
                         try:
                             result = func(*args, **kwargs)
-                            
+
                             # Set output attributes
                             span.set_attribute(result_attr, _serialize_value(result))
                             span.set_attribute(HeimdallAttributes.STATUS, SpanStatus.OK.value)
                             span.set_status(Status(StatusCode.OK))
-                            
+
                             return result
                         except Exception as e:
                             _record_error(span, e)
@@ -125,11 +273,11 @@ def _create_span_decorator(
                         finally:
                             duration_ms = (time.perf_counter() - start_time) * 1000
                             span.set_attribute(HeimdallAttributes.DURATION_MS, duration_ms)
-                
+
                 return sync_wrapped  # type: ignore
-        
+
         return wrapper
-    
+
     return decorator
 
 
@@ -160,6 +308,15 @@ trace_mcp_tool = _create_span_decorator(
 trace_mcp_tool.__doc__ = """
 Decorator to trace MCP tool calls.
 
+Args:
+    name: Custom name for the span (defaults to function name)
+    user_extractor: Function to extract user ID from (args, kwargs).
+        Useful for extracting user info from MCP Context.
+        Returns user ID string or None to use default from client.
+    session_extractor: Function to extract session ID from (args, kwargs).
+        Useful for extracting session info from MCP Context.
+        Returns session ID string or None to use default from client.
+
 Example:
     >>> @trace_mcp_tool()
     ... def my_tool(arg1: str, arg2: int) -> str:
@@ -168,6 +325,14 @@ Example:
     >>> @trace_mcp_tool("custom-tool-name")
     ... async def async_tool(data: dict) -> dict:
     ...     return {"processed": data}
+
+    # Extract user and session from MCP Context (first argument)
+    >>> @trace_mcp_tool(
+    ...     user_extractor=lambda args, kwargs: getattr(args[0], 'user_id', None) if args else None,
+    ...     session_extractor=lambda args, kwargs: getattr(args[0], 'session_id', None) if args else None,
+    ... )
+    ... def my_tool_with_ctx(ctx, query: str) -> str:
+    ...     return f"Query: {query}"
 """
 
 trace_mcp_resource = _create_span_decorator(

{hmdl-0.0.1.dist-info → hmdl-0.0.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hmdl
-Version: 0.0.1
+Version: 0.0.2
 Summary: Observability SDK for MCP (Model Context Protocol) servers - Heimdall Platform
 Project-URL: Homepage, https://tryheimdall.com
 Project-URL: Documentation, https://docs.tryheimdall.com
@@ -141,6 +141,8 @@ async def async_search(query: str) -> list:
 | `HEIMDALL_DEBUG` | Enable debug logging | `false` |
 | `HEIMDALL_BATCH_SIZE` | Spans per batch | `100` |
 | `HEIMDALL_FLUSH_INTERVAL_MS` | Flush interval (ms) | `5000` |
+| `HEIMDALL_SESSION_ID` | Default session ID | - |
+| `HEIMDALL_USER_ID` | Default user ID | - |
 
 ### Local Development
 
@@ -155,6 +157,45 @@ export HEIMDALL_ENABLED="true"
 
 ## Advanced Usage
 
+### Session and User Tracking
+
+`trace_mcp_tool` automatically includes session and user IDs in spans. You just need to provide them via one of these methods:
+
+#### Option 1: HTTP Headers (Recommended for MCP servers)
+
+Pass HTTP headers directly to `trace_mcp_tool`. Session ID is extracted from the `Mcp-Session-Id` header, and user ID from the JWT token in the `Authorization` header:
+
+```python
+from hmdl import trace_mcp_tool
+
+@app.post("/mcp")
+def handle_request():
+    @trace_mcp_tool(headers=dict(request.headers))
+    def search_tool(query: str):
+        return results
+
+    return search_tool("test")  # Session/user included in span
+```
+
+#### Option 2: Extractors (Per-tool extraction)
+
+```python
+from typing import Optional
+
+@trace_mcp_tool(
+    session_extractor=lambda args, kwargs: kwargs.get('session_id'),
+    user_extractor=lambda args, kwargs: kwargs.get('user_id'),
+)
+def my_tool(query: str, session_id: Optional[str] = None, user_id: Optional[str] = None):
+    return f"Query: {query}"
+```
+
+#### Resolution Priority
+
+1. Extractor callback → 2. HTTP headers → 3. Client value (initialized from environment variables)
+
+> **Note**: If no user ID is found through any of these methods, `"anonymous"` is used as the default.
+
 ### Custom span names
 
 ```python

hmdl-0.0.2.dist-info/RECORD

@@ -0,0 +1,10 @@
+hmdl/__init__.py,sha256=2OCWYFkHkNNRhYawHGVd9tUVfZ8WMMUuxhvJvgHsvFQ,648
+hmdl/client.py,sha256=XItMSgojsDwz-NAjFzWF2dXHDo7o6rXKBHppbCgsSz0,8633
+hmdl/config.py,sha256=crHQu_b2KO6WJ6vJ8GTtdCbuizAQopI5-WQtgaCHfj0,3520
+hmdl/decorators.py,sha256=HupF7DUZRIqPPR8mvw2jCTv4US7fAQgAOuqPMlCGBNY,17561
+hmdl/py.typed,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+hmdl/types.py,sha256=C_QDl4YTJjnrhvMUdnDnTG35jKmovLFqMrHXL4LHOG0,3130
+hmdl-0.0.2.dist-info/METADATA,sha256=pxcj7lxXZd1Bnu7aLnvOo6zhgWDhjM9Ldx9sQwIzTR4,8083
+hmdl-0.0.2.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+hmdl-0.0.2.dist-info/licenses/LICENSE,sha256=b8jAb5oXJiKCT9GmhRp2uDLqZXIA63QnLT4_3JvzxhE,1064
+hmdl-0.0.2.dist-info/RECORD,,

hmdl-0.0.1.dist-info/RECORD

@@ -1,10 +0,0 @@
-hmdl/__init__.py,sha256=EA49ssIg0WJTzi1-Oi0J8GWV6k55R-oav45xbufd4WU,570
-hmdl/client.py,sha256=_VHVIlfq8JV_FucLN9rjgj_0fE_C0h5DZWQCUPXSSCw,6782
-hmdl/config.py,sha256=jXo0XC962JM3-N2uUl6o8Dt0hUZjIWvDJcthm04Ux0U,3028
-hmdl/decorators.py,sha256=yUywG_AMCA-tqT-w4fO1bSpOrVn3TDk7Y981R2UmJjQ,11697
-hmdl/py.typed,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-hmdl/types.py,sha256=C_QDl4YTJjnrhvMUdnDnTG35jKmovLFqMrHXL4LHOG0,3130
-hmdl-0.0.1.dist-info/METADATA,sha256=gbak-CcmYZkcegGv0GqSvMaC41EtCo99IgWC4or3UE8,6743
-hmdl-0.0.1.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-hmdl-0.0.1.dist-info/licenses/LICENSE,sha256=b8jAb5oXJiKCT9GmhRp2uDLqZXIA63QnLT4_3JvzxhE,1064
-hmdl-0.0.1.dist-info/RECORD,,