set_message_context

def set_message_context(
    func: Callable[
        [StreamMessage[MsgType]],
        Awaitable[WrappedReturn[T_HandlerReturn]],
    ],
) -> Callable[[StreamMessage[MsgType]], Awaitable[WrappedReturn[T_HandlerReturn]]]:
    """Sets the message context for a function.

    Args:
        func: The function to set the message context for.

    Returns:
        The function with the message context set.

    """

    @wraps(func)
    async def set_message_wrapper(
        message: StreamMessage[MsgType],
    ) -> WrappedReturn[T_HandlerReturn]:
        """Wraps a function that handles a stream message.

        Args:
            message: The stream message to be handled.

        Returns:
            The wrapped return value of the handler function.

        """
        with context.scope("message", message):
            return await func(message)

    return set_message_wrapper