fix(event_handler): isolate local ASGI request state

Author: wuodar

aws_lambda_powertools/event_handler/http_resolver.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 import base64
+from contextvars import ContextVar
+from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Callable
 from urllib.parse import parse_qs
 
@@ -13,6 +15,8 @@
 from aws_lambda_powertools.utilities.data_classes.common import BaseProxyEvent
 
 if TYPE_CHECKING:
+    from collections.abc import Mapping, MutableMapping
+
     from aws_lambda_powertools.shared.cookies import Cookie
 
 
@@ -95,7 +99,7 @@ def _from_dict(cls, data: dict[str, Any]) -> HttpProxyEvent:
         return instance
 
     @classmethod
-    def from_asgi(cls, scope: dict[str, Any], body: bytes | None = None) -> HttpProxyEvent:
+    def from_asgi(cls, scope: Mapping[str, Any], body: bytes | None = None) -> HttpProxyEvent:
         """
         Create an HttpProxyEvent from an ASGI scope dict.
 
@@ -159,6 +163,14 @@ def get_remaining_time_in_millis(self) -> int:  # pragma: no cover
         return 300000  # 5 minutes
 
 
+@dataclass
+class _RequestState:
+    event: BaseProxyEvent | None = None
+    lambda_context: Any = None
+    context: dict = field(default_factory=dict)
+    processed_stack_frames: list[str] = field(default_factory=list)
+
+
 class HttpResolverLocal(ApiGatewayResolver):
     """
     ASGI-compatible HTTP resolver.
@@ -204,6 +216,8 @@ def __init__(
         strip_prefixes: list[str | Any] | None = None,
         enable_validation: bool = False,
     ):
