How to modify the content-length post modification of response in a middleware in FastAPI?

Question

Example Code:

# Here is a minimal reproducible example

import json
from starlette.datastructures import MutableHeaders
from starlette.types import ASGIApp, Receive, Scope, Send, Message
import datetime
import socket
import uvicorn
from fastapi import FastAPI

class MetaDataAdderMiddleware:
    application_generic_urls = ['/openapi.json', '/docs', '/docs/oauth2-redirect', '/redoc']
    
    def __init__(self, app: ASGIApp) -> None:
        self.app = app

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:

        start_time = datetime.datetime.utcnow()

        async def send_wrapper(message: Message) -> None:   
         
            if message["type"] == "http.response.body" and len(message["body"]) and not any([scope["path"].startswith(endpoint) for endpoint in MetaDataAdderMiddleware.application_generic_urls]):

                response_body = json.loads(message["body"].decode())

                end_time = datetime.datetime.utcnow()
                response_processing_time_seconds = end_time - start_time
                data = {}
                data["data"] = response_body
                data['metadata'] = {
                    'request_timestamp_utc': start_time,
                    'response_timestamp_utc': end_time,
                    'processing_time_seconds': response_processing_time_seconds,
                    'service_host': socket.gethostname()
                }

                data_to_be_sent_to_user = json.dumps(data, default=str).encode("utf-8")
                message["body"] = data_to_be_sent_to_user

            await send(message)
        await self.app(scope, receive, send_wrapper)


app = FastAPI(
    title="MY DUMMY APP",
)

app.add_middleware(MetaDataAdderMiddleware)

@app.get("/")
async def root():
    return {"message": "Hello World"}

Description: So here is my usecase: All of my endpoints in FastAPI APP, whatever response they are sending, I need to wrap that response, with some metadata. Let's say, some endpoint is sending me this: {"data_key": "data_value"}. But, the users should see, this as the final output:

{
    "data": {"data_key": "data_value"}
     "metadata": {
          "request_timestamp_utc": "somevalue",
          ...and so on
     }
}

I have a big application, and numerous routers. We have achieved the functionality of adding Request ID, Authentication and Authorization, so far by writing middlewares.

However, when I hit APIs of my app, after adding the abovementioned MetaDataAdderMiddleware, I am greeted with this following error:

ERROR:    Exception in ASGI application
Traceback (most recent call last):
  File "<MY PYTHON PATH>/lib/python3.6/site-packages/uvicorn/protocols/http/httptools_impl.py", line 521, in send
    raise RuntimeError("Response content longer than Content-Length")
RuntimeError: Response content longer than Content-Length

This error is logical enough, since I have modified the Response body but not changed the content-length headers.

Here is snapshot of message and scope bodies in the send_wrapper function level, and as well as header values I have printed from the httptools_impl.py level: (I have edited out some fields, to mask org-specific things)

send_wrapper called
message: {'type': 'http.response.start', 'status': 200, 'headers': [(b'content-length', b'58'), (b'content-type', b'application/json')]}
scope: {'type': 'http', 'asgi': {'version': '3.0', 'spec_version': '2.1'}, 'http_version': '1.1', 'scheme': 'http', 'method': 'POST', 'root_path': '', 'query_string': b'', 'headers': [(b'content-type', b'application/json'), (b'accept', b'*/*'), (b'cache-control', b'no-cache'), (b'accept-encoding', b'gzip, deflate'), (b'content-length', b'238'), (b'connection', b'keep-alive')], 'app': <fastapi.applications.FastAPI object at >, 'fastapi_astack': <contextlib2.AsyncExitStack object at >, 'router': <fastapi.routing.APIRouter object at >, 'endpoint': <function initiate_playbook_execution at >, 'path_params': {}, 'route': <fastapi.routing.APIRoute object at >}
INFO:     - "POST /MYAPI" 200 OK

INSIDE httptools_impl
name: b'content-length' | value: b'58'
self.expected_content_length: 58

