feat: auto-detect ASGI mode for @aio decorated functions

Functions decorated with @aio.http or @aio.cloud_event now automatically
run in ASGI mode without requiring the --asgi flag. This improves the
developer experience by removing the need to remember to pass the flag
when using async decorators.

Author: taeold

src/functions_framework/_cli.py

@@ -16,7 +16,7 @@
 
 import click
 
-from functions_framework import create_app
+from functions_framework import create_app, _function_registry
 from functions_framework._http import create_server
 
 
@@ -39,6 +39,9 @@
     help="Use ASGI server for function execution",
 )
 def _cli(target, source, signature_type, host, port, debug, asgi):
+    if not asgi and target in _function_registry.ASGI_FUNCTIONS:
+        asgi = True
+
     if asgi:  # pragma: no cover
         from functions_framework.aio import create_asgi_app
 

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

@@ -69,6 +69,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)
@@ -89,6 +90,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):
 

tests/test_cli.py

@@ -20,6 +20,8 @@
 from click.testing import CliRunner
 
 import functions_framework
+import functions_framework._function_registry as _function_registry
+import functions_framework.aio
 
 from functions_framework._cli import _cli
 
@@ -124,3 +126,49 @@ 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_auto_asgi_for_aio_decorated_functions(monkeypatch):
+    original_asgi_functions = _function_registry.ASGI_FUNCTIONS.copy()
+    _function_registry.ASGI_FUNCTIONS.clear()
+    _function_registry.ASGI_FUNCTIONS.add("my_aio_func")
+
+    asgi_app = pretend.stub()
+    create_asgi_app = pretend.call_recorder(lambda *a, **k: asgi_app)
+    aio_module = pretend.stub(create_asgi_app=create_asgi_app)
+    monkeypatch.setitem(sys.modules, "functions_framework.aio", aio_module)
+
+    asgi_server = pretend.stub(run=pretend.call_recorder(lambda host, port: None))
+    create_server = pretend.call_recorder(lambda app, debug: asgi_server)
+    monkeypatch.setattr(functions_framework._cli, "create_server", create_server)
+
+    runner = CliRunner()
+    result = runner.invoke(_cli, ["--target", "my_aio_func"])
+
+    assert create_asgi_app.calls == [pretend.call("my_aio_func", None, "http")]
+    assert asgi_server.run.calls == [pretend.call("0.0.0.0", 8080)]
+    
+    _function_registry.ASGI_FUNCTIONS.clear()
+    _function_registry.ASGI_FUNCTIONS.update(original_asgi_functions)
+
+
+def test_no_auto_asgi_for_regular_functions(monkeypatch):
+    original_asgi_functions = _function_registry.ASGI_FUNCTIONS.copy()
+    _function_registry.ASGI_FUNCTIONS.clear()
+
+    app = pretend.stub()
+    create_app = pretend.call_recorder(lambda *a, **k: app)
+    monkeypatch.setattr(functions_framework._cli, "create_app", create_app)
+
+    flask_server = pretend.stub(run=pretend.call_recorder(lambda host, port: None))
+    create_server = pretend.call_recorder(lambda app, debug: flask_server)
+    monkeypatch.setattr(functions_framework._cli, "create_server", create_server)
+
+    runner = CliRunner()
+    result = runner.invoke(_cli, ["--target", "regular_func"])
+
+    assert create_app.calls == [pretend.call("regular_func", None, "http")]
+    assert flask_server.run.calls == [pretend.call("0.0.0.0", 8080)]
+    
+    _function_registry.ASGI_FUNCTIONS.clear()
+    _function_registry.ASGI_FUNCTIONS.update(original_asgi_functions)

tests/test_decorator_functions.py

@@ -18,6 +18,7 @@
 
 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):
@@ -128,3 +129,43 @@ 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."""
+    original_registry_map = registry.REGISTRY_MAP.copy()
+    original_asgi_functions = registry.ASGI_FUNCTIONS.copy()
+    registry.REGISTRY_MAP.clear()
+    registry.ASGI_FUNCTIONS.clear()
+
+    from functions_framework.aio import http, cloud_event
+
+    @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
+    
+    registry.REGISTRY_MAP.clear()
+    registry.REGISTRY_MAP.update(original_registry_map)
+    registry.ASGI_FUNCTIONS.clear()
+    registry.ASGI_FUNCTIONS.update(original_asgi_functions)