feat: add execution_id support for async stack (#377)

* feat: add execution_id support for async stack

- Add contextvars support to execution_id.py for async-safe context storage
- Create AsgiMiddleware class to inject execution_id into ASGI requests
- Add set_execution_context_async decorator for both sync and async functions
- Update LoggingHandlerAddExecutionId to support both Flask g and contextvars
- Integrate execution_id support in aio/__init__.py with proper exception handling
- Add comprehensive async tests matching sync test functionality
- Follow Starlette best practices for exception handling

The implementation enables automatic execution_id injection and logging for async
functions when LOG_EXECUTION_ID=true, matching the existing sync stack behavior.

* refactor: move exception logging to crash handler for cleaner code

- Remove try/catch blocks from wrapper functions
- Centralize exception logging in _crash_handler
- Extract execution_id directly from request headers in crash handler
- Temporarily set context when logging exceptions to ensure execution_id is included
- This approach is cleaner and more similar to Flask's centralized exception handling

* refactor: improve code organization based on feedback

- Move imports to top of file instead of inside functions
- Extract common header parsing logic into _extract_context_from_headers helper
- Reduce code duplication between sync and async decorators
- Add comment explaining why crash handler needs to extract context from headers
- This addresses the context reset issue where decorators clean up before exception handlers run

* fix: preserve execution context for exception handlers

- Don't reset context on exception, only on successful completion
- This allows exception handlers to access execution_id naturally
- Simplify crash handler since context is now available
- Rely on Python's automatic contextvar cleanup when task completes
- Each request runs in its own task, so no risk of context leakage

This is more correct and follows the principle that context should be
available throughout the entire request lifecycle, including error handling.

* style: apply black and isort formatting

- Format code with black for consistent style
- Sort imports with isort for better organization
- All linting checks now pass

* refactor: clean up async tests and remove redundant comments

* chore: remove uv.lock from version control

* style: fix black formatting

* fix: skip async execution_id tests on Python 3.7

* refactor: reuse _enable_execution_id_logging from main module

* chore: more cleanup.

* test: remove unnecessary pragma no cover for sync_wrapper

* test: improve coverage by removing unnecessary pragma no cover annotations

* style: fix black formatting

* style: fix isort import ordering

* test: add back pragma no cover for genuinely hard-to-test edge cases

* refactor: simplify async decorator by removing dead code branch

* fix: exclude async-specific code from py37 coverage

The AsgiMiddleware class and set_execution_context_async function in
execution_id.py require Python 3.8+ due to async dependencies. These
are now excluded from coverage calculations in Python 3.7 environments.

* fix: improve async execution ID context propagation using contextvars

- Use contextvars.copy_context() to properly propagate execution context in async functions
- Implement AsyncExecutionIdHandler to handle JSON logging with execution_id
- Redirect logging output from stderr to stdout for consistency
- Add build dependency to dev dependencies
- Update tests to reflect new logging output location

* feat: Add execution ID logging for async functions

Refactors the async logging implementation to align with the sync version, ensuring consistent execution ID logging across both stacks.

* chore: clean up impl.

* refactor: define custom exception handling middleware to avoid duplicate log of traceback.

* style: run black

* chore: clean up code a little more.

* fix: propagate context in ce fns.

* style: more nits.

* chore: remove unncessary debug flag.

* fix: respond to PR comments

* fix: respond to more PR comments

Author: taeold

.coveragerc-py37

@@ -18,4 +18,8 @@ exclude_lines =
     # Don't complain about async-specific imports and code
     from functions_framework.aio import
     from functions_framework._http.asgi import
-    from functions_framework._http.gunicorn import UvicornApplication
+    from functions_framework._http.gunicorn import UvicornApplication
+    
+    # Exclude async-specific classes and functions in execution_id.py
+    class AsgiMiddleware:
+    def set_execution_context_async

conftest.py

@@ -50,8 +50,12 @@ def pytest_ignore_collect(collection_path, config):
     if sys.version_info >= (3, 8):
         return None
 
-    # Skip test_aio.py and test_asgi.py entirely on Python 3.7
-    if collection_path.name in ["test_aio.py", "test_asgi.py"]:
+    # Skip test_aio.py, test_asgi.py, and test_execution_id_async.py entirely on Python 3.7
+    if collection_path.name in [
+        "test_aio.py",
+        "test_asgi.py",
+        "test_execution_id_async.py",
+    ]:
         return True
 
     return None

pyproject.toml

@@ -61,3 +61,13 @@ functions_framework = ["py.typed"]
 
 [tool.setuptools.package-dir]
 "" = "src"
+
+[dependency-groups]
+dev = [
+    "black>=23.3.0",
+    "build>=1.1.1",
+    "isort>=5.11.5",
+    "pretend>=1.0.9",
+    "pytest>=7.4.4",
+    "pytest-asyncio>=0.21.2",
+]

src/functions_framework/aio/__init__.py

@@ -13,16 +13,24 @@
 # limitations under the License.
 
 import asyncio
+import contextvars
 import functools
 import inspect
+import logging
+import logging.config
 import os
+import traceback
 
 from typing import Any, Awaitable, Callable, Dict, Tuple, Union
 
 from cloudevents.http import from_http
 from cloudevents.http.event import CloudEvent
 
-from functions_framework import _function_registry
+from functions_framework import (
+    _enable_execution_id_logging,
+    _function_registry,
+    execution_id,
+)
 from functions_framework.exceptions import (
     FunctionsFrameworkException,
     MissingSourceException,
@@ -31,6 +39,7 @@
 try:
     from starlette.applications import Starlette
     from starlette.exceptions import HTTPException
+    from starlette.middleware import Middleware
     from starlette.requests import Request
     from starlette.responses import JSONResponse, Response
     from starlette.routing import Route
@@ -96,29 +105,27 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
-async def _crash_handler(request, exc):
-    headers = {_FUNCTION_STATUS_HEADER_FIELD: _CRASH}
-    return Response(f"Internal Server Error: {exc}", status_code=500, headers=headers)
-
-
-def _http_func_wrapper(function, is_async):
+def _http_func_wrapper(function, is_async, enable_id_logging=False):
+    @execution_id.set_execution_context_async(enable_id_logging)
     @functools.wraps(function)
     async def handler(request):
         if is_async:
             result = await function(request)
         else:
             # TODO: Use asyncio.to_thread when we drop Python 3.8 support
-            # Python 3.8 compatible version of asyncio.to_thread
             loop = asyncio.get_event_loop()
-            result = await loop.run_in_executor(None, function, request)
+            ctx = contextvars.copy_context()
+            result = await loop.run_in_executor(None, ctx.run, function, request)
         if isinstance(result, str):
             return Response(result)
         elif isinstance(result, dict):
             return JSONResponse(result)
         elif isinstance(result, tuple) and len(result) == 2:
-            # Support Flask-style tuple response
             content, status_code = result
-            return Response(content, status_code=status_code)
+            if isinstance(content, dict):
+                return JSONResponse(content, status_code=status_code)
+            else:
+                return Response(content, status_code=status_code)
         elif result is None:
             raise HTTPException(status_code=500, detail="No response returned")
         else:
@@ -127,7 +134,8 @@ async def handler(request):
     return handler
 
 
-def _cloudevent_func_wrapper(function, is_async):
+def _cloudevent_func_wrapper(function, is_async, enable_id_logging=False):
+    @execution_id.set_execution_context_async(enable_id_logging)
     @functools.wraps(function)
     async def handler(request):
         data = await request.body()
@@ -142,9 +150,9 @@ async def handler(request):
             await function(event)
         else:
             # TODO: Use asyncio.to_thread when we drop Python 3.8 support
-            # Python 3.8 compatible version of asyncio.to_thread
             loop = asyncio.get_event_loop()
-            await loop.run_in_executor(None, function, event)
+            ctx = contextvars.copy_context()
+            await loop.run_in_executor(None, ctx.run, function, event)
         return Response("OK")
 
     return handler
@@ -154,6 +162,64 @@ async def _handle_not_found(request: Request):
     raise HTTPException(status_code=404, detail="Not Found")
 
 
+def _configure_app_execution_id_logging():
+    logging.config.dictConfig(
+        {
+            "version": 1,
+            "handlers": {
+                "asgi": {
+                    "class": "logging.StreamHandler",
+                    "stream": "ext://functions_framework.execution_id.logging_stream",
+                },
+            },
+            "root": {"level": "INFO", "handlers": ["asgi"]},
+        }
+    )
+
+
+class ExceptionHandlerMiddleware:
+    def __init__(self, app):
+        self.app = app
+
+    async def __call__(self, scope, receive, send):
+        if scope["type"] != "http":  # pragma: no cover
+            await self.app(scope, receive, send)
+            return
+
+        try:
+            await self.app(scope, receive, send)
+        except Exception as exc:
+            logger = logging.getLogger()
+            tb_lines = traceback.format_exception(type(exc), exc, exc.__traceback__)
+            tb_text = "".join(tb_lines)
+
+            path = scope.get("path", "/")
+            method = scope.get("method", "GET")
+            error_msg = f"Exception on {path} [{method}]\n{tb_text}".rstrip()
+
+            logger.error(error_msg)
+
+            headers = [
+                [b"content-type", b"text/plain"],
+                [_FUNCTION_STATUS_HEADER_FIELD.encode(), _CRASH.encode()],
+            ]
+
+            await send(
+                {
+                    "type": "http.response.start",
+                    "status": 500,
+                    "headers": headers,
+                }
+            )
+            await send(
+                {
+                    "type": "http.response.body",
+                    "body": b"Internal Server Error",
+                }
+            )
+            # Don't re-raise to prevent starlette from printing traceback again
+
+
 def create_asgi_app(target=None, source=None, signature_type=None):
     """Create an ASGI application for the function.
 
@@ -175,14 +241,19 @@ def create_asgi_app(target=None, source=None, signature_type=None):
         )
 
     source_module, spec = _function_registry.load_function_module(source)
+
+    enable_id_logging = _enable_execution_id_logging()
+    if enable_id_logging:
+        _configure_app_execution_id_logging()
+
     spec.loader.exec_module(source_module)
     function = _function_registry.get_user_function(source, source_module, target)
     signature_type = _function_registry.get_func_signature_type(target, signature_type)
 
     is_async = inspect.iscoroutinefunction(function)
     routes = []
     if signature_type == _function_registry.HTTP_SIGNATURE_TYPE:
-        http_handler = _http_func_wrapper(function, is_async)
+        http_handler = _http_func_wrapper(function, is_async, enable_id_logging)
         routes.append(
             Route(
                 "/",
@@ -202,7 +273,9 @@ def create_asgi_app(target=None, source=None, signature_type=None):
             )
         )
     elif signature_type == _function_registry.CLOUDEVENT_SIGNATURE_TYPE:
-        cloudevent_handler = _cloudevent_func_wrapper(function, is_async)
+        cloudevent_handler = _cloudevent_func_wrapper(
+            function, is_async, enable_id_logging
+        )
         routes.append(
             Route("/{path:path}", endpoint=cloudevent_handler, methods=["POST"])
         )
@@ -221,10 +294,14 @@ def create_asgi_app(target=None, source=None, signature_type=None):
             f"Unsupported signature type for ASGI server: {signature_type}"
         )
 
-    exception_handlers = {
-        500: _crash_handler,
-    }
-    app = Starlette(routes=routes, exception_handlers=exception_handlers)
+    app = Starlette(
+        routes=routes,
+        middleware=[
+            Middleware(ExceptionHandlerMiddleware),
+            Middleware(execution_id.AsgiMiddleware),
+        ],
+    )
+
     return app