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

asap/transport/middleware.py

@@ -32,8 +32,10 @@ Example:
     >>> middleware = AuthenticationMiddleware(manifest, validator)
 """
 
+import asyncio
+import inspect
 import uuid
-from typing import Any, Awaitable, Callable, Protocol
+from typing import Any, Awaitable, Callable, Protocol, cast
 from collections.abc import Sequence
 
 from fastapi import HTTPException, Request
@@ -54,7 +56,9 @@ logger = get_logger(__name__)
 AUTH_SCHEME_BEARER = "bearer"
 
 # Rate limiting default configuration
-DEFAULT_RATE_LIMIT = "100/minute"
+# Uses token bucket pattern: burst limit (per second) + sustained limit (per minute)
+# This allows short bursts while preventing sustained abuse
+DEFAULT_RATE_LIMIT = "10/second;100/minute"
 
 
 def _get_sender_from_envelope(request: Request) -> str:
@@ -80,9 +84,7 @@ def _get_sender_from_envelope(request: Request) -> str:
         >>> sender = _get_sender_from_envelope(request)
         >>> # Returns "192.168.1.1" (IP address, not sender URN)
     """
-    # Try to extract sender from envelope if already parsed (early returns reduce complexity)
     try:
-        # Check if envelope is stored in request state (after parsing)
         if hasattr(request.state, "envelope") and request.state.envelope:
             envelope = request.state.envelope
             if hasattr(envelope, "sender") and isinstance(envelope.sender, str):
@@ -105,17 +107,12 @@ def _get_sender_from_envelope(request: Request) -> str:
         # Envelope not available, fall back to IP
         pass
 
-    # Fallback to client IP address
     remote_addr = get_remote_address(request)
-    # Type narrowing: get_remote_address returns str, but mypy may see it as Any
     if isinstance(remote_addr, str):
         return remote_addr
     return str(remote_addr)
 
 
-# Create rate limiter instance with IP-based key function
-# Note: The key function attempts to extract sender but always falls back to IP
-# because rate limiting executes before request body parsing
 limiter = Limiter(
     key_func=_get_sender_from_envelope,
     default_limits=[DEFAULT_RATE_LIMIT],
@@ -154,8 +151,15 @@ def create_test_limiter(limits: Sequence[str] | None = None) -> Limiter:
 def create_limiter(limits: Sequence[str] | None = None) -> Limiter:
     """Create a new limiter instance for production use.
 
-    Creates an isolated limiter instance with its own storage, allowing
-    multiple FastAPI app instances to have independent rate limiters.
+    Creates an isolated limiter instance with its own in-memory storage
+    (``memory://``), allowing multiple FastAPI app instances to have
+    independent rate limiters.
+
+    **Multi-worker warning:** ``memory://`` storage is per-process. In
+    multi-worker deployments (e.g., Gunicorn with 4 workers), each worker
+    has isolated limits—effective rate = configured limit × number of workers
+    (e.g., 10/s → 40/s across workers). For shared limits in production,
+    use Redis-backed storage via slowapi's ``storage_uri``.
 
     Args:
         limits: Optional list of rate limit strings (e.g., ["100/minute"]).
@@ -173,10 +177,16 @@ def create_limiter(limits: Sequence[str] | None = None) -> Limiter:
 
     # Use unique storage URI to ensure isolation between app instances
     unique_storage_id = str(uuid.uuid4())
+    storage_uri = f"memory://{unique_storage_id}"
+    logger.warning(
+        "asap.rate_limit.memory_storage",
+        message="memory:// storage is per-process; in multi-worker deployments "
+        "(e.g., Gunicorn), effective rate = limit × workers. Use Redis for shared limits.",
+    )
     return Limiter(
         key_func=_get_sender_from_envelope,
         default_limits=list(limits),
-        storage_uri=f"memory://{unique_storage_id}",
+        storage_uri=storage_uri,
     )
 
 
@@ -197,7 +207,6 @@ def rate_limit_handler(request: Request, exc: Exception) -> JSONResponse:
         >>> response = rate_limit_handler(request, exc)
         >>> # Returns JSONResponse with status_code=429 and JSON-RPC error
     """
-    # Type narrowing: FastAPI passes RateLimitExceeded but handler signature uses Exception
     if not isinstance(exc, RateLimitExceeded):
         # Fallback for unexpected exception types
         logger.warning("asap.rate_limit.unexpected_exception", exc_type=type(exc).__name__)
@@ -269,25 +278,32 @@ class TokenValidator(Protocol):
     """Protocol for token validation implementations.
 
     Custom validators must implement this interface to integrate
-    with the authentication middleware.
+    with the authentication middleware. Validators may be sync or async.
+    Sync validators that perform I/O (e.g., DB/Redis lookups) are run
+    in a thread pool to avoid blocking the event loop.
 
-    Example:
-        >>> class MyTokenValidator:
+    Example (sync):
+        >>> class MySyncValidator:
         ...     def __call__(self, token: str) -> str | None:
-        ...         # Validate token against database, JWT, etc.
         ...         if is_valid(token):
         ...             return extract_agent_id(token)
         ...         return None
+
+    Example (async):
+        >>> class MyAsyncValidator:
+        ...     async def __call__(self, token: str) -> str | None:
+        ...         return await db.lookup_agent(token)
     """
 
-    def __call__(self, token: str) -> str | None:
+    def __call__(self, token: str) -> str | None | Awaitable[str | None]:
         """Validate a token and return the authenticated agent ID.
 
         Args:
             token: The authentication token to validate
 
         Returns:
-            The agent ID (URN) if token is valid, None otherwise
+            The agent ID (URN) if token is valid, None otherwise.
+            May return a coroutine for async validators.
 
         Example:
             >>> validator = BearerTokenValidator(my_validate_func)
@@ -301,15 +317,14 @@ class BearerTokenValidator:
     """Default Bearer token validator implementation.
 
     Wraps a validation function to conform to the TokenValidator protocol.
-    The validation function should take a token string and return an agent ID
-    if valid, or None if invalid.
+    The validation function may be sync or async. Sync validators that perform
+    I/O are run in a thread pool by the middleware to avoid blocking the loop.
 
     Attributes:
         validate_func: Function that validates tokens and returns agent IDs
 
     Example:
         >>> def my_validator(token: str) -> str | None:
-        ...     # Check token in database
         ...     if token in valid_tokens:
         ...         return valid_tokens[token]["agent_id"]
         ...     return None
@@ -318,22 +333,28 @@ class BearerTokenValidator:
         >>> agent_id = validator("abc123")
     """
 
-    def __init__(self, validate_func: Callable[[str], str | None]) -> None:
+    def __init__(
+        self,
+        validate_func: Callable[[str], str | None | Awaitable[str | None]],
+    ) -> None:
         """Initialize the Bearer token validator.
 
         Args:
-            validate_func: Function that validates tokens and returns agent IDs
+            validate_func: Sync or async function that validates tokens
+                and returns agent IDs. Sync functions with I/O are run
+                in a thread pool to avoid blocking the event loop.
         """
         self.validate_func = validate_func
 
-    def __call__(self, token: str) -> str | None:
+    def __call__(self, token: str) -> str | None | Awaitable[str | None]:
         """Validate a token and return the authenticated agent ID.
 
         Args:
             token: The authentication token to validate
 
         Returns:
-            The agent ID (URN) if token is valid, None otherwise
+            The agent ID (URN) if token is valid, None otherwise.
+            May return a coroutine for async validate_func.
         """
         return self.validate_func(token)
 
@@ -388,7 +409,6 @@ class AuthenticationMiddleware:
         self.validator = validator
         self.security = HTTPBearer(auto_error=False)
 
-        # Validate configuration
         if self._is_auth_required() and validator is None:
             raise ValueError(
                 "Token validator required when authentication is configured in manifest"
@@ -461,7 +481,6 @@ class AuthenticationMiddleware:
                 headers={"WWW-Authenticate": "Bearer"},
             )
 
-        # Validate Authorization scheme (case-insensitive)
         if credentials.scheme.lower() != AUTH_SCHEME_BEARER:
             logger.warning(
                 "asap.auth.invalid_scheme",
@@ -475,7 +494,6 @@ class AuthenticationMiddleware:
                 headers={"WWW-Authenticate": "Bearer"},
             )
 
-        # Validate Bearer token support in manifest
         if not self._supports_bearer_auth():
             logger.warning(
                 "asap.auth.scheme_not_supported",
@@ -488,15 +506,23 @@ class AuthenticationMiddleware:
                 headers={"WWW-Authenticate": "Bearer"},
             )
 
-        # Validate token and get agent ID
         token = credentials.credentials
-        # Type narrowing: validator is not None when auth is required (validated in __init__)
         if self.validator is None:
             raise RuntimeError(
                 "Token validator is None but authentication is required. "
                 "This should not happen if middleware was initialized correctly."
             )
-        agent_id = self.validator(token)
+        if inspect.iscoroutinefunction(self.validator.__call__):
+            result = self.validator(token)
+            agent_id = await cast(Awaitable[str | None], result)
+        else:
+            # Cast: to_thread expects Callable[..., T]; TokenValidator may return Awaitable
+            # but we handle that below (BearerTokenValidator wrapping async func)
+            agent_id = await asyncio.to_thread(
+                cast("Callable[[str], str | None]", self.validator), token
+            )
+            if inspect.isawaitable(agent_id):
+                agent_id = await cast(Awaitable[str | None], agent_id)
 
         if agent_id is None:
             # Log sanitized token to avoid exposing full token data
@@ -615,7 +641,6 @@ class SizeLimitMiddleware(BaseHTTPMiddleware):
         Returns:
             Response from next handler or error response if size exceeded
         """
-        # Check Content-Length header if present
         content_length = request.headers.get("content-length")
         if content_length:
             try:
@@ -634,7 +659,6 @@ class SizeLimitMiddleware(BaseHTTPMiddleware):
                         },
                     )
             except ValueError:
-                # Invalid Content-Length header, let route handler validate actual body size
                 pass
 
         # Continue to next middleware or route handler