feat: auto-detect ASGI mode for @aio decorated functions (GoogleCloudPlatform#387)

## Summary
- Implement automatic ASGI mode detection for functions decorated with `@aio.http` or `@aio.cloud_event`
- The approach creates a Flask app first, loads the module within its context, then checks if ASGI is needed
- This results in an unused Flask app for ASGI functions, but we accept this memory overhead as a trade-off
- The `--asgi` CLI flag still works and skips the Flask app creation for optimization

## Implementation Details
- Added `ASGI_FUNCTIONS` set to `_function_registry.py` to track functions that require ASGI
- Updated `@aio.http` and `@aio.cloud_event` decorators to register functions in `ASGI_FUNCTIONS`
- Modified `create_app()` to auto-detect ASGI requirements after module loading:
  1. Always creates a Flask app first
  2. Loads the user module within Flask app context
  3. Checks if target function is in `ASGI_FUNCTIONS` registry
  4. If ASGI is needed, delegates to `create_asgi_app_from_module()`
- The `--asgi` CLI flag continues to work, bypassing Flask app creation entirely for performance

## Trade-offs
- **Memory overhead**: ASGI functions will have an unused Flask app instance created during auto-detection
- **Accepted trade-off**: This avoids loading modules twice which could cause side effects
- **Optimization available**: Users can still use `--asgi` flag to skip Flask app creation entirely

## Test plan
- [x] Added tests to verify decorators register functions in `ASGI_FUNCTIONS`
- [x] Added CLI tests to verify auto-detection works for `@aio` decorated functions
- [x] Added CLI tests to verify regular functions still use Flask/WSGI mode
- [x] Added proper test isolation with registry cleanup fixtures
- [x] All existing tests pass
- [x] Linting passes

Author: taeold

.coveragerc-py37

@@ -12,14 +12,11 @@ omit =
 
 [report]
 exclude_lines =
-    # Have to re-enable the standard pragma
     pragma: no cover
-    
-    # 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
-    
-    # Exclude async-specific classes and functions in execution_id.py
     class AsgiMiddleware:
-    def set_execution_context_async
+    def set_execution_context_async
+    return create_asgi_app_from_module
+    app = create_asgi_app\(target, source, signature_type\)

src/functions_framework/__init__.py

@@ -327,6 +327,16 @@ def crash_handler(e):
 
 
 def create_app(target=None, source=None, signature_type=None):
+    """Create an app for the function.
+
+    Args:
+        target: The name of the target function to invoke
+        source: The source file containing the function
+        signature_type: The signature type of the function
+
+    Returns:
+        A Flask WSGI app or Starlette ASGI app depending on function decorators
+    """
     target = _function_registry.get_function_target(target)
     source = _function_registry.get_function_source(source)
 
@@ -370,6 +380,7 @@ def handle_none(rv):
         setup_logging()
 
     _app.wsgi_app = execution_id.WsgiMiddleware(_app.wsgi_app)
+
     # Execute the module, within the application context
     with _app.app_context():
         try:
@@ -394,6 +405,23 @@ def function(*_args, **_kwargs):
                 # command fails.
                 raise e from None
 
+    use_asgi = target in _function_registry.ASGI_FUNCTIONS
+    if use_asgi:
+        # This function needs ASGI, delegate to create_asgi_app
+        # Note: @aio decorators only register functions in ASGI_FUNCTIONS when the
+        # module is imported. We can't know if a function uses @aio until after
+        # we load the module.
+        #
+        # To avoid loading modules twice, we always create a Flask app first, load the
+        # module within its context, then check if ASGI is needed. This results in an
+        # unused Flask app for ASGI functions, but we accept this memory overhead as a
+        # trade-off.
+        from functions_framework.aio import create_asgi_app_from_module
+
+        return create_asgi_app_from_module(
+            target, source, signature_type, source_module, spec
+        )
+
     # Get the configured function signature type
     signature_type = _function_registry.get_func_signature_type(target, signature_type)
 

src/functions_framework/_cli.py

@@ -16,7 +16,7 @@
 
 import click
 
-from functions_framework import create_app
+from functions_framework import _function_registry, create_app
 from functions_framework._http import create_server
 
 
@@ -39,11 +39,10 @@
     help="Use ASGI server for function execution",
 )
 def _cli(target, source, signature_type, host, port, debug, asgi):
-    if asgi:  # pragma: no cover
+    if asgi:
         from functions_framework.aio import create_asgi_app
 
         app = create_asgi_app(target, source, signature_type)
     else:
         app = create_app(target, source, signature_type)
-
     create_server(app, debug).run(host, port)

src/functions_framework/_function_registry.py

@@ -40,6 +40,10 @@
 # Keys are the user function name, values are the type of the function input
 INPUT_TYPE_MAP = {}
 
+# ASGI_FUNCTIONS stores function names that require ASGI mode.
+# Functions decorated with @aio.http or @aio.cloud_event are added here.
+ASGI_FUNCTIONS = set()
+
 
 def get_user_function(source, source_module, target):
     """Returns user function, raises exception for invalid function."""

src/functions_framework/aio/__init__.py

@@ -62,6 +62,7 @@ def cloud_event(func: CloudEventFunction) -> CloudEventFunction:
     _function_registry.REGISTRY_MAP[func.__name__] = (
         _function_registry.CLOUDEVENT_SIGNATURE_TYPE
     )
+    _function_registry.ASGI_FUNCTIONS.add(func.__name__)
     if inspect.iscoroutinefunction(func):
 
         @functools.wraps(func)
@@ -82,6 +83,7 @@ def http(func: HTTPFunction) -> HTTPFunction:
     _function_registry.REGISTRY_MAP[func.__name__] = (
         _function_registry.HTTP_SIGNATURE_TYPE
     )
+    _function_registry.ASGI_FUNCTIONS.add(func.__name__)
 
     if inspect.iscoroutinefunction(func):
 
@@ -213,6 +215,29 @@ async def __call__(self, scope, receive, send):
             # Don't re-raise to prevent starlette from printing traceback again
 
 
+def create_asgi_app_from_module(target, source, signature_type, source_module, spec):
+    """Create an ASGI application from an already-loaded module.
+
+    Args:
+        target: The name of the target function to invoke
+        source: The source file containing the function
+        signature_type: The signature type of the function
+        source_module: The already-loaded module
+        spec: The module spec
+
+    Returns:
+        A Starlette ASGI application instance
+    """
+    enable_id_logging = _enable_execution_id_logging()
+    if enable_id_logging:  # pragma: no cover
+        _configure_app_execution_id_logging()
+
+    function = _function_registry.get_user_function(source, source_module, target)
+    signature_type = _function_registry.get_func_signature_type(target, signature_type)
+
+    return _create_asgi_app_with_function(function, signature_type, enable_id_logging)
+
+
 def create_asgi_app(target=None, source=None, signature_type=None):
     """Create an ASGI application for the function.
 
@@ -243,6 +268,11 @@ def create_asgi_app(target=None, source=None, signature_type=None):
     function = _function_registry.get_user_function(source, source_module, target)
     signature_type = _function_registry.get_func_signature_type(target, signature_type)
 
+    return _create_asgi_app_with_function(function, signature_type, enable_id_logging)
+
+
+def _create_asgi_app_with_function(function, signature_type, enable_id_logging):
+    """Create an ASGI app with the given function and signature type."""
     is_async = inspect.iscoroutinefunction(function)
     routes = []
     if signature_type == _function_registry.HTTP_SIGNATURE_TYPE:

tests/test_cli.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
+import pathlib
 import sys
 
 import pretend
@@ -20,9 +22,30 @@
 from click.testing import CliRunner
 
 import functions_framework
+import functions_framework._function_registry as _function_registry
 
 from functions_framework._cli import _cli
 
+# Conditional import for Starlette (Python 3.8+)
+if sys.version_info >= (3, 8):
+    from starlette.applications import Starlette
+else:
+    Starlette = None
+
+
+@pytest.fixture(autouse=True)
+def clean_registries():
+    """Clean up both REGISTRY_MAP and ASGI_FUNCTIONS registries."""
+    original_registry_map = _function_registry.REGISTRY_MAP.copy()
+    original_asgi = _function_registry.ASGI_FUNCTIONS.copy()
+    _function_registry.REGISTRY_MAP.clear()
+    _function_registry.ASGI_FUNCTIONS.clear()
+    yield
+    _function_registry.REGISTRY_MAP.clear()
+    _function_registry.REGISTRY_MAP.update(original_registry_map)
+    _function_registry.ASGI_FUNCTIONS.clear()
+    _function_registry.ASGI_FUNCTIONS.update(original_asgi)
+
 
 def test_cli_no_arguments():
     runner = CliRunner()
@@ -124,3 +147,19 @@ def test_asgi_cli(monkeypatch):
     assert result.exit_code == 0
     assert create_asgi_app.calls == [pretend.call("foo", None, "http")]
     assert asgi_server.run.calls == [pretend.call("0.0.0.0", 8080)]
+
+
+def test_cli_auto_detects_asgi_decorator():
+    """Test that CLI auto-detects @aio decorated functions without --asgi flag."""
+    # Use the actual async_decorator.py test file which has @aio.http decorated functions
+    test_functions_dir = pathlib.Path(__file__).parent / "test_functions" / "decorators"
+    source = test_functions_dir / "async_decorator.py"
+
+    # Call create_app without any asgi flag - should auto-detect
+    app = functions_framework.create_app(target="function_http", source=str(source))
+
+    # Verify it created a Starlette app (ASGI)
+    assert isinstance(app, Starlette)
+
+    # Verify the function was registered in ASGI_FUNCTIONS
+    assert "function_http" in _function_registry.ASGI_FUNCTIONS

tests/test_decorator_functions.py

@@ -19,6 +19,8 @@
 from cloudevents import conversion as ce_conversion
 from cloudevents.http import CloudEvent
 
+import functions_framework._function_registry as registry
+
 # Conditional import for Starlette
 if sys.version_info >= (3, 8):
     from starlette.testclient import TestClient as StarletteTestClient
@@ -35,6 +37,21 @@
 
 TEST_FUNCTIONS_DIR = pathlib.Path(__file__).resolve().parent / "test_functions"
 
+
+@pytest.fixture(autouse=True)
+def clean_registries():
+    """Clean up both REGISTRY_MAP and ASGI_FUNCTIONS registries."""
+    original_registry_map = registry.REGISTRY_MAP.copy()
+    original_asgi = registry.ASGI_FUNCTIONS.copy()
+    registry.REGISTRY_MAP.clear()
+    registry.ASGI_FUNCTIONS.clear()
+    yield
+    registry.REGISTRY_MAP.clear()
+    registry.REGISTRY_MAP.update(original_registry_map)
+    registry.ASGI_FUNCTIONS.clear()
+    registry.ASGI_FUNCTIONS.update(original_asgi)
+
+
 # Python 3.5: ModuleNotFoundError does not exist
 try:
     _ModuleNotFoundError = ModuleNotFoundError
@@ -128,3 +145,33 @@ def test_aio_http_dict_response():
     resp = client.post("/")
     assert resp.status_code == 200
     assert resp.json() == {"message": "hello", "count": 42, "success": True}
+
+
+def test_aio_decorators_register_asgi_functions():
+    """Test that @aio decorators add function names to ASGI_FUNCTIONS registry."""
+    from functions_framework.aio import cloud_event, http
+
+    @http
+    async def test_http_func(request):
+        return "test"
+
+    @cloud_event
+    async def test_cloud_event_func(event):
+        pass
+
+    assert "test_http_func" in registry.ASGI_FUNCTIONS
+    assert "test_cloud_event_func" in registry.ASGI_FUNCTIONS
+
+    assert registry.REGISTRY_MAP["test_http_func"] == "http"
+    assert registry.REGISTRY_MAP["test_cloud_event_func"] == "cloudevent"
+
+    @http
+    def test_http_sync(request):
+        return "sync"
+
+    @cloud_event
+    def test_cloud_event_sync(event):
+        pass
+
+    assert "test_http_sync" in registry.ASGI_FUNCTIONS
+    assert "test_cloud_event_sync" in registry.ASGI_FUNCTIONS

tests/test_function_registry.py

@@ -13,9 +13,25 @@
 # limitations under the License.
 import os
 
+import pytest
+
 from functions_framework import _function_registry
 
 
+@pytest.fixture(autouse=True)
+def clean_registries():
+    """Clean up both REGISTRY_MAP and ASGI_FUNCTIONS registries."""
+    original_registry_map = _function_registry.REGISTRY_MAP.copy()
+    original_asgi = _function_registry.ASGI_FUNCTIONS.copy()
+    _function_registry.REGISTRY_MAP.clear()
+    _function_registry.ASGI_FUNCTIONS.clear()
+    yield
+    _function_registry.REGISTRY_MAP.clear()
+    _function_registry.REGISTRY_MAP.update(original_registry_map)
+    _function_registry.ASGI_FUNCTIONS.clear()
+    _function_registry.ASGI_FUNCTIONS.update(original_asgi)
+
+
 def test_get_function_signature():
     test_cases = [
         {

tests/test_functions.py

@@ -495,7 +495,7 @@ def test_error_paths(http_trigger_client, path):
 def test_lazy_wsgi_app(monkeypatch, target, source, signature_type):
     actual_app_stub = pretend.stub()
     wsgi_app = pretend.call_recorder(lambda *a, **kw: actual_app_stub)
-    create_app = pretend.call_recorder(lambda *a: wsgi_app)
+    create_app = pretend.call_recorder(lambda *a, **kw: wsgi_app)
     monkeypatch.setattr(functions_framework, "create_app", create_app)
 
     # Test that it's lazy