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.

Author: taeold

src/functions_framework/aio/__init__.py

@@ -64,23 +64,11 @@ async def _crash_handler(request, exc):
     tb_text = ''.join(tb_lines)
     error_msg = f"Exception on {request.url.path} [{request.method}]\n{tb_text}".rstrip()
 
-    # When execution ID logging is enabled, we need to extract execution_id from headers
-    # because the decorator resets the context before exception handlers run
+    # Context should still be available since we don't reset on exception
     if _enable_execution_id_logging():
-        context = execution_id._extract_context_from_headers(request.headers)
-        if context.execution_id:
-            # Temporarily set context for logging
-            token = execution_id.execution_context_var.set(context)
-            try:
-                # Output as JSON so LoggingHandlerAddExecutionId can process it
-                log_entry = {"message": error_msg, "levelname": "ERROR"}
-                logger.error(json.dumps(log_entry))
-            finally:
-                execution_id.execution_context_var.reset(token)
-        else:  # pragma: no cover
-            # No execution ID in headers
-            log_entry = {"message": error_msg, "levelname": "ERROR"}
-            logger.error(json.dumps(log_entry))
+        # Output as JSON so LoggingHandlerAddExecutionId can process it
+        log_entry = {"message": error_msg, "levelname": "ERROR"}
+        logger.error(json.dumps(log_entry))
     else:
         # Execution ID logging not enabled, log plain text
         logger.error(error_msg)

src/functions_framework/execution_id.py

@@ -166,7 +166,11 @@ def wrapper(*args, **kwargs):
             _set_current_context(context)
 
             with stderr_redirect, stdout_redirect:
-                return view_function(*args, **kwargs)
+                result = view_function(*args, **kwargs)
+                
+                # Context cleanup happens automatically via Flask's request context
+                # No need to manually clean up flask.g
+                return result
 
         return wrapper
 
@@ -195,16 +199,17 @@ async def async_wrapper(request, *args, **kwargs):
             # Set context using contextvars
             token = execution_context_var.set(context)
 
-            try:
-                with stderr_redirect, stdout_redirect:
-                    # Handle both sync and async functions
-                    if inspect.iscoroutinefunction(view_function):
-                        return await view_function(request, *args, **kwargs)
-                    else:
-                        return view_function(request, *args, **kwargs)  # pragma: no cover
-            finally:
-                # Reset context
+            with stderr_redirect, stdout_redirect:
+                # Handle both sync and async functions
+                if inspect.iscoroutinefunction(view_function):
+                    result = await view_function(request, *args, **kwargs)
+                else:
+                    result = view_function(request, *args, **kwargs)  # pragma: no cover
+                
+                # Only reset context on successful completion
+                # On exception, leave context available for exception handlers
                 execution_context_var.reset(token)
+                return result
 
         @functools.wraps(view_function)
         def sync_wrapper(request, *args, **kwargs):  # pragma: no cover
@@ -214,12 +219,13 @@ def sync_wrapper(request, *args, **kwargs):  # pragma: no cover
             # Set context using contextvars
             token = execution_context_var.set(context)
 
-            try:
-                with stderr_redirect, stdout_redirect:
-                    return view_function(request, *args, **kwargs)
-            finally:
-                # Reset context
+            with stderr_redirect, stdout_redirect:
+                result = view_function(request, *args, **kwargs)
+                
+                # Only reset context on successful completion
+                # On exception, leave context available for exception handlers
                 execution_context_var.reset(token)
+                return result
 
         # Return appropriate wrapper based on whether the function is async
         if inspect.iscoroutinefunction(view_function):