+        self._startup_state = _RequestState()
+        self._request_state: ContextVar[_RequestState | None] = ContextVar("local_http_request", default=None)
         super().__init__(
             proxy_type=ProxyEventType.APIGatewayProxyEvent,  # Use REST API format internally
             cors=cors,
@@ -212,7 +226,46 @@ def __init__(
             strip_prefixes=strip_prefixes,
             enable_validation=enable_validation,
         )
-        self._is_async_mode = False
+
+    @property
+    def _state(self) -> _RequestState:
+        return self._request_state.get() or self._startup_state
+
+    # Powertools declares these as mutable attributes. Properties preserve that
+    # interface while directing each task to its own state. asyncio.to_thread
+    # propagates the ContextVar, so middleware sees the same request dictionary.
+    @property
+    def current_event(self) -> BaseProxyEvent:
+        # Preserve the inherited synchronous resolve() path outside ASGI calls.
+        return self._state.event or BaseRouter.current_event
+
+    @current_event.setter
+    def current_event(self, value: BaseProxyEvent) -> None:
+        self._state.event = value
+
+    @property
+    def lambda_context(self) -> Any:
+        return self._state.lambda_context or BaseRouter.lambda_context
+
+    @lambda_context.setter
+    def lambda_context(self, value: Any) -> None:
+        self._state.lambda_context = value
+
+    @property
+    def context(self) -> dict:
+        return self._state.context
+
+    @context.setter
+    def context(self, value: dict) -> None:
+        self._state.context = value
+
+    @property
+    def processed_stack_frames(self) -> list[str]:
+        return self._state.processed_stack_frames
+
+    @processed_stack_frames.setter
+    def processed_stack_frames(self, value: list[str]) -> None:
+        self._state.processed_stack_frames = value
 
     def _to_proxy_event(self, event: dict) -> BaseProxyEvent:
         """Convert event dict to HttpProxyEvent."""
@@ -234,7 +287,7 @@ async def _resolve_async(self) -> dict:  # type: ignore[override]
         response_builder = await super()._resolve_async()
         return response_builder.build(self.current_event, self._cors)
 
-    async def asgi_handler(self, scope: dict, receive: Callable, send: Callable) -> None:
+    async def asgi_handler(self, scope: MutableMapping[str, Any], receive: Callable, send: Callable) -> None:
         """
         ASGI interface - allows running with uvicorn/hypercorn/etc.
 
@@ -274,25 +327,27 @@ async def asgi_handler(self, scope: dict, receive: Callable, send: Callable) ->
         # Create mock Lambda context
         context: Any = MockLambdaContext()
 
-        # Set up resolver state (similar to resolve())
-        BaseRouter.current_event = self._to_proxy_event(event._data)
-        BaseRouter.lambda_context = context
-
-        self._is_async_mode = True
-
+        # Never write BaseRouter's class attributes: another ASGI request may
+        # enter while validation or the handler is awaiting I/O.
+        state = _RequestState(
+            event=self._to_proxy_event(event._data),
+            lambda_context=context,
+            context=self._startup_state.context.copy(),
+        )
+        token = self._request_state.set(state)
         try:
-            # Use async resolve
             response = await self._resolve_async()
         finally:
-            self._is_async_mode = False
-            self.clear_context()
+            # Reset only this task's binding. Middleware threads may still be
+            # unwinding after cancellation and retain their request's state.
+            self._request_state.reset(token)
 
         # Send HTTP response
         await self._send_response(send, response)
 
     async def __call__(  # type: ignore[override]
         self,
-        scope: dict,
+        scope: MutableMapping[str, Any],
         receive: Callable,
         send: Callable,
     ) -> None:

tests/functional/event_handler/_pydantic/test_http_resolver_pydantic.py

@@ -365,3 +365,85 @@ def create_user(user: UserModel) -> UserResponse:
     # THEN schema includes 422 response
     post_operation = schema.paths["/users"].post
     assert 422 in post_operation.responses
+
+
+async def _post_concurrent_request(app, name: str | None):
+    payload = {} if name is None else {"name": name, "age": 30}
+    scope = {
+        "type": "http",
+        "method": "POST",
+        "path": "/concurrent",
+        "headers": [(b"content-type", b"application/json"), (b"x-request-name", (name or "invalid").encode())],
+        "query_string": b"",
+    }
+    send, captured = make_asgi_send()
+    await asyncio.wait_for(app(scope, make_asgi_receive(json.dumps(payload).encode()), send), timeout=5)
+    return captured["status_code"], json.loads(captured["body"])
+
+
+def test_concurrent_asgi_validation_preserves_each_request_body():
+    # GIVEN one local application using native request validation
+    app = HttpResolverLocal(enable_validation=True)
+
+    @app.post("/concurrent")
+    async def echo(user: UserModel) -> dict:
+        await asyncio.sleep(0)
+        return {"name": user.name}
+
+    async def scenario():
+        # WHEN distinct bodies are submitted concurrently
+        return await asyncio.gather(*(_post_concurrent_request(app, str(i)) for i in range(6)))
+
+    # THEN each caller receives its own input
+    assert asyncio.run(scenario()) == [(200, {"name": str(i)}) for i in range(6)]
+
+
+@pytest.mark.parametrize("interruption", ["invalid", "cancelled"])
+def test_interrupted_asgi_request_does_not_clear_another_requests_state(interruption):
+    async def scenario():
+        # GIVEN an active request that reads its context after awaiting I/O
+        app = HttpResolverLocal(enable_validation=True)
+        entered = {name: asyncio.Event() for name in ("first", "second", "later")}
+        release = {name: asyncio.Event() for name in entered}
+        pending = []
+
+        @app.post("/concurrent")
+        async def echo(user: UserModel) -> dict:
+            app.append_context(name=user.name)
+            entered[user.name].set()
+            await release[user.name].wait()
+            return {"name": app.context["name"], "header": app.current_event.headers["x-request-name"]}
+
+        async def start(name):
+            task = asyncio.create_task(_post_concurrent_request(app, name))
+            pending.append(task)
+            await asyncio.wait_for(entered[name].wait(), timeout=5)
+            return task
+
+        try:
+            first = await start("first")
+            # WHEN another request fails validation or an overlapping request is cancelled
+            if interruption == "invalid":
+                status, _ = await _post_concurrent_request(app, None)
+                assert status == 422
+                survivor, name = first, "first"
+            else:
+                survivor, name = await start("second"), "second"
+                first.cancel()
+                with pytest.raises(asyncio.CancelledError):
+                    await first
+
+            # THEN the surviving request retains both context and headers
+            release[name].set()
+            assert await survivor == (200, {"name": name, "header": name})
+            release["later"].set()
+            assert await _post_concurrent_request(app, "later") == (200, {"name": "later", "header": "later"})
+        finally:
+            for event in release.values():
+                event.set()
+            for task in pending:
+                if not task.done():
+                    task.cancel()
+            await asyncio.gather(*pending, return_exceptions=True)
+
+    asyncio.run(scenario())