Round

async def capture_last_snapshot(self, sub_round_id: Optional[int] = None) -> None:
    """
    Capture the last snapshot of the application, including the screenshot and the XML file if configured.
    :param sub_round_id: The id of the sub-round, default is None.
    """

    # Capture the final screenshot
    if sub_round_id is None:
        screenshot_save_path = self.log_path + f"action_round_{self.id}_final.png"
    else:
        screenshot_save_path = (
            self.log_path
            + f"action_round_{self.id}_sub_round_{sub_round_id}_final.png"
        )

    if (
        self.application_window is not None
        or self.application_window_info is not None
    ):

        try:

            result = await self.context.command_dispatcher.execute_commands(
                [
                    Command(
                        tool_name="capture_window_screenshot",
                        parameters={},
                        tool_type="data_collection",
                    )
                ]
            )

            image = result[0].result
            utils.save_image_string(image, screenshot_save_path)
            self.logger.info(
                f"Captured application window screenshot at final: {screenshot_save_path}"
            )

        except Exception as e:
            self.logger.warning(
                f"The last snapshot capture failed, due to the error: {e}"
            )
        if ufo_config.system.save_ui_tree:
            # Get session data manager from context

            ui_tree_path = os.path.join(self.log_path, "ui_trees")
            ui_tree_file_name = (
                f"ui_tree_round_{self.id}_final.json"
                if sub_round_id is None
                else f"ui_tree_round_{self.id}_sub_round_{sub_round_id}_final.json"
            )
            ui_tree_save_path = os.path.join(ui_tree_path, ui_tree_file_name)

            await self.save_ui_tree(ui_tree_save_path)

        if ufo_config.system.save_full_screen:

            desktop_save_path = (
                self.log_path
                + f"desktop_round_{self.id}_sub_round_{sub_round_id}_final.png"
            )

            result = await self.context.command_dispatcher.execute_commands(
                [
                    Command(
                        tool_name="capture_desktop_screenshot",
                        parameters={"all_screens": True},
                        tool_type="data_collection",
                    )
                ]
            )

            desktop_screenshot_url = result[0].result
            utils.save_image_string(desktop_screenshot_url, desktop_save_path)
            self.logger.info(f"Desktop screenshot saved to {desktop_save_path}")

def evaluation(self) -> None:
    """
    TODO: Evaluate the round.
    """
    pass

def is_finished(self) -> bool:
    """
    Check if the round is finished.
    return: True if the round is finished, otherwise False.
    """
    return (
        self.state.is_round_end()
        or self.context.get(ContextNames.SESSION_STEP) >= ufo_config.system.max_step
    )

def print_cost(self) -> None:
    """
    Print the total cost of the round.
    """

    total_cost = self.cost
    if isinstance(total_cost, float):
        formatted_cost = "${:.2f}".format(total_cost)
        console.print(
            f"💰 Request total cost for current round is {formatted_cost}",
            style="yellow",
        )

async def run(self) -> None:
    """
    Run the round.
    """

    while not self.is_finished():

        await self.agent.handle(self.context)

        # Take action

        self.state = self.agent.state.next_state(self.agent)
        self.agent = self.agent.state.next_agent(self.agent)

        self.logger.info(
            f"Agent {self.agent.name} transitioned to state {self.state.name()}"
        )

        self.agent.set_state(self.state)

        # If the subtask ends, capture the last snapshot of the application.
        if self.state.is_subtask_end():
            time.sleep(ufo_config.system.sleep_time)
            await self.capture_last_snapshot(sub_round_id=self.subtask_amount)
            self.subtask_amount += 1

    self.agent.blackboard.add_requests(
        {"request_{i}".format(i=self.id): self.request}
    )

    await self.capture_last_snapshot()

    if self._should_evaluate:
        self.evaluation()

    return self.context.get(ContextNames.ROUND_RESULT)

async def save_ui_tree(self, save_path: str):
    """
    Save the UI tree of the current application window.
    """
    if self.application_window is not None:
        result = await self.context.command_dispatcher.execute_commands(
            [
                Command(
                    tool_name="get_ui_tree",
                    parameters={},
                    tool_type="data_collection",
                )
            ]
        )
        step_ui_tree = result[0].result

        if step_ui_tree:

            save_dir = os.path.dirname(save_path)
            if not os.path.exists(save_dir):
                os.makedirs(save_dir)

            with open(save_path, "w") as file:
                json.dump(step_ui_tree, file, indent=4)
                self.logger.info(f"UI tree saved to {save_path}")