Enhancement: Support for returning AsyncGenerator[ServerSentEventMessage, None] at route handler

[winstxnhdw]

Summary

Currently, Litestar forces you to create an intermediary function that yields ServerSentEventMessage then wrap the generator in ServerSentEvent.

class SomeController(Controller):
    async def intermediary_function(self) -> AsyncGenerator[ServerSentEventMessage, None]:
        yield ServerSentEventMessage(data='Doing this..')
        yield ServerSentEventMessage(data='Doing that..')
        yield ServerSentEventMessage(data='Finished doing this and that..')

    @get()
    async def actual_server_sent_event_handler(self) -> ServerSentEvent:
        return ServerSentEvent(self.intermediary_function())   

This is not an ideal pattern as you now have a method that isn't a route handler in your controller. If you were to abstract it out as a function, you'll now have to find an elegant place to define this function. In any case, this is an annoying problem.

Basic Example

Ideally, this is what you'd want. It has good code locality and you don't have to jump through functions or files to understand the full context of the code.

class SomeController(Controller):
    @get()
    async def actual_server_sent_event_handler(self) -> AsyncGenerator[ServerSentEventMessage, None]:
        yield ServerSentEventMessage(data='Doing this..')
        yield ServerSentEventMessage(data='Doing that..')
        yield ServerSentEventMessage(data='Finished doing this and that..')

Drawbacks and Impact

I have a pseudo solution to this using decorators that has zero runtime cost aside from the initial injection at startup. Ideally, the actual solution should not have to depend on any decorators.

from collections.abc import AsyncIterator, Awaitable, Callable, Iterator
from functools import partial, update_wrapper
from inspect import isasyncgenfunction

from litestar.response import ServerSentEvent, ServerSentEventMessage


async def async_handler[**P](
    handler: Callable[P, Iterator[ServerSentEventMessage] | AsyncIterator[ServerSentEventMessage]],
    *args: P.args,
    **kwargs: P.kwargs,
) -> ServerSentEvent:
    return ServerSentEvent(handler(*args, **kwargs))


def server_sent_event[**P](
    handler: Callable[P, Iterator[ServerSentEventMessage] | AsyncIterator[ServerSentEventMessage]],
) -> Callable[P, ServerSentEvent | Awaitable[ServerSentEvent]]:
    new_handler: Callable[P, ServerSentEvent | Awaitable[ServerSentEvent]] = (
        partial(async_handler, handler)
        if isasyncgenfunction(handler)
        else lambda *args, **kwargs: ServerSentEvent(handler(*args, **kwargs))
    )

    new_handler = update_wrapper(new_handler, handler)
    new_handler.__annotations__['return'] = ServerSentEvent

    return new_handler


class SomeController(Controller):
    @get()
    @server_sent_event
    async def actual_server_sent_event_handler(self) -> AsyncGenerator[ServerSentEventMessage, None]:
        yield ServerSentEventMessage(data='Doing this..')
        yield ServerSentEventMessage(data='Doing that..')
        yield ServerSentEventMessage(data='Finished doing this and that..')

With this solution, I am still able to correctly generate OpenAPI docs and have no issues with DI. I am quite confident that this can be implemented without issues. If you know of any potential drawbacks, I'd like to hear it.

Unresolved questions

No response

---

Note

While we are open for sponsoring on GitHub Sponsors and
OpenCollective, we also utilize Polar.sh to engage in pledge-based sponsorship.

Check out all issues funded or available for funding on our Polar.sh dashboard

• If you would like to see an issue prioritized, make a pledge towards it!
• We receive the pledge once the issue is completed & verified
• This, along with engagement in the community, helps us know which features are a priority to our users.