class AppAgentStatus(Enum):
    """Store the status of the app agent."""

    CONTINUE = "CONTINUE"      # Main execution state
    SCREENSHOT = "SCREENSHOT"  # Re-annotation state
    FINISH = "FINISH"          # Subtask completed successfully
    FAIL = "FAIL"              # Subtask failed but recoverable
    PENDING = "PENDING"        # Awaiting user input
    CONFIRM = "CONFIRM"        # Safety confirmation required
    ERROR = "ERROR"            # Critical failure

@AppAgentStateManager.register
class ContinueAppAgentState(AppAgentState):
    """The class for the continue app agent state."""

    async def handle(
        self, agent: "AppAgent", context: Optional["Context"] = None
    ) -> None:
        """
        Handle the agent for the current step.
        :param agent: The agent for the current step.
        :param context: The context for the agent and session.
        """
        await agent.process(context)

    def is_subtask_end(self) -> bool:
        """Check if the subtask ends."""
        return False

    @classmethod
    def name(cls) -> str:
        """The class name of the state."""
        return AppAgentStatus.CONTINUE.value

CONTINUE → Capture UI → LLM selects "Export [12]" → Click control 12 
→ LLM returns Status: "SCREENSHOT" → Transition to SCREENSHOT

@AppAgentStateManager.register
class ScreenshotAppAgentState(ContinueAppAgentState):
    """The class for the screenshot app agent state."""

    @classmethod
    def name(cls) -> str:
        """The class name of the state."""
        return AppAgentStatus.SCREENSHOT.value

    def next_state(self, agent: BasicAgent) -> AgentState:
        """Determine next state based on control_reannotate."""
        agent_processor = agent.processor

        if agent_processor is None:
            agent.status = AppAgentStatus.CONTINUE.value
            return ContinueAppAgentState()

        control_reannotate = agent_processor.control_reannotate

        if control_reannotate is None or len(control_reannotate) == 0:
            agent.status = AppAgentStatus.CONTINUE.value
            return ContinueAppAgentState()
        else:
            return super().next_state(agent)

    def is_subtask_end(self) -> bool:
        """Check if the subtask ends."""
        return False

Action: Click "Export" button [12]
→ Dialog opens with new controls
→ LLM sets Status: "SCREENSHOT"
→ SCREENSHOT state re-annotates dialog controls as [1], [2], [3]...
→ Transitions to CONTINUE with fresh annotations

@AppAgentStateManager.register
class FinishAppAgentState(AppAgentState):
    """The class for the finish app agent state."""

    async def handle(
        self, agent: "AppAgent", context: Optional["Context"] = None
    ) -> None:
        """Archive subtask result."""
        if agent.processor:
            result = agent.processor.processing_context.get_local("result")
        else:
            result = None

        await self.archive_subtask(context, result)

    def next_agent(self, agent: "AppAgent") -> HostAgent:
        """Get the agent for the next step."""
        return agent.host

    def next_state(self, agent: "AppAgent") -> HostAgentState:
        """Get the next state of the agent."""
        if agent.mode == "follower":
            return FinishHostAgentState()
        else:
            return ContinueHostAgentState()

@AppAgentStateManager.register
class PendingAppAgentState(AppAgentState):
    """The class for the pending app agent state."""

    async def handle(
        self, agent: "AppAgent", context: Optional["Context"] = None
    ) -> None:
        """Ask the user questions to help the agent proceed."""
        agent.process_asker(ask_user=ufo_config.system.ask_question)

    def next_state(self, agent: AppAgent) -> AppAgentState:
        """Get the next state of the agent."""
        agent.status = AppAgentStatus.CONTINUE.value
        return ContinueAppAgentState()

    def is_subtask_end(self) -> bool:
        """Check if the subtask ends."""
        return False

    @classmethod
    def name(cls) -> str:
        """The class name of the state."""
        return AppAgentStatus.PENDING.value

@AppAgentStateManager.register
class ConfirmAppAgentState(AppAgentState):
    """The class for the confirm app agent state."""

    def __init__(self) -> None:
        """Initialize the confirm state."""
        self._confirm = None

    async def handle(
        self, agent: "AppAgent", context: Optional["Context"] = None
    ) -> None:
        """Request user confirmation for the action."""
        # If safe guard disabled, proceed automatically
        if not ufo_config.system.safe_guard:
            await agent.process_resume()
            self._confirm = True
            return

        # Ask user for confirmation
        self._confirm = agent.process_confirmation()

        # If user confirms, resume the task
        if self._confirm:
            await agent.process_resume()

    def next_state(self, agent: AppAgent) -> AppAgentState:
        """Get the next state based on user decision."""
        if self._confirm:
            agent.status = AppAgentStatus.CONTINUE.value
            return ContinueAppAgentState()
        else:
            agent.status = AppAgentStatus.FINISH.value
            return FinishAppAgentState()

    def is_subtask_end(self) -> bool:
        """Check if the subtask ends."""
        return False

    @classmethod
    def name(cls) -> str:
        """The class name of the state."""
        return AppAgentStatus.CONFIRM.value

@AppAgentStateManager.register
class ErrorAppAgentState(AppAgentState):
    """The class for the error app agent state."""

    async def handle(
        self, agent: "AppAgent", context: Optional["Context"] = None
    ) -> None:
        """Archive subtask with error result."""
        if agent.processor:
            result = agent.processor.processing_context.get_local("result")
        else:
            result = None

        await self.archive_subtask(context, result)

    def next_agent(self, agent: "AppAgent") -> HostAgent:
        """Get the agent for the next step."""
        return agent.host

    def next_state(self, agent: "AppAgent") -> HostAgentState:
        """Get the next state of the agent."""
        return FinishHostAgentState()

    def is_round_end(self) -> bool:
        """Check if the round ends."""
        return True

    def is_subtask_end(self) -> bool:
        """Check if the subtask ends."""
        return True

    @classmethod
    def name(cls) -> str:
        """The class name of the state."""
        return AppAgentStatus.ERROR.value

@AppAgentStateManager.register
class FailAppAgentState(AppAgentState):
    """The class for the fail app agent state."""

    async def handle(
        self, agent: "AppAgent", context: Optional["Context"] = None
    ) -> None:
        """Archive subtask with failure result."""
        if agent.processor:
            result = agent.processor.processing_context.get_local("result")
        else:
            result = None

        await self.archive_subtask(context, result)

    def next_agent(self, agent: "AppAgent") -> HostAgent:
        """Get the agent for the next step."""
        return agent.host

    def next_state(self, agent: "AppAgent") -> HostAgentState:
        """Get the next state of the agent."""
        return FinishHostAgentState()

    def is_round_end(self) -> bool:
        """Check if the round ends."""
        return False

    def is_subtask_end(self) -> bool:
        """Check if the subtask ends."""
        return True

    @classmethod
    def name(cls) -> str:
        """The class name of the state."""
        return AppAgentStatus.FAIL.value

{
  "Observation": "Word document with Export button [12] visible",
  "Thought": "I should click the Export button to extract table data",
  "ControlLabel": "12",
  "ControlText": "Export",
  "Function": "click_input",
  "Args": {"button": "left"},
  "Status": "SCREENSHOT",
  "Comment": "Clicking Export will open a dialog"
}

# Exception handling in processor
try:
    result = await processor.process(agent, context)
except Exception as e:
    agent.status = AppAgentStatus.ERROR.value
    # Transitions to ERROR state

# Screenshot re-annotation check
if control_reannotate and len(control_reannotate) > 0:
    # Stay in SCREENSHOT state
    return ScreenshotAppAgentState()
else:
    # Transition to CONTINUE
    agent.status = AppAgentStatus.CONTINUE.value
    return ContinueAppAgentState()

class AppAgentStateManager(AgentStateManager):
    """State manager for AppAgent with registration system."""

    _state_mapping: Dict[str, Type[AppAgentState]] = {}

    @property
    def none_state(self) -> AgentState:
        """The none state of the state manager."""
        return NoneAppAgentState()

# States are registered via decorator
@AppAgentStateManager.register
class ContinueAppAgentState(AppAgentState):
    ...