from judy_core.messages import GetTaskState, TaskState
from tongcell.core import MessageContext, RoutedActor, message_handler


class TaskStateManager(RoutedActor):
    def __init__(self):
        super().__init__()
        self._task_state = TaskState()

    @message_handler()
    async def get_task_state(self, _: GetTaskState, ctx: MessageContext) -> TaskState:
        return self._task_state