send_wrapper called
message: {'type': 'http.response.body', 'body': b'{"status":true,"stdout":null,"stderr":null,"message":null}'}
scope: {'type': 'http', 'asgi': {'version': '3.0', 'spec_version': '2.1'}, 'http_version': '1.1', 'scheme': 'http', 'method': 'POST', 'root_path': '', 'query_string': b'', 'headers': [(b'content-type', b'application/json'), (b'accept', b'*/*'), (b'cache-control', b'no-cache'), (b'accept-encoding', b'gzip, deflate'), (b'content-length', b'238'), (b'connection', b'keep-alive')], 'app': <fastapi.applications.FastAPI object at >, 'fastapi_astack': <contextlib2.AsyncExitStack object at >, 'router': <fastapi.routing.APIRouter object at >, 'endpoint': <function initiate_playbook_execution at >, 'path_params': {}, 'route': <fastapi.routing.APIRoute object at >}

INSIDE httptools_impl
body: b'{"data": {"status": true, "stdout": null, "stderr": null, "message": null}, "metadata": {"request_timestamp_utc": "BLAH", "response_timestamp_utc": "BLAH", "processing_time_seconds": "0:00:00.469472", "some_field": "some_value"}}'
num_bytes: 286

Here are the attempts that I have made to update the content-length:

• In the send wrapper function just after I update the response body, I have tried doing the following:

  data_to_be_sent_to_user = json.dumps(data, default=str).encode("utf-8")
  message["body"] = data_to_be_sent_to_user
  headers = MutableHeaders(scope=scope)
  headers["content-length"] = str(len(data_to_be_sent_to_user))
  
  # But this hasn't worked, no change in situation!
  

How can I proceed forward?

Answer 1

Thanks to @MatsLindh comment, I referred to Starlette's GZipMiddleware codebase here: https://github.com/encode/starlette/blob/fcc4c705ff69182ebd663bc686cb55c242d32683/starlette/middleware/gzip.py#L60

So the idea is, the problematic content-length value is in header present in http.response.start message. So, how GZipMiddleware has been written is, they have simply not sent this first http.response.start message instantly. Instead, they also capture http.response.body, then modify the response, then find its length, then update the length in http.response.start message, and then send both these messages in the correct order.

The working implementation that I was able to write, borrowing heavily from GZipMiddleware is here:

from starlette.types import ASGIApp, Receive, Scope, Send, Message
from starlette.requests import Request
import json
from starlette.datastructures import MutableHeaders


class MetaDataAdderMiddleware:
    application_generic_urls = ['/openapi.json', '/docs', '/docs/oauth2-redirect', '/redoc']
    
    def __init__(
            self, 
            app: ASGIApp
    ) -> None:
        self.app = app

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] == "http" and not any([scope["path"].startswith(endpoint) for endpoint in MetaDataAdderMiddleware.application_generic_urls]):
            responder = MetaDataAdderMiddlewareResponder(self.app, self.standard_meta_data, self.additional_custom_information)
            await responder(scope, receive, send)
            return
        await self.app(scope, receive, send)


class MetaDataAdderMiddlewareResponder:

    def __init__(
            self,
            app: ASGIApp,
    ) -> None:
        """
        """
        self.app = app
        self.initial_message: Message = {}

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        self.send = send
        await self.app(scope, receive, self.send_with_meta_response)

    async def send_with_meta_response(self, message: Message):

        message_type = message["type"]
        if message_type == "http.response.start":
            # Don't send the initial message until we've determined how to
            # modify the outgoing headers correctly.
            self.initial_message = message

        elif message_type == "http.response.body":
            response_body = json.loads(message["body"].decode())

            data = {}
            data["data"] = response_body
            data['metadata'] = {
                'field_1': 'value_1',
                'field_2': 'value_2'
            }

            data_to_be_sent_to_user = json.dumps(data, default=str).encode("utf-8")

            headers = MutableHeaders(raw=self.initial_message["headers"])
            headers["Content-Length"] = str(len(data_to_be_sent_to_user))
            message["body"] = data_to_be_sent_to_user

            await self.send(self.initial_message)
            await self.send(message)