asap-protocol 0.3.0__py3-none-any.whl → 1.0.0__py3-none-any.whl

asap/transport/compression.py

@@ -0,0 +1,389 @@
+"""Compression support for ASAP protocol transport layer.
+
+This module provides compression and decompression utilities for reducing
+bandwidth usage in agent-to-agent communication.
+
+Supported compression algorithms:
+- gzip: Standard compression, widely supported (default)
+- brotli: Better compression ratio (optional, requires brotli package)
+
+Compression is applied automatically when:
+- Request body exceeds COMPRESSION_THRESHOLD (default: 1KB)
+- Server indicates support via Accept-Encoding header
+
+Example:
+    >>> from asap.transport.compression import compress_payload, decompress_payload
+    >>>
+    >>> # Compress large payload
+    >>> data = b'{"large": "payload..."}' * 100
+    >>> compressed, encoding = compress_payload(data)
+    >>> # encoding = "gzip" or "br" depending on configuration
+    >>>
+    >>> # Decompress received payload
+    >>> original = decompress_payload(compressed, encoding)
+"""
+
+import gzip
+from enum import Enum
+
+from asap.observability import get_logger
+
+# Module logger
+logger = get_logger(__name__)
+
+# Compression threshold: only compress payloads larger than this size (bytes)
+COMPRESSION_THRESHOLD = 1024  # 1KB
+
+# Gzip compression level (1-9, higher = better compression but slower)
+GZIP_COMPRESSION_LEVEL = 6
+
+
+class CompressionAlgorithm(str, Enum):
+    """Supported compression algorithms."""
+
+    GZIP = "gzip"
+    BROTLI = "br"
+    IDENTITY = "identity"  # No compression
+
+
+def is_brotli_available() -> bool:
+    """Check if brotli compression is available.
+
+    Brotli is an optional dependency that provides better compression
+    ratios than gzip, especially for JSON payloads.
+
+    Returns:
+        True if brotli package is installed, False otherwise
+    """
+    try:
+        import brotli  # noqa: F401
+
+        return True
+    except ImportError:
+        return False
+
+
+def get_supported_encodings() -> list[str]:
+    """Get list of supported Content-Encoding values.
+
+    Returns:
+        List of supported encoding names (e.g., ["gzip", "br"])
+    """
+    encodings = ["gzip"]
+    if is_brotli_available():
+        encodings.insert(0, "br")  # Prefer brotli when available
+    return encodings
+
+
+def get_accept_encoding_header() -> str:
+    """Get Accept-Encoding header value for client requests.
+
+    Returns a comma-separated list of supported encodings with quality
+    values to indicate preference.
+
+    Returns:
+        Accept-Encoding header value (e.g., "br, gzip, identity")
+    """
+    encodings = get_supported_encodings()
+    # Add identity as fallback (always supported)
+    encodings.append("identity")
+    return ", ".join(encodings)
+
+
+def compress_gzip(data: bytes) -> bytes:
+    """Compress data using gzip algorithm.
+
+    Args:
+        data: Raw bytes to compress
+
+    Returns:
+        Gzip compressed bytes
+    """
+    return gzip.compress(data, compresslevel=GZIP_COMPRESSION_LEVEL)
+
+
+def decompress_gzip(data: bytes) -> bytes:
+    """Decompress gzip data.
+
+    Args:
+        data: Gzip compressed bytes
+
+    Returns:
+        Decompressed bytes
+
+    Raises:
+        OSError: If data is not valid gzip format
+    """
+    return gzip.decompress(data)
+
+
+def compress_brotli(data: bytes) -> bytes:
+    """Compress data using brotli algorithm.
+
+    Args:
+        data: Raw bytes to compress
+
+    Returns:
+        Brotli compressed bytes
+
+    Raises:
+        ImportError: If brotli package is not installed
+    """
+    import brotli
+
+    # Quality 4 is a good balance of speed and compression ratio
+    result: bytes = brotli.compress(data, quality=4)
+    return result
+
+
+def decompress_brotli(data: bytes) -> bytes:
+    """Decompress brotli data.
+
+    Args:
+        data: Brotli compressed bytes
+
+    Returns:
+        Decompressed bytes
+
+    Raises:
+        ImportError: If brotli package is not installed
+        brotli.error: If data is not valid brotli format
+    """
+    import brotli
+
+    try:
+        result: bytes = brotli.decompress(data)
+        return result
+    except brotli.error as e:
+        raise OSError(f"Brotli decompression failed: {e}") from e
+
+
+def compress_payload(
+    data: bytes,
+    preferred_algorithm: CompressionAlgorithm | None = None,
+    threshold: int = COMPRESSION_THRESHOLD,
+) -> tuple[bytes, CompressionAlgorithm]:
+    """Compress payload if it exceeds threshold.
+
+    Compresses the payload using the preferred algorithm if specified,
+    otherwise uses the best available algorithm (brotli > gzip).
+
+    Args:
+        data: Raw bytes to compress
+        preferred_algorithm: Optional preferred compression algorithm.
+            If None, uses brotli if available, otherwise gzip.
+        threshold: Minimum payload size to trigger compression (default: 1KB).
+            Payloads smaller than this are returned as-is with IDENTITY encoding.
+
+    Returns:
+        Tuple of (compressed_data, algorithm_used)
+        If compression is skipped, returns (original_data, IDENTITY)
+
+        >>> compressed, algorithm = compress_payload(data)
+        >>> print(f"Compressed {len(data)} -> {len(compressed)} bytes using {algorithm.value}")
+
+    Note:
+        Compression adds CPU overhead which may increase latency for very small payloads.
+        For extremely latency-sensitive scenarios, consider increasing the threshold
+        or disabling compression if payloads are typically small.
+    """
+    # Skip compression for small payloads
+    if len(data) < threshold:
+        logger.debug(
+            "asap.compression.skipped",
+            size=len(data),
+            threshold=threshold,
+            reason="payload_below_threshold",
+        )
+        return data, CompressionAlgorithm.IDENTITY
+
+    # Determine algorithm to use
+    if preferred_algorithm is not None:
+        algorithm = preferred_algorithm
+    elif is_brotli_available():
+        # TODO: Add prefer_fast_compression option to prefer gzip even when brotli is available
+        algorithm = CompressionAlgorithm.BROTLI
+    else:
+        algorithm = CompressionAlgorithm.GZIP
+
+    # Skip if identity is requested
+    if algorithm == CompressionAlgorithm.IDENTITY:
+        return data, CompressionAlgorithm.IDENTITY
+
+    # Compress using selected algorithm
+    original_size = len(data)
+    try:
+        if algorithm == CompressionAlgorithm.BROTLI:
+            compressed = compress_brotli(data)
+        else:
+            compressed = compress_gzip(data)
+
+        compressed_size = len(compressed)
+        reduction_pct = (1 - compressed_size / original_size) * 100
+
+        logger.debug(
+            "asap.compression.applied",
+            algorithm=algorithm.value,
+            original_size=original_size,
+            compressed_size=compressed_size,
+            reduction_percent=round(reduction_pct, 1),
+        )
+
+        # Only use compression if it actually reduces size
+        if compressed_size >= original_size:
+            logger.debug(
+                "asap.compression.ineffective",
+                algorithm=algorithm.value,
+                original_size=original_size,
+                compressed_size=compressed_size,
+                reason="compression_increased_size",
+            )
+            return data, CompressionAlgorithm.IDENTITY
+
+        return compressed, algorithm
+
+    except ImportError:
+        # Brotli not available, fallback to gzip
+        logger.warning(
+            "asap.compression.fallback",
+            requested=algorithm.value,
+            fallback="gzip",
+            reason="brotli_not_installed",
+        )
+        return compress_payload(data, CompressionAlgorithm.GZIP, threshold)
+    except Exception as e:
+        # Compression failed, return original
+        logger.warning(
+            "asap.compression.failed",
+            algorithm=algorithm.value,
+            error=str(e),
+            error_type=type(e).__name__,
+        )
+        return data, CompressionAlgorithm.IDENTITY
+
+
+def decompress_payload(
+    data: bytes,
+    encoding: str,
+) -> bytes:
+    """Decompress payload based on Content-Encoding header.
+
+    Args:
+        data: Compressed bytes
+        encoding: Content-Encoding header value (e.g., "gzip", "br", "identity")
+
+    Returns:
+        Decompressed bytes
+
+    Raises:
+        ValueError: If encoding is not supported
+        OSError: If decompression fails (invalid compressed data)
+
+    Example:
+        >>> compressed_data = gzip.compress(b'{"message": "hello"}')
+        >>> original = decompress_payload(compressed_data, "gzip")
+    """
+    # Normalize encoding name
+    encoding_lower = encoding.lower().strip()
+
+    # Handle identity (no compression)
+    if encoding_lower in ("identity", ""):
+        return data
+
+    original_size = len(data)
+
+    try:
+        if encoding_lower == "gzip":
+            decompressed = decompress_gzip(data)
+        elif encoding_lower == "br":
+            if not is_brotli_available():
+                raise ValueError(
+                    "Brotli decompression requested but brotli package is not installed. "
+                    "Install with: pip install brotli"
+                )
+            decompressed = decompress_brotli(data)
+        else:
+            raise ValueError(
+                f"Unsupported Content-Encoding: {encoding}. Supported: gzip, br, identity"
+            )
+
+        decompressed_size = len(decompressed)
+        logger.debug(
+            "asap.decompression.applied",
+            encoding=encoding_lower,
+            compressed_size=original_size,
+            decompressed_size=decompressed_size,
+        )
+
+        return decompressed
+
+    except ImportError as e:
+        raise ValueError(
+            f"Cannot decompress {encoding}: required package not installed. {e}"
+        ) from e
+
+
+def select_best_encoding(accept_encoding: str | None) -> CompressionAlgorithm:
+    """Select best compression algorithm based on Accept-Encoding header.
+
+    Parses the Accept-Encoding header and selects the best supported algorithm
+    based on client preferences and availability.
+
+    Args:
+        accept_encoding: Accept-Encoding header value (e.g., "gzip, br;q=0.9")
+            If None, returns IDENTITY (no compression).
+
+    Returns:
+        Best available compression algorithm
+
+    Example:
+        >>> algorithm = select_best_encoding("br, gzip;q=0.9, identity;q=0.5")
+        >>> print(algorithm.value)  # "br" if brotli available, else "gzip"
+    """
+    if not accept_encoding:
+        return CompressionAlgorithm.IDENTITY
+
+    # Parse Accept-Encoding header (simplified parser)
+    # Format: encoding[;q=weight], encoding[;q=weight], ...
+    encodings: dict[str, float] = {}
+    for part in accept_encoding.split(","):
+        part = part.strip()
+        if not part:
+            continue
+
+        # Split encoding and quality
+        if ";q=" in part.lower():
+            encoding, q_str = part.lower().split(";q=", 1)
+            try:
+                quality = float(q_str)
+            except ValueError:
+                quality = 1.0
+        else:
+            encoding = part.lower()
+            quality = 1.0
+
+        encoding = encoding.strip()
+        if encoding:
+            encodings[encoding] = quality
+
+    # Sort by quality (highest first) and filter to supported encodings
+    supported = get_supported_encodings()
+    candidates: list[tuple[str, float]] = []
+
+    for enc, quality in encodings.items():
+        if enc in supported or enc == "identity":
+            candidates.append((enc, quality))
+
+    # Sort by quality descending
+    candidates.sort(key=lambda x: x[1], reverse=True)
+
+    if not candidates:
+        return CompressionAlgorithm.IDENTITY
+
+    best_encoding = candidates[0][0]
+
+    if best_encoding == "br" and is_brotli_available():
+        return CompressionAlgorithm.BROTLI
+    if best_encoding == "gzip":
+        return CompressionAlgorithm.GZIP
+    return CompressionAlgorithm.IDENTITY

asap/transport/handlers.py

@@ -32,7 +32,7 @@ import time
 from collections.abc import Awaitable
 from concurrent.futures import Executor
 from threading import RLock
-from typing import Protocol, cast
+from typing import Callable, Protocol, TypeAlias, cast
 
 from asap.errors import ASAPError
 from asap.models.entities import Manifest
@@ -40,7 +40,8 @@ from asap.models.enums import TaskStatus
 from asap.models.envelope import Envelope
 from asap.models.ids import generate_id
 from asap.models.payloads import TaskRequest, TaskResponse
-from asap.observability import get_logger
+from asap.observability import get_logger, get_metrics
+from asap.observability.tracing import handler_span_context
 
 # Module logger
 logger = get_logger(__name__)
@@ -94,6 +95,49 @@ Returns:
     Response envelope to send back (sync) or awaitable (async)
 """
 
+# Type alias for factories that return a sync handler (useful in tests)
+SyncHandlerFactory: TypeAlias = Callable[[], SyncHandler]
+"""Type alias for callables that return a SyncHandler (e.g. create_echo_handler)."""
+
+
+def validate_handler(handler: Handler) -> None:
+    """Validate that a handler has the required signature (envelope, manifest).
+
+    Checks that the handler is callable and accepts exactly two parameters
+    (envelope and manifest), matching the Handler protocol. Use when
+    registering custom handlers to fail fast on invalid signatures.
+
+    Args:
+        handler: The handler callable to validate (sync or async).
+
+    Raises:
+        TypeError: If handler is not callable or does not have the required
+            signature (two parameters for a function, or three for a bound
+            callable with self/cls).
+
+    Example:
+        >>> validate_handler(create_echo_handler())
+        >>> def bad(x, y, z): ...
+        >>> validate_handler(bad)
+        Traceback (most recent call last):
+            ...
+        TypeError: Handler must accept (envelope, manifest); got 3 parameters
+    """
+    if not callable(handler):
+        raise TypeError("Handler must be callable")
+    try:
+        sig = inspect.signature(handler)
+    except (ValueError, TypeError):
+        raise TypeError("Handler signature could not be inspected") from None
+    params = list(sig.parameters)
+    # Allow (envelope, manifest) or (self, envelope, manifest) / (cls, envelope, manifest)
+    if len(params) == 2 or len(params) == 3 and params[0] in ("self", "cls"):
+        pass
+    else:
+        raise TypeError(
+            f"Handler must accept (envelope, manifest); got {len(params)} parameters: {params}"
+        )
+
 
 class HandlerNotFoundError(ASAPError):
     """Raised when no handler is registered for a payload type.
@@ -181,6 +225,7 @@ class HandlerRegistry:
             >>> registry = HandlerRegistry()
             >>> registry.register("task.request", create_echo_handler())
         """
+        validate_handler(handler)
         with self._lock:
             is_override = payload_type in self._handlers
             self._handlers[payload_type] = handler
@@ -246,7 +291,6 @@ class HandlerRegistry:
                 raise HandlerNotFoundError(payload_type)
             handler = self._handlers[payload_type]
 
-        # Log dispatch start
         logger.debug(
             "asap.handler.dispatch",
             payload_type=payload_type,
@@ -254,13 +298,8 @@ class HandlerRegistry:
             handler_name=handler.__name__ if hasattr(handler, "__name__") else str(handler),
         )
 
-        # Execute handler outside the lock to allow concurrent dispatches
         try:
-            # Note: dispatch() only works with sync handlers that return Envelope directly
-            # For async handlers, use dispatch_async() instead
-            # Type narrowing: we expect sync handlers here
             result = handler(envelope, manifest)
-            # For sync handlers, result is Envelope directly
             if inspect.isawaitable(result):
                 raise TypeError(
                     f"Handler {handler} returned awaitable in sync dispatch(). "
@@ -323,7 +362,6 @@ class HandlerRegistry:
                 raise HandlerNotFoundError(payload_type)
             handler = self._handlers[payload_type]
 
-        # Log dispatch start
         logger.debug(
             "asap.handler.dispatch",
             payload_type=payload_type,
@@ -331,48 +369,62 @@ class HandlerRegistry:
             handler_name=handler.__name__ if hasattr(handler, "__name__") else str(handler),
         )
 
-        # Execute handler outside the lock to allow concurrent dispatches
-        try:
-            # Support both sync and async handlers
-            response: Envelope
-            if inspect.iscoroutinefunction(handler):
-                # Async handler - await it directly
-                response = await handler(envelope, manifest)
-            else:
-                # Sync handler - run in thread pool to avoid blocking event loop
-                # Also handle async callable objects that return awaitables
-                loop = asyncio.get_event_loop()
-                # Use bounded executor if provided, otherwise use default (unbounded)
-                executor = self._executor if self._executor is not None else None
-                result: object = await loop.run_in_executor(executor, handler, envelope, manifest)
-                # Check if result is awaitable (handles async __call__ methods)
-                if inspect.isawaitable(result):
-                    response = await result
+        agent_urn = manifest.id
+        with handler_span_context(
+            payload_type=payload_type,
+            agent_urn=agent_urn,
+            envelope_id=envelope.id,
+        ):
+            try:
+                response: Envelope
+                if inspect.iscoroutinefunction(handler):
+                    response = await handler(envelope, manifest)
                 else:
-                    # Type narrowing: result is Envelope for sync handlers
-                    # After checking it's not awaitable, we know it's Envelope
-                    response = cast(Envelope, result)
-
-            duration_ms = (time.perf_counter() - start_time) * 1000
-            logger.debug(
-                "asap.handler.completed",
-                payload_type=payload_type,
-                envelope_id=envelope.id,
-                response_id=response.id,
-                duration_ms=round(duration_ms, 2),
-            )
-            return response
-        except Exception as e:
-            duration_ms = (time.perf_counter() - start_time) * 1000
-            logger.exception(
-                "asap.handler.error",
-                payload_type=payload_type,
-                envelope_id=envelope.id,
-                error=str(e),
-                error_type=type(e).__name__,
-                duration_ms=round(duration_ms, 2),
-            )
-            raise
+                    loop = asyncio.get_running_loop()
+                    executor = self._executor if self._executor is not None else None
+                    result: object = await loop.run_in_executor(
+                        executor, handler, envelope, manifest
+                    )
+                    if inspect.isawaitable(result):
+                        response = await result
+                    else:
+                        response = cast(Envelope, result)
+
+                duration_ms = (time.perf_counter() - start_time) * 1000
+                duration_seconds = duration_ms / 1000.0
+                logger.debug(
+                    "asap.handler.completed",
+                    payload_type=payload_type,
+                    envelope_id=envelope.id,
+                    response_id=response.id,
+                    duration_ms=round(duration_ms, 2),
+                )
+                metrics = get_metrics()
+                metrics.increment_counter(
+                    "asap_handler_executions_total",
+                    {"payload_type": payload_type},
+                )
+                metrics.observe_histogram(
+                    "asap_handler_duration_seconds",
+                    duration_seconds,
+                    {"payload_type": payload_type},
+                )
+                return response
+            except Exception as e:
+                duration_ms = (time.perf_counter() - start_time) * 1000
+                get_metrics().increment_counter(
+                    "asap_handler_errors_total",
+                    {"payload_type": payload_type},
+                )
+                logger.exception(
+                    "asap.handler.error",
+                    payload_type=payload_type,
+                    envelope_id=envelope.id,
+                    error=str(e),
+                    error_type=type(e).__name__,
+                    duration_ms=round(duration_ms, 2),
+                )
+                raise
 
     def list_handlers(self) -> list[str]:
         """List all registered payload types.
@@ -391,18 +443,19 @@ class HandlerRegistry:
             return list(self._handlers.keys())
 
 
-def create_echo_handler() -> Handler:
-    """Create an echo handler that echoes TaskRequest input.
+def create_echo_handler() -> SyncHandler:
+    """Create a synchronous echo handler that echoes TaskRequest input.
 
     The echo handler is a simple implementation that:
     - Receives a TaskRequest envelope
     - Returns a TaskResponse with the input echoed back
     - Preserves trace_id and sets correlation_id
 
+    Returns SyncHandler (not Handler) so tests can use it without casting.
     This is useful for testing and as a base for custom handlers.
 
     Returns:
-        Handler function that echoes TaskRequest input
+        SyncHandler that echoes TaskRequest input
 
     Example:
         >>> handler = create_echo_handler()