Add websocket support to middleware

Author: tpoliaw

src/blueapi/service/main.py

@@ -36,11 +36,14 @@
 from starlette.responses import JSONResponse
 from super_state_machine.errors import TransitionError
 
-from blueapi import __version__
 from blueapi.config import ApplicationConfig, OIDCConfig, Tag
 from blueapi.core.bluesky_types import DataEvent
 from blueapi.service import interface
 from blueapi.service.authentication import CommonHttpOAuth
+from blueapi.service.middleware import (
+    ObservabilityContextPropagator,
+    VersionHeaders,
+)
 from blueapi.worker import TrackableTask, WorkerState
 from blueapi.worker.event import ProgressEvent, TaskStatusEnum, WorkerEvent
 from blueapi.worker.worker_errors import WorkerBusyError
@@ -132,8 +135,9 @@ def get_app(config: ApplicationConfig):
     app.include_router(secure_router, dependencies=dependencies)
     app.add_exception_handler(KeyError, on_key_error_404)
     app.add_exception_handler(jwt.PyJWTError, on_token_error_401)
-    app.middleware("http")(add_version_headers)
-    app.middleware("http")(inject_propagated_observability_context)
+
+    app.add_middleware(ObservabilityContextPropagator)
+    app.add_middleware(VersionHeaders)
     app.middleware("http")(log_request_details)
     if config.api.cors:
         app.add_middleware(
@@ -625,15 +629,6 @@ def start(config: ApplicationConfig):
     )
 
 
-async def add_version_headers(
-    request: Request, call_next: Callable[[Request], Awaitable[Response]]
-):
-    response = await call_next(request)
-    response.headers["X-API-Version"] = ApplicationConfig.REST_API_VERSION
-    response.headers["X-BlueAPI-Version"] = __version__
-    return response
-
-
 async def log_request_details(
     request: Request, call_next: Callable[[Request], Awaitable[StreamingResponse]]
 ) -> Response:

src/blueapi/service/middleware.py

@@ -0,0 +1,60 @@
+import logging
+
+from opentelemetry.context import attach
+from opentelemetry.propagate import get_global_textmap
+from starlette.types import ASGIApp, Message, Receive, Scope, Send
+
+from blueapi import __version__
+from blueapi.config import ApplicationConfig
+
+OBS_LOGGER = logging.getLogger("blueapi.service.middleware.obs")
+VER_LOGGER = logging.getLogger("blueapi.service.middleware.version")
+
+CONTEXT_HEADER = ApplicationConfig.CONTEXT_HEADER.encode()
+VENDOR_CONTEXT_HEADER = ApplicationConfig.VENDOR_CONTEXT_HEADER.encode()
+
+API_VERSION = (b"x-api-version", ApplicationConfig.REST_API_VERSION.encode("utf-8"))
+VERSION = (b"x-blueapi-version", __version__.encode("utf-8"))
+
+
+class VersionHeaders:
+    def __init__(self, app: ASGIApp):
+        self.app = app
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send):
+        if scope.get("type") not in ("websocket", "http"):
+            return await self.app(scope, receive, send)
+
+        async def local_send(message: Message):
+            VER_LOGGER.info("message: %s", message)
+            if message["type"] in ("websocket.accept", "http.response.start"):
+                message["headers"].append(VERSION)
+                message["headers"].append(API_VERSION)
+            await send(message)
+
+        await self.app(scope, receive, local_send)
+
+
+class ObservabilityContextPropagator:
+    def __init__(self, app: ASGIApp):
+        self.app = app
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send):
+        if scope["type"] not in ("http", "websocket"):
+            return await self.app(scope, receive, send)
+
+        ctx = None
+        v_ctx = None
+        for key, val in scope.get("headers", ()):
+            if key == CONTEXT_HEADER:
+                ctx = val.decode()
+            elif key == VENDOR_CONTEXT_HEADER:
+                v_ctx = val.decode()
+        if ctx:
+            OBS_LOGGER.debug("Propagating observability context: %s, %s", ctx, v_ctx)
+            carrier = {ApplicationConfig.CONTEXT_HEADER: ctx}
+            if v_ctx:
+                carrier[ApplicationConfig.VENDOR_CONTEXT_HEADER] = v_ctx
+            attach(get_global_textmap().extract(carrier))
+
+        await self.app(scope, receive, send)

tests/unit_tests/service/test_main.py

@@ -8,15 +8,15 @@
 from blueapi import __version__
 from blueapi.config import ApplicationConfig
 from blueapi.service.main import (
-    add_version_headers,
     get_passthrough_headers,
     log_request_details,
 )
+from blueapi.service.middleware import VersionHeaders
 
 
 async def test_add_version_header():
     app = FastAPI()
-    app.middleware("http")(add_version_headers)
+    app.add_middleware(VersionHeaders)
 
     @app.get("/")
     async def root():