Fix flask and fastapi get logs (#42)

Co-authored-by: José Valim <[email protected]>

Author: jonatanklosko

src/tidewave/django/__init__.py

@@ -69,15 +69,11 @@ def __init__(self, get_response: Callable):
         self.base_middleware = BaseMiddleware(self._dummy_wsgi_app, self.mcp_handler, self.config)
 
     def _setup_logging(self):
+        file_handler.addFilter(lambda record: record.name != "django.utils.autoreload")
         file_handler.addFilter(
-            CallbackFilter(lambda record: record.name != "django.utils.autoreload")
-        )
-        file_handler.addFilter(
-            CallbackFilter(
-                lambda record: not (
-                    (record.name == "django.server" or record.name == "django.request")
-                    and "/tidewave" in record.getMessage()
-                )
+            lambda record: not (
+                (record.name == "django.server" or record.name == "django.request")
+                and " /tidewave" in record.getMessage()
             )
         )
 

src/tidewave/fastapi/__init__.py

@@ -2,6 +2,7 @@
 FastAPI-specific integration for Tidewave
 """
 
+import logging
 import os
 from typing import Any, Optional
 
@@ -14,6 +15,7 @@
 from tidewave.fastapi.middleware import Middleware
 from tidewave.mcp_handler import MCPHandler
 from tidewave.middleware import Middleware as MCPMiddleware
+from tidewave.tools.get_logs import file_handler
 
 
 class Tidewave:
@@ -74,3 +76,21 @@ def wsgi_app(environ, start_response):
         mcp_middleware = MCPMiddleware(wsgi_app, mcp_handler, config)
         app.mount("/tidewave", WSGIMiddleware(mcp_middleware))
         app.add_middleware(Middleware)
+
+        self._setup_logging()
+
+    def _setup_logging(self):
+        file_handler.addFilter(
+            lambda record: not (
+                (record.name == "uvicorn.access") and " /tidewave" in record.getMessage()
+            )
+        )
+
+        # We set a global logger handler. The "uvicorn.error" logger
+        # propagates [1], so it will also invoke that handler, on the
+        # other hand, "uvicorn.access" logger does not propagate, so
+        # we need to add the handler separately.
+        #
+        # [1]: https://github.com/Kludex/uvicorn/blob/0.37.0/uvicorn/config.py#L94-L98
+        logging.getLogger().addHandler(file_handler)
+        logging.getLogger("uvicorn.access").addHandler(file_handler)

src/tidewave/flask/__init__.py

@@ -2,9 +2,10 @@
 Flask-specific integration for Tidewave
 """
 
+import logging
 from typing import Any, Optional
 
-from flask import current_app
+from flask import current_app, got_request_exception, request
 
 from tidewave import tools
 from tidewave.flask.middleware import Middleware
@@ -69,3 +70,22 @@ def init_app(self, app):
 
         app.wsgi_app = MCPMiddleware(Middleware(app.wsgi_app), mcp_handler, middleware_config)
         app.jinja_env.add_extension(Extension)
+
+        # In debug mode, exceptions are not logged [1], they propagate,
+        # the debugger shows an error page and only writes to stderr [2].
+        # We want the exceptions to appear in get_logs, so we log them
+        # explicitly.
+        #
+        # [1]: https://flask.palletsprojects.com/en/stable/api/#flask.Flask.handle_exception
+        # [2]: https://github.com/pallets/werkzeug/blob/3.1.3/src/werkzeug/debug/__init__.py#L378
+        got_request_exception.connect(app_exception_handler, app)
+
+
+def app_exception_handler(sender, exception, **extra):
+    # We log to the default logger. This does not log to the terminal,
+    # because terminal output comes from terminal handlers in more
+    # specific flask/werkzeug loggers.
+    logging.getLogger().exception(
+        f"Exception on {request.path} [{request.method}]",
+        exc_info=exception,
+    )

src/tidewave/flask/middleware.py

@@ -2,22 +2,49 @@
 Flask-specific middleware for dealing with CSP/frame headers.
 """
 
+import logging
 from typing import Any, Callable
 
 from tidewave.middleware import modify_csp
+from tidewave.tools.get_logs import file_handler
 
 
 class Middleware:
     def __init__(self, app: Callable):
         self.app = app
+        self.logging_initialized = False
 
     def __call__(self, environ: dict[str, Any], start_response: Callable):
+        self._maybe_initialize_logging()
+
         def handle_response(status, headers):
             modified_headers = self._process_response(headers)
             return start_response(status, modified_headers)
 
         return self.app(environ, handle_response)
 
+    def _maybe_initialize_logging(self):
+        # Ideally we would add the logger handler as part of
+        # tidewave.flask.Tidewave, however that is too soon. Both
+        # werkzeug [1] and flask [2] have conditional logic that adds
+        # a default terminal handler, only if no other handler is
+        # present, so if we add our handler too soon, no logs would
+        # show up in the terminal. Instead, we add our handler on the
+        # first request, by which point the default handler is already
+        # added.
+        #
+        # [1]: https://github.com/pallets/werkzeug/blob/3.1.3/src/werkzeug/_internal.py#L94-L95
+        # [2]: https://github.com/pallets/flask/blob/3.1.2/src/flask/logging.py#L76-L77
+
+        if not self.logging_initialized:
+            file_handler.addFilter(
+                lambda record: not (
+                    record.name == "werkzeug" and " /tidewave" in record.getMessage()
+                )
+            )
+            logging.getLogger().addHandler(file_handler)
+            self.logging_initialized = True
+
     def _process_response(self, headers):
         """
         Modify headers to allow embedding in Tidewave: