def METHOD_NAME(self, unit: "Unit", agent: "Agent") -> None:
    """
    Running a task with access to remote queries
    """
    start_time = time.time()
    while (
        not agent.await_submit(timeout=None)
        and unit.db_id in self.running_units
        and time.time() - start_time < self.assignment_duration_in_seconds
    ):
        self._run_server_timestep_for_agent(agent)
    remaining_time = self.assignment_duration_in_seconds - (time.time() - start_time)
    agent.await_submit(timeout=remaining_time)