diff --git a/droidrun/tools/ui/provider.py b/droidrun/tools/ui/provider.py
index ca4e50f..cfeafa0 100644
--- a/droidrun/tools/ui/provider.py
+++ b/droidrun/tools/ui/provider.py
@@ -49,12 +49,13 @@ class AndroidStateProvider(StateProvider):
         tree_formatter: "TreeFormatter",
         use_normalized: bool = False,
         stealth: bool = False,
+        ui_cls: "type[UIState] | None" = None,
     ) -> None:
         super().__init__(driver)
         self.tree_filter = tree_filter
         self.tree_formatter = tree_formatter
         self.use_normalized = use_normalized
-        self._ui_cls = StealthUIState if stealth else UIState
+        self._ui_cls = ui_cls or (StealthUIState if stealth else UIState)
 
     async def get_state(self) -> UIState:
         max_retries = 3
diff --git a/pyproject.toml b/pyproject.toml
index 0447cbf..9552975 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "droidrun"
-version = "0.5.0.dev3"
+version = "0.5.0.dev4"
 description = "A framework for controlling Android devices through LLM agents"
 authors = [{ name = "Niels Schmidt", email = "niels.schmidt@droidrun.ai" }]
 dependencies = [
diff --git a/uv.lock b/uv.lock
index d80e3a5..197af09 100644
--- a/uv.lock
+++ b/uv.lock
@@ -760,7 +760,7 @@ wheels = [
 
 [[package]]
 name = "droidrun"
-version = "0.5.0.dev3"
+version = "0.5.0.dev4"
 source = { editable = "." }
 dependencies = [
     { name = "aiofiles" },