diff --git a/.bazelignore b/.bazelignore index 19915887..0ecf67c7 100644 --- a/.bazelignore +++ b/.bazelignore @@ -1 +1,2 @@ android_env/apps +android_env/playground/frontend/node_modules diff --git a/.bazelrc b/.bazelrc index 331e296e..3bd889b5 100644 --- a/.bazelrc +++ b/.bazelrc @@ -2,3 +2,6 @@ build --cxxopt=-std=c++17 build --host_cxxopt=-std=c++17 build --incompatible_enable_proto_toolchain_resolution + +# Required by aspect_rules_ts (Playground frontend). +common --@aspect_rules_ts//ts:skipLibCheck=always diff --git a/MODULE.bazel b/MODULE.bazel index ebb152d7..4db5cf9b 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -50,3 +50,11 @@ register_toolchains("@nodejs_toolchains//:all") rules_ts_ext = use_extension("@aspect_rules_ts//ts:extensions.bzl", "ext") rules_ts_ext.deps() use_repo(rules_ts_ext, "npm_typescript") +# npm dependencies for the Playground frontend. +npm = use_extension("@aspect_rules_js//npm:extensions.bzl", "npm") +npm.npm_translate_lock( + name = "npm", + pnpm_lock = "//android_env/playground/frontend:pnpm-lock.yaml", + verify_node_modules_ignored = "//:.bazelignore", +) +use_repo(npm, "npm") diff --git a/android_env/BUILD.bazel b/android_env/BUILD.bazel index ebe91c3f..4d5cd88d 100644 --- a/android_env/BUILD.bazel +++ b/android_env/BUILD.bazel @@ -94,6 +94,10 @@ py_package( packages = ["android_env"], deps = [ ":loader", + "//android_env/components:android_device_loader", + "//android_env/mcp:android_mcp_tools", + "//android_env/playground:playground_lib", + "//android_env/rl:environment", "//android_env/wrappers:a11y_grpc_wrapper", "//android_env/wrappers:discrete_action_wrapper", "//android_env/wrappers:flat_interface_wrapper", @@ -104,6 +108,7 @@ py_package( "//android_env/wrappers:rate_limit_wrapper", "//android_env/wrappers:swipe_action_wrapper", "//android_env/wrappers:tap_action_wrapper", + "//android_env/wrappers:task_wrapper", ], ) @@ -111,6 +116,10 @@ py_wheel( name = "android_env_wheel", author = "DeepMind", distribution = "android_env", + extra_requires = { + "mcp": ["mcp>=1.26,<2"], + "playground": ["ml-collections"], + }, license = "Apache License, Version 2.0", requires = [ "absl-py>=0.1.0", @@ -122,6 +131,8 @@ py_wheel( "gymnasium>=0.29.1", "protobuf>=5.27.2", "grpcio>=1.64.2", + "opencv-python-headless>=4.10.0", + "websocket-client>=1.8.0", ], summary = "AndroidEnv: A Gym/Gymnasium Environment for Android", version = "1.3.0", diff --git a/android_env/apps/java/com/google/androidenv/core/Action.kt b/android_env/apps/java/com/google/androidenv/core/Action.kt new file mode 100644 index 00000000..f61f8411 --- /dev/null +++ b/android_env/apps/java/com/google/androidenv/core/Action.kt @@ -0,0 +1,39 @@ +// Copyright 2026 DeepMind Technologies Limited. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package com.google.androidenv.core + +/** Types of actions that can be injected. */ +enum class ActionType { + TOUCH_DOWN, + TOUCH_UP, + TOUCH_MOVE, + KEY_EVENT, + SET_TEXT, + SET_SELECTION, + CUT, + COPY, + PASTE, +} + +/** Represents an action to be injected into the device. */ +data class Action( + val actionType: ActionType, + val touchPosition: List?, // [x, y] normalized 0.0-1.0 + val keycode: Int?, + val pointerId: Int? = null, + val text: String? = null, + val selectionStart: Int? = null, + val selectionEnd: Int? = null, +) diff --git a/android_env/apps/java/com/google/androidenv/core/ActionHandler.kt b/android_env/apps/java/com/google/androidenv/core/ActionHandler.kt new file mode 100644 index 00000000..d1ee4a34 --- /dev/null +++ b/android_env/apps/java/com/google/androidenv/core/ActionHandler.kt @@ -0,0 +1,390 @@ +// Copyright 2026 DeepMind Technologies Limited. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package com.google.androidenv.core + +import android.accessibilityservice.AccessibilityService +import android.accessibilityservice.GestureDescription +import android.graphics.Path +import android.util.Log +import android.view.accessibility.AccessibilityNodeInfo + +/** + * Handles touch gestures, key events, and accessibility actions for AndroidEnv. + * + * Note on WebContents / Chrome HTML inputs: Native Android View text controls (e.g. EditText, URL + * bar) fully support text manipulation via ACTION_SET_TEXT, ACTION_SET_SELECTION, and clipboard + * actions. WebContents rendered inside Chrome / WebViews use Chromium's virtual DOM accessibility + * tree, which may not consistently receive or process AccessibilityNodeInfo text actions across all + * web pages. + */ +class ActionHandler(private val gestureInjector: GestureInjector) { + + private val gestureLock = Any() + private var latestStrokeDescription: GestureDescription.StrokeDescription? = null + private var isTouching: Boolean = false + private var latestX: Float = 0f + private var latestY: Float = 0f + + fun injectAction(action: Action, width: Int, height: Int) { + try { + if (action.actionType == ActionType.KEY_EVENT) { + action.keycode?.let { keycode -> + when (keycode) { + android.view.KeyEvent.KEYCODE_BACK -> { + gestureInjector.performGlobalAction(AccessibilityService.GLOBAL_ACTION_BACK) + } + android.view.KeyEvent.KEYCODE_HOME -> { + gestureInjector.performGlobalAction(AccessibilityService.GLOBAL_ACTION_HOME) + } + android.view.KeyEvent.KEYCODE_DEL, + android.view.KeyEvent.KEYCODE_FORWARD_DEL -> { + val targetNode = getTargetNode() + if (targetNode != null) { + val currentText = getActualNodeText(targetNode) + if (currentText.isNotEmpty()) { + val selStart = targetNode.textSelectionStart + val selEnd = targetNode.textSelectionEnd + val newText = + if (selStart in 0..selEnd && selStart != selEnd) { + currentText.removeRange( + selStart.coerceAtMost(currentText.length), + selEnd.coerceAtMost(currentText.length), + ) + } else if (keycode == android.view.KeyEvent.KEYCODE_DEL) { + val pos = + if (selStart > 0 && selStart <= currentText.length) selStart + else currentText.length + if (pos > 0) currentText.removeRange(pos - 1, pos) else currentText + } else { + val pos = if (selStart >= 0) selStart else 0 + if (pos < currentText.length) currentText.removeRange(pos, pos + 1) + else currentText + } + val arguments = android.os.Bundle() + arguments.putCharSequence( + AccessibilityNodeInfo.ACTION_ARGUMENT_SET_TEXT_CHARSEQUENCE, + newText, + ) + targetNode.performAction(AccessibilityNodeInfo.ACTION_SET_TEXT, arguments) + } + } + } + android.view.KeyEvent.KEYCODE_ENTER -> { + val targetNode = getTargetNode() + if (targetNode != null) { + var handled = false + if (android.os.Build.VERSION.SDK_INT >= android.os.Build.VERSION_CODES.R) { + handled = + targetNode.performAction( + AccessibilityNodeInfo.AccessibilityAction.ACTION_IME_ENTER.id + ) + } + if (!handled) { + targetNode.performAction(AccessibilityNodeInfo.ACTION_CLICK) + } + } + } + android.view.KeyEvent.KEYCODE_DPAD_LEFT -> { + val targetNode = getTargetNode() + if (targetNode != null) { + val start = targetNode.textSelectionStart + val newPos = (start - 1).coerceAtLeast(0) + val args = + android.os.Bundle().apply { + putInt(AccessibilityNodeInfo.ACTION_ARGUMENT_SELECTION_START_INT, newPos) + putInt(AccessibilityNodeInfo.ACTION_ARGUMENT_SELECTION_END_INT, newPos) + } + targetNode.performAction(AccessibilityNodeInfo.ACTION_SET_SELECTION, args) + } + } + android.view.KeyEvent.KEYCODE_DPAD_RIGHT -> { + val targetNode = getTargetNode() + if (targetNode != null) { + val start = targetNode.textSelectionStart + val textLen = getActualNodeText(targetNode).length + val newPos = (start + 1).coerceAtMost(textLen) + val args = + android.os.Bundle().apply { + putInt(AccessibilityNodeInfo.ACTION_ARGUMENT_SELECTION_START_INT, newPos) + putInt(AccessibilityNodeInfo.ACTION_ARGUMENT_SELECTION_END_INT, newPos) + } + targetNode.performAction(AccessibilityNodeInfo.ACTION_SET_SELECTION, args) + } + } + android.view.KeyEvent.KEYCODE_MOVE_HOME -> { + val targetNode = getTargetNode() + targetNode?.performAction( + AccessibilityNodeInfo.ACTION_SET_SELECTION, + android.os.Bundle().apply { + putInt(AccessibilityNodeInfo.ACTION_ARGUMENT_SELECTION_START_INT, 0) + putInt(AccessibilityNodeInfo.ACTION_ARGUMENT_SELECTION_END_INT, 0) + }, + ) + } + android.view.KeyEvent.KEYCODE_MOVE_END -> { + val targetNode = getTargetNode() + if (targetNode != null) { + val len = getActualNodeText(targetNode).length + targetNode.performAction( + AccessibilityNodeInfo.ACTION_SET_SELECTION, + android.os.Bundle().apply { + putInt(AccessibilityNodeInfo.ACTION_ARGUMENT_SELECTION_START_INT, len) + putInt(AccessibilityNodeInfo.ACTION_ARGUMENT_SELECTION_END_INT, len) + }, + ) + } + } + android.view.KeyEvent.KEYCODE_PAGE_UP -> { + val targetNode = getTargetNode() + targetNode?.performAction(AccessibilityNodeInfo.ACTION_SCROLL_BACKWARD) + } + android.view.KeyEvent.KEYCODE_PAGE_DOWN -> { + val targetNode = getTargetNode() + targetNode?.performAction(AccessibilityNodeInfo.ACTION_SCROLL_FORWARD) + } + else -> { + Log.w(TAG, "Keycode $keycode injection not supported on-device without root.") + } + } + } + } else if (action.actionType == ActionType.SET_TEXT) { + action.text?.let { textToSet -> + val targetNode = getTargetNode() + if (targetNode != null) { + val currentText = getActualNodeText(targetNode) + val newText = currentText + textToSet + val arguments = android.os.Bundle() + arguments.putCharSequence( + AccessibilityNodeInfo.ACTION_ARGUMENT_SET_TEXT_CHARSEQUENCE, + newText, + ) + targetNode.performAction(AccessibilityNodeInfo.ACTION_SET_TEXT, arguments) + } else { + Log.w(TAG, "Could not find active focus node for SET_TEXT action.") + } + } + } else if (action.actionType == ActionType.SET_SELECTION) { + val targetNode = getTargetNode() + if (targetNode != null) { + val textLen = getActualNodeText(targetNode).length + val selStart = action.selectionStart ?: 0 + val rawEnd = action.selectionEnd ?: selStart + val selEnd = if (rawEnd < 0) textLen else rawEnd.coerceAtMost(textLen) + val args = + android.os.Bundle().apply { + putInt(AccessibilityNodeInfo.ACTION_ARGUMENT_SELECTION_START_INT, selStart) + putInt(AccessibilityNodeInfo.ACTION_ARGUMENT_SELECTION_END_INT, selEnd) + } + targetNode.performAction(AccessibilityNodeInfo.ACTION_SET_SELECTION, args) + } + } else if ( + action.actionType == ActionType.CUT || + action.actionType == ActionType.COPY || + action.actionType == ActionType.PASTE + ) { + val targetNode = getTargetNode() + if (targetNode != null) { + val a11yAction = + when (action.actionType) { + ActionType.CUT -> AccessibilityNodeInfo.ACTION_CUT + ActionType.COPY -> AccessibilityNodeInfo.ACTION_COPY + else -> AccessibilityNodeInfo.ACTION_PASTE + } + targetNode.performAction(a11yAction) + } + } else { + // Touch action + val gesture = buildGesture(action, width, height) + if (gesture != null) { + gestureInjector.injectGesture( + gesture, + object : AccessibilityService.GestureResultCallback() { + override fun onCompleted(gestureDescription: GestureDescription?) { + Log.i(TAG, "Gesture completed successfully") + } + + override fun onCancelled(gestureDescription: GestureDescription?) { + Log.w(TAG, "Gesture cancelled") + } + }, + ) + } + } + } catch (e: Exception) { + Log.e(TAG, "Failed to inject action: ${e.message}", e) + } + } + + private fun buildGesture(action: Action, width: Int, height: Int): GestureDescription? { + synchronized(gestureLock) { + val pos = action.touchPosition ?: return null + if (pos.size < 2) return null + val x = pos[0] * width + val y = pos[1] * height + + val builder = GestureDescription.Builder() + val segmentDuration: Long = 20 + val path = Path() + + when (action.actionType) { + ActionType.TOUCH_DOWN -> { + path.moveTo(x, y) + path.lineTo(x, y) + val stroke = GestureDescription.StrokeDescription(path, 0, segmentDuration, true) + builder.addStroke(stroke) + latestStrokeDescription = stroke + isTouching = true + } + ActionType.TOUCH_MOVE -> { + if (isTouching && latestStrokeDescription != null) { + path.moveTo(latestX, latestY) + path.lineTo(x, y) + val stroke = latestStrokeDescription!!.continueStroke(path, 0, segmentDuration, true) + builder.addStroke(stroke) + latestStrokeDescription = stroke + } else { + path.moveTo(x, y) + val stroke = GestureDescription.StrokeDescription(path, 0, segmentDuration, true) + builder.addStroke(stroke) + latestStrokeDescription = stroke + isTouching = true + } + } + ActionType.TOUCH_UP -> { + if (isTouching && latestStrokeDescription != null) { + path.moveTo(latestX, latestY) + path.lineTo(x, y) + val stroke = latestStrokeDescription!!.continueStroke(path, 0, segmentDuration, false) + builder.addStroke(stroke) + } + latestStrokeDescription = null + isTouching = false + } + else -> return null + } + latestX = x + latestY = y + if (path.isEmpty) return null + return builder.build() + } + } + + private fun getActiveRootNode(): AccessibilityNodeInfo? { + return gestureInjector.getRootNode() + ?: try { + InputService.getInstance().getRootNode() + } catch (e: Exception) { + null + } + } + + private fun getAllRootNodes(): List { + val roots = gestureInjector.getRootNodes() + if (roots.isNotEmpty()) return roots + val singleRoot = gestureInjector.getRootNode() + if (singleRoot != null) return listOf(singleRoot) + return try { + InputService.getInstance().getRootNodes() + } catch (e: Exception) { + emptyList() + } + } + + private fun getTargetNode(): AccessibilityNodeInfo? { + return findTargetNodeFromRoots(getActiveRootNode(), getAllRootNodes()) + } + + companion object { + private const val TAG = "ActionHandler" + + private fun findTargetNodeFromRoots( + activeRoot: AccessibilityNodeInfo?, + allRoots: List, + ): AccessibilityNodeInfo? { + // 1. Try input focus in active window first (e.g. Chrome URL bar or native EditText) + val activeInputFocused = activeRoot?.findFocus(AccessibilityNodeInfo.FOCUS_INPUT) + if (activeInputFocused != null) return activeInputFocused + + // 2. Try accessibility focus in active window + val activeA11yFocused = activeRoot?.findFocus(AccessibilityNodeInfo.FOCUS_ACCESSIBILITY) + if ( + activeA11yFocused != null && (activeA11yFocused.isEditable || activeA11yFocused.isFocused) + ) { + return activeA11yFocused + } + + // 3. Try input focus across all windows + for (root in allRoots) { + val inputFocused = root.findFocus(AccessibilityNodeInfo.FOCUS_INPUT) + if (inputFocused != null) return inputFocused + } + + // 4. Try accessibility focus across all windows + for (root in allRoots) { + val a11yFocused = root.findFocus(AccessibilityNodeInfo.FOCUS_ACCESSIBILITY) + if (a11yFocused != null && (a11yFocused.isEditable || a11yFocused.isFocused)) { + return a11yFocused + } + } + + // 5. Search for a node that is explicitly focused AND editable (e.g. focused HTML input) + for (root in allRoots) { + val focusedEditable = findFocusedEditableNode(root) + if (focusedEditable != null) return focusedEditable + } + + // 6. Search for any editable node in all roots + for (root in allRoots) { + val editable = findAnyEditableNode(root) + if (editable != null) return editable + } + + return activeRoot ?: allRoots.firstOrNull() + } + + private fun findFocusedEditableNode(node: AccessibilityNodeInfo?): AccessibilityNodeInfo? { + if (node == null) return null + if (node.isEditable && (node.isFocused || node.isAccessibilityFocused)) { + return node + } + for (i in 0 until node.childCount) { + val found = findFocusedEditableNode(node.getChild(i)) + if (found != null) return found + } + return null + } + + private fun findAnyEditableNode(node: AccessibilityNodeInfo?): AccessibilityNodeInfo? { + if (node == null) return null + if (node.isEditable) { + return node + } + for (i in 0 until node.childCount) { + val found = findAnyEditableNode(node.getChild(i)) + if (found != null) return found + } + return null + } + + private fun getActualNodeText(node: AccessibilityNodeInfo): String { + val text = node.text?.toString() ?: return "" + val hint = node.hintText?.toString() + if (node.isShowingHintText || (hint != null && text == hint)) { + return "" + } + return text + } + } +} diff --git a/android_env/apps/java/com/google/androidenv/core/ActionHandlerTest.kt b/android_env/apps/java/com/google/androidenv/core/ActionHandlerTest.kt new file mode 100644 index 00000000..1e3701b3 --- /dev/null +++ b/android_env/apps/java/com/google/androidenv/core/ActionHandlerTest.kt @@ -0,0 +1,217 @@ +// Copyright 2026 DeepMind Technologies Limited. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package com.google.androidenv.core + +import android.accessibilityservice.GestureDescription +import android.view.accessibility.AccessibilityNodeInfo +import com.google.common.truth.Truth.assertThat +import org.junit.Before +import org.junit.Test +import org.junit.runner.RunWith +import org.mockito.kotlin.any +import org.mockito.kotlin.argumentCaptor +import org.mockito.kotlin.doReturn +import org.mockito.kotlin.mock +import org.mockito.kotlin.verify +import org.mockito.kotlin.whenever + +@RunWith(androidx.test.ext.junit.runners.AndroidJUnit4::class) +class ActionHandlerTest { + + private val mockGestureInjector = mock() + private lateinit var handler: ActionHandler + + @Before + fun setUp() { + handler = ActionHandler(mockGestureInjector) + } + + @Test + fun injectAction_touch_callsInputServiceWithScaledCoordinates() { + val action = + Action( + actionType = ActionType.TOUCH_DOWN, + touchPosition = listOf(0.5f, 0.25f), + keycode = null, + ) + + handler.injectAction(action, 1000, 2000) + + val gestureCaptor = argumentCaptor() + verify(mockGestureInjector).injectGesture(gestureCaptor.capture(), any()) + + val gesture = gestureCaptor.firstValue + assertThat(gesture).isNotNull() + } + + @Test + fun injectAction_keyEvent_callsPerformGlobalAction() { + val action = + Action( + actionType = ActionType.KEY_EVENT, + touchPosition = null, + keycode = 4, // BACK -> GLOBAL_ACTION_BACK = 1 + ) + + handler.injectAction(action, 1000, 2000) + + verify(mockGestureInjector).performGlobalAction(1) + } + + @Test + fun injectAction_setText_appendsToExistingNodeText() { + val mockNode = mock() + whenever(mockNode.findFocus(AccessibilityNodeInfo.FOCUS_INPUT)).doReturn(mockNode) + whenever(mockNode.text).doReturn("hello ") + whenever(mockGestureInjector.getRootNode()).doReturn(mockNode) + + val action = + Action(actionType = ActionType.SET_TEXT, touchPosition = null, keycode = null, text = "world") + + handler.injectAction(action, 1000, 2000) + + val argsCaptor = argumentCaptor() + verify(mockNode) + .performAction( + org.mockito.kotlin.eq(AccessibilityNodeInfo.ACTION_SET_TEXT), + argsCaptor.capture(), + ) + val charSeq = + argsCaptor.firstValue.getCharSequence( + AccessibilityNodeInfo.ACTION_ARGUMENT_SET_TEXT_CHARSEQUENCE + ) + assertThat(charSeq.toString()).isEqualTo("hello world") + } + + @Test + fun injectAction_setText_ignoresHintText() { + val mockNode = mock() + whenever(mockNode.findFocus(AccessibilityNodeInfo.FOCUS_INPUT)).doReturn(mockNode) + whenever(mockNode.text).doReturn("Search or type URL") + whenever(mockNode.hintText).doReturn("Search or type URL") + whenever(mockGestureInjector.getRootNode()).doReturn(mockNode) + + val action = + Action(actionType = ActionType.SET_TEXT, touchPosition = null, keycode = null, text = "g") + + handler.injectAction(action, 1000, 2000) + + val argsCaptor = argumentCaptor() + verify(mockNode) + .performAction( + org.mockito.kotlin.eq(AccessibilityNodeInfo.ACTION_SET_TEXT), + argsCaptor.capture(), + ) + val charSeq = + argsCaptor.firstValue.getCharSequence( + AccessibilityNodeInfo.ACTION_ARGUMENT_SET_TEXT_CHARSEQUENCE + ) + assertThat(charSeq.toString()).isEqualTo("g") + } + + @Test + fun injectAction_keyEvent_del_removesLastChar() { + val mockNode = mock() + whenever(mockNode.findFocus(AccessibilityNodeInfo.FOCUS_INPUT)).doReturn(mockNode) + whenever(mockNode.text).doReturn("hello") + whenever(mockGestureInjector.getRootNode()).doReturn(mockNode) + + val action = + Action( + actionType = ActionType.KEY_EVENT, + touchPosition = null, + keycode = 67, // KEYCODE_DEL + ) + + handler.injectAction(action, 1000, 2000) + + val argsCaptor = argumentCaptor() + verify(mockNode) + .performAction( + org.mockito.kotlin.eq(AccessibilityNodeInfo.ACTION_SET_TEXT), + argsCaptor.capture(), + ) + val charSeq = + argsCaptor.firstValue.getCharSequence( + AccessibilityNodeInfo.ACTION_ARGUMENT_SET_TEXT_CHARSEQUENCE + ) + assertThat(charSeq.toString()).isEqualTo("hell") + } + + @Test + fun injectAction_setText_targetsEditableHtmlNodeInChrome() { + val mockRoot = mock() + val mockHtmlInput = mock() + whenever(mockRoot.findFocus(AccessibilityNodeInfo.FOCUS_INPUT)).doReturn(null) + whenever(mockRoot.findFocus(AccessibilityNodeInfo.FOCUS_ACCESSIBILITY)).doReturn(null) + whenever(mockRoot.childCount).doReturn(1) + whenever(mockRoot.getChild(0)).doReturn(mockHtmlInput) + whenever(mockHtmlInput.isEditable).doReturn(true) + whenever(mockHtmlInput.text).doReturn("test") + whenever(mockGestureInjector.getRootNode()).doReturn(mockRoot) + + val action = + Action(actionType = ActionType.SET_TEXT, touchPosition = null, keycode = null, text = "ing") + + handler.injectAction(action, 1000, 2000) + + val argsCaptor = argumentCaptor() + verify(mockHtmlInput) + .performAction( + org.mockito.kotlin.eq(AccessibilityNodeInfo.ACTION_SET_TEXT), + argsCaptor.capture(), + ) + val charSeq = + argsCaptor.firstValue.getCharSequence( + AccessibilityNodeInfo.ACTION_ARGUMENT_SET_TEXT_CHARSEQUENCE + ) + assertThat(charSeq.toString()).isEqualTo("testing") + } + + @Test + fun injectAction_setText_findsEditableNodeInMultiWindowAppWhenKeyboardIsUp() { + val keyboardRoot = mock() + whenever(keyboardRoot.findFocus(AccessibilityNodeInfo.FOCUS_INPUT)).doReturn(null) + whenever(keyboardRoot.findFocus(AccessibilityNodeInfo.FOCUS_ACCESSIBILITY)).doReturn(null) + + val chromeAppRoot = mock() + val htmlInputNode = mock() + whenever(chromeAppRoot.findFocus(AccessibilityNodeInfo.FOCUS_INPUT)).doReturn(null) + whenever(chromeAppRoot.findFocus(AccessibilityNodeInfo.FOCUS_ACCESSIBILITY)).doReturn(null) + whenever(chromeAppRoot.childCount).doReturn(1) + whenever(chromeAppRoot.getChild(0)).doReturn(htmlInputNode) + whenever(htmlInputNode.isEditable).doReturn(true) + whenever(htmlInputNode.text).doReturn("google") + + whenever(mockGestureInjector.getRootNodes()).doReturn(listOf(keyboardRoot, chromeAppRoot)) + + val action = + Action(actionType = ActionType.SET_TEXT, touchPosition = null, keycode = null, text = ".com") + + handler.injectAction(action, 1000, 2000) + + val argsCaptor = argumentCaptor() + verify(htmlInputNode) + .performAction( + org.mockito.kotlin.eq(AccessibilityNodeInfo.ACTION_SET_TEXT), + argsCaptor.capture(), + ) + val charSeq = + argsCaptor.firstValue.getCharSequence( + AccessibilityNodeInfo.ACTION_ARGUMENT_SET_TEXT_CHARSEQUENCE + ) + assertThat(charSeq.toString()).isEqualTo("google.com") + } +} diff --git a/android_env/apps/java/com/google/androidenv/core/AndroidDeviceEnv.kt b/android_env/apps/java/com/google/androidenv/core/AndroidDeviceEnv.kt new file mode 100644 index 00000000..cccb4c2e --- /dev/null +++ b/android_env/apps/java/com/google/androidenv/core/AndroidDeviceEnv.kt @@ -0,0 +1,39 @@ +// Copyright 2026 DeepMind Technologies Limited. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package com.google.androidenv.core + +/** + * The core, raw Android Device Environment interface. Provides direct, non-blocking state access + * and asynchronous action injection. + */ +interface AndroidDeviceEnv { + /** Dynamically updates the active set of signals. Stops unneeded services immediately. */ + fun updateSubscriptions(activeSignals: Set) + + /** + * Non-blocking pull of the absolute latest device state. Unsubscribed signals in the returned + * DeviceState will have StreamedSignal.value = null. + */ + fun getDeviceState(): DeviceState + + /** + * Asynchronously injects an action (touches, keycodes) into the device. Consumes resources only + * during the event injection (event-driven). + */ + fun injectAction(action: Action) + + /** Frees up resources (e.g. stops background logcat readers, sensor listeners). */ + fun close() {} +} diff --git a/android_env/apps/java/com/google/androidenv/core/AndroidDeviceEnvImpl.kt b/android_env/apps/java/com/google/androidenv/core/AndroidDeviceEnvImpl.kt new file mode 100644 index 00000000..4a699bda --- /dev/null +++ b/android_env/apps/java/com/google/androidenv/core/AndroidDeviceEnvImpl.kt @@ -0,0 +1,565 @@ +// Copyright 2026 DeepMind Technologies Limited. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package com.google.androidenv.core + +import android.content.Context +import android.graphics.Bitmap +import android.hardware.Sensor +import android.hardware.SensorEvent +import android.hardware.SensorEventListener +import android.hardware.SensorManager +import android.util.DisplayMetrics +import android.util.Log +import android.view.WindowManager +import android.view.accessibility.AccessibilityNodeInfo +import java.io.BufferedReader +import java.io.ByteArrayOutputStream +import java.io.InputStreamReader +import java.util.concurrent.ConcurrentLinkedQueue + +/** + * Concrete implementation of AndroidDeviceEnv that interacts with Android APIs. Uses + * ScreenShareService and InputService under the hood. + */ +class AndroidDeviceEnvImpl( + private val context: Context, + private val gestureInjector: GestureInjector = SystemGestureInjector(), + private val screenShareProvider: ScreenShareProvider = SystemScreenShareProvider(), +) : AndroidDeviceEnv { + + init { + screenShareProvider.setAudioListener( + object : ScreenShareService.AudioCaptureListener { + override fun onAudioCaptured( + rawBytes: ByteArray, + sampleRateHz: Int, + channelCount: Int, + sampleFormat: Int, + ) { + if (synchronized(this@AndroidDeviceEnvImpl) { activeSignals.contains(DeviceSignal.AUDIO_OUTPUT) }) { + val formatStr = + when (sampleFormat) { + android.media.AudioFormat.ENCODING_PCM_16BIT -> "pcm16" + else -> "unknown" + } + val frame = AudioFrame(rawBytes, sampleRateHz, channelCount, formatStr) + audioTracker.update(frame, System.nanoTime() / 1000) + } + } + } + ) + } + + private val sensorManager = context.getSystemService(Context.SENSOR_SERVICE) as SensorManager + private val activeSignals = mutableSetOf() + + // Trackers for streamed signals (thread-safe updates) + private val screenshotTracker = StreamTracker() + private val a11yTracker = StreamTracker() + private val a11yEventsTracker = + StreamTracker>() + private val activePackageTracker = StreamTracker() + private val activeActivityTracker = StreamTracker() + private val systemLogsTracker = StreamTracker>() + private val audioTracker = AudioStreamTracker() + private val batteryPowerTracker = StreamTracker() + private val batteryPowerCollector = BatteryPowerCollector(context) + + private var logcatStreamer: LogcatStreamer? = null + + // Sensor trackers + private val accelTracker = StreamTracker>() + private val gyroTracker = StreamTracker>() + private val magTracker = StreamTracker>() + private val baroTracker = StreamTracker() + private val tempTracker = StreamTracker() + private val lightTracker = StreamTracker() + private val proxTracker = StreamTracker() + + // Screenshot cache to detect changes + private var lastScreenshotTimestampUs: Long = 0 + + // Touch gesture tracking state + private var isTouching = false + private var latestX = 0.0f + private var latestY = 0.0f + private var latestStrokeDescription: + android.accessibilityservice.GestureDescription.StrokeDescription? = + null + + private val gestureLock = Any() + + private val sensorListener = + object : SensorEventListener { + override fun onSensorChanged(event: SensorEvent) { + val eventTimestampUs = event.timestamp / 1000 + synchronized(this@AndroidDeviceEnvImpl) { + when (event.sensor.type) { + Sensor.TYPE_ACCELEROMETER -> { + if (activeSignals.contains(DeviceSignal.SENSOR_ACCELEROMETER)) { + accelTracker.update(event.values.toList(), eventTimestampUs) + } + } + Sensor.TYPE_GYROSCOPE -> { + if (activeSignals.contains(DeviceSignal.SENSOR_GYROSCOPE)) { + gyroTracker.update(event.values.toList(), eventTimestampUs) + } + } + Sensor.TYPE_MAGNETIC_FIELD -> { + if (activeSignals.contains(DeviceSignal.SENSOR_MAGNETOMETER)) { + magTracker.update(event.values.toList(), eventTimestampUs) + } + } + Sensor.TYPE_PRESSURE -> { + if (activeSignals.contains(DeviceSignal.SENSOR_BAROMETER)) { + baroTracker.update(event.values[0], eventTimestampUs) + } + } + Sensor.TYPE_AMBIENT_TEMPERATURE -> { + if (activeSignals.contains(DeviceSignal.SENSOR_AMBIENT_TEMP)) { + tempTracker.update(event.values[0], eventTimestampUs) + } + } + Sensor.TYPE_LIGHT -> { + if (activeSignals.contains(DeviceSignal.SENSOR_LIGHT)) { + lightTracker.update(event.values[0], eventTimestampUs) + } + } + Sensor.TYPE_PROXIMITY -> { + if (activeSignals.contains(DeviceSignal.SENSOR_PROXIMITY)) { + proxTracker.update(event.values[0], eventTimestampUs) + } + } + } + } + } + + override fun onAccuracyChanged(sensor: Sensor, accuracy: Int) {} + } + + override fun updateSubscriptions(activeSignals: Set) { + synchronized(this) { + this.activeSignals.clear() + this.activeSignals.addAll(activeSignals) + Log.i(TAG, "updateSubscriptions: $activeSignals") + + // Manage Logcat Streamer + if (activeSignals.contains(DeviceSignal.SYSTEM_LOGS)) { + if (logcatStreamer == null) { + logcatStreamer = LogcatStreamer().apply { start() } + } + } else { + logcatStreamer?.stop() + logcatStreamer = null + } + + if (!activeSignals.contains(DeviceSignal.ACCESSIBILITY_TREE)) { + a11yTracker.recycle() + } + + if (activeSignals.contains(DeviceSignal.ACCESSIBILITY_EVENTS)) { + gestureInjector.setEventTrackingEnabled(true) + } else { + a11yEventsTracker.recycle() + gestureInjector.setEventTrackingEnabled(false) + gestureInjector.clearEvents() + } + + // Re-register sensor listeners + + sensorManager.unregisterListener(sensorListener) + + if (activeSignals.contains(DeviceSignal.SENSOR_ACCELEROMETER)) { + sensorManager.getDefaultSensor(Sensor.TYPE_ACCELEROMETER)?.let { + sensorManager.registerListener(sensorListener, it, SensorManager.SENSOR_DELAY_GAME) + } + } + if (activeSignals.contains(DeviceSignal.SENSOR_GYROSCOPE)) { + sensorManager.getDefaultSensor(Sensor.TYPE_GYROSCOPE)?.let { + sensorManager.registerListener(sensorListener, it, SensorManager.SENSOR_DELAY_GAME) + } + } + if (activeSignals.contains(DeviceSignal.SENSOR_MAGNETOMETER)) { + sensorManager.getDefaultSensor(Sensor.TYPE_MAGNETIC_FIELD)?.let { + sensorManager.registerListener(sensorListener, it, SensorManager.SENSOR_DELAY_GAME) + } + } + if (activeSignals.contains(DeviceSignal.SENSOR_BAROMETER)) { + sensorManager.getDefaultSensor(Sensor.TYPE_PRESSURE)?.let { + sensorManager.registerListener(sensorListener, it, SensorManager.SENSOR_DELAY_GAME) + } + } + if (activeSignals.contains(DeviceSignal.SENSOR_AMBIENT_TEMP)) { + sensorManager.getDefaultSensor(Sensor.TYPE_AMBIENT_TEMPERATURE)?.let { + sensorManager.registerListener(sensorListener, it, SensorManager.SENSOR_DELAY_GAME) + } + } + if (activeSignals.contains(DeviceSignal.SENSOR_LIGHT)) { + sensorManager.getDefaultSensor(Sensor.TYPE_LIGHT)?.let { + sensorManager.registerListener(sensorListener, it, SensorManager.SENSOR_DELAY_GAME) + } + } + if (activeSignals.contains(DeviceSignal.SENSOR_PROXIMITY)) { + sensorManager.getDefaultSensor(Sensor.TYPE_PROXIMITY)?.let { + sensorManager.registerListener(sensorListener, it, SensorManager.SENSOR_DELAY_GAME) + } + } + + if (activeSignals.contains(DeviceSignal.BATTERY_POWER)) { + batteryPowerCollector.start() + } else { + batteryPowerCollector.stop() + } + } + } + + private val actionHandler = ActionHandler(gestureInjector) + + override fun injectAction(action: Action) { + try { + val wm = context.getSystemService(Context.WINDOW_SERVICE) as? WindowManager + val metrics = DisplayMetrics() + wm?.defaultDisplay?.getRealMetrics(metrics) + val width = + if (screenShareProvider.width() > 0) screenShareProvider.width() else metrics.widthPixels + val height = + if (screenShareProvider.height() > 0) screenShareProvider.height() else metrics.heightPixels + + actionHandler.injectAction(action, width, height) + } catch (e: Exception) { + Log.e(TAG, "Failed to inject action: ${e.message}", e) + } + } + + override fun close() { + synchronized(this) { + logcatStreamer?.stop() + logcatStreamer = null + sensorManager.unregisterListener(sensorListener) + batteryPowerCollector.stop() + a11yTracker.recycle() + } + } + + override fun getDeviceState(): DeviceState { + val nowUs = System.currentTimeMillis() * 1000L + val wm = context.getSystemService(Context.WINDOW_SERVICE) as? android.view.WindowManager + val rotation = wm?.defaultDisplay?.rotation ?: 0 + + // 1. Gather ScreenShare updates + try { + val imgTs = screenShareProvider.latestScreenshotTimestampUs() + if (imgTs > lastScreenshotTimestampUs) { + val rawBytes = screenShareProvider.latestScreenshot() + if (rawBytes != null) { + val bmp = + rgbaToBitmap( + rawBytes, + screenShareProvider.width(), + screenShareProvider.height(), + screenShareProvider.rowStride(), + ) + screenshotTracker.update(bmp, imgTs) + lastScreenshotTimestampUs = imgTs + } + } + } catch (e: Exception) { + Log.w(TAG, "Screenshot not available: ${e.message}") + } + + // 2. Gather Accessibility updates + try { + if (synchronized(this) { activeSignals.contains(DeviceSignal.ACCESSIBILITY_TREE) }) { + gestureInjector.getRootNode()?.let { root -> a11yTracker.update(root, nowUs) } + } + if (synchronized(this) { activeSignals.contains(DeviceSignal.ACCESSIBILITY_EVENTS) }) { + val newEvents = gestureInjector.popEvents() + a11yEventsTracker.update(newEvents, nowUs) + } + if (synchronized(this) { activeSignals.contains(DeviceSignal.ACTIVE_PACKAGE) }) { + val pkg = gestureInjector.getActivePackage() + if (pkg.isNotEmpty()) activePackageTracker.update(pkg, nowUs) + } + if (synchronized(this) { activeSignals.contains(DeviceSignal.ACTIVE_ACTIVITY) }) { + val act = gestureInjector.getActiveActivity() + if (act.isNotEmpty()) activeActivityTracker.update(act, nowUs) + } + } catch (e: Exception) { + Log.w(TAG, "Accessibility not available: ${e.message}") + } + + // 3. Gather Logcat updates + try { + logcatStreamer?.let { streamer -> + if (synchronized(this) { activeSignals.contains(DeviceSignal.SYSTEM_LOGS) }) { + val newLogs = streamer.getNewLogs() + // Update the tracker unconditionally so value is updated to the latest batch + // (empty list when no logs occurred) and sequenceId is incremented for downstream delta + // tracking. + systemLogsTracker.update(newLogs, nowUs) + } + } + } catch (e: Exception) { + Log.w(TAG, "Logcat not available: ${e.message}") + } + + // 4. Gather Battery & Power updates + try { + if (synchronized(this) { activeSignals.contains(DeviceSignal.BATTERY_POWER) }) { + val batteryMetrics = batteryPowerCollector.getLatestMetrics() + batteryPowerTracker.update(batteryMetrics, nowUs) + } + } catch (e: Exception) { + Log.w(TAG, "BatteryPower metrics not available: ${e.message}") + } + + // Assemble and return the state + synchronized(this) { + return DeviceState( + screenshot = + if (activeSignals.contains(DeviceSignal.SCREENSHOT)) screenshotTracker.toStreamedSignal() + else StreamedSignal.empty(), + accessibilityRoot = + if (activeSignals.contains(DeviceSignal.ACCESSIBILITY_TREE)) + a11yTracker.toStreamedSignal() + else StreamedSignal.empty(), + accessibilityEvents = + if (activeSignals.contains(DeviceSignal.ACCESSIBILITY_EVENTS)) + a11yEventsTracker.toStreamedSignal() + else StreamedSignal.empty(), + activePackage = + if (activeSignals.contains(DeviceSignal.ACTIVE_PACKAGE)) + activePackageTracker.toStreamedSignal() + else StreamedSignal.empty(), + activeActivity = + if (activeSignals.contains(DeviceSignal.ACTIVE_ACTIVITY)) + activeActivityTracker.toStreamedSignal() + else StreamedSignal.empty(), + systemLogs = + if (activeSignals.contains(DeviceSignal.SYSTEM_LOGS)) systemLogsTracker.toStreamedSignal() + else StreamedSignal.empty(), + runningProcesses = StreamedSignal.empty(), // Stubbed + deviceSettings = StreamedSignal.empty(), // Stubbed + accelerometer = + if (activeSignals.contains(DeviceSignal.SENSOR_ACCELEROMETER)) + accelTracker.toStreamedSignal() + else StreamedSignal.empty(), + gyroscope = + if (activeSignals.contains(DeviceSignal.SENSOR_GYROSCOPE)) gyroTracker.toStreamedSignal() + else StreamedSignal.empty(), + magnetometer = + if (activeSignals.contains(DeviceSignal.SENSOR_MAGNETOMETER)) + magTracker.toStreamedSignal() + else StreamedSignal.empty(), + barometer = + if (activeSignals.contains(DeviceSignal.SENSOR_BAROMETER)) baroTracker.toStreamedSignal() + else StreamedSignal.empty(), + ambientTemp = + if (activeSignals.contains(DeviceSignal.SENSOR_AMBIENT_TEMP)) + tempTracker.toStreamedSignal() + else StreamedSignal.empty(), + lightLevel = + if (activeSignals.contains(DeviceSignal.SENSOR_LIGHT)) lightTracker.toStreamedSignal() + else StreamedSignal.empty(), + proximity = + if (activeSignals.contains(DeviceSignal.SENSOR_PROXIMITY)) proxTracker.toStreamedSignal() + else StreamedSignal.empty(), + locationGps = StreamedSignal.empty(), // Stubbed + heartRate = StreamedSignal.empty(), // Stubbed + audioOutput = + if (activeSignals.contains(DeviceSignal.AUDIO_OUTPUT)) audioTracker.toStreamedSignal() + else StreamedSignal.empty(), + batteryPower = + if (activeSignals.contains(DeviceSignal.BATTERY_POWER)) batteryPowerTracker.toStreamedSignal() + else StreamedSignal.empty(), + snapshotTimestampUs = nowUs, + orientation = rotation, + ) + } + } + + @Suppress("UseKtx") + private fun rgbaToBitmap(rgba: ByteArray, width: Int, height: Int, rowStride: Int): Bitmap { + val bitmap = Bitmap.createBitmap(width, height, Bitmap.Config.ARGB_8888) + if (rowStride == width * 4 || rowStride <= 0) { + bitmap.copyPixelsFromBuffer(java.nio.ByteBuffer.wrap(rgba)) + } else { + val buffer = java.nio.ByteBuffer.wrap(rgba) + val cleanBuffer = java.nio.ByteBuffer.allocate(width * height * 4) + val rowBytes = width * 4 + val row = ByteArray(rowBytes) + for (i in 0 until height) { + buffer.position(i * rowStride) + buffer.get(row, 0, rowBytes) + cleanBuffer.put(row) + } + cleanBuffer.rewind() + bitmap.copyPixelsFromBuffer(cleanBuffer) + } + return bitmap + } + + private class StreamTracker { + var value: T? = null + private set + + var timestampUs: Long = 0 + private set + + var sequenceId: Long = 0 + private set + + @Synchronized + fun update(newValue: T, newTimestampUs: Long) { + @Suppress("DEPRECATION") (value as? AccessibilityNodeInfo)?.recycle() + value = newValue + timestampUs = newTimestampUs + sequenceId++ + } + + @Synchronized + fun recycle() { + @Suppress("DEPRECATION") (value as? AccessibilityNodeInfo)?.recycle() + value = null + } + + @Synchronized + fun toStreamedSignal(): StreamedSignal { + return StreamedSignal(value, timestampUs, sequenceId) + } + } + + private class AudioStreamTracker { + private val byteBuffer = ByteArrayOutputStream() + private var sampleRateHz: Int = 44100 + private var channelCount: Int = 2 + private var sampleFormat: String? = null + private var timestampUs: Long = 0 + private var sequenceId: Long = 0 + + @Synchronized + fun update(frame: AudioFrame, newTimestampUs: Long) { + if (byteBuffer.size() == 0) { + timestampUs = newTimestampUs + sampleRateHz = frame.sampleRateHz + channelCount = frame.channelCount + sampleFormat = frame.sampleFormat + } + try { + byteBuffer.write(frame.rawBytes) + } catch (e: Exception) { + // Ignored + } + sequenceId++ + } + + @Synchronized + fun toStreamedSignal(): StreamedSignal { + if (byteBuffer.size() == 0) { + return StreamedSignal(null, timestampUs, sequenceId) + } + val accumulatedBytes = byteBuffer.toByteArray() + byteBuffer.reset() + val frame = + AudioFrame( + rawBytes = accumulatedBytes, + sampleRateHz = sampleRateHz, + channelCount = channelCount, + sampleFormat = sampleFormat, + ) + return StreamedSignal(frame, timestampUs, sequenceId) + } + } + + private class LogcatStreamer { + private var process: Process? = null + private var reader: BufferedReader? = null + private val logQueue = ConcurrentLinkedQueue() + private var thread: Thread? = null + private val maxQueueSize = 10000 + + @Synchronized + fun start() { + if (thread != null) return + Log.i(TAG, "Starting LogcatStreamer.") + thread = + Thread { + try { + process = ProcessBuilder("logcat", "-T", "1").start() + reader = BufferedReader(InputStreamReader(process!!.inputStream)) + var line: String? + while (!Thread.currentThread().isInterrupted) { + line = reader?.readLine() + if (line == null) break + if (line.isNotEmpty()) { + if (logQueue.size >= maxQueueSize) { + logQueue.poll() + } + logQueue.add(line) + } + } + Log.i(TAG, "Logcat reading thread finished.") + } catch (e: Exception) { + Log.e(TAG, "Error in LogcatStreamer thread", e) + } finally { + cleanup() + } + } + .apply { start() } + } + + private fun cleanup() { + try { + reader?.close() + } catch (e: Exception) {} + reader = null + try { + process?.destroy() + } catch (e: Exception) {} + process = null + } + + @Synchronized + fun stop() { + if (thread == null) return + Log.i(TAG, "Stopping LogcatStreamer.") + + cleanup() + thread?.interrupt() + + try { + thread?.join(5000) + } catch (e: InterruptedException) { + Log.w(TAG, "Interrupted while joining LogcatStreamer thread.") + Thread.currentThread().interrupt() + } finally { + thread = null + } + } + + fun getNewLogs(maxLogs: Int = 1000): List { + val logs = mutableListOf() + while (logs.size < maxLogs) { + val line = logQueue.poll() ?: break + logs.add(line) + } + return logs + } + } + + companion object { + private const val TAG = "AndroidDeviceEnvImpl" + } +} diff --git a/android_env/apps/java/com/google/androidenv/core/AndroidDeviceEnvImplTest.kt b/android_env/apps/java/com/google/androidenv/core/AndroidDeviceEnvImplTest.kt new file mode 100644 index 00000000..d27a62fa --- /dev/null +++ b/android_env/apps/java/com/google/androidenv/core/AndroidDeviceEnvImplTest.kt @@ -0,0 +1,225 @@ +// Copyright 2026 DeepMind Technologies Limited. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package com.google.androidenv.core + +import android.accessibilityservice.GestureDescription +import android.content.Context +import androidx.test.core.app.ApplicationProvider +import androidx.test.ext.junit.runners.AndroidJUnit4 +import com.google.common.truth.Truth.assertThat +import org.junit.After +import org.junit.Before +import org.junit.Test +import org.junit.runner.RunWith +import org.mockito.kotlin.any +import org.mockito.kotlin.argumentCaptor +import org.mockito.kotlin.doReturn +import org.mockito.kotlin.mock +import org.mockito.kotlin.verify +import org.mockito.kotlin.whenever + +@RunWith(AndroidJUnit4::class) +class AndroidDeviceEnvImplTest { + + private val context = ApplicationProvider.getApplicationContext() + private val mockGestureInjector = mock() + private val mockScreenShareProvider = mock() + private lateinit var env: AndroidDeviceEnv + + @Before + fun setUp() { + env = + AndroidDeviceEnvImpl( + context = context, + gestureInjector = mockGestureInjector, + screenShareProvider = mockScreenShareProvider, + ) + } + + @After + fun tearDown() { + env.close() + } + + @Test + fun injectAction_touch_callsInputServiceWithScaledCoordinates() { + whenever(mockScreenShareProvider.width()).doReturn(1000) + whenever(mockScreenShareProvider.height()).doReturn(2000) + + val action = + Action( + actionType = ActionType.TOUCH_DOWN, + touchPosition = listOf(0.5f, 0.25f), // x=50%, y=25% + keycode = null, + ) + + env.injectAction(action) + + val gestureCaptor = argumentCaptor() + verify(mockGestureInjector).injectGesture(gestureCaptor.capture(), any()) + + val gesture = gestureCaptor.firstValue + assertThat(gesture).isNotNull() + } + + @Test + fun getDeviceState_returnsStateWithScreenshot() { + val width = 100 + val height = 100 + val rowStride = 400 // 100 * 4 bytes + val rawBytes = ByteArray(width * height * 4) // RGBA + rawBytes[0] = 255.toByte() + rawBytes[1] = 255.toByte() + rawBytes[2] = 255.toByte() + rawBytes[3] = 255.toByte() + + whenever(mockScreenShareProvider.width()).doReturn(width) + whenever(mockScreenShareProvider.height()).doReturn(height) + whenever(mockScreenShareProvider.rowStride()).doReturn(rowStride) + whenever(mockScreenShareProvider.latestScreenshot()).doReturn(rawBytes) + whenever(mockScreenShareProvider.latestScreenshotTimestampUs()).doReturn(12345L) + + env.updateSubscriptions(setOf(DeviceSignal.SCREENSHOT)) + val state = env.getDeviceState() + + assertThat(state.snapshotTimestampUs).isGreaterThan(0L) + assertThat(state.screenshot.timestampUs).isEqualTo(12345L) + assertThat(state.screenshot.value).isNotNull() + val bmp = state.screenshot.value!! + assertThat(bmp.width).isEqualTo(width) + assertThat(bmp.height).isEqualTo(height) + val pixel = bmp.getPixel(0, 0) + assertThat(android.graphics.Color.red(pixel)).isEqualTo(255) + } + + @Test + fun getDeviceState_systemLogsDeltaStreaming() { + env.updateSubscriptions(setOf(DeviceSignal.SYSTEM_LOGS)) + val state1 = env.getDeviceState() + + assertThat(state1.systemLogs).isNotNull() + val seq1 = state1.systemLogs.sequenceId + + // On subsequent call when no new logs exist, systemLogs.value is empty list (not stale logs) + // and sequenceId is incremented for delta tracking. + val state2 = env.getDeviceState() + assertThat(state2.systemLogs).isNotNull() + assertThat(state2.systemLogs.sequenceId).isEqualTo(seq1 + 1) + assertThat(state2.systemLogs.value).isEmpty() + } + + @Test + fun injectAction_keyEvent_callsPerformGlobalAction() { + val action = + Action( + actionType = ActionType.KEY_EVENT, + touchPosition = null, + keycode = 4, // BACK -> GLOBAL_ACTION_BACK = 1 + ) + + env.injectAction(action) + + verify(mockGestureInjector).performGlobalAction(1) + } + + @Test + fun getDeviceState_returnsAccessibilityTreeAndEvents() { + val mockNode = mock() + whenever(mockGestureInjector.getRootNode()).doReturn(mockNode) + whenever(mockGestureInjector.getActivePackage()).doReturn("com.example.app") + whenever(mockGestureInjector.getActiveActivity()).doReturn("com.example.app.MainActivity") + + val mockEvent = + com.google.androidenv.accessibilityforwarder.EventRequest.newBuilder() + .putEvent("event_type", "TYPE_VIEW_CLICKED") + .putEvent("package_name", "com.example.app") + .build() + whenever(mockGestureInjector.popEvents()).doReturn(listOf(mockEvent)) + + env.updateSubscriptions( + setOf( + DeviceSignal.ACCESSIBILITY_TREE, + DeviceSignal.ACCESSIBILITY_EVENTS, + DeviceSignal.ACTIVE_PACKAGE, + DeviceSignal.ACTIVE_ACTIVITY, + ) + ) + + val state = env.getDeviceState() + + assertThat(state.accessibilityRoot.value).isEqualTo(mockNode) + assertThat(state.activePackage.value).isEqualTo("com.example.app") + assertThat(state.activeActivity.value).isEqualTo("com.example.app.MainActivity") + assertThat(state.accessibilityEvents.value).containsExactly(mockEvent) + } + + @Test + fun updateSubscriptions_enablesAndDisablesEventTracking() { + env.updateSubscriptions(setOf(DeviceSignal.ACCESSIBILITY_EVENTS)) + verify(mockGestureInjector).setEventTrackingEnabled(true) + + env.updateSubscriptions(emptySet()) + verify(mockGestureInjector).setEventTrackingEnabled(false) + verify(mockGestureInjector).clearEvents() + } + + @Test + fun updateSubscriptions_registersSensors() { + env.updateSubscriptions( + setOf( + DeviceSignal.SENSOR_ACCELEROMETER, + DeviceSignal.SENSOR_GYROSCOPE, + DeviceSignal.SENSOR_MAGNETOMETER, + DeviceSignal.SENSOR_BAROMETER, + DeviceSignal.SENSOR_AMBIENT_TEMP, + DeviceSignal.SENSOR_LIGHT, + DeviceSignal.SENSOR_PROXIMITY, + DeviceSignal.SENSOR_HEART_RATE, + DeviceSignal.LOCATION_GPS, + ) + ) + val state = env.getDeviceState() + assertThat(state).isNotNull() + } + + @Test + fun updateSubscriptions_enablesAndDisablesBatteryCollector() { + env.updateSubscriptions(setOf(DeviceSignal.BATTERY_POWER)) + val stateWithBattery = env.getDeviceState() + assertThat(stateWithBattery.batteryPower.value).isNotNull() + + env.updateSubscriptions(emptySet()) + val stateWithoutBattery = env.getDeviceState() + assertThat(stateWithoutBattery.batteryPower.value).isNull() + } + + @Test + fun getDeviceState_returnsBatteryPowerMetrics() { + env.updateSubscriptions(setOf(DeviceSignal.BATTERY_POWER)) + val state = env.getDeviceState() + + assertThat(state.batteryPower.value).isNotNull() + val metrics = state.batteryPower.value!! + assertThat(metrics.levelPercent).isAtLeast(0) + assertThat(metrics.levelPercent).isAtMost(100) + assertThat(metrics.voltageMv).isGreaterThan(0) + assertThat(metrics.powerMw).isFinite() + assertThat(metrics.chargeCounterUah).isAtLeast(0L) + assertThat(metrics.energyCounterNwh).isAtLeast(0L) + assertThat(metrics.status).isNotEmpty() + assertThat(metrics.pluggedSource).isNotEmpty() + } +} + diff --git a/android_env/apps/java/com/google/androidenv/core/AndroidDeviceEnvWebSocketServer.kt b/android_env/apps/java/com/google/androidenv/core/AndroidDeviceEnvWebSocketServer.kt new file mode 100644 index 00000000..a1d2fb4f --- /dev/null +++ b/android_env/apps/java/com/google/androidenv/core/AndroidDeviceEnvWebSocketServer.kt @@ -0,0 +1,171 @@ +// Copyright 2026 DeepMind Technologies Limited. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package com.google.androidenv.core + +import android.util.Log +import com.androidenv.proto.ClientMessage +import com.androidenv.proto.ErrorResponse +import com.androidenv.proto.ServerMessage +import com.google.androidenv.core.ProtoConversions.toCore +import com.google.androidenv.core.ProtoConversions.toProto +import java.io.IOException +import java.net.InetSocketAddress +import java.nio.ByteBuffer +import org.java_websocket.WebSocket +import org.java_websocket.exceptions.WebsocketNotConnectedException +import org.java_websocket.handshake.ClientHandshake +import org.java_websocket.server.WebSocketServer + +/** + * WebSocket server that exposes AndroidDeviceEnv over the network. Communicates using serialized + * Protobuf messages and continuous background state streaming. + */ +class AndroidDeviceEnvWebSocketServer(private val env: AndroidDeviceEnv, port: Int) : + WebSocketServer(InetSocketAddress("127.0.0.1", port)) { + + private var streamerThread: Thread? = null + @Volatile private var isStreaming = false + + init { + isReuseAddr = true + } + + override fun onOpen(conn: WebSocket, handshake: ClientHandshake) { + // We only support one active session for control, but we don't enforce it strictly here. + } + + override fun onClose(conn: WebSocket, code: Int, reason: String, remote: Boolean) {} + + override fun onMessage(conn: WebSocket, message: String) { + sendError(conn, "Text messages are not supported. Use binary protobuf messages.") + } + + override fun onMessage(conn: WebSocket, message: ByteBuffer) { + try { + // Convert ByteBuffer to ByteString for Protobuf parsing + val byteString = com.google.protobuf.ByteString.copyFrom(message) + val clientMessage = ClientMessage.parseFrom(byteString) + handleClientMessage(conn, clientMessage) + } catch (e: Exception) { + sendError(conn, "Failed to parse message: ${e.message}") + } + } + + override fun onError(conn: WebSocket?, ex: Exception) { + // Log or handle server errors + } + + override fun onStart() { + startStreaming() + } + + @Synchronized + private fun startStreaming() { + if (isStreaming) return + isStreaming = true + streamerThread = + Thread { runStreamingLoop() } + .apply { + isDaemon = true + start() + } + } + + private fun runStreamingLoop() { + val frameIntervalMs = 33L // ~30 FPS + while (isStreaming && !Thread.currentThread().isInterrupted) { + try { + broadcastLatestDeviceState() + Thread.sleep(frameIntervalMs) + } catch (e: InterruptedException) { + break + } catch (e: WebsocketNotConnectedException) { + // Client disconnected during frame broadcast. + } catch (e: IOException) { + Log.w(TAG, "IO error streaming device state: ${e.message}") + } + } + } + + private fun broadcastLatestDeviceState() { + val activeConns = connections + if (activeConns.isEmpty()) return + + val coreState = env.getDeviceState() + val protoState = coreState.toProto() + val serverMsg = ServerMessage.newBuilder().setState(protoState).build() + val msgBytes = serverMsg.toByteArray() + + for (conn in activeConns) { + if (conn.isOpen) { + conn.send(msgBytes) + } + } + } + + @Synchronized + private fun stopStreaming() { + isStreaming = false + streamerThread?.interrupt() + streamerThread = null + } + + private fun handleClientMessage(conn: WebSocket, msg: ClientMessage) { + try { + when (msg.payloadCase) { + ClientMessage.PayloadCase.UPDATE_SUBSCRIPTIONS -> { + val protoSignals = msg.updateSubscriptions.activeSignalsList + val coreSignals = protoSignals.mapNotNull { it.toCore() }.toSet() + env.updateSubscriptions(coreSignals) + } + ClientMessage.PayloadCase.INJECT_ACTION -> { + val protoAction = msg.injectAction.action + env.injectAction(protoAction.toCore()) + } + ClientMessage.PayloadCase.GET_STATE -> { + val coreState = env.getDeviceState() + val protoState = coreState.toProto() + val serverMsg = ServerMessage.newBuilder().setState(protoState).build() + conn.send(serverMsg.toByteArray()) + } + else -> { + sendError(conn, "Unsupported payload: ${msg.payloadCase}") + } + } + } catch (e: Exception) { + sendError(conn, "Error handling command: ${e.message}") + } + } + + private fun sendError(conn: WebSocket, errorMsg: String) { + try { + val serverMsg = + ServerMessage.newBuilder().setError(ErrorResponse.newBuilder().setMessage(errorMsg)).build() + conn.send(serverMsg.toByteArray()) + } catch (e: Exception) { + // Connection might be closed, ignore + } + } + + override fun stop() { + stopStreaming() + super.stop() + env.close() + } + + companion object { + private const val TAG = "AndroidDeviceEnvWS" + } +} diff --git a/android_env/apps/java/com/google/androidenv/core/AndroidDeviceEnvWebSocketServerTest.kt b/android_env/apps/java/com/google/androidenv/core/AndroidDeviceEnvWebSocketServerTest.kt new file mode 100644 index 00000000..5ef7de63 --- /dev/null +++ b/android_env/apps/java/com/google/androidenv/core/AndroidDeviceEnvWebSocketServerTest.kt @@ -0,0 +1,165 @@ +// Copyright 2026 DeepMind Technologies Limited. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package com.google.androidenv.core + +import androidx.test.ext.junit.runners.AndroidJUnit4 +import com.androidenv.proto.ClientMessage +import com.androidenv.proto.DeviceSignal +import com.androidenv.proto.InjectActionRequest +import com.androidenv.proto.UpdateSubscriptionsRequest +import com.google.common.truth.Truth.assertThat +import java.net.URI +import org.java_websocket.client.WebSocketClient +import org.java_websocket.handshake.ServerHandshake +import org.junit.After +import org.junit.Before +import org.junit.Test +import org.junit.runner.RunWith +import org.mockito.kotlin.any +import org.mockito.kotlin.doReturn +import org.mockito.kotlin.mock +import org.mockito.kotlin.verify +import org.mockito.kotlin.whenever + +@RunWith(AndroidJUnit4::class) +class AndroidDeviceEnvWebSocketServerTest { + + private val mockEnv = mock() + private var serverPort = 0 + private var server: AndroidDeviceEnvWebSocketServer? = null + + @Before + fun setUp() { + val dummyState = + DeviceState( + screenshot = StreamedSignal.empty(), + accessibilityRoot = StreamedSignal.empty(), + activePackage = StreamedSignal.empty(), + activeActivity = StreamedSignal.empty(), + systemLogs = StreamedSignal.empty(), + runningProcesses = StreamedSignal.empty(), + deviceSettings = StreamedSignal.empty(), + accelerometer = StreamedSignal.empty(), + gyroscope = StreamedSignal.empty(), + magnetometer = StreamedSignal.empty(), + barometer = StreamedSignal.empty(), + ambientTemp = StreamedSignal.empty(), + lightLevel = StreamedSignal.empty(), + proximity = StreamedSignal.empty(), + locationGps = StreamedSignal.empty(), + heartRate = StreamedSignal.empty(), + audioOutput = StreamedSignal.empty(), + snapshotTimestampUs = 12345L, + ) + whenever(mockEnv.getDeviceState()).doReturn(dummyState) + + // Find an available port and start server + val socket = java.net.ServerSocket(0) + serverPort = socket.localPort + socket.close() + + server = AndroidDeviceEnvWebSocketServer(mockEnv, serverPort) + server?.start() + // Wait for server to start + Thread.sleep(200) + } + + @After + fun tearDown() { + server?.stop() + } + + @Test + fun testWebSocketServer_startAndConnect() { + var connected = false + val client = + object : WebSocketClient(URI("ws://127.0.0.1:$serverPort")) { + override fun onOpen(handshakedata: ServerHandshake?) { + connected = true + } + + override fun onMessage(message: String?) {} + + override fun onMessage(bytes: java.nio.ByteBuffer?) {} + + override fun onClose(code: Int, reason: String?, remote: Boolean) {} + + override fun onError(ex: Exception?) {} + } + + client.connectBlocking() + assertThat(connected).isTrue() + client.close() + } + + @Test + fun testWebSocketServer_handleUpdateSubscriptions() { + val client = + object : WebSocketClient(URI("ws://127.0.0.1:$serverPort")) { + override fun onOpen(handshakedata: ServerHandshake?) {} + + override fun onMessage(message: String?) {} + + override fun onMessage(bytes: java.nio.ByteBuffer?) {} + + override fun onClose(code: Int, reason: String?, remote: Boolean) {} + + override fun onError(ex: Exception?) {} + } + + client.connectBlocking() + + val req = + UpdateSubscriptionsRequest.newBuilder() + .addActiveSignals(DeviceSignal.DEVICE_SIGNAL_SCREENSHOT) + .build() + val msg = ClientMessage.newBuilder().setUpdateSubscriptions(req).build() + client.send(msg.toByteArray()) + + Thread.sleep(150) + verify(mockEnv).updateSubscriptions(setOf(com.google.androidenv.core.DeviceSignal.SCREENSHOT)) + client.close() + } + + @Test + fun testWebSocketServer_handleInjectAction() { + val client = + object : WebSocketClient(URI("ws://127.0.0.1:$serverPort")) { + override fun onOpen(handshakedata: ServerHandshake?) {} + + override fun onMessage(message: String?) {} + + override fun onMessage(bytes: java.nio.ByteBuffer?) {} + + override fun onClose(code: Int, reason: String?, remote: Boolean) {} + + override fun onError(ex: Exception?) {} + } + + client.connectBlocking() + + val action = + com.androidenv.proto.Action.newBuilder() + .setActionType(com.androidenv.proto.ActionType.ACTION_TYPE_TOUCH_DOWN) + .build() + val req = InjectActionRequest.newBuilder().setAction(action).build() + val msg = ClientMessage.newBuilder().setInjectAction(req).build() + client.send(msg.toByteArray()) + + Thread.sleep(150) + verify(mockEnv).injectAction(any()) + client.close() + } +} diff --git a/android_env/apps/java/com/google/androidenv/core/AndroidManifest.xml b/android_env/apps/java/com/google/androidenv/core/AndroidManifest.xml new file mode 100644 index 00000000..7a069190 --- /dev/null +++ b/android_env/apps/java/com/google/androidenv/core/AndroidManifest.xml @@ -0,0 +1,66 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/android_env/apps/java/com/google/androidenv/core/BatteryPowerCollector.kt b/android_env/apps/java/com/google/androidenv/core/BatteryPowerCollector.kt new file mode 100644 index 00000000..a550e8c9 --- /dev/null +++ b/android_env/apps/java/com/google/androidenv/core/BatteryPowerCollector.kt @@ -0,0 +1,281 @@ +// Copyright 2026 DeepMind Technologies Limited. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package com.google.androidenv.core + +import android.content.BroadcastReceiver +import android.content.Context +import android.content.Intent +import android.content.IntentFilter +import android.os.BatteryManager +import android.os.Build +import android.os.HardwarePropertiesManager +import android.os.PowerManager +import android.util.Log +import kotlin.math.abs + +/** + * APK-native Battery and Power metric extraction collector for AndroidDeviceEnv. + * Collects real-time battery voltage, current draw, charge counter, thermal status, + * and calculated power consumption without relying on ADB commands or sysfs file access. + * + * Implements hardened safeguards against integer overflow, OEM unit scaling differences, + * and Goldfish/QEMU emulator virtualization quirks. + */ +class BatteryPowerCollector( + private val context: Context, + private val batteryManager: BatteryManager? = + context.getSystemService(Context.BATTERY_SERVICE) as? BatteryManager, + private val powerManager: PowerManager? = + context.getSystemService(Context.POWER_SERVICE) as? PowerManager, + private val hardwarePropertiesManager: HardwarePropertiesManager? = + context.getSystemService(Context.HARDWARE_PROPERTIES_SERVICE) as? HardwarePropertiesManager, +) { + + private var _isStarted = false + val isStarted: Boolean + @Synchronized get() = _isStarted + + private var lastKnownVoltageMv: Int = 4000 // Default 3.7V - 4.0V nominal battery + private var lastKnownTempCelsius: Float = 25.0f + + // Possible values: "CHARGING", "DISCHARGING", "FULL", "NOT_CHARGING", "UNKNOWN" + private var lastKnownStatus: String = "UNKNOWN" + + // Possible values: "AC", "USB", "WIRELESS", "UNPLUGGED", "DOCK_OR_OTHER" + private var lastKnownPlugged: String = "UNPLUGGED" + private var lastKnownLevelPercent: Int = 100 + + private val batteryReceiver = + object : BroadcastReceiver() { + override fun onReceive(context: Context?, intent: Intent?) { + if (intent != null) { + parseBatteryIntent(intent) + } + } + } + + @Synchronized + @Suppress("UnprotectedReceiver") + fun start() { + if (_isStarted) return + val filter = + IntentFilter().apply { + addAction(Intent.ACTION_BATTERY_CHANGED) + addAction(Intent.ACTION_POWER_CONNECTED) + addAction(Intent.ACTION_POWER_DISCONNECTED) + } + val stickyIntent = + if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.TIRAMISU) { + context.registerReceiver(batteryReceiver, filter, Context.RECEIVER_NOT_EXPORTED) + } else { + context.registerReceiver(batteryReceiver, filter) + } + if (stickyIntent != null) { + parseBatteryIntent(stickyIntent) + } + _isStarted = true + Log.i(TAG, "BatteryPowerCollector started") + } + + @Synchronized + fun stop() { + if (!_isStarted) return + try { + context.unregisterReceiver(batteryReceiver) + } catch (e: IllegalArgumentException) { + Log.w(TAG, "Receiver was not registered", e) + } + _isStarted = false + Log.i(TAG, "BatteryPowerCollector stopped") + } + + /** Captures the latest instantaneous battery and power telemetry. */ + @Synchronized + @Suppress("UnprotectedReceiver") + fun getLatestMetrics(): BatteryPowerMetrics { + // Re-query sticky broadcast if not started or to ensure freshest voltage/temp + if (!_isStarted) { + val stickyIntent = + context.registerReceiver(null, IntentFilter(Intent.ACTION_BATTERY_CHANGED)) + if (stickyIntent != null) { + parseBatteryIntent(stickyIntent) + } + } + + val rawCurrent = + batteryManager?.getLongProperty(BatteryManager.BATTERY_PROPERTY_CURRENT_NOW) + ?: Long.MIN_VALUE + val chargeCounter = + batteryManager?.getLongProperty(BatteryManager.BATTERY_PROPERTY_CHARGE_COUNTER) + ?: Long.MIN_VALUE + val capacityPercent = + batteryManager?.getIntProperty(BatteryManager.BATTERY_PROPERTY_CAPACITY) ?: -1 + + val levelPercent = if (capacityPercent in 1..100) capacityPercent else lastKnownLevelPercent + + // Normalization & Defensive Safeguards: + // Detect sentinels (Long.MIN_VALUE, Long.MAX_VALUE, 0 on emulators, -1 or Int.MIN_VALUE on older HALs) + val isCurrentSentinel = + rawCurrent == Long.MIN_VALUE || + rawCurrent == Long.MAX_VALUE || + rawCurrent == 0L || + rawCurrent == -1L || + rawCurrent == Int.MIN_VALUE.toLong() + + val currentNowUa: Long = + if (!isCurrentSentinel) { + // Auto-detect mA vs uA scaling across OEMs + // If absolute value is smaller than 10000, OEM likely returned mA instead of uA + if (abs(rawCurrent) > 0 && abs(rawCurrent) < 10000) { + rawCurrent * 1000L + } else { + rawCurrent + } + } else { + 0L + } + + // 64-bit floating point arithmetic to prevent 32-bit Math.abs(Integer.MIN_VALUE) overflow trap + val powerMw = (abs(currentNowUa.toDouble()) * lastKnownVoltageMv.toDouble()) / 1_000_000.0 + + val thermalStatus = + if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.Q && powerManager != null) { + try { + powerManager.currentThermalStatus + } catch (e: Exception) { + 0 + } + } else { + 0 + } + + val rawEnergyNwh = + try { + batteryManager?.getLongProperty(BatteryManager.BATTERY_PROPERTY_ENERGY_COUNTER) + ?: Long.MIN_VALUE + } catch (e: Exception) { + Long.MIN_VALUE + } + + val validChargeCounter = + if (chargeCounter <= 0L || chargeCounter == Long.MAX_VALUE) { + 0L + } else { + chargeCounter + } + + val validEnergyCounterNwh = + if (rawEnergyNwh <= 0L || rawEnergyNwh == Long.MAX_VALUE) { + 0L + } else { + rawEnergyNwh + } + + val bestTempCelsius = getBestBatteryTemperatureCelsius() + + return BatteryPowerMetrics( + levelPercent = levelPercent, + currentNowUa = currentNowUa, + voltageMv = lastKnownVoltageMv, + powerMw = powerMw, + temperatureCelsius = bestTempCelsius, + status = lastKnownStatus, + pluggedSource = lastKnownPlugged, + chargeCounterUah = validChargeCounter, + thermalStatus = thermalStatus, + energyCounterNwh = validEnergyCounterNwh, + ) + } + + private fun getBestBatteryTemperatureCelsius(): Float { + if (hardwarePropertiesManager != null) { + try { + val temps = + hardwarePropertiesManager.getDeviceTemperatures( + HardwarePropertiesManager.DEVICE_TEMPERATURE_BATTERY, + HardwarePropertiesManager.TEMPERATURE_CURRENT + ) + if (temps.isNotEmpty() && temps[0] > -100.0f) { + return temps[0] + } + } catch (e: Exception) { + // Fallback to EXTRA_TEMPERATURE when unprivileged app or SecurityException + } + } + return lastKnownTempCelsius + } + + @Synchronized + private fun parseBatteryIntent(intent: Intent) { + when (intent.action) { + Intent.ACTION_POWER_CONNECTED -> { + lastKnownStatus = "CHARGING" + if (lastKnownPlugged == "UNPLUGGED") { + lastKnownPlugged = "AC" + } + return + } + Intent.ACTION_POWER_DISCONNECTED -> { + lastKnownStatus = "DISCHARGING" + lastKnownPlugged = "UNPLUGGED" + return + } + } + + val voltage = intent.getIntExtra(BatteryManager.EXTRA_VOLTAGE, -1) + if (voltage > 0) { + // Android EXTRA_VOLTAGE is typically reported in millivolts (mV). + // Some old hardware reports in volts or decivolts; ensure proper mV range + lastKnownVoltageMv = if (voltage < 100) voltage * 1000 else voltage + } + + val tempTenths = intent.getIntExtra(BatteryManager.EXTRA_TEMPERATURE, Int.MIN_VALUE) + if (tempTenths > -1000) { + lastKnownTempCelsius = tempTenths / 10.0f + } + + val statusInt = intent.getIntExtra(BatteryManager.EXTRA_STATUS, -1) + lastKnownStatus = + when (statusInt) { + BatteryManager.BATTERY_STATUS_CHARGING -> "CHARGING" + BatteryManager.BATTERY_STATUS_DISCHARGING -> "DISCHARGING" + BatteryManager.BATTERY_STATUS_FULL -> "FULL" + BatteryManager.BATTERY_STATUS_NOT_CHARGING -> "NOT_CHARGING" + else -> "UNKNOWN" + } + + if (intent.hasExtra(BatteryManager.EXTRA_PLUGGED)) { + val pluggedInt = intent.getIntExtra(BatteryManager.EXTRA_PLUGGED, 0) + lastKnownPlugged = + when { + (pluggedInt and BatteryManager.BATTERY_PLUGGED_AC) != 0 -> "AC" + (pluggedInt and BatteryManager.BATTERY_PLUGGED_USB) != 0 -> "USB" + (pluggedInt and BatteryManager.BATTERY_PLUGGED_WIRELESS) != 0 -> "WIRELESS" + pluggedInt == 0 -> "UNPLUGGED" + else -> "DOCK_OR_OTHER" + } + } + + val level = intent.getIntExtra(BatteryManager.EXTRA_LEVEL, -1) + val scale = intent.getIntExtra(BatteryManager.EXTRA_SCALE, -1) + if (level >= 0 && scale > 0) { + lastKnownLevelPercent = ((level.toFloat() / scale.toFloat()) * 100.0f).toInt() + } + } + + companion object { + private const val TAG = "BatteryPowerCollector" + } +} diff --git a/android_env/apps/java/com/google/androidenv/core/BatteryPowerCollectorTest.kt b/android_env/apps/java/com/google/androidenv/core/BatteryPowerCollectorTest.kt new file mode 100644 index 00000000..0eecc679 --- /dev/null +++ b/android_env/apps/java/com/google/androidenv/core/BatteryPowerCollectorTest.kt @@ -0,0 +1,574 @@ +// Copyright 2026 DeepMind Technologies Limited. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package com.google.androidenv.core + +import androidx.test.core.app.ApplicationProvider +import androidx.test.ext.junit.runners.AndroidJUnit4 +import com.google.androidenv.core.ProtoConversions.toCore +import com.google.androidenv.core.ProtoConversions.toProto +import com.google.common.truth.Truth.assertThat +import org.junit.Test +import org.junit.runner.RunWith +import org.robolectric.shadows.ShadowLooper + +@RunWith(AndroidJUnit4::class) +class BatteryPowerCollectorTest { + + @Test + fun testGetLatestMetrics_returnsValidDefaultsOrSimulatedMetrics() { + val collector = BatteryPowerCollector(ApplicationProvider.getApplicationContext()) + val metrics = collector.getLatestMetrics() + + assertThat(metrics.levelPercent).isAtLeast(0) + assertThat(metrics.levelPercent).isAtMost(100) + assertThat(metrics.voltageMv).isGreaterThan(0) + assertThat(metrics.temperatureCelsius).isNotNull() + assertThat(metrics.status).isNotEmpty() + assertThat(metrics.pluggedSource).isNotEmpty() + // Verify powerMw calculation did not overflow or produce invalid negative megawatt spikes + assertThat(metrics.powerMw).isFinite() + } + + @Test + fun testStartAndStop_lifecycleSucceedsWithoutExceptions() { + val collector = BatteryPowerCollector(ApplicationProvider.getApplicationContext()) + collector.start() + val metricsAfterStart = collector.getLatestMetrics() + assertThat(metricsAfterStart.voltageMv).isGreaterThan(0) + + collector.stop() + // Stopping twice should be a no-op and not throw + collector.stop() + } + + @Test + fun testGetLatestMetrics_defensiveSafeguardsAndSentinels() { + val collector = BatteryPowerCollector(ApplicationProvider.getApplicationContext()) + val metrics = collector.getLatestMetrics() + + // Verify temperature is not corrupted by missing extra default (-1 / 10 = -0.1) + assertThat(metrics.temperatureCelsius).isAtLeast(-20.0f) + assertThat(metrics.temperatureCelsius).isAtMost(80.0f) + + // Verify pluggedSource is valid and not overridden by bitwise two's-complement trap + assertThat(metrics.pluggedSource).isAnyOf("UNPLUGGED", "AC", "USB", "WIRELESS", "DOCK_OR_OTHER") + + // Verify chargeCounter is non-negative (no -1L sentinel leak) + assertThat(metrics.chargeCounterUah).isAtLeast(0L) + // Verify energyCounterNwh is non-negative (no Long.MIN_VALUE leak) + assertThat(metrics.energyCounterNwh).isAtLeast(0L) + } + + @Test + fun testGetLatestMetrics_multipleSequentialCallsSucceed() { + val collector = BatteryPowerCollector(ApplicationProvider.getApplicationContext()) + collector.start() + for (i in 1..10) { + val metrics = collector.getLatestMetrics() + assertThat(metrics.levelPercent).isAtLeast(0) + assertThat(metrics.voltageMv).isGreaterThan(0) + assertThat(metrics.powerMw).isFinite() + } + collector.stop() + } + + @Test + fun testDeviceSignalEnum_batteryPowerConversion() { + val signal = DeviceSignal.BATTERY_POWER + val protoSignal = signal.toProto() + assertThat(protoSignal).isEqualTo(com.androidenv.proto.DeviceSignal.DEVICE_SIGNAL_BATTERY_POWER) + assertThat(protoSignal.toCore()).isEqualTo(DeviceSignal.BATTERY_POWER) + } + + @Test + fun testBatteryPowerMetrics_protoConversionInDeviceState() { + val metrics = + BatteryPowerMetrics( + levelPercent = 85, + currentNowUa = 450_000L, + voltageMv = 4200, + powerMw = 1890.0, + temperatureCelsius = 28.5f, + status = "DISCHARGING", + pluggedSource = "UNPLUGGED", + chargeCounterUah = 3_000_000L, + thermalStatus = 0, + energyCounterNwh = 15_000_000_000L, + ) + + val deviceState = + DeviceState( + screenshot = StreamedSignal.empty(), + accessibilityRoot = StreamedSignal.empty(), + activePackage = StreamedSignal.empty(), + activeActivity = StreamedSignal.empty(), + systemLogs = StreamedSignal.empty(), + runningProcesses = StreamedSignal.empty(), + deviceSettings = StreamedSignal.empty(), + accelerometer = StreamedSignal.empty(), + gyroscope = StreamedSignal.empty(), + magnetometer = StreamedSignal.empty(), + barometer = StreamedSignal.empty(), + ambientTemp = StreamedSignal.empty(), + lightLevel = StreamedSignal.empty(), + proximity = StreamedSignal.empty(), + locationGps = StreamedSignal.empty(), + heartRate = StreamedSignal.empty(), + audioOutput = StreamedSignal.empty(), + batteryPower = StreamedSignal(metrics, 123456789L, 1L), + snapshotTimestampUs = 123456789L, + orientation = 0, + ) + + val protoState = deviceState.toProto() + val batteryProtoSignal = + protoState.signalsList.find { + it.type == com.androidenv.proto.DeviceSignal.DEVICE_SIGNAL_BATTERY_POWER + } + + assertThat(batteryProtoSignal).isNotNull() + assertThat(batteryProtoSignal!!.hasBatteryPower()).isTrue() + + val protoMetrics = batteryProtoSignal.batteryPower + assertThat(protoMetrics.levelPercent).isEqualTo(85) + assertThat(protoMetrics.currentNowUa).isEqualTo(450_000L) + assertThat(protoMetrics.voltageMv).isEqualTo(4200) + assertThat(protoMetrics.powerMw).isEqualTo(1890.0) + assertThat(protoMetrics.temperatureCelsius).isEqualTo(28.5f) + assertThat(protoMetrics.status).isEqualTo("DISCHARGING") + assertThat(protoMetrics.pluggedSource).isEqualTo("UNPLUGGED") + assertThat(protoMetrics.chargeCounterUah).isEqualTo(3_000_000L) + assertThat(protoMetrics.thermalStatus).isEqualTo(0) + assertThat(protoMetrics.energyCounterNwh).isEqualTo(15_000_000_000L) + } + + @Test + fun testParseBatteryIntent_powerConnectedAndDisconnected() { + val collector = BatteryPowerCollector(ApplicationProvider.getApplicationContext()) + collector.start() + val context = ApplicationProvider.getApplicationContext() + val connectedIntent = android.content.Intent(android.content.Intent.ACTION_POWER_CONNECTED) + context.sendBroadcast(connectedIntent) + ShadowLooper.idleMainLooper() + + val metricsConnected = collector.getLatestMetrics() + assertThat(metricsConnected.status).isEqualTo("CHARGING") + assertThat(metricsConnected.pluggedSource).isEqualTo("AC") + + val disconnectedIntent = + android.content.Intent(android.content.Intent.ACTION_POWER_DISCONNECTED) + context.sendBroadcast(disconnectedIntent) + ShadowLooper.idleMainLooper() + + val metricsDisconnected = collector.getLatestMetrics() + assertThat(metricsDisconnected.status).isEqualTo("DISCHARGING") + assertThat(metricsDisconnected.pluggedSource).isEqualTo("UNPLUGGED") + collector.stop() + } + + @Test + fun testParseBatteryIntent_powerConnectedWhenAlreadyPluggedUsb_preservesUsbSource() { + val collector = BatteryPowerCollector(ApplicationProvider.getApplicationContext()) + collector.start() + val context = ApplicationProvider.getApplicationContext() + + val intentUsb = + android.content.Intent(android.content.Intent.ACTION_BATTERY_CHANGED).apply { + putExtra( + android.os.BatteryManager.EXTRA_PLUGGED, + android.os.BatteryManager.BATTERY_PLUGGED_USB, + ) + putExtra( + android.os.BatteryManager.EXTRA_STATUS, + android.os.BatteryManager.BATTERY_STATUS_CHARGING, + ) + } + context.sendBroadcast(intentUsb) + ShadowLooper.idleMainLooper() + + assertThat(collector.getLatestMetrics().pluggedSource).isEqualTo("USB") + + val connectedIntent = android.content.Intent(android.content.Intent.ACTION_POWER_CONNECTED) + context.sendBroadcast(connectedIntent) + ShadowLooper.idleMainLooper() + + val metrics = collector.getLatestMetrics() + assertThat(metrics.status).isEqualTo("CHARGING") + assertThat(metrics.pluggedSource).isEqualTo("USB") + collector.stop() + } + + @Test + fun testParseBatteryIntent_powerConnectedWhenAlreadyPluggedWireless_preservesWirelessSource() { + val collector = BatteryPowerCollector(ApplicationProvider.getApplicationContext()) + collector.start() + val context = ApplicationProvider.getApplicationContext() + + val intentWireless = + android.content.Intent(android.content.Intent.ACTION_BATTERY_CHANGED).apply { + putExtra( + android.os.BatteryManager.EXTRA_PLUGGED, + android.os.BatteryManager.BATTERY_PLUGGED_WIRELESS, + ) + putExtra( + android.os.BatteryManager.EXTRA_STATUS, + android.os.BatteryManager.BATTERY_STATUS_CHARGING, + ) + } + context.sendBroadcast(intentWireless) + ShadowLooper.idleMainLooper() + + assertThat(collector.getLatestMetrics().pluggedSource).isEqualTo("WIRELESS") + + val connectedIntent = android.content.Intent(android.content.Intent.ACTION_POWER_CONNECTED) + context.sendBroadcast(connectedIntent) + ShadowLooper.idleMainLooper() + + val metrics = collector.getLatestMetrics() + assertThat(metrics.status).isEqualTo("CHARGING") + assertThat(metrics.pluggedSource).isEqualTo("WIRELESS") + collector.stop() + } + + @Test + fun testParseBatteryIntent_batteryChangedVariations() { + val collector = BatteryPowerCollector(ApplicationProvider.getApplicationContext()) + collector.start() + val context = ApplicationProvider.getApplicationContext() + + // Test charging status, AC plugged, low voltage scaling (<100 decivolts/volts -> mV + // conversion), and level/scale + val intentAc = + android.content.Intent(android.content.Intent.ACTION_BATTERY_CHANGED).apply { + putExtra(android.os.BatteryManager.EXTRA_VOLTAGE, 4) // 4V -> 4000mV + putExtra(android.os.BatteryManager.EXTRA_TEMPERATURE, 312) // 31.2 C + putExtra( + android.os.BatteryManager.EXTRA_STATUS, + android.os.BatteryManager.BATTERY_STATUS_CHARGING, + ) + putExtra( + android.os.BatteryManager.EXTRA_PLUGGED, + android.os.BatteryManager.BATTERY_PLUGGED_AC, + ) + putExtra(android.os.BatteryManager.EXTRA_LEVEL, 45) + putExtra(android.os.BatteryManager.EXTRA_SCALE, 50) // 90% + } + context.sendBroadcast(intentAc) + ShadowLooper.idleMainLooper() + var metrics = collector.getLatestMetrics() + + assertThat(metrics.voltageMv).isEqualTo(4000) + assertThat(metrics.temperatureCelsius).isEqualTo(31.2f) + assertThat(metrics.status).isEqualTo("CHARGING") + assertThat(metrics.pluggedSource).isEqualTo("AC") + assertThat(metrics.levelPercent).isEqualTo(90) + + // Test USB plugged, FULL status + val intentUsb = + android.content.Intent(android.content.Intent.ACTION_BATTERY_CHANGED).apply { + putExtra(android.os.BatteryManager.EXTRA_VOLTAGE, 4100) + putExtra( + android.os.BatteryManager.EXTRA_STATUS, + android.os.BatteryManager.BATTERY_STATUS_FULL, + ) + putExtra( + android.os.BatteryManager.EXTRA_PLUGGED, + android.os.BatteryManager.BATTERY_PLUGGED_USB, + ) + } + context.sendBroadcast(intentUsb) + ShadowLooper.idleMainLooper() + metrics = collector.getLatestMetrics() + + assertThat(metrics.status).isEqualTo("FULL") + assertThat(metrics.pluggedSource).isEqualTo("USB") + + // Test WIRELESS plugged, NOT_CHARGING status (no EXTRA_VOLTAGE in intent, voltageMv remains + // last known) + val intentWireless = + android.content.Intent(android.content.Intent.ACTION_BATTERY_CHANGED).apply { + putExtra( + android.os.BatteryManager.EXTRA_STATUS, + android.os.BatteryManager.BATTERY_STATUS_NOT_CHARGING, + ) + putExtra( + android.os.BatteryManager.EXTRA_PLUGGED, + android.os.BatteryManager.BATTERY_PLUGGED_WIRELESS, + ) + } + context.sendBroadcast(intentWireless) + ShadowLooper.idleMainLooper() + metrics = collector.getLatestMetrics() + + assertThat(metrics.status).isEqualTo("NOT_CHARGING") + assertThat(metrics.pluggedSource).isEqualTo("WIRELESS") + assertThat(metrics.voltageMv).isEqualTo(4100) + + // Test UNKNOWN status and DOCK_OR_OTHER plugged + val intentOther = + android.content.Intent(android.content.Intent.ACTION_BATTERY_CHANGED).apply { + putExtra(android.os.BatteryManager.EXTRA_STATUS, 999) + putExtra(android.os.BatteryManager.EXTRA_PLUGGED, 8) // Non-standard bit flag + } + context.sendBroadcast(intentOther) + ShadowLooper.idleMainLooper() + metrics = collector.getLatestMetrics() + + assertThat(metrics.status).isEqualTo("UNKNOWN") + assertThat(metrics.pluggedSource).isEqualTo("DOCK_OR_OTHER") + assertThat(metrics.voltageMv).isEqualTo(4100) + + // Test EXTRA_PLUGGED = 0 -> UNPLUGGED + val intentUnplugged = + android.content.Intent(android.content.Intent.ACTION_BATTERY_CHANGED).apply { + putExtra(android.os.BatteryManager.EXTRA_PLUGGED, 0) + } + context.sendBroadcast(intentUnplugged) + ShadowLooper.idleMainLooper() + metrics = collector.getLatestMetrics() + + assertThat(metrics.pluggedSource).isEqualTo("UNPLUGGED") + collector.stop() + } + + @Test + fun testMockBatteryManager_nonSentinelCurrentAndCounters() { + val context = ApplicationProvider.getApplicationContext() + val mockBatteryManager = org.mockito.Mockito.mock(android.os.BatteryManager::class.java) + + org.mockito.Mockito.`when`( + mockBatteryManager.getLongProperty(android.os.BatteryManager.BATTERY_PROPERTY_CURRENT_NOW) + ) + .thenReturn(450_000L) + org.mockito.Mockito.`when`( + mockBatteryManager.getLongProperty( + android.os.BatteryManager.BATTERY_PROPERTY_CHARGE_COUNTER + ) + ) + .thenReturn(3_000_000L) + org.mockito.Mockito.`when`( + mockBatteryManager.getLongProperty( + android.os.BatteryManager.BATTERY_PROPERTY_ENERGY_COUNTER + ) + ) + .thenReturn(15_000_000_000L) + org.mockito.Mockito.`when`( + mockBatteryManager.getIntProperty(android.os.BatteryManager.BATTERY_PROPERTY_CAPACITY) + ) + .thenReturn(85) + + val injectedCollector = BatteryPowerCollector(context, batteryManager = mockBatteryManager) + val metrics = injectedCollector.getLatestMetrics() + + assertThat(metrics.currentNowUa).isEqualTo(450_000L) + assertThat(metrics.chargeCounterUah).isEqualTo(3_000_000L) + assertThat(metrics.energyCounterNwh).isEqualTo(15_000_000_000L) + assertThat(metrics.levelPercent).isEqualTo(85) + } + + @Test + fun testMockBatteryManager_mAUnitAutoScaling() { + val context = ApplicationProvider.getApplicationContext() + val mockBatteryManager = org.mockito.Mockito.mock(android.os.BatteryManager::class.java) + + // Current < 10000 microamps (e.g. 450) triggers OEM mA -> uA auto-scaling (*1000) + org.mockito.Mockito.`when`( + mockBatteryManager.getLongProperty(android.os.BatteryManager.BATTERY_PROPERTY_CURRENT_NOW) + ) + .thenReturn(450L) + + val injectedCollector = BatteryPowerCollector(context, batteryManager = mockBatteryManager) + val metrics = injectedCollector.getLatestMetrics() + + assertThat(metrics.currentNowUa).isEqualTo(450_000L) + } + + @Test + fun testMockBatteryManager_allCurrentSentinels() { + val context = ApplicationProvider.getApplicationContext() + val mockBatteryManager = org.mockito.Mockito.mock(android.os.BatteryManager::class.java) + + val sentinels = listOf(Long.MIN_VALUE, Long.MAX_VALUE, 0L, -1L, Int.MIN_VALUE.toLong()) + + for (sentinel in sentinels) { + org.mockito.Mockito.`when`( + mockBatteryManager.getLongProperty(android.os.BatteryManager.BATTERY_PROPERTY_CURRENT_NOW) + ) + .thenReturn(sentinel) + + val injectedCollector = BatteryPowerCollector(context, batteryManager = mockBatteryManager) + val metrics = injectedCollector.getLatestMetrics() + + assertThat(metrics.currentNowUa).isEqualTo(0L) + } + } + + @Test + fun testRealPowerManager_thermalStatus() { + val context = ApplicationProvider.getApplicationContext() + val powerManager = + context.getSystemService(android.content.Context.POWER_SERVICE) as android.os.PowerManager + val injectedCollector = BatteryPowerCollector(context, powerManager = powerManager) + val metrics = injectedCollector.getLatestMetrics() + + assertThat(metrics.thermalStatus).isEqualTo(0) + } + + @Test + fun testParseBatteryIntent_invalidOrMissingTemperatureExtra() { + val context = ApplicationProvider.getApplicationContext() + val testCollector = BatteryPowerCollector(context) + testCollector.start() + + // Intent missing EXTRA_TEMPERATURE + val intentMissingTemp = + android.content.Intent(android.content.Intent.ACTION_BATTERY_CHANGED).apply { + putExtra(android.os.BatteryManager.EXTRA_VOLTAGE, 4000) + } + context.sendBroadcast(intentMissingTemp) + ShadowLooper.idleMainLooper() + var metrics = testCollector.getLatestMetrics() + assertThat(metrics.temperatureCelsius).isEqualTo(25.0f) + + // Intent with invalid extra temperature (-2000 = -200°C) + val intentInvalidTemp = + android.content.Intent(android.content.Intent.ACTION_BATTERY_CHANGED).apply { + putExtra(android.os.BatteryManager.EXTRA_TEMPERATURE, -2000) + } + context.sendBroadcast(intentInvalidTemp) + ShadowLooper.idleMainLooper() + metrics = testCollector.getLatestMetrics() + assertThat(metrics.temperatureCelsius).isEqualTo(25.0f) + testCollector.stop() + } + + @Test + fun testMockHardwarePropertiesManager_batteryTemperature() { + val context = ApplicationProvider.getApplicationContext() + val mockHwManager = org.mockito.Mockito.mock(android.os.HardwarePropertiesManager::class.java) + + org.mockito.Mockito.`when`( + mockHwManager.getDeviceTemperatures( + android.os.HardwarePropertiesManager.DEVICE_TEMPERATURE_BATTERY, + android.os.HardwarePropertiesManager.TEMPERATURE_CURRENT, + ) + ) + .thenReturn(floatArrayOf(35.5f)) + + val injectedCollector = + BatteryPowerCollector(context, hardwarePropertiesManager = mockHwManager) + val metrics = injectedCollector.getLatestMetrics() + + assertThat(metrics.temperatureCelsius).isEqualTo(35.5f) + } + + @Test + @Suppress("DEPRECATION") + fun testGetLatestMetrics_withoutStart_queriesStickyIntent() { + val context = ApplicationProvider.getApplicationContext() + val unstartedCollector = BatteryPowerCollector(context) + assertThat(unstartedCollector.isStarted).isFalse() + + // Send sticky battery changed broadcast + val stickyIntent = + android.content.Intent(android.content.Intent.ACTION_BATTERY_CHANGED).apply { + putExtra(android.os.BatteryManager.EXTRA_VOLTAGE, 3800) + putExtra(android.os.BatteryManager.EXTRA_LEVEL, 77) + putExtra(android.os.BatteryManager.EXTRA_SCALE, 100) + } + context.sendStickyBroadcast(stickyIntent) + ShadowLooper.idleMainLooper() + + // getLatestMetrics without calling start() re-queries sticky broadcast + val metrics = unstartedCollector.getLatestMetrics() + assertThat(metrics.voltageMv).isEqualTo(3800) + assertThat(metrics.levelPercent).isEqualTo(77) + } + + @Test + fun testStartAndStop_lifecycle() { + val context = ApplicationProvider.getApplicationContext() + val lifecycleCollector = BatteryPowerCollector(context) + assertThat(lifecycleCollector.isStarted).isFalse() + + // Verify start and stop sequence + lifecycleCollector.start() + assertThat(lifecycleCollector.isStarted).isTrue() + + lifecycleCollector.start() // Repeated start is no-op + assertThat(lifecycleCollector.isStarted).isTrue() + + lifecycleCollector.stop() + assertThat(lifecycleCollector.isStarted).isFalse() + + lifecycleCollector.stop() // Repeated stop catches unregister IllegalArgumentException cleanly + assertThat(lifecycleCollector.isStarted).isFalse() + } + + private class TestContextWrapper(base: android.content.Context) : + android.content.ContextWrapper(base) { + var registerCount = 0 + var unregisterCount = 0 + var throwOnUnregister = false + + override fun registerReceiver( + receiver: android.content.BroadcastReceiver?, + filter: android.content.IntentFilter?, + ): android.content.Intent? { + registerCount++ + return super.registerReceiver(receiver, filter) + } + + override fun registerReceiver( + receiver: android.content.BroadcastReceiver?, + filter: android.content.IntentFilter?, + flags: Int, + ): android.content.Intent? { + registerCount++ + return super.registerReceiver(receiver, filter, flags) + } + + override fun unregisterReceiver(receiver: android.content.BroadcastReceiver?) { + unregisterCount++ + if (throwOnUnregister) { + throw IllegalArgumentException("Simulated unregister exception") + } + super.unregisterReceiver(receiver) + } + } + + @Test + fun testStartAndStop_verifiesRegisterUnregisterGuardAndExceptionHandling() { + val context = ApplicationProvider.getApplicationContext() + val testContext = TestContextWrapper(context) + val testCollector = BatteryPowerCollector(testContext) + + // 1. Calling start() twice must call context.registerReceiver EXACTLY once + testCollector.start() + testCollector.start() + assertThat(testContext.registerCount).isEqualTo(1) + + // 2. Calling stop() twice must call context.unregisterReceiver EXACTLY once + testCollector.stop() + testCollector.stop() + assertThat(testContext.unregisterCount).isEqualTo(1) + + // 3. Exception handling when unregisterReceiver throws IllegalArgumentException + testContext.throwOnUnregister = true + testCollector.start() + testCollector.stop() + assertThat(testCollector.isStarted).isFalse() + } +} diff --git a/android_env/apps/java/com/google/androidenv/core/CoreModelsTest.kt b/android_env/apps/java/com/google/androidenv/core/CoreModelsTest.kt new file mode 100644 index 00000000..37cd09ec --- /dev/null +++ b/android_env/apps/java/com/google/androidenv/core/CoreModelsTest.kt @@ -0,0 +1,232 @@ +// Copyright 2026 DeepMind Technologies Limited. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package com.google.androidenv.core + +import android.graphics.Bitmap +import android.view.accessibility.AccessibilityNodeInfo +import androidx.test.ext.junit.runners.AndroidJUnit4 +import com.google.common.truth.Truth.assertThat +import org.junit.Test +import org.junit.runner.RunWith +import org.robolectric.annotation.Config + +@RunWith(AndroidJUnit4::class) +@Config(manifest = Config.NONE) +class CoreModelsTest { + + @Test + fun testActionEqualsAndHashCode() { + val pos1 = listOf(0.1f, 0.2f) + val pos2 = listOf(0.1f, 0.2f) + val pos3 = listOf(0.3f, 0.4f) + + val action1 = Action(ActionType.TOUCH_DOWN, pos1, null) + val action2 = Action(ActionType.TOUCH_DOWN, pos2, null) + val action3 = Action(ActionType.TOUCH_DOWN, pos3, null) + val actionDifferentType = Action(ActionType.TOUCH_UP, pos1, null) + val actionNullPos = Action(ActionType.TOUCH_DOWN, null, null) + + // Structural equality check (automatically supported by List in data class) + assertThat(action1).isEqualTo(action2) + assertThat(action1.hashCode()).isEqualTo(action2.hashCode()) + + assertThat(action1).isNotEqualTo(action3) + assertThat(action1).isNotEqualTo(actionDifferentType) + assertThat(action1).isNotEqualTo(actionNullPos) + } + + @Test + fun testAudioFrameEqualsAndHashCode() { + val bytes1 = byteArrayOf(-31) + val bytes2 = byteArrayOf(-31) + val bytes3 = byteArrayOf(4, 5, 6) + + val frame1 = AudioFrame(bytes1, 44100, 2) + val frame2 = AudioFrame(bytes2, 44100, 2) + val frame3 = AudioFrame(bytes3, 44100, 2) + val frameDifferentRate = AudioFrame(bytes1, 22050, 2) + val frameDifferentChannels = AudioFrame(bytes1, 44100, 1) + val frameDifferentFormat = AudioFrame(bytes1, 44100, 2, "PCM_16") + + assertThat(frame1).isEqualTo(frame1) + assertThat(frame1).isEqualTo(frame2) + assertThat(frame1.hashCode()).isEqualTo(frame2.hashCode()) + + assertThat(frame1).isNotEqualTo(frame3) + assertNotEqualObjects(frame1, "different_type") + assertThat(frame1).isNotEqualTo(null) + assertThat(frame1).isNotEqualTo(frameDifferentRate) + assertThat(frame1).isNotEqualTo(frameDifferentChannels) + assertThat(frame1).isNotEqualTo(frameDifferentFormat) + + // Test hash code inequality to prevent mutations + assertThat(frame1.hashCode()).isNotEqualTo(frame3.hashCode()) + assertThat(frame1.hashCode()).isNotEqualTo(frameDifferentRate.hashCode()) + assertThat(frame1.hashCode()).isNotEqualTo(frameDifferentFormat.hashCode()) + + // Specific test for channelCount mutation (result = 31 * result + channelCount) + // We use channelCount = 0 for both frames to make them equal under mutation (result * 0 = 0) + // but they differ in sampleRateHz so they should have different hash codes originally. + val frameChannels0Rate1 = AudioFrame(bytes1, 44100, 0) + val frameChannels0Rate2 = AudioFrame(bytes1, 22050, 0) + assertThat(frameChannels0Rate1.hashCode()).isNotEqualTo(frameChannels0Rate2.hashCode()) + + // Specific test for sampleFormat hashCode mutation (result = 31 * result + + // (sampleFormat?.hashCode() ?: 0)) + // We choose strings "foo?" and "foo " which have hash codes differing by 31. + // We choose channelCount 2 and 1 respectively, so that BaseA - BaseB = 1. + // Under mutation (+ is changed to -), they will have the same hashCode. + val frameFormatMutation1 = AudioFrame(bytes1, 44100, 2, "foo?") + val frameFormatMutation2 = AudioFrame(bytes1, 44100, 1, "foo ") + assertThat(frameFormatMutation1.hashCode()).isNotEqualTo(frameFormatMutation2.hashCode()) + } + + @Test + fun testGpsCoordinates() { + val gps1 = GpsCoordinates(1.0, 2.0, 3.0) + val gps2 = GpsCoordinates(1.0, 2.0, 3.0) + val gps3 = GpsCoordinates(4.0, 5.0, 6.0) + + assertThat(gps1).isEqualTo(gps2) + assertThat(gps1.hashCode()).isEqualTo(gps2.hashCode()) + assertThat(gps1).isNotEqualTo(gps3) + assertThat(gps1.latitude).isEqualTo(1.0) + assertThat(gps1.longitude).isEqualTo(2.0) + assertThat(gps1.altitude).isEqualTo(3.0) + } + + @Test + fun testStreamedSignal() { + val signal1 = StreamedSignal("test", 100L, 1L) + val signal2 = StreamedSignal("test", 100L, 1L) + val signal3 = StreamedSignal("other", 200L, 2L) + val emptySignal = StreamedSignal.empty() + + assertThat(signal1).isEqualTo(signal2) + assertThat(signal1.hashCode()).isEqualTo(signal2.hashCode()) + assertThat(signal1).isNotEqualTo(signal3) + + assertThat(signal1.value).isEqualTo("test") + assertThat(signal1.timestampUs).isEqualTo(100L) + assertThat(signal1.sequenceId).isEqualTo(1L) + + assertThat(emptySignal.value).isNull() + assertThat(emptySignal.timestampUs).isEqualTo(0L) + assertThat(emptySignal.sequenceId).isEqualTo(0L) + } + + @Test + fun testDeviceState() { + val screenshot = StreamedSignal.empty() + val accessibilityRoot = StreamedSignal.empty() + val activePackage = StreamedSignal.empty() + val activeActivity = StreamedSignal.empty() + val systemLogs = StreamedSignal.empty>() + val runningProcesses = StreamedSignal.empty>() + val deviceSettings = StreamedSignal.empty>() + val accelerometer = StreamedSignal.empty>() + val gyroscope = StreamedSignal.empty>() + val magnetometer = StreamedSignal.empty>() + val barometer = StreamedSignal.empty() + val ambientTemp = StreamedSignal.empty() + val lightLevel = StreamedSignal.empty() + val proximity = StreamedSignal.empty() + val locationGps = StreamedSignal.empty() + val heartRate = StreamedSignal.empty() + val audioOutput = StreamedSignal.empty() + val timestamp = 123456L + + val state1 = + DeviceState( + screenshot = screenshot, + accessibilityRoot = accessibilityRoot, + activePackage = activePackage, + activeActivity = activeActivity, + systemLogs = systemLogs, + runningProcesses = runningProcesses, + deviceSettings = deviceSettings, + accelerometer = accelerometer, + gyroscope = gyroscope, + magnetometer = magnetometer, + barometer = barometer, + ambientTemp = ambientTemp, + lightLevel = lightLevel, + proximity = proximity, + locationGps = locationGps, + heartRate = heartRate, + audioOutput = audioOutput, + snapshotTimestampUs = timestamp, + ) + + val state2 = + DeviceState( + screenshot = screenshot, + accessibilityRoot = accessibilityRoot, + activePackage = activePackage, + activeActivity = activeActivity, + systemLogs = systemLogs, + runningProcesses = runningProcesses, + deviceSettings = deviceSettings, + accelerometer = accelerometer, + gyroscope = gyroscope, + magnetometer = magnetometer, + barometer = barometer, + ambientTemp = ambientTemp, + lightLevel = lightLevel, + proximity = proximity, + locationGps = locationGps, + heartRate = heartRate, + audioOutput = audioOutput, + snapshotTimestampUs = timestamp, + ) + + assertThat(state1).isEqualTo(state2) + assertThat(state1.hashCode()).isEqualTo(state2.hashCode()) + assertThat(state1.snapshotTimestampUs).isEqualTo(timestamp) + } + + @Test + fun testAndroidDeviceEnvDefaultClose() { + val env = + object : AndroidDeviceEnv { + override fun updateSubscriptions(activeSignals: Set) {} + + override fun getDeviceState(): DeviceState { + throw UnsupportedOperationException() + } + + override fun injectAction(action: Action) {} + } + env.close() + env.updateSubscriptions(emptySet()) + try { + val unused = env.getDeviceState() + } catch (e: UnsupportedOperationException) { + assertThat(e).isNotNull() + } + env.injectAction(Action(ActionType.TOUCH_DOWN, null, null)) + } + + @Test + fun testDeviceSignalEnum() { + val signals = DeviceSignal.values() + assertThat(signals).isNotEmpty() + assertThat(DeviceSignal.valueOf("SCREENSHOT")).isEqualTo(DeviceSignal.SCREENSHOT) + } + + private fun assertNotEqualObjects(actual: Any?, expected: Any?) { + assertThat(actual).isNotEqualTo(expected) + } +} diff --git a/android_env/apps/java/com/google/androidenv/core/CoreModelsTestManifest.xml b/android_env/apps/java/com/google/androidenv/core/CoreModelsTestManifest.xml new file mode 100644 index 00000000..7c1fadf5 --- /dev/null +++ b/android_env/apps/java/com/google/androidenv/core/CoreModelsTestManifest.xml @@ -0,0 +1,24 @@ + + + + + + + + + diff --git a/android_env/apps/java/com/google/androidenv/core/DeviceSignal.kt b/android_env/apps/java/com/google/androidenv/core/DeviceSignal.kt new file mode 100644 index 00000000..dbd7df21 --- /dev/null +++ b/android_env/apps/java/com/google/androidenv/core/DeviceSignal.kt @@ -0,0 +1,38 @@ +// Copyright 2026 DeepMind Technologies Limited. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package com.google.androidenv.core + +/** Signals that can be dynamically subscribed to. */ +enum class DeviceSignal { + SCREENSHOT, + ACCESSIBILITY_TREE, + ACCESSIBILITY_EVENTS, + ACTIVE_PACKAGE, + ACTIVE_ACTIVITY, + SYSTEM_LOGS, + RUNNING_PROCESSES, + DEVICE_SETTINGS, + SENSOR_ACCELEROMETER, + SENSOR_GYROSCOPE, + SENSOR_MAGNETOMETER, + SENSOR_BAROMETER, + SENSOR_AMBIENT_TEMP, + SENSOR_LIGHT, + SENSOR_PROXIMITY, + SENSOR_HEART_RATE, + LOCATION_GPS, + AUDIO_OUTPUT, + BATTERY_POWER, +} diff --git a/android_env/apps/java/com/google/androidenv/core/DeviceState.kt b/android_env/apps/java/com/google/androidenv/core/DeviceState.kt new file mode 100644 index 00000000..13240d68 --- /dev/null +++ b/android_env/apps/java/com/google/androidenv/core/DeviceState.kt @@ -0,0 +1,106 @@ +// Copyright 2026 DeepMind Technologies Limited. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package com.google.androidenv.core + +import android.graphics.Bitmap +import android.view.accessibility.AccessibilityNodeInfo + +/** GPS coordinates. */ +data class GpsCoordinates(val latitude: Double, val longitude: Double, val altitude: Double) + +/** Represents a single audio frame. */ +data class AudioFrame( + val rawBytes: ByteArray, + val sampleRateHz: Int, + val channelCount: Int, + val sampleFormat: String? = null, +) { + override fun equals(other: Any?): Boolean { + if (javaClass != other?.javaClass) return false + + other as AudioFrame + + if (!rawBytes.contentEquals(other.rawBytes)) return false + if (sampleRateHz != other.sampleRateHz) return false + if (channelCount != other.channelCount) return false + if (sampleFormat != other.sampleFormat) return false + + return true + } + + override fun hashCode(): Int { + var result = rawBytes.contentHashCode() + result = 31 * result + sampleRateHz + result = 31 * result + channelCount + result = 31 * result + (sampleFormat?.hashCode() ?: 0) + return result + } +} + +/** Telemetry metrics for device battery state, current draw, voltage, and calculated power. */ +data class BatteryPowerMetrics( + val levelPercent: Int, + val currentNowUa: Long, + val voltageMv: Int, + val powerMw: Double, + val temperatureCelsius: Float, + val status: String, + val pluggedSource: String, + val chargeCounterUah: Long, + val thermalStatus: Int, + val energyCounterNwh: Long = 0L, +) + +/** Represents the complete, raw state of the Android device at a point in time. */ +data class DeviceState( + val screenshot: StreamedSignal, + val accessibilityRoot: StreamedSignal, + val accessibilityEvents: + StreamedSignal> = + StreamedSignal.empty(), + val activePackage: StreamedSignal, + val activeActivity: StreamedSignal, + val systemLogs: StreamedSignal>, + val runningProcesses: StreamedSignal>, + val deviceSettings: StreamedSignal>, + + // Fine-grained optional sensor streams + // [x, y, z] in m/s^2 (linear acceleration) + val accelerometer: StreamedSignal>, + // [x, y, z] in rad/s (angular speed) + val gyroscope: StreamedSignal>, + // [x, y, z] in uT (magnetic field / compass) + val magnetometer: StreamedSignal>, + // Atmospheric pressure in hPa + val barometer: StreamedSignal, + // Ambient temperature in degrees Celsius (°C) + val ambientTemp: StreamedSignal, + // Illuminance in lx (lux) + val lightLevel: StreamedSignal, + // Proximity distance in centimeters (cm) + val proximity: StreamedSignal, + // Latitude, Longitude, Altitude + val locationGps: StreamedSignal, + + // Smartwatch & Audio Wearables + // Heart rate in beats per minute (BPM) + val heartRate: StreamedSignal, + // Audio frequencies streamed from speaker + val audioOutput: StreamedSignal, + // Battery & power telemetry metrics + val batteryPower: StreamedSignal = StreamedSignal.empty(), + val snapshotTimestampUs: Long, + val orientation: Int = 0, +) diff --git a/android_env/apps/java/com/google/androidenv/core/GestureInjector.kt b/android_env/apps/java/com/google/androidenv/core/GestureInjector.kt new file mode 100644 index 00000000..b87a257b --- /dev/null +++ b/android_env/apps/java/com/google/androidenv/core/GestureInjector.kt @@ -0,0 +1,104 @@ +// Copyright 2026 DeepMind Technologies Limited. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package com.google.androidenv.core + +import android.accessibilityservice.GestureDescription +import android.view.accessibility.AccessibilityNodeInfo + +/** Interface for injecting gestures and querying accessibility node info. */ +interface GestureInjector { + fun injectGesture( + gesture: GestureDescription, + callback: android.accessibilityservice.AccessibilityService.GestureResultCallback? = null, + ) + + fun getRootNode(): AccessibilityNodeInfo? + + fun getRootNodes(): List { + return getRootNode()?.let { listOf(it) } ?: emptyList() + } + + fun getActivePackage(): String + + fun getActiveActivity(): String + + fun popEvents(): List + + fun setEventTrackingEnabled(enabled: Boolean) + + fun clearEvents() + + @com.google.errorprone.annotations.CanIgnoreReturnValue + fun performGlobalAction(action: Int): Boolean +} + +/** Production implementation that delegates to the InputService singleton. */ +class SystemGestureInjector : GestureInjector { + override fun injectGesture( + gesture: GestureDescription, + callback: android.accessibilityservice.AccessibilityService.GestureResultCallback?, + ) { + InputService.getInstance().injectGesture(gesture, callback) + } + + override fun getRootNode(): AccessibilityNodeInfo? { + return InputService.getInstance().getRootNode() + } + + override fun getRootNodes(): List { + return try { + InputService.getInstance().getRootNodes() + } catch (e: Exception) { + val root = getRootNode() + if (root != null) listOf(root) else emptyList() + } + } + + override fun getActivePackage(): String { + return InputService.activePackage + } + + override fun getActiveActivity(): String { + return InputService.activeActivity + } + + override fun popEvents(): List { + return try { + InputService.getInstance().popEvents() + } catch (e: Exception) { + emptyList() + } + } + + override fun setEventTrackingEnabled(enabled: Boolean) { + try { + InputService.getInstance().isEventTrackingEnabled = enabled + } catch (e: Exception) { + // Ignored if service not bound + } + } + + override fun clearEvents() { + try { + InputService.getInstance().clearEvents() + } catch (e: Exception) { + // Ignored if service not bound + } + } + + override fun performGlobalAction(action: Int): Boolean { + return InputService.getInstance().performGlobalAction(action) + } +} diff --git a/android_env/apps/java/com/google/androidenv/core/InputService.kt b/android_env/apps/java/com/google/androidenv/core/InputService.kt new file mode 100644 index 00000000..20561a0c --- /dev/null +++ b/android_env/apps/java/com/google/androidenv/core/InputService.kt @@ -0,0 +1,205 @@ +// Copyright 2026 DeepMind Technologies Limited. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package com.google.androidenv.core + +import android.accessibilityservice.AccessibilityService +import android.accessibilityservice.GestureDescription +import android.util.Log +import android.view.accessibility.AccessibilityEvent +import android.view.accessibility.AccessibilityNodeInfo + +/** + * AccessibilityService that handles action injection (gestures) and accessibility tree gathering. + * Must be manually enabled in Android Settings. + */ +class InputService : AccessibilityService() { + + override fun onCreate() { + Log.i(TAG, "InputService::onCreate()") + super.onCreate() + instance = this + } + + override fun onServiceConnected() { + Log.i(TAG, "InputService::onServiceConnected()") + super.onServiceConnected() + instance = this + } + + @Volatile var isEventTrackingEnabled: Boolean = false + private val eventBuffer = + mutableListOf() + + override fun onAccessibilityEvent(event: AccessibilityEvent) { + if (isEventTrackingEnabled) { + val map = mutableMapOf() + map["event_type"] = AccessibilityEvent.eventTypeToString(event.eventType) + map["package_name"] = event.packageName?.toString() ?: "" + map["class_name"] = event.className?.toString() ?: "" + map["text"] = event.text.filterNotNull().joinToString(" ") + map["content_description"] = event.contentDescription?.toString() ?: "" + map["event_time"] = + ((System.currentTimeMillis() - android.os.SystemClock.uptimeMillis() + event.eventTime) * + 1000L) + .toString() + + val eventProto = + com.google.androidenv.accessibilityforwarder.EventRequest.newBuilder() + .putAllEvent(map) + .build() + + synchronized(eventBuffer) { + if (eventBuffer.size < 1000) { + eventBuffer.add(eventProto) + } + } + } + + if (event.eventType == AccessibilityEvent.TYPE_WINDOW_STATE_CHANGED) { + val pkg = event.packageName?.toString() + val cls = event.className?.toString() + if (pkg != null) { + activePackage = pkg + } + if (cls != null) { + activeActivity = cls + } + if ( + pkg == "com.android.systemui" && cls?.contains("MediaProjectionPermissionActivity") == true + ) { + autoAcceptMediaProjection() + } + } + } + + fun clearEvents() { + synchronized(eventBuffer) { eventBuffer.clear() } + } + + fun popEvents(): List { + synchronized(eventBuffer) { + val copy = ArrayList(eventBuffer) + eventBuffer.clear() + return copy + } + } + + private fun autoAcceptMediaProjection() { + Log.i(TAG, "Attempting to auto-accept MediaProjection permission dialog...") + val rootNode = rootInActiveWindow + if (rootNode == null) { + Log.w(TAG, "rootInActiveWindow is null") + return + } + + // Search for "Start now" (English) + var nodes = rootNode.findAccessibilityNodeInfosByText("Start now") + if (nodes.isEmpty()) { + // Try "Start recording" + nodes = rootNode.findAccessibilityNodeInfosByText("Start recording") + } + if (nodes.isEmpty()) { + Log.w(TAG, "Could not find 'Start now' or 'Start recording' button") + dumpNodeTree(rootNode, 0) + return + } + + for (node in nodes) { + if (node.isClickable) { + val success = node.performAction(AccessibilityNodeInfo.ACTION_CLICK) + Log.i(TAG, "Clicked auto-accept button: ${node.text}, success=$success") + if (success) return + } else { + var parent = node.parent + while (parent != null) { + if (parent.isClickable) { + val success = parent.performAction(AccessibilityNodeInfo.ACTION_CLICK) + Log.i(TAG, "Clicked parent auto-accept button: success=$success") + if (success) return + } + parent = parent.parent + } + } + } + Log.w(TAG, "Failed to click any matching button") + } + + private fun dumpNodeTree(node: AccessibilityNodeInfo?, depth: Int) { + if (node == null) return + val indent = " ".repeat(depth) + Log.d( + TAG, + "$indent Node: text=${node.text}, class=${node.className}, id=${node.viewIdResourceName}, clickable=${node.isClickable}", + ) + for (i in 0 until node.childCount) { + dumpNodeTree(node.getChild(i), depth + 1) + } + } + + override fun onInterrupt() {} + + fun getRootNode(): AccessibilityNodeInfo? { + return rootInActiveWindow + } + + fun getRootNodes(): List { + val roots = mutableListOf() + val allWindows = + try { + windows + } catch (e: Exception) { + null + } + if (!allWindows.isNullOrEmpty()) { + for (w in allWindows) { + val root = w.root + if (root != null) { + roots.add(root) + } + } + } + if (roots.isEmpty()) { + val activeRoot = rootInActiveWindow + if (activeRoot != null) { + roots.add(activeRoot) + } + } + return roots + } + + fun injectGesture(gesture: GestureDescription, callback: GestureResultCallback? = null) { + dispatchGesture(gesture, callback, /* handler= */ null) + } + + companion object { + private const val TAG = "AndroidEnvInputService" + + @Volatile private var instance: InputService? = null + + @com.google.errorprone.annotations.CanIgnoreReturnValue + fun getInstance(): InputService { + return instance + ?: throw IllegalStateException( + "InputService not connected. Enable it in Accessibility Settings." + ) + } + + var activePackage: String = "" + private set + + var activeActivity: String = "" + private set + } +} diff --git a/android_env/apps/java/com/google/androidenv/core/InputServiceTest.kt b/android_env/apps/java/com/google/androidenv/core/InputServiceTest.kt new file mode 100644 index 00000000..b9731459 --- /dev/null +++ b/android_env/apps/java/com/google/androidenv/core/InputServiceTest.kt @@ -0,0 +1,252 @@ +// Copyright 2026 DeepMind Technologies Limited. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package com.google.androidenv.core + +import android.accessibilityservice.GestureDescription +import android.graphics.Path +import android.view.accessibility.AccessibilityEvent +import android.view.accessibility.AccessibilityNodeInfo +import androidx.test.ext.junit.runners.AndroidJUnit4 +import com.google.common.truth.Truth.assertThat +import org.junit.After +import org.junit.Before +import org.junit.Test +import org.junit.runner.RunWith +import org.robolectric.Robolectric +import org.robolectric.Shadows.shadowOf +import org.robolectric.android.controller.ServiceController +import org.robolectric.shadows.ShadowAccessibilityService + +@RunWith(AndroidJUnit4::class) +class InputServiceTest { + + private lateinit var serviceController: ServiceController + private lateinit var service: InputService + private lateinit var shadowService: ShadowAccessibilityService + + @Before + fun setUp() { + serviceController = Robolectric.buildService(InputService::class.java) + service = serviceController.create().get() + shadowService = shadowOf(service) + } + + @After + fun tearDown() { + serviceController.destroy() + // Clean up singleton + val field = InputService::class.java.getDeclaredField("instance") + field.isAccessible = true + field.set(null, null) + + // Clean up companion variables + val activePackageField = InputService::class.java.getDeclaredField("activePackage") + activePackageField.isAccessible = true + activePackageField.set(null, "") + + val activeActivityField = InputService::class.java.getDeclaredField("activeActivity") + activeActivityField.isAccessible = true + activeActivityField.set(null, "") + } + + @Test + fun onCreate_setsSingletonInstance() { + assertThat(InputService.getInstance()).isEqualTo(service) + } + + @Test + fun injectGesture_dispatchesGestureToSystem() { + val path = + Path().apply { + moveTo(100f, 200f) + lineTo(300f, 400f) + } + val stroke = GestureDescription.StrokeDescription(path, 0, 100) + val gesture = GestureDescription.Builder().addStroke(stroke).build() + + service.injectGesture(gesture) + + val dispatched = shadowService.gesturesDispatched + assertThat(dispatched).hasSize(1) + val dispatchedGesture = dispatched[0].description() + assertThat(dispatchedGesture.strokeCount).isEqualTo(1) + + val dispatchedStroke = dispatchedGesture.getStroke(0) + assertThat(dispatchedStroke.duration).isEqualTo(100) + } + + @Test + fun getRootNode_returnsRootFromActiveWindow() { + val rootNode = + AccessibilityNodeInfo().apply { + packageName = "com.test.package" + className = "android.widget.FrameLayout" + } + + // Direct setup using shadow + shadowService.setRootInActiveWindow(rootNode) + + val root = service.getRootNode() + assertThat(root).isNotNull() + assertThat(root?.packageName).isEqualTo("com.test.package") + assertThat(root?.className).isEqualTo("android.widget.FrameLayout") + } + + @Test + fun onAccessibilityEvent_windowStateChanged_updatesActivePackageAndActivity() { + val event = + AccessibilityEvent().apply { + eventType = AccessibilityEvent.TYPE_WINDOW_STATE_CHANGED + packageName = "com.example.app" + className = "com.example.app.MainActivity" + } + + service.onAccessibilityEvent(event) + + assertThat(InputService.activePackage).isEqualTo("com.example.app") + assertThat(InputService.activeActivity).isEqualTo("com.example.app.MainActivity") + } + + @Test + fun onServiceConnected_setsSingletonInstance() { + val field = InputService::class.java.getDeclaredField("instance") + field.isAccessible = true + field.set(null, null) + + val method = InputService::class.java.getDeclaredMethod("onServiceConnected") + method.isAccessible = true + method.invoke(service) + + assertThat(InputService.getInstance()).isEqualTo(service) + } + + @Test + fun onInterrupt_doesNothing() { + service.onInterrupt() + } + + @Test(expected = IllegalStateException::class) + @Suppress("CheckReturnValue") + fun getInstance_notConnected_throwsException() { + val field = InputService::class.java.getDeclaredField("instance") + field.isAccessible = true + field.set(null, null) + + InputService.getInstance() + } + + @Test + fun onAccessibilityEvent_whenTrackingEnabled_buffersEvent() { + service.isEventTrackingEnabled = true + val event = + AccessibilityEvent().apply { + eventType = AccessibilityEvent.TYPE_VIEW_CLICKED + packageName = "com.test.app" + className = "android.widget.Button" + eventTime = 12345L + } + + service.onAccessibilityEvent(event) + + val events = service.popEvents() + assertThat(events).hasSize(1) + assertThat(events[0].eventMap["event_type"]).contains("TYPE_VIEW_CLICKED") + assertThat(events[0].eventMap["package_name"]).isEqualTo("com.test.app") + assertThat(events[0].eventMap["class_name"]).isEqualTo("android.widget.Button") + + // Buffer cleared after popEvents + assertThat(service.popEvents()).isEmpty() + } + + @Test + fun onAccessibilityEvent_whenTrackingDisabled_doesNotBufferEvent() { + service.isEventTrackingEnabled = false + val event = + AccessibilityEvent().apply { + eventType = AccessibilityEvent.TYPE_VIEW_CLICKED + packageName = "com.test.app" + className = "android.widget.Button" + eventTime = 12345L + } + + service.onAccessibilityEvent(event) + + val events = service.popEvents() + assertThat(events).isEmpty() + } + + @Test + fun clearEvents_clearsBufferedEvents() { + service.isEventTrackingEnabled = true + val event = + AccessibilityEvent().apply { + eventType = AccessibilityEvent.TYPE_VIEW_CLICKED + packageName = "com.test.app" + className = "android.widget.Button" + } + service.onAccessibilityEvent(event) + + service.clearEvents() + assertThat(service.popEvents()).isEmpty() + } + + @Test + fun systemGestureInjector_allMethods_whenBoundAndUnbound() { + val injector = SystemGestureInjector() + + // 1. When bound + assertThat(injector.getRootNode()).isNull() + assertThat(injector.getActivePackage()).isEqualTo("") + assertThat(injector.getActiveActivity()).isEqualTo("") + assertThat(injector.performGlobalAction(1)).isTrue() + try { + val stroke = + GestureDescription.StrokeDescription( + android.graphics.Path().apply { lineTo(10f, 10f) }, + 0L, + 100L, + ) + val gesture = GestureDescription.Builder().addStroke(stroke).build() + injector.injectGesture(gesture, null) + } catch (e: Exception) { + // Ignored if service.dispatchGesture fails in unit test environment + } + injector.setEventTrackingEnabled(true) + val event = + AccessibilityEvent().apply { + eventType = AccessibilityEvent.TYPE_VIEW_CLICKED + packageName = "com.test.app" + className = "android.widget.Button" + } + service.onAccessibilityEvent(event) + + val events = injector.popEvents() + assertThat(events).hasSize(1) + + injector.setEventTrackingEnabled(true) + service.onAccessibilityEvent(event) + injector.clearEvents() + assertThat(injector.popEvents()).isEmpty() + + // 2. When unbound (test catch blocks) + val field = InputService::class.java.getDeclaredField("instance") + field.isAccessible = true + field.set(null, null) + + assertThat(injector.popEvents()).isEmpty() + injector.setEventTrackingEnabled(true) + injector.clearEvents() + } +} diff --git a/android_env/apps/java/com/google/androidenv/core/MainActivity.kt b/android_env/apps/java/com/google/androidenv/core/MainActivity.kt new file mode 100644 index 00000000..ff04768b --- /dev/null +++ b/android_env/apps/java/com/google/androidenv/core/MainActivity.kt @@ -0,0 +1,267 @@ +// Copyright 2026 DeepMind Technologies Limited. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package com.google.androidenv.core + +import android.accessibilityservice.AccessibilityService +import android.annotation.SuppressLint +import android.app.Activity +import android.content.Context +import android.content.Intent +import android.media.projection.MediaProjectionManager +import android.os.Bundle +import android.provider.Settings +import android.util.Log +import android.widget.Button +import android.widget.LinearLayout +import android.widget.TextView +import androidx.activity.ComponentActivity +import androidx.activity.result.ActivityResult +import androidx.activity.result.contract.ActivityResultContracts + +/** + * Main activity that orchestrates permissions (Accessibility and MediaProjection) and starts either + * the WebSocket server (PoC 1) or the In-Process Agent (PoC 2). Can be controlled via Intents for + * headless automation. + */ +@SuppressLint("SetTextI18n") +class MainActivity : ComponentActivity() { + + private lateinit var statusTextView: TextView + private lateinit var serverButton: Button + + private enum class PendingAction { + NONE, + START_SERVER, + } + + private var pendingAction = PendingAction.NONE + private var isResumed = false + private var pendingCommand: String? = null + + private val screenCaptureLauncher = + registerForActivityResult(ActivityResultContracts.StartActivityForResult()) { result -> + onScreenCaptureResult(result) + } + + override fun onCreate(savedInstanceState: Bundle?) { + Log.e("CRITICAL_DEBUG", "MainActivity onCreate called!!!") + super.onCreate(savedInstanceState) + + if (savedInstanceState != null) { + val actionName = savedInstanceState.getString("PENDING_ACTION") + if (actionName != null) { + pendingAction = PendingAction.valueOf(actionName) + Log.e("CRITICAL_DEBUG", "Restored pendingAction: $pendingAction") + } + } + + // Programmatic UI + val layout = + LinearLayout(this).apply { + orientation = LinearLayout.VERTICAL + setPadding(50, 50, 50, 50) + } + statusTextView = + TextView(this).apply { + text = "Status: Idle" + textSize = 18f + layoutParams = + LinearLayout.LayoutParams( + LinearLayout.LayoutParams.MATCH_PARENT, + LinearLayout.LayoutParams.WRAP_CONTENT, + ) + .apply { setMargins(0, 0, 0, 50) } + } + serverButton = + Button(this).apply { + text = "Start WebSocket Server" + layoutParams = + LinearLayout.LayoutParams( + LinearLayout.LayoutParams.MATCH_PARENT, + LinearLayout.LayoutParams.WRAP_CONTENT, + ) + .apply { setMargins(0, 0, 0, 20) } + setOnClickListener { triggerAction(PendingAction.START_SERVER) } + } + layout.addView(statusTextView) + layout.addView(serverButton) + setContentView(layout) + + handleIntent(intent) + } + + override fun onResume() { + super.onResume() + isResumed = true + updateAccessibilityStatus() + updateServerUIStatus() + Log.e("CRITICAL_DEBUG", "onResume: pendingCommand=$pendingCommand") + pendingCommand?.let { + executeCommand(it) + pendingCommand = null + } + } + + override fun onPause() { + super.onPause() + isResumed = false + } + + override fun onNewIntent(intent: Intent) { + super.onNewIntent(intent) + handleIntent(intent) + } + + override fun onSaveInstanceState(outState: Bundle) { + super.onSaveInstanceState(outState) + outState.putString("PENDING_ACTION", pendingAction.name) + Log.e("CRITICAL_DEBUG", "Saved pendingAction: ${pendingAction.name}") + } + + private fun handleIntent(intent: Intent?) { + val command = intent?.getStringExtra("COMMAND") + Log.i(TAG, "Received command: $command, isResumed=$isResumed") + if (command != null) { + intent.removeExtra("COMMAND") + if (isResumed) { + executeCommand(command) + } else { + pendingCommand = command + Log.e("CRITICAL_DEBUG", "Set pendingCommand to $command") + } + } + } + + private fun executeCommand(command: String) { + Log.e("CRITICAL_DEBUG", "executeCommand: $command") + when (command) { + "START_SERVER" -> triggerAction(PendingAction.START_SERVER) + "STOP" -> stopWebSocketServer() + } + } + + private fun updateAccessibilityStatus() { + val enabled = isAccessibilityServiceEnabled(this, InputService::class.java) + if (!enabled) { + statusTextView.text = + "Status: Accessibility Service disabled.\nPlease enable 'AndroidEnv Input Service' in Settings." + serverButton.isEnabled = false + } else { + statusTextView.text = "Status: Ready" + serverButton.isEnabled = true + } + } + + private fun updateServerUIStatus() { + if (!isAccessibilityServiceEnabled(this, InputService::class.java)) { + return + } + val isRunning = + try { + ScreenShareService.getInstance().isServerRunning() + } catch (e: IllegalStateException) { + false + } + if (isRunning) { + statusTextView.text = + "Server: Running on ws://localhost:8000\nAccessibility: ACTIVE\nScreenCapture: ACTIVE" + serverButton.text = "Stop WebSocket Server" + serverButton.setOnClickListener { stopWebSocketServer() } + } else { + statusTextView.text = "Status: Ready" + serverButton.text = "Start WebSocket Server" + serverButton.setOnClickListener { triggerAction(PendingAction.START_SERVER) } + } + } + + private fun triggerAction(action: PendingAction) { + if (!isAccessibilityServiceEnabled(this, InputService::class.java)) { + Log.w(TAG, "Cannot trigger action: Accessibility disabled.") + return + } + + Log.e("CRITICAL_DEBUG", "triggerAction: action=$action, current pendingAction=$pendingAction") + pendingAction = action + val manager = getSystemService(Context.MEDIA_PROJECTION_SERVICE) as MediaProjectionManager + screenCaptureLauncher.launch(manager.createScreenCaptureIntent()) + } + + private fun onScreenCaptureResult(result: ActivityResult) { + Log.e( + "CRITICAL_DEBUG", + "onScreenCaptureResult: result=${result.resultCode}, pendingAction=$pendingAction", + ) + if (result.resultCode != Activity.RESULT_OK) { + statusTextView.text = "Status: Screen Capture permission denied." + pendingAction = PendingAction.NONE + return + } + val data = result.data ?: return + + val serviceIntent = + Intent(this, ScreenShareService::class.java).apply { + putExtra("ACTION", "START") + putExtra("DATA", data) + putExtra("RESULT_CODE", result.resultCode) + } + try { + startForegroundService(serviceIntent) + statusTextView.post { updateServerUIStatus() } + } catch (e: Exception) { + statusTextView.text = "Error: Failed to start screen capture: ${e.message}" + Log.e(TAG, "Failed to start ScreenShareService", e) + } finally { + pendingAction = PendingAction.NONE + } + } + + private fun stopWebSocketServer() { + stopScreenShareService() + statusTextView.text = "Status: Ready" + serverButton.text = "Start WebSocket Server" + serverButton.setOnClickListener { triggerAction(PendingAction.START_SERVER) } + } + + private fun stopScreenShareService() { + val serviceIntent = + Intent(this, ScreenShareService::class.java).apply { putExtra("ACTION", "STOP") } + startService(serviceIntent) + } + + private fun isAccessibilityServiceEnabled( + context: Context, + service: Class, + ): Boolean { + val expectedId = "${context.packageName}/${service.name}" + val enabledServicesSetting = + Settings.Secure.getString( + context.contentResolver, + Settings.Secure.ENABLED_ACCESSIBILITY_SERVICES, + ) ?: return false + val colonSplitter = android.text.TextUtils.SimpleStringSplitter(':') + colonSplitter.setString(enabledServicesSetting) + while (colonSplitter.hasNext()) { + val componentNameString = colonSplitter.next() + if (componentNameString.equals(expectedId, ignoreCase = true)) { + return true + } + } + return false + } + + companion object { + private const val TAG = "MainActivity" + } +} diff --git a/android_env/apps/java/com/google/androidenv/core/MainActivityTest.kt b/android_env/apps/java/com/google/androidenv/core/MainActivityTest.kt new file mode 100644 index 00000000..40e09070 --- /dev/null +++ b/android_env/apps/java/com/google/androidenv/core/MainActivityTest.kt @@ -0,0 +1,214 @@ +// Copyright 2026 DeepMind Technologies Limited. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package com.google.androidenv.core + +import android.app.Activity +import android.content.Context +import android.content.Intent +import android.provider.Settings +import android.view.ViewGroup +import android.widget.Button +import android.widget.LinearLayout +import android.widget.TextView +import androidx.test.core.app.ActivityScenario +import androidx.test.core.app.ApplicationProvider +import androidx.test.ext.junit.runners.AndroidJUnit4 +import com.google.common.truth.Truth.assertThat +import org.junit.After +import org.junit.Before +import org.junit.Test +import org.junit.runner.RunWith +import org.robolectric.Robolectric +import org.robolectric.Shadows.shadowOf +import org.robolectric.annotation.Config + +@RunWith(AndroidJUnit4::class) +@Config(sdk = [26]) +class MainActivityTest { + + private val context = ApplicationProvider.getApplicationContext() + + @Before + fun setUp() { + setAccessibilityEnabled(false) + } + + @After + fun tearDown() { + setAccessibilityEnabled(false) + } + + private fun setAccessibilityEnabled(enabled: Boolean) { + val value = + if (enabled) { + "${context.packageName}/${InputService::class.java.name}" + } else { + "" + } + Settings.Secure.putString( + context.contentResolver, + Settings.Secure.ENABLED_ACCESSIBILITY_SERVICES, + value, + ) + } + + @Test + fun onCreate_accessibilityDisabled_showsDisabledStatusAndDisablesButton() { + ActivityScenario.launch(MainActivity::class.java).use { scenario -> + scenario.onActivity { activity -> + val content = activity.findViewById(android.R.id.content) + val layout = content.getChildAt(0) as LinearLayout + val statusTextView = layout.getChildAt(0) as TextView + val serverButton = layout.getChildAt(1) as Button + + assertThat(statusTextView.text.toString()).contains("Accessibility Service disabled") + assertThat(serverButton.isEnabled).isFalse() + } + } + } + + @Test + fun onCreate_accessibilityEnabled_showsReadyStatusAndEnablesButton() { + setAccessibilityEnabled(true) + ActivityScenario.launch(MainActivity::class.java).use { scenario -> + scenario.onActivity { activity -> + val content = activity.findViewById(android.R.id.content) + val layout = content.getChildAt(0) as LinearLayout + val statusTextView = layout.getChildAt(0) as TextView + val serverButton = layout.getChildAt(1) as Button + + assertThat(statusTextView.text.toString()).isEqualTo("Status: Ready") + assertThat(serverButton.isEnabled).isTrue() + assertThat(serverButton.text.toString()).isEqualTo("Start WebSocket Server") + } + } + } + + @Test + fun clickStartServer_launchesScreenCaptureIntent() { + setAccessibilityEnabled(true) + ActivityScenario.launch(MainActivity::class.java).use { scenario -> + scenario.onActivity { activity -> + val content = activity.findViewById(android.R.id.content) + val layout = content.getChildAt(0) as LinearLayout + val serverButton = layout.getChildAt(1) as Button + + serverButton.performClick() + + val shadowActivity = shadowOf(activity) + val startedIntent = shadowActivity.nextStartedActivityForResult + assertThat(startedIntent).isNotNull() + } + } + } + + @Test + fun screenCaptureResult_ok_startsScreenShareService() { + setAccessibilityEnabled(true) + ActivityScenario.launch(MainActivity::class.java).use { scenario -> + scenario.onActivity { activity -> + val content = activity.findViewById(android.R.id.content) + val layout = content.getChildAt(0) as LinearLayout + val serverButton = layout.getChildAt(1) as Button + + serverButton.performClick() + + val shadowActivity = shadowOf(activity) + val startedIntent = shadowActivity.nextStartedActivityForResult + + val resultData = Intent().apply { putExtra("EXTRA_DATA", "dummy") } + shadowActivity.receiveResult(startedIntent.intent, Activity.RESULT_OK, resultData) + + val serviceIntent = shadowActivity.nextStartedService + assertThat(serviceIntent).isNotNull() + assertThat(serviceIntent.component?.className) + .isEqualTo(ScreenShareService::class.java.name) + assertThat(serviceIntent.getStringExtra("ACTION")).isEqualTo("START") + assertThat(serviceIntent.getIntExtra("RESULT_CODE", Activity.RESULT_CANCELED)) + .isEqualTo(Activity.RESULT_OK) + } + } + } + + @Test + fun screenCaptureResult_canceled_doesNotStartServiceAndShowsError() { + setAccessibilityEnabled(true) + ActivityScenario.launch(MainActivity::class.java).use { scenario -> + scenario.onActivity { activity -> + val content = activity.findViewById(android.R.id.content) + val layout = content.getChildAt(0) as LinearLayout + val statusTextView = layout.getChildAt(0) as TextView + val serverButton = layout.getChildAt(1) as Button + + serverButton.performClick() + + val shadowActivity = shadowOf(activity) + val startedIntent = shadowActivity.nextStartedActivityForResult + + shadowActivity.receiveResult(startedIntent.intent, Activity.RESULT_CANCELED, null) + + val serviceIntent = shadowActivity.nextStartedService + assertThat(serviceIntent).isNull() + assertThat(statusTextView.text.toString()).contains("Screen Capture permission denied") + } + } + } + + @Test + fun intentCommand_startServer_launchesScreenCapture() { + setAccessibilityEnabled(true) + val intent = + Intent(context, MainActivity::class.java).apply { putExtra("COMMAND", "START_SERVER") } + ActivityScenario.launch(intent).use { scenario -> + scenario.onActivity { activity -> + val shadowActivity = shadowOf(activity) + val startedIntent = shadowActivity.nextStartedActivityForResult + assertThat(startedIntent).isNotNull() + } + } + } + + @Test + fun intentCommand_stop_stopsService() { + setAccessibilityEnabled(true) + val intent = Intent(context, MainActivity::class.java).apply { putExtra("COMMAND", "STOP") } + ActivityScenario.launch(intent).use { scenario -> + scenario.onActivity { activity -> + val shadowActivity = shadowOf(activity) + val serviceIntent = shadowActivity.nextStartedService + assertThat(serviceIntent).isNotNull() + assertThat(serviceIntent.component?.className) + .isEqualTo(ScreenShareService::class.java.name) + assertThat(serviceIntent.getStringExtra("ACTION")).isEqualTo("STOP") + } + } + } + + @Test + fun onNewIntent_handlesCommand() { + setAccessibilityEnabled(true) + val controller = Robolectric.buildActivity(MainActivity::class.java) + val activity = controller.create().resume().get() + + val newIntent = Intent(activity, MainActivity::class.java).apply { putExtra("COMMAND", "STOP") } + controller.newIntent(newIntent) + + val shadowActivity = shadowOf(activity) + val serviceIntent = shadowActivity.nextStartedService + assertThat(serviceIntent).isNotNull() + assertThat(serviceIntent.getStringExtra("ACTION")).isEqualTo("STOP") + controller.destroy() + } +} diff --git a/android_env/apps/java/com/google/androidenv/core/MediaProjectionWrapper.kt b/android_env/apps/java/com/google/androidenv/core/MediaProjectionWrapper.kt new file mode 100644 index 00000000..ceb663c0 --- /dev/null +++ b/android_env/apps/java/com/google/androidenv/core/MediaProjectionWrapper.kt @@ -0,0 +1,107 @@ +// Copyright 2026 DeepMind Technologies Limited. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package com.google.androidenv.core + +import android.content.Intent +import android.hardware.display.VirtualDisplay +import android.media.projection.MediaProjection +import android.media.projection.MediaProjectionManager +import android.os.Handler +import android.view.Surface + +/** Wrapper for [VirtualDisplay] to facilitate testing. */ +interface VirtualDisplayWrapper { + fun release() +} + +/** Real implementation of [VirtualDisplayWrapper]. */ +class RealVirtualDisplayWrapper(private val virtualDisplay: VirtualDisplay) : + VirtualDisplayWrapper { + override fun release() { + virtualDisplay.release() + } +} + +/** Wrapper for [MediaProjection] to facilitate testing. */ +interface MediaProjectionWrapper { + fun registerCallback(callback: MediaProjection.Callback, handler: Handler?) + + fun createVirtualDisplay( + name: String, + width: Int, + height: Int, + dpi: Int, + flags: Int, + surface: Surface?, + callback: VirtualDisplay.Callback?, + handler: Handler?, + ): VirtualDisplayWrapper? + + fun stop() + + fun getActualMediaProjection(): MediaProjection +} + +/** Real implementation of [MediaProjectionWrapper]. */ +class RealMediaProjectionWrapper(private val mediaProjection: MediaProjection) : + MediaProjectionWrapper { + override fun registerCallback(callback: MediaProjection.Callback, handler: Handler?) { + mediaProjection.registerCallback(callback, handler) + } + + override fun createVirtualDisplay( + name: String, + width: Int, + height: Int, + dpi: Int, + flags: Int, + surface: Surface?, + callback: VirtualDisplay.Callback?, + handler: Handler?, + ): VirtualDisplayWrapper? { + val vd = + mediaProjection.createVirtualDisplay( + name, + width, + height, + dpi, + flags, + surface, + callback, + handler, + ) + return vd?.let { RealVirtualDisplayWrapper(it) } + } + + override fun stop() { + mediaProjection.stop() + } + + override fun getActualMediaProjection(): MediaProjection = mediaProjection +} + +/** Wrapper for [MediaProjectionManager] to facilitate testing. */ +interface MediaProjectionManagerWrapper { + fun getMediaProjection(resultCode: Int, resultData: Intent): MediaProjectionWrapper? +} + +/** Real implementation of [MediaProjectionManagerWrapper]. */ +class RealMediaProjectionManagerWrapper(private val manager: MediaProjectionManager) : + MediaProjectionManagerWrapper { + override fun getMediaProjection(resultCode: Int, resultData: Intent): MediaProjectionWrapper? { + val mp = manager.getMediaProjection(resultCode, resultData) + return mp?.let { RealMediaProjectionWrapper(it) } + } +} diff --git a/android_env/apps/java/com/google/androidenv/core/ProtoConversions.kt b/android_env/apps/java/com/google/androidenv/core/ProtoConversions.kt new file mode 100644 index 00000000..73d63d04 --- /dev/null +++ b/android_env/apps/java/com/google/androidenv/core/ProtoConversions.kt @@ -0,0 +1,372 @@ +// Copyright 2026 DeepMind Technologies Limited. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package com.google.androidenv.core + +import android.media.MediaFormat +import android.view.accessibility.AccessibilityNodeInfo +import com.androidenv.proto.Action as ProtoAction +import com.androidenv.proto.ActionType as ProtoActionType +import com.androidenv.proto.AudioFrame as ProtoAudioFrame +import com.androidenv.proto.DeviceSignal as ProtoDeviceSignal +import com.androidenv.proto.DeviceState as ProtoDeviceState +import com.androidenv.proto.FloatArray as ProtoFloatArray +import com.androidenv.proto.GpsCoordinates as ProtoGpsCoordinates +import com.androidenv.proto.StreamedSignalMetadata as ProtoMetadata +import com.androidenv.proto.StringList as ProtoStringList +import com.androidenv.proto.StringMap as ProtoStringMap + +/** Utility functions to convert between Core model classes and Proto classes. */ +object ProtoConversions { + + fun ProtoAction.toCore(): Action { + val coreType = + when (actionType) { + ProtoActionType.ACTION_TYPE_TOUCH_DOWN -> ActionType.TOUCH_DOWN + ProtoActionType.ACTION_TYPE_TOUCH_UP -> ActionType.TOUCH_UP + ProtoActionType.ACTION_TYPE_TOUCH_MOVE -> ActionType.TOUCH_MOVE + ProtoActionType.ACTION_TYPE_KEY_EVENT -> ActionType.KEY_EVENT + ProtoActionType.ACTION_TYPE_SET_TEXT -> ActionType.SET_TEXT + ProtoActionType.ACTION_TYPE_SET_SELECTION -> ActionType.SET_SELECTION + ProtoActionType.ACTION_TYPE_CUT -> ActionType.CUT + ProtoActionType.ACTION_TYPE_COPY -> ActionType.COPY + ProtoActionType.ACTION_TYPE_PASTE -> ActionType.PASTE + else -> ActionType.KEY_EVENT // Default fallback + } + val pos = + if (hasTouchPosition()) { + listOf(touchPosition.x, touchPosition.y) + } else { + null + } + val key = if (hasKeycode()) keycode else null + val pointer = if (hasPointerId()) pointerId else null + val textVal = if (hasText()) text else null + val selStart = if (hasSelectionStart()) selectionStart else null + val selEnd = if (hasSelectionEnd()) selectionEnd else null + return Action(coreType, pos, key, pointer, textVal, selStart, selEnd) + } + + fun ProtoDeviceSignal.toCore(): DeviceSignal? { + return when (this) { + ProtoDeviceSignal.DEVICE_SIGNAL_SCREENSHOT -> DeviceSignal.SCREENSHOT + ProtoDeviceSignal.DEVICE_SIGNAL_ACCESSIBILITY_TREE -> DeviceSignal.ACCESSIBILITY_TREE + ProtoDeviceSignal.DEVICE_SIGNAL_ACCESSIBILITY_EVENTS -> DeviceSignal.ACCESSIBILITY_EVENTS + ProtoDeviceSignal.DEVICE_SIGNAL_ACTIVE_PACKAGE -> DeviceSignal.ACTIVE_PACKAGE + ProtoDeviceSignal.DEVICE_SIGNAL_ACTIVE_ACTIVITY -> DeviceSignal.ACTIVE_ACTIVITY + ProtoDeviceSignal.DEVICE_SIGNAL_SYSTEM_LOGS -> DeviceSignal.SYSTEM_LOGS + ProtoDeviceSignal.DEVICE_SIGNAL_RUNNING_PROCESSES -> DeviceSignal.RUNNING_PROCESSES + ProtoDeviceSignal.DEVICE_SIGNAL_DEVICE_SETTINGS -> DeviceSignal.DEVICE_SETTINGS + ProtoDeviceSignal.DEVICE_SIGNAL_SENSOR_ACCELEROMETER -> DeviceSignal.SENSOR_ACCELEROMETER + ProtoDeviceSignal.DEVICE_SIGNAL_SENSOR_GYROSCOPE -> DeviceSignal.SENSOR_GYROSCOPE + ProtoDeviceSignal.DEVICE_SIGNAL_SENSOR_MAGNETOMETER -> DeviceSignal.SENSOR_MAGNETOMETER + ProtoDeviceSignal.DEVICE_SIGNAL_SENSOR_BAROMETER -> DeviceSignal.SENSOR_BAROMETER + ProtoDeviceSignal.DEVICE_SIGNAL_SENSOR_AMBIENT_TEMP -> DeviceSignal.SENSOR_AMBIENT_TEMP + ProtoDeviceSignal.DEVICE_SIGNAL_SENSOR_LIGHT -> DeviceSignal.SENSOR_LIGHT + ProtoDeviceSignal.DEVICE_SIGNAL_SENSOR_PROXIMITY -> DeviceSignal.SENSOR_PROXIMITY + ProtoDeviceSignal.DEVICE_SIGNAL_SENSOR_HEART_RATE -> DeviceSignal.SENSOR_HEART_RATE + ProtoDeviceSignal.DEVICE_SIGNAL_LOCATION_GPS -> DeviceSignal.LOCATION_GPS + ProtoDeviceSignal.DEVICE_SIGNAL_AUDIO_OUTPUT -> DeviceSignal.AUDIO_OUTPUT + ProtoDeviceSignal.DEVICE_SIGNAL_BATTERY_POWER -> DeviceSignal.BATTERY_POWER + else -> null + } + } + + fun DeviceSignal.toProto(): ProtoDeviceSignal { + return when (this) { + DeviceSignal.SCREENSHOT -> ProtoDeviceSignal.DEVICE_SIGNAL_SCREENSHOT + DeviceSignal.ACCESSIBILITY_TREE -> ProtoDeviceSignal.DEVICE_SIGNAL_ACCESSIBILITY_TREE + DeviceSignal.ACCESSIBILITY_EVENTS -> ProtoDeviceSignal.DEVICE_SIGNAL_ACCESSIBILITY_EVENTS + DeviceSignal.ACTIVE_PACKAGE -> ProtoDeviceSignal.DEVICE_SIGNAL_ACTIVE_PACKAGE + DeviceSignal.ACTIVE_ACTIVITY -> ProtoDeviceSignal.DEVICE_SIGNAL_ACTIVE_ACTIVITY + DeviceSignal.SYSTEM_LOGS -> ProtoDeviceSignal.DEVICE_SIGNAL_SYSTEM_LOGS + DeviceSignal.RUNNING_PROCESSES -> ProtoDeviceSignal.DEVICE_SIGNAL_RUNNING_PROCESSES + DeviceSignal.DEVICE_SETTINGS -> ProtoDeviceSignal.DEVICE_SIGNAL_DEVICE_SETTINGS + DeviceSignal.SENSOR_ACCELEROMETER -> ProtoDeviceSignal.DEVICE_SIGNAL_SENSOR_ACCELEROMETER + DeviceSignal.SENSOR_GYROSCOPE -> ProtoDeviceSignal.DEVICE_SIGNAL_SENSOR_GYROSCOPE + DeviceSignal.SENSOR_MAGNETOMETER -> ProtoDeviceSignal.DEVICE_SIGNAL_SENSOR_MAGNETOMETER + DeviceSignal.SENSOR_BAROMETER -> ProtoDeviceSignal.DEVICE_SIGNAL_SENSOR_BAROMETER + DeviceSignal.SENSOR_AMBIENT_TEMP -> ProtoDeviceSignal.DEVICE_SIGNAL_SENSOR_AMBIENT_TEMP + DeviceSignal.SENSOR_LIGHT -> ProtoDeviceSignal.DEVICE_SIGNAL_SENSOR_LIGHT + DeviceSignal.SENSOR_PROXIMITY -> ProtoDeviceSignal.DEVICE_SIGNAL_SENSOR_PROXIMITY + DeviceSignal.SENSOR_HEART_RATE -> ProtoDeviceSignal.DEVICE_SIGNAL_SENSOR_HEART_RATE + DeviceSignal.LOCATION_GPS -> ProtoDeviceSignal.DEVICE_SIGNAL_LOCATION_GPS + DeviceSignal.AUDIO_OUTPUT -> ProtoDeviceSignal.DEVICE_SIGNAL_AUDIO_OUTPUT + DeviceSignal.BATTERY_POWER -> ProtoDeviceSignal.DEVICE_SIGNAL_BATTERY_POWER + } + } + + fun DeviceState.toProto(): ProtoDeviceState { + val builder = + ProtoDeviceState.newBuilder() + .setSnapshotTimestampUs(snapshotTimestampUs) + .setOrientation(orientation) + + fun addSignalIfPresent( + signal: StreamedSignal, + signalType: ProtoDeviceSignal, + valueBuilder: (T, com.androidenv.proto.StreamedSignal.Builder) -> Unit, + ) { + val rawValue = signal.value + if (rawValue != null) { + val signalBuilder = + com.androidenv.proto.StreamedSignal.newBuilder() + .setType(signalType) + .setMetadata( + ProtoMetadata.newBuilder() + .setTimestampUs(signal.timestampUs) + .setSequenceId(signal.sequenceId) + .build() + ) + valueBuilder(rawValue, signalBuilder) + builder.addSignals(signalBuilder.build()) + } + } + + // 1. Screenshot (Compressed JPEG bytes included directly in StreamedSignal over WebSocket) + val rawScreenshot = screenshot.value + if (rawScreenshot != null) { + val stream = java.io.ByteArrayOutputStream() + rawScreenshot.compress(android.graphics.Bitmap.CompressFormat.JPEG, 80, stream) + val jpegBytes = stream.toByteArray() + + val screenshotVal = + com.androidenv.proto.ScreenshotValue.newBuilder() + .setEncodedBytes(com.google.protobuf.ByteString.copyFrom(jpegBytes)) + .build() + + val signalBuilder = + com.androidenv.proto.StreamedSignal.newBuilder() + .setType(ProtoDeviceSignal.DEVICE_SIGNAL_SCREENSHOT) + .setMetadata( + ProtoMetadata.newBuilder() + .setTimestampUs(screenshot.timestampUs) + .setSequenceId(screenshot.sequenceId) + .build() + ) + .setScreenshot(screenshotVal) + builder.addSignals(signalBuilder.build()) + } + + // 2. Accessibility Tree + addSignalIfPresent(accessibilityRoot, ProtoDeviceSignal.DEVICE_SIGNAL_ACCESSIBILITY_TREE) { + value, + signalBuilder -> + val treeProto = + com.google.androidenv.accessibilityforwarder.AndroidAccessibilityTree.newBuilder() + .addNodes(value.toA11yNodeInfoProto()) + .build() + val windowInfo = + com.google.androidenv.accessibilityforwarder.AndroidAccessibilityWindowInfo.newBuilder() + .setTree(treeProto) + .build() + signalBuilder.setAccessibilityForest( + com.google.androidenv.accessibilityforwarder.AndroidAccessibilityForest.newBuilder() + .addWindows(windowInfo) + .build() + ) + } + + // 3. Accessibility Events + addSignalIfPresent(accessibilityEvents, ProtoDeviceSignal.DEVICE_SIGNAL_ACCESSIBILITY_EVENTS) { + value, + signalBuilder -> + signalBuilder.setAccessibilityEvents( + com.google.androidenv.accessibilityforwarder.EventRequestList.newBuilder() + .addAllEvents(value) + .build() + ) + } + + // 3. Active Package + addSignalIfPresent(activePackage, ProtoDeviceSignal.DEVICE_SIGNAL_ACTIVE_PACKAGE) { + value, + signalBuilder -> + signalBuilder.setActivePackage(value) + } + + // 4. Active Activity + addSignalIfPresent(activeActivity, ProtoDeviceSignal.DEVICE_SIGNAL_ACTIVE_ACTIVITY) { + value, + signalBuilder -> + signalBuilder.setActiveActivity(value) + } + + // 5. System Logs + addSignalIfPresent(systemLogs, ProtoDeviceSignal.DEVICE_SIGNAL_SYSTEM_LOGS) { + value, + signalBuilder -> + signalBuilder.setSystemLogs(ProtoStringList.newBuilder().addAllValues(value).build()) + } + + // 6. Running Processes + addSignalIfPresent(runningProcesses, ProtoDeviceSignal.DEVICE_SIGNAL_RUNNING_PROCESSES) { + value, + signalBuilder -> + signalBuilder.setRunningProcesses(ProtoStringList.newBuilder().addAllValues(value).build()) + } + + // 7. Device Settings + val settingsMap = deviceSettings.value?.toMutableMap() ?: mutableMapOf() + settingsMap["device_time_ms"] = System.currentTimeMillis().toString() + val screenService = + try { + ScreenShareService.getInstance() + } catch (e: IllegalStateException) { + null + } + if (screenService != null) { + val mime = screenService.activeVideoMime + val codec = + when (mime) { + MediaFormat.MIMETYPE_VIDEO_HEVC -> "h265" + "image/jpeg" -> "mjpeg" + else -> "h264" + } + settingsMap["video_codec"] = codec + val width = if (screenService.activeVideoWidth > 0) screenService.activeVideoWidth else 1080 + val height = + if (screenService.activeVideoHeight > 0) screenService.activeVideoHeight else 1920 + settingsMap["video_width"] = width.toString() + settingsMap["video_height"] = height.toString() + } + val signalBuilder = + com.androidenv.proto.StreamedSignal.newBuilder() + .setType(ProtoDeviceSignal.DEVICE_SIGNAL_DEVICE_SETTINGS) + .setMetadata( + ProtoMetadata.newBuilder() + .setTimestampUs( + if (deviceSettings.timestampUs > 0) deviceSettings.timestampUs + else snapshotTimestampUs + ) + .setSequenceId(deviceSettings.sequenceId) + .build() + ) + .setDeviceSettings(ProtoStringMap.newBuilder().putAllEntries(settingsMap).build()) + builder.addSignals(signalBuilder.build()) + + // Sensors + addSignalIfPresent(accelerometer, ProtoDeviceSignal.DEVICE_SIGNAL_SENSOR_ACCELEROMETER) { + value, + signalBuilder -> + signalBuilder.setAccelerometer(ProtoFloatArray.newBuilder().addAllValues(value).build()) + } + addSignalIfPresent(gyroscope, ProtoDeviceSignal.DEVICE_SIGNAL_SENSOR_GYROSCOPE) { + value, + signalBuilder -> + signalBuilder.setGyroscope(ProtoFloatArray.newBuilder().addAllValues(value).build()) + } + addSignalIfPresent(magnetometer, ProtoDeviceSignal.DEVICE_SIGNAL_SENSOR_MAGNETOMETER) { + value, + signalBuilder -> + signalBuilder.setMagnetometer(ProtoFloatArray.newBuilder().addAllValues(value).build()) + } + addSignalIfPresent(barometer, ProtoDeviceSignal.DEVICE_SIGNAL_SENSOR_BAROMETER) { + value, + signalBuilder -> + signalBuilder.setBarometer(value) + } + addSignalIfPresent(ambientTemp, ProtoDeviceSignal.DEVICE_SIGNAL_SENSOR_AMBIENT_TEMP) { + value, + signalBuilder -> + signalBuilder.setAmbientTemp(value) + } + addSignalIfPresent(lightLevel, ProtoDeviceSignal.DEVICE_SIGNAL_SENSOR_LIGHT) { + value, + signalBuilder -> + signalBuilder.setLightLevel(value) + } + addSignalIfPresent(proximity, ProtoDeviceSignal.DEVICE_SIGNAL_SENSOR_PROXIMITY) { + value, + signalBuilder -> + signalBuilder.setProximity(value) + } + + addSignalIfPresent(locationGps, ProtoDeviceSignal.DEVICE_SIGNAL_LOCATION_GPS) { + value, + signalBuilder -> + signalBuilder.setLocationGps( + ProtoGpsCoordinates.newBuilder() + .setLatitude(value.latitude) + .setLongitude(value.longitude) + .setAltitude(value.altitude) + .build() + ) + } + + addSignalIfPresent(heartRate, ProtoDeviceSignal.DEVICE_SIGNAL_SENSOR_HEART_RATE) { + value, + signalBuilder -> + signalBuilder.setHeartRate(value) + } + addSignalIfPresent(audioOutput, ProtoDeviceSignal.DEVICE_SIGNAL_AUDIO_OUTPUT) { + value, + signalBuilder -> + val audioBuilder = + ProtoAudioFrame.newBuilder() + .setRawBytes(com.google.protobuf.ByteString.copyFrom(value.rawBytes)) + .setSampleRateHz(value.sampleRateHz) + .setChannelCount(value.channelCount) + value.sampleFormat?.let { audioBuilder.setSampleFormat(it) } + signalBuilder.setAudioOutput(audioBuilder.build()) + } + addSignalIfPresent(batteryPower, ProtoDeviceSignal.DEVICE_SIGNAL_BATTERY_POWER) { + value, + signalBuilder -> + signalBuilder.setBatteryPower( + com.androidenv.proto.BatteryPowerMetrics.newBuilder() + .setLevelPercent(value.levelPercent) + .setCurrentNowUa(value.currentNowUa) + .setVoltageMv(value.voltageMv) + .setPowerMw(value.powerMw) + .setTemperatureCelsius(value.temperatureCelsius) + .setStatus(value.status) + .setPluggedSource(value.pluggedSource) + .setChargeCounterUah(value.chargeCounterUah) + .setThermalStatus(value.thermalStatus) + .setEnergyCounterNwh(value.energyCounterNwh) + .build() + ) + } + + return builder.build() + } + + @Suppress("DEPRECATION") + private fun AccessibilityNodeInfo.toA11yNodeInfoProto(): + com.google.androidenv.accessibilityforwarder.AndroidAccessibilityNodeInfo { + val builder = + com.google.androidenv.accessibilityforwarder.AndroidAccessibilityNodeInfo.newBuilder() + className?.let { builder.className = it.toString() } + text?.let { builder.text = it.toString() } + contentDescription?.let { builder.contentDescription = it.toString() } + viewIdResourceName?.let { builder.viewIdResourceName = it } + + val rect = android.graphics.Rect() + getBoundsInScreen(rect) + builder.boundsInScreen = + com.google.androidenv.accessibilityforwarder.ProtoRect.newBuilder() + .setLeft(rect.left) + .setTop(rect.top) + .setRight(rect.right) + .setBottom(rect.bottom) + .build() + + builder.isClickable = isClickable + builder.isEnabled = isEnabled + builder.isFocusable = isFocusable + builder.isFocused = isFocused + return builder.build() + } +} diff --git a/android_env/apps/java/com/google/androidenv/core/ProtoConversionsTest.kt b/android_env/apps/java/com/google/androidenv/core/ProtoConversionsTest.kt new file mode 100644 index 00000000..dab28d99 --- /dev/null +++ b/android_env/apps/java/com/google/androidenv/core/ProtoConversionsTest.kt @@ -0,0 +1,340 @@ +// Copyright 2026 DeepMind Technologies Limited. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package com.google.androidenv.core + +import androidx.test.ext.junit.runners.AndroidJUnit4 +import com.androidenv.proto.Action as ProtoAction +import com.androidenv.proto.ActionType as ProtoActionType +import com.androidenv.proto.DeviceSignal as ProtoDeviceSignal +import com.google.common.truth.Truth.assertThat +import org.junit.Test +import org.junit.runner.RunWith +import org.robolectric.annotation.Config + +@RunWith(AndroidJUnit4::class) +@Config(manifest = Config.NONE) +class ProtoConversionsTest { + + @Test + fun testProtoActionToCore_touchDown() { + val protoAction = + ProtoAction.newBuilder() + .setActionType(ProtoActionType.ACTION_TYPE_TOUCH_DOWN) + .setTouchPosition(com.androidenv.proto.Point.newBuilder().setX(0.5f).setY(0.8f)) + .setPointerId(3) + .build() + + val coreAction: Action = with(ProtoConversions) { protoAction.toCore() } + + assertThat(coreAction.actionType).isEqualTo(ActionType.TOUCH_DOWN) + val touchPosition = coreAction.touchPosition + requireNotNull(touchPosition) + assertThat(touchPosition[0]).isEqualTo(0.5f) + assertThat(touchPosition[1]).isEqualTo(0.8f) + assertThat(coreAction.pointerId).isEqualTo(3) + } + + @Test + fun testProtoActionToCore_keyEvent() { + val protoAction = + ProtoAction.newBuilder() + .setActionType(ProtoActionType.ACTION_TYPE_KEY_EVENT) + .setKeycode(42) + .build() + + val coreAction: Action = with(ProtoConversions) { protoAction.toCore() } + + assertThat(coreAction.actionType).isEqualTo(ActionType.KEY_EVENT) + assertThat(coreAction.keycode).isEqualTo(42) + assertThat(coreAction.pointerId).isNull() + } + + @Test + fun testProtoActionToCore_setText() { + val protoAction = + ProtoAction.newBuilder() + .setActionType(ProtoActionType.ACTION_TYPE_SET_TEXT) + .setText("hello world") + .build() + + val coreAction: Action = with(ProtoConversions) { protoAction.toCore() } + + assertThat(coreAction.actionType).isEqualTo(ActionType.SET_TEXT) + assertThat(coreAction.text).isEqualTo("hello world") + } + + @Test + fun testProtoActionToCore_setSelection() { + val protoAction = + ProtoAction.newBuilder() + .setActionType(ProtoActionType.ACTION_TYPE_SET_SELECTION) + .setSelectionStart(2) + .setSelectionEnd(5) + .build() + + val coreAction: Action = with(ProtoConversions) { protoAction.toCore() } + + assertThat(coreAction.actionType).isEqualTo(ActionType.SET_SELECTION) + assertThat(coreAction.selectionStart).isEqualTo(2) + assertThat(coreAction.selectionEnd).isEqualTo(5) + } + + @Test + fun testProtoActionToCore_cutCopyPaste() { + val cutAction = ProtoAction.newBuilder().setActionType(ProtoActionType.ACTION_TYPE_CUT).build() + val copyAction = + ProtoAction.newBuilder().setActionType(ProtoActionType.ACTION_TYPE_COPY).build() + val pasteAction = + ProtoAction.newBuilder().setActionType(ProtoActionType.ACTION_TYPE_PASTE).build() + + with(ProtoConversions) { + assertThat(cutAction.toCore().actionType).isEqualTo(ActionType.CUT) + assertThat(copyAction.toCore().actionType).isEqualTo(ActionType.COPY) + assertThat(pasteAction.toCore().actionType).isEqualTo(ActionType.PASTE) + } + } + + @Test + fun testProtoSignalToCore() { + with(ProtoConversions) { + assertThat(ProtoDeviceSignal.DEVICE_SIGNAL_SCREENSHOT.toCore()) + .isEqualTo(DeviceSignal.SCREENSHOT) + assertThat(ProtoDeviceSignal.DEVICE_SIGNAL_ACTIVE_PACKAGE.toCore()) + .isEqualTo(DeviceSignal.ACTIVE_PACKAGE) + assertThat(ProtoDeviceSignal.DEVICE_SIGNAL_BATTERY_POWER.toCore()) + .isEqualTo(DeviceSignal.BATTERY_POWER) + } + } + + @Test + fun testCoreSignalToProto() { + with(ProtoConversions) { + assertThat(DeviceSignal.SCREENSHOT.toProto()) + .isEqualTo(ProtoDeviceSignal.DEVICE_SIGNAL_SCREENSHOT) + assertThat(DeviceSignal.ACTIVE_PACKAGE.toProto()) + .isEqualTo(ProtoDeviceSignal.DEVICE_SIGNAL_ACTIVE_PACKAGE) + assertThat(DeviceSignal.BATTERY_POWER.toProto()) + .isEqualTo(ProtoDeviceSignal.DEVICE_SIGNAL_BATTERY_POWER) + } + } + + @Test + fun testCoreDeviceStateToProto_packsSignalsCorrectly() { + val dummyBitmap = + android.graphics.Bitmap.createBitmap(1, 1, android.graphics.Bitmap.Config.ARGB_8888) + val screenshotSignal = StreamedSignal(value = dummyBitmap, timestampUs = 1000L, sequenceId = 1L) + val activePkgSignal = + StreamedSignal(value = "com.android.chrome", timestampUs = 2000L, sequenceId = 2L) + + val coreState = + DeviceState( + screenshot = screenshotSignal, + accessibilityRoot = StreamedSignal(null, 0L, 0L), + accessibilityEvents = StreamedSignal(null, 0L, 0L), + activePackage = activePkgSignal, + activeActivity = StreamedSignal(null, 0L, 0L), + systemLogs = StreamedSignal(null, 0L, 0L), + runningProcesses = StreamedSignal(null, 0L, 0L), + deviceSettings = StreamedSignal(null, 0L, 0L), + accelerometer = StreamedSignal(null, 0L, 0L), + gyroscope = StreamedSignal(null, 0L, 0L), + magnetometer = StreamedSignal(null, 0L, 0L), + barometer = StreamedSignal(null, 0L, 0L), + ambientTemp = StreamedSignal(null, 0L, 0L), + lightLevel = StreamedSignal(null, 0L, 0L), + proximity = StreamedSignal(null, 0L, 0L), + locationGps = StreamedSignal(null, 0L, 0L), + heartRate = StreamedSignal(null, 0L, 0L), + audioOutput = StreamedSignal(null, 0L, 0L), + snapshotTimestampUs = 9999L, + ) + + val protoState = with(ProtoConversions) { coreState.toProto() } + + assertThat(protoState.snapshotTimestampUs).isEqualTo(9999L) + assertThat(protoState.signalsCount).isEqualTo(3) + + // Verify device settings signal (always present in WIP) + val protoSettings = + protoState.signalsList.find { it.type == ProtoDeviceSignal.DEVICE_SIGNAL_DEVICE_SETTINGS } + assertThat(protoSettings).isNotNull() + assertThat(protoSettings!!.deviceSettings.entriesMap).containsKey("device_time_ms") + + // Verify screenshot signal with embedded screenshot value over WebSocket + val protoScreenshot = + protoState.signalsList.find { it.type == ProtoDeviceSignal.DEVICE_SIGNAL_SCREENSHOT } + assertThat(protoScreenshot).isNotNull() + assertThat(protoScreenshot!!.metadata.timestampUs).isEqualTo(1000L) + assertThat(protoScreenshot.metadata.sequenceId).isEqualTo(1L) + assertThat(protoScreenshot.hasScreenshot()).isTrue() + + // Verify active package signal (value present) + val protoPkg = + protoState.signalsList.find { it.type == ProtoDeviceSignal.DEVICE_SIGNAL_ACTIVE_PACKAGE } + assertThat(protoPkg).isNotNull() + assertThat(protoPkg!!.metadata.timestampUs).isEqualTo(2000L) + assertThat(protoPkg.metadata.sequenceId).isEqualTo(2L) + assertThat(protoPkg.activePackage).isEqualTo("com.android.chrome") + } + + @Test + fun testCoreDeviceStateToProto_packsAccessibilitySignals() { + val mockNode = + android.view.accessibility.AccessibilityNodeInfo.obtain().apply { + className = "android.widget.Button" + text = "Click Me" + } + val event = + com.google.androidenv.accessibilityforwarder.EventRequest.newBuilder() + .putEvent("event_type", "TYPE_VIEW_CLICKED") + .build() + + val coreState = + DeviceState( + screenshot = StreamedSignal.empty(), + accessibilityRoot = StreamedSignal(value = mockNode, timestampUs = 3000L, sequenceId = 3L), + accessibilityEvents = + StreamedSignal(value = listOf(event), timestampUs = 4000L, sequenceId = 4L), + activePackage = StreamedSignal.empty(), + activeActivity = StreamedSignal.empty(), + systemLogs = StreamedSignal.empty(), + runningProcesses = StreamedSignal.empty(), + deviceSettings = StreamedSignal.empty(), + accelerometer = StreamedSignal.empty(), + gyroscope = StreamedSignal.empty(), + magnetometer = StreamedSignal.empty(), + barometer = StreamedSignal.empty(), + ambientTemp = StreamedSignal.empty(), + lightLevel = StreamedSignal.empty(), + proximity = StreamedSignal.empty(), + locationGps = StreamedSignal.empty(), + heartRate = StreamedSignal.empty(), + audioOutput = StreamedSignal.empty(), + snapshotTimestampUs = 9999L, + ) + + val protoState = with(ProtoConversions) { coreState.toProto() } + + val protoTree = + protoState.signalsList.find { it.type == ProtoDeviceSignal.DEVICE_SIGNAL_ACCESSIBILITY_TREE } + assertThat(protoTree).isNotNull() + assertThat(protoTree!!.hasAccessibilityForest()).isTrue() + assertThat(protoTree.accessibilityForest.windowsCount).isEqualTo(1) + + val protoEvents = + protoState.signalsList.find { + it.type == ProtoDeviceSignal.DEVICE_SIGNAL_ACCESSIBILITY_EVENTS + } + assertThat(protoEvents).isNotNull() + assertThat(protoEvents!!.hasAccessibilityEvents()).isTrue() + assertThat(protoEvents.accessibilityEvents.eventsCount).isEqualTo(1) + } + + @Test + fun testCoreDeviceStateToProto_withScreenShareServiceSettings() { + val controller = + org.robolectric.Robolectric.buildService(ScreenShareService::class.java).create() + val service = controller.get() + service.activeVideoMime = android.media.MediaFormat.MIMETYPE_VIDEO_HEVC + service.activeVideoWidth = 1920 + service.activeVideoHeight = 1080 + + val coreState = + DeviceState( + screenshot = StreamedSignal.empty(), + accessibilityRoot = StreamedSignal.empty(), + accessibilityEvents = StreamedSignal.empty(), + activePackage = StreamedSignal.empty(), + activeActivity = StreamedSignal.empty(), + systemLogs = StreamedSignal.empty(), + runningProcesses = StreamedSignal.empty(), + deviceSettings = StreamedSignal.empty(), + accelerometer = StreamedSignal.empty(), + gyroscope = StreamedSignal.empty(), + magnetometer = StreamedSignal.empty(), + barometer = StreamedSignal.empty(), + ambientTemp = StreamedSignal.empty(), + lightLevel = StreamedSignal.empty(), + proximity = StreamedSignal.empty(), + locationGps = StreamedSignal.empty(), + heartRate = StreamedSignal.empty(), + audioOutput = StreamedSignal.empty(), + snapshotTimestampUs = 9999L, + ) + + val protoState = with(ProtoConversions) { coreState.toProto() } + + val protoSettings = + protoState.signalsList.find { it.type == ProtoDeviceSignal.DEVICE_SIGNAL_DEVICE_SETTINGS } + assertThat(protoSettings).isNotNull() + val entries = protoSettings!!.deviceSettings.entriesMap + assertThat(entries["video_codec"]).isEqualTo("h265") + assertThat(entries["video_width"]).isEqualTo("1920") + assertThat(entries["video_height"]).isEqualTo("1080") + } + + @Test + fun testCoreDeviceStateToProto_packsBatteryPowerSignal() { + val metrics = + BatteryPowerMetrics( + levelPercent = 90, + currentNowUa = 300_000L, + voltageMv = 4100, + powerMw = 1230.0, + temperatureCelsius = 27.0f, + status = "CHARGING", + pluggedSource = "USB", + chargeCounterUah = 2_800_000L, + thermalStatus = 1, + energyCounterNwh = 12_000_000_000L, + ) + + val coreState = + DeviceState( + screenshot = StreamedSignal.empty(), + accessibilityRoot = StreamedSignal.empty(), + accessibilityEvents = StreamedSignal.empty(), + activePackage = StreamedSignal.empty(), + activeActivity = StreamedSignal.empty(), + systemLogs = StreamedSignal.empty(), + runningProcesses = StreamedSignal.empty(), + deviceSettings = StreamedSignal.empty(), + accelerometer = StreamedSignal.empty(), + gyroscope = StreamedSignal.empty(), + magnetometer = StreamedSignal.empty(), + barometer = StreamedSignal.empty(), + ambientTemp = StreamedSignal.empty(), + lightLevel = StreamedSignal.empty(), + proximity = StreamedSignal.empty(), + locationGps = StreamedSignal.empty(), + heartRate = StreamedSignal.empty(), + audioOutput = StreamedSignal.empty(), + batteryPower = StreamedSignal(value = metrics, timestampUs = 5000L, sequenceId = 5L), + snapshotTimestampUs = 9999L, + ) + + val protoState = with(ProtoConversions) { coreState.toProto() } + + val protoBattery = + protoState.signalsList.find { it.type == ProtoDeviceSignal.DEVICE_SIGNAL_BATTERY_POWER } + assertThat(protoBattery).isNotNull() + assertThat(protoBattery!!.metadata.timestampUs).isEqualTo(5000L) + assertThat(protoBattery.metadata.sequenceId).isEqualTo(5L) + assertThat(protoBattery.hasBatteryPower()).isTrue() + assertThat(protoBattery.batteryPower.levelPercent).isEqualTo(90) + assertThat(protoBattery.batteryPower.energyCounterNwh).isEqualTo(12_000_000_000L) + } +} + diff --git a/android_env/apps/java/com/google/androidenv/core/ScreenShareProvider.kt b/android_env/apps/java/com/google/androidenv/core/ScreenShareProvider.kt new file mode 100644 index 00000000..05fd5211 --- /dev/null +++ b/android_env/apps/java/com/google/androidenv/core/ScreenShareProvider.kt @@ -0,0 +1,48 @@ +// Copyright 2026 DeepMind Technologies Limited. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package com.google.androidenv.core + +/** Interface for providing screen dimensions, screenshot bytes, and audio capture. */ +interface ScreenShareProvider { + fun width(): Int + + fun height(): Int + + fun rowStride(): Int + + fun latestScreenshot(): ByteArray? + + fun latestScreenshotTimestampUs(): Long + + fun setAudioListener(listener: ScreenShareService.AudioCaptureListener?) +} + +/** Production implementation that delegates to the ScreenShareService singleton. */ +class SystemScreenShareProvider : ScreenShareProvider { + override fun width(): Int = ScreenShareService.getInstance().width() + + override fun height(): Int = ScreenShareService.getInstance().height() + + override fun rowStride(): Int = ScreenShareService.getInstance().rowStride() + + override fun latestScreenshot(): ByteArray? = ScreenShareService.getInstance().latestScreenshot() + + override fun latestScreenshotTimestampUs(): Long = + ScreenShareService.getInstance().latestScreenshotTimestampUs() + + override fun setAudioListener(listener: ScreenShareService.AudioCaptureListener?) { + ScreenShareService.audioListener = listener + } +} diff --git a/android_env/apps/java/com/google/androidenv/core/ScreenShareService.kt b/android_env/apps/java/com/google/androidenv/core/ScreenShareService.kt new file mode 100644 index 00000000..8311a0f1 --- /dev/null +++ b/android_env/apps/java/com/google/androidenv/core/ScreenShareService.kt @@ -0,0 +1,392 @@ +// Copyright 2026 DeepMind Technologies Limited. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package com.google.androidenv.core + +import android.app.Activity +import android.app.Notification +import android.app.NotificationChannel +import android.app.NotificationManager +import android.app.Service +import android.content.Context +import android.content.Intent +import android.content.res.Resources +import android.graphics.PixelFormat +import android.hardware.display.DisplayManager +import android.media.AudioAttributes +import android.media.AudioFormat +import android.media.AudioPlaybackCaptureConfiguration +import android.media.AudioRecord +import android.media.Image +import android.media.ImageReader +import android.media.projection.MediaProjection +import android.media.projection.MediaProjectionManager +import android.os.IBinder +import android.util.DisplayMetrics +import android.util.Log +import android.view.WindowManager +import androidx.annotation.RequiresApi +import androidx.core.content.IntentCompat + +/** A Service to project the screen and capture screenshots in the background. */ +class ScreenShareService : Service() { + + private var mediaProjection: MediaProjectionWrapper? = null + private var virtualDisplay: VirtualDisplayWrapper? = null + private var displayDensity: Int = -1 + private var displayWidth: Int = -1 + private var displayHeight: Int = -1 + private var displayRowStride: Int = -1 + private var imageReader: ImageReader? = null + private var server: AndroidDeviceEnvWebSocketServer? = null + + lateinit var mediaProjectionManagerWrapper: MediaProjectionManagerWrapper + + private var audioRecord: AudioRecord? = null + private var audioThread: Thread? = null + @Volatile private var isRecording = false + + @Volatile private var latestImage: ByteArray? = null + @Volatile private var latestImageTimestampUs: Long = 0 + + @Volatile var activeVideoMime: String? = null + @Volatile var activeVideoWidth: Int = -1 + @Volatile var activeVideoHeight: Int = -1 + + interface ScreenCaptureListener { + fun onFrameCaptured( + rgbaBytes: ByteArray, + width: Int, + height: Int, + rowStride: Int, + timestampUs: Long, + ) + } + + @Volatile var frameListener: ScreenCaptureListener? = null + + interface AudioCaptureListener { + fun onAudioCaptured( + rawBytes: ByteArray, + sampleRateHz: Int, + channelCount: Int, + sampleFormat: Int, + ) + } + + override fun onBind(intent: Intent): IBinder? = null + + override fun onCreate() { + Log.i(TAG, "ScreenShareService::onCreate()") + instance = this + if (!::mediaProjectionManagerWrapper.isInitialized) { + mediaProjectionManagerWrapper = + RealMediaProjectionManagerWrapper( + getSystemService(Context.MEDIA_PROJECTION_SERVICE) as MediaProjectionManager + ) + } + } + + override fun onStartCommand(intent: Intent, flags: Int, startId: Int): Int { + Log.i(TAG, "onStartCommand()") + + val action = intent.getStringExtra("ACTION") + if (action == "START") { + Log.i(TAG, "Starting MediaProjection") + val notification = getNotification(this) + startForeground(notification.first, notification.second) + + val data: Intent? = IntentCompat.getParcelableExtra(intent, "DATA", Intent::class.java) + val resultCode = intent.getIntExtra("RESULT_CODE", Activity.RESULT_CANCELED) + startProjection(data, resultCode) + startWebSocketServer() + } else if (action == "STOP") { + Log.i(TAG, "Stopping MediaProjection") + stopWebSocketServer() + stopProjection() + stopSelf() + } else { + stopSelf() + } + + return START_NOT_STICKY + } + + fun latestScreenshot(): ByteArray? = latestImage + + fun latestScreenshotTimestampUs(): Long = latestImageTimestampUs + + fun getMediaProjection(): MediaProjectionWrapper? = mediaProjection + + fun width(): Int = displayWidth + + fun height(): Int = displayHeight + + fun density(): Int = displayDensity + + fun rowStride(): Int = displayRowStride + + private fun startProjection(data: Intent?, resultCode: Int) { + if (data == null) return + if (mediaProjection != null) return + + val proj = mediaProjectionManagerWrapper.getMediaProjection(resultCode, data) ?: return + mediaProjection = proj + + val wm = getSystemService(Context.WINDOW_SERVICE) as? WindowManager + val metrics = DisplayMetrics() + wm?.defaultDisplay?.getRealMetrics(metrics) + + displayDensity = + if (metrics.densityDpi > 0) metrics.densityDpi + else Resources.getSystem().displayMetrics.densityDpi + displayWidth = + if (metrics.widthPixels > 0) metrics.widthPixels + else Resources.getSystem().displayMetrics.widthPixels + displayHeight = + if (metrics.heightPixels > 0) metrics.heightPixels + else Resources.getSystem().displayMetrics.heightPixels + + proj.registerCallback( + object : MediaProjection.Callback() { + override fun onStop() { + Log.i(TAG, "MediaProjection stopped") + virtualDisplay?.release() + imageReader?.setOnImageAvailableListener(null, null) + stopAudioRecordingInternal() + mediaProjection = null + } + }, + handler = null, + ) + createVirtualDisplay() + if (android.os.Build.VERSION.SDK_INT >= 29) { + startAudioRecording(proj) + } + } + + private fun stopProjection() { + isRecording = false + audioThread?.interrupt() + mediaProjection?.stop() + mediaProjection = null + } + + @RequiresApi(29) + private fun startAudioRecording(proj: MediaProjectionWrapper) { + if (isRecording) return + isRecording = true + audioThread = + Thread { + val config = + AudioPlaybackCaptureConfiguration.Builder(proj.getActualMediaProjection()) + .addMatchingUsage(AudioAttributes.USAGE_MEDIA) + .addMatchingUsage(AudioAttributes.USAGE_GAME) + .build() + + val sampleRate = 44100 + val channelConfig = AudioFormat.CHANNEL_IN_STEREO + val audioFormat = AudioFormat.ENCODING_PCM_16BIT + val bufferSize = AudioRecord.getMinBufferSize(sampleRate, channelConfig, audioFormat) + + if (bufferSize == AudioRecord.ERROR_BAD_VALUE || bufferSize == AudioRecord.ERROR) { + Log.e(TAG, "Invalid buffer size: $bufferSize") + return@Thread + } + + try { + val record = + AudioRecord.Builder() + .setAudioFormat( + AudioFormat.Builder() + .setEncoding(audioFormat) + .setSampleRate(sampleRate) + .setChannelMask(channelConfig) + .build() + ) + .setAudioPlaybackCaptureConfig(config) + .setBufferSizeInBytes(bufferSize) + .build() + + audioRecord = record + record.startRecording() + Log.i(TAG, "Audio recording started.") + + val audioBuffer = ByteArray(bufferSize) + while (isRecording && !Thread.currentThread().isInterrupted) { + val read = record.read(audioBuffer, 0, bufferSize) + if (read > 0) { + val data = audioBuffer.copyOf(read) + audioListener?.onAudioCaptured(data, sampleRate, 2, audioFormat) + } else if (read < 0) { + Log.e(TAG, "Error reading audio data: $read") + break + } + } + } catch (e: SecurityException) { + Log.e(TAG, "SecurityException starting audio recording", e) + } catch (e: Exception) { + Log.e(TAG, "Exception in audio recording thread", e) + } finally { + stopAudioRecordingInternal() + } + } + .apply { start() } + } + + private fun stopAudioRecordingInternal() { + isRecording = false + try { + audioRecord?.stop() + } catch (e: Exception) { + // Ignore + } + try { + audioRecord?.release() + } catch (e: Exception) { + // Ignore + } + audioRecord = null + Log.i(TAG, "Audio recording stopped.") + } + + private fun createVirtualDisplay() { + imageReader = + ImageReader.newInstance( + displayWidth, + displayHeight, + PixelFormat.RGBA_8888, + /* maxImages = */ 2, + ) + val flags = + DisplayManager.VIRTUAL_DISPLAY_FLAG_OWN_CONTENT_ONLY or + DisplayManager.VIRTUAL_DISPLAY_FLAG_PUBLIC + virtualDisplay = + mediaProjection?.createVirtualDisplay( + "androidenv_screen", + displayWidth, + displayHeight, + displayDensity, + flags, + imageReader?.surface, + null, + null, + ) + imageReader?.setOnImageAvailableListener(onImageAvailableListener, null) + } + + private val onImageAvailableListener = ImageReader.OnImageAvailableListener { reader -> + val image: Image = reader?.acquireLatestImage() ?: return@OnImageAvailableListener + try { + val plane = image.planes[0] + displayRowStride = plane.rowStride + val buffer = plane.buffer + val bytes = ByteArray(buffer.capacity()) + buffer.get(bytes) + latestImage = bytes + latestImageTimestampUs = System.nanoTime() / 1000 + frameListener?.onFrameCaptured( + bytes, + displayWidth, + displayHeight, + displayRowStride, + latestImageTimestampUs, + ) + } finally { + image.close() + } + } + + private fun getNotification(context: Context): Pair { + createNotificationChannel(context) + val notification = createNotification(context) + val manager = context.getSystemService(Context.NOTIFICATION_SERVICE) as NotificationManager + manager.notify(NOTIFICATION_ID, notification) + return Pair(NOTIFICATION_ID, notification) + } + + private fun createNotificationChannel(context: Context) { + val channel = + NotificationChannel( + NOTIFICATION_CHANNEL_ID, + NOTIFICATION_CHANNEL_NAME, + NotificationManager.IMPORTANCE_LOW, + ) + channel.lockscreenVisibility = Notification.VISIBILITY_PRIVATE + val manager = context.getSystemService(Context.NOTIFICATION_SERVICE) as NotificationManager + manager.createNotificationChannel(channel) + } + + private fun createNotification(context: Context): Notification { + val builder = Notification.Builder(context, NOTIFICATION_CHANNEL_ID) + // Use system drawable to avoid custom resource dependency + builder.setSmallIcon(android.R.drawable.ic_menu_camera) + builder.setContentTitle("AndroidEnv") + builder.setContentText("Screen Recording Active") + builder.setOngoing(true) + builder.setCategory(Notification.CATEGORY_SERVICE) + builder.setShowWhen(true) + return builder.build() + } + + fun isServerRunning(): Boolean = server != null + + private fun startWebSocketServer() { + if (server != null) return + try { + Log.i(TAG, "Starting WebSocket server on port 8000...") + val env = AndroidDeviceEnvImpl(applicationContext) + val srv = AndroidDeviceEnvWebSocketServer(env, 8000) + srv.start() + server = srv + Log.i(TAG, "WebSocket server started successfully.") + } catch (e: Exception) { + Log.e(TAG, "Failed to start WebSocket server", e) + } + } + + private fun stopWebSocketServer() { + Log.i(TAG, "Stopping WebSocket server...") + server?.stop() + server = null + Log.i(TAG, "WebSocket server stopped.") + } + + override fun onDestroy() { + Log.i(TAG, "ScreenShareService::onDestroy()") + super.onDestroy() + stopWebSocketServer() + instance = null + } + + companion object { + private const val TAG = "ScreenShareService" + private const val NOTIFICATION_ID = 492459162 + private const val NOTIFICATION_CHANNEL_ID = "com.androidenv" + private const val NOTIFICATION_CHANNEL_NAME = "com.androidenv" + + @Volatile private var instance: ScreenShareService? = null + + @Volatile + var audioListener: AudioCaptureListener? = null + set(value) { + field = value + Log.i("ScreenShareService", "audioListener set to $value") + } + + fun getInstance(): ScreenShareService { + return instance ?: throw IllegalStateException("ScreenShareService not started.") + } + } +} diff --git a/android_env/apps/java/com/google/androidenv/core/ScreenShareServiceTest.kt b/android_env/apps/java/com/google/androidenv/core/ScreenShareServiceTest.kt new file mode 100644 index 00000000..4f53a8f2 --- /dev/null +++ b/android_env/apps/java/com/google/androidenv/core/ScreenShareServiceTest.kt @@ -0,0 +1,296 @@ +// Copyright 2026 DeepMind Technologies Limited. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package com.google.androidenv.core + +import android.app.Activity +import android.app.Application +import android.app.NotificationManager +import android.content.Context +import android.content.Intent +import android.media.projection.MediaProjection +import android.media.projection.MediaProjectionManager +import android.os.Bundle +import android.os.IBinder +import android.os.IInterface +import androidx.test.core.app.ApplicationProvider +import androidx.test.ext.junit.runners.AndroidJUnit4 +import com.google.common.truth.Truth.assertThat +import org.junit.After +import org.junit.Before +import org.junit.Test +import org.junit.runner.RunWith +import org.mockito.kotlin.any +import org.mockito.kotlin.doReturn +import org.mockito.kotlin.eq +import org.mockito.kotlin.mock +import org.mockito.kotlin.verify +import org.mockito.kotlin.whenever +import org.robolectric.Robolectric +import org.robolectric.Shadows.shadowOf +import org.robolectric.android.controller.ServiceController +import org.robolectric.annotation.Config +import org.robolectric.shadows.ShadowNotificationManager + +@RunWith(AndroidJUnit4::class) +@Config(sdk = [26]) // ScreenShareService is @RequiresApi(26) +class ScreenShareServiceTest { + + private val context = ApplicationProvider.getApplicationContext() + private lateinit var shadowNotificationManager: ShadowNotificationManager + private var serviceController: ServiceController? = null + + @Before + fun setUp() { + val notificationManager = + context.getSystemService(Context.NOTIFICATION_SERVICE) as NotificationManager + shadowNotificationManager = shadowOf(notificationManager) + } + + @After + fun tearDown() { + serviceController?.destroy() + // Clean up singleton + val field = ScreenShareService::class.java.getDeclaredField("instance") + field.isAccessible = true + field.set(null, null) + } + + @Test + fun onCreate_setsSingletonInstance() { + serviceController = Robolectric.buildService(ScreenShareService::class.java) + val service = serviceController!!.create().get() + assertThat(ScreenShareService.getInstance()).isEqualTo(service) + } + + @Test + fun initialDimensions_areMinusOne() { + serviceController = Robolectric.buildService(ScreenShareService::class.java) + val service = serviceController!!.create().get() + assertThat(service.width()).isEqualTo(-1) + assertThat(service.height()).isEqualTo(-1) + assertThat(service.rowStride()).isEqualTo(-1) + } + + @Test + fun onStartCommand_START_withoutData_stopsSelf() { + val intent = + Intent(context, ScreenShareService::class.java).apply { putExtra("ACTION", "START") } + serviceController = Robolectric.buildService(ScreenShareService::class.java, intent) + serviceController!!.create().startCommand(0, 1) + + // It should have posted foreground notification + assertThat(shadowNotificationManager.size()).isEqualTo(1) + val notification = shadowNotificationManager.allNotifications[0] + assertThat(notification).isNotNull() + } + + @Test + fun onStartCommand_STOP_stopsForegroundAndSelf() { + val stopIntent = + Intent(context, ScreenShareService::class.java).apply { putExtra("ACTION", "STOP") } + val controller = Robolectric.buildService(ScreenShareService::class.java, stopIntent) + serviceController = controller // Store for tearDown + controller.create().startCommand(0, 1) + + // Verifies it runs without crash. stopSelf() is called internally. + } + + @Test + fun onStartCommand_START_withData_startsProjection() { + val mockMediaProjection = mock() + val mockMediaProjectionManager = mock() + val resultCode = Activity.RESULT_OK + val dataIntent = Intent("data") + + whenever(mockMediaProjectionManager.getMediaProjection(eq(resultCode), any())) + .doReturn(mockMediaProjection) + + val intent = + Intent(context, ScreenShareService::class.java).apply { + putExtra("ACTION", "START") + putExtra("RESULT_CODE", resultCode) + putExtra("DATA", dataIntent) + } + val controller = Robolectric.buildService(ScreenShareService::class.java, intent) + val service = controller.get() + service.mediaProjectionManagerWrapper = mockMediaProjectionManager + serviceController = controller // Store for tearDown + + controller.create().startCommand(0, 1) + + assertThat(service.getMediaProjection()).isEqualTo(mockMediaProjection) + } + + @Test + fun onStartCommand_STOP_stopsProjection() { + val mockMediaProjection = mock() + val mockMediaProjectionManager = mock() + val resultCode = Activity.RESULT_OK + val dataIntent = Intent("data") + + whenever(mockMediaProjectionManager.getMediaProjection(eq(resultCode), any())) + .doReturn(mockMediaProjection) + + val intent = + Intent(context, ScreenShareService::class.java).apply { + putExtra("ACTION", "START") + putExtra("RESULT_CODE", resultCode) + putExtra("DATA", dataIntent) + } + val controller = Robolectric.buildService(ScreenShareService::class.java, intent) + val service = controller.get() + service.mediaProjectionManagerWrapper = mockMediaProjectionManager + serviceController = controller + + controller.create().startCommand(0, 1) + + assertThat(service.getMediaProjection()).isEqualTo(mockMediaProjection) + + // Change action to STOP and call startCommand again + intent.putExtra("ACTION", "STOP") + controller.startCommand(0, 2) + + verify(mockMediaProjection).stop() + assertThat(service.getMediaProjection()).isNull() + } + + @Test + fun realMediaProjectionManagerWrapper_returnsNull_whenResultCancelled() { + val manager = + context.getSystemService(Context.MEDIA_PROJECTION_SERVICE) as MediaProjectionManager + val wrapper = RealMediaProjectionManagerWrapper(manager) + val result = wrapper.getMediaProjection(Activity.RESULT_CANCELED, Intent()) + assertThat(result).isNull() + } + + @Test + fun realMediaProjectionManagerWrapper_returnsWrapper_whenValid() { + val manager = + context.getSystemService(Context.MEDIA_PROJECTION_SERVICE) as MediaProjectionManager + val wrapper = RealMediaProjectionManagerWrapper(manager) + + val imClass = Class.forName("android.media.projection.IMediaProjection") + val mockImpl = org.mockito.Mockito.mock(imClass) + + val mockBinder = mock() + whenever(mockBinder.queryLocalInterface(eq("android.media.projection.IMediaProjection"))) + .doReturn(mockImpl as IInterface) + + val bundle = + Bundle().apply { + putBinder("android.media.projection.extra.EXTRA_MEDIA_PROJECTION", mockBinder) + } + val intent = Intent().putExtras(bundle) + + val result = wrapper.getMediaProjection(Activity.RESULT_OK, intent) + assertThat(result).isNotNull() + assertThat(result).isInstanceOf(RealMediaProjectionWrapper::class.java) + } + + @Test + fun realMediaProjectionWrapper_stop_delegatesToActual() { + val imClass = Class.forName("android.media.projection.IMediaProjection") + val mockImpl = org.mockito.Mockito.mock(imClass) + val constructor = MediaProjection::class.java.getConstructor(Context::class.java, imClass) + val realMp = constructor.newInstance(context, mockImpl) as MediaProjection + val wrapper = RealMediaProjectionWrapper(realMp) + + wrapper.stop() + + val invocations = org.mockito.Mockito.mockingDetails(mockImpl).invocations + val methodNames = invocations.map { it.method.name } + assertThat(methodNames).contains("stop") + } + + @Test + fun realMediaProjectionWrapper_createVirtualDisplay_delegatesToActual() { + val imClass = Class.forName("android.media.projection.IMediaProjection") + val mockImpl = org.mockito.Mockito.mock(imClass) + + val constructor = MediaProjection::class.java.getConstructor(Context::class.java, imClass) + val realMp = constructor.newInstance(context, mockImpl) as MediaProjection + val wrapper = RealMediaProjectionWrapper(realMp) + + val vdWrapper = + wrapper.createVirtualDisplay( + name = "test_display", + width = 640, + height = 480, + dpi = 160, + flags = 0, + surface = null, + callback = null, + handler = null, + ) + + assertThat(vdWrapper).isNotNull() + vdWrapper!!.release() + } + + @Test + fun realMediaProjectionWrapper_registerCallback_delegatesToActual() { + val imClass = Class.forName("android.media.projection.IMediaProjection") + val mockImpl = org.mockito.Mockito.mock(imClass) + val constructor = MediaProjection::class.java.getConstructor(Context::class.java, imClass) + val realMp = constructor.newInstance(context, mockImpl) as MediaProjection + val wrapper = RealMediaProjectionWrapper(realMp) + + val callback = object : MediaProjection.Callback() {} + wrapper.registerCallback(callback, null) + } + + @Test + fun realMediaProjectionWrapper_getActualMediaProjection_returnsActual() { + val imClass = Class.forName("android.media.projection.IMediaProjection") + val mockImpl = org.mockito.Mockito.mock(imClass) + val constructor = MediaProjection::class.java.getConstructor(Context::class.java, imClass) + val realMp = constructor.newInstance(context, mockImpl) as MediaProjection + val wrapper = RealMediaProjectionWrapper(realMp) + + assertThat(wrapper.getActualMediaProjection()).isEqualTo(realMp) + } + + @Test + @Config(sdk = [29]) + fun onStartCommand_START_withData_startsAudioRecording_atSdk29() { + val intent = + Intent(context, ScreenShareService::class.java).apply { + putExtra("ACTION", "START") + putExtra("RESULT_CODE", Activity.RESULT_OK) + putExtra("DATA", Intent()) + } + + val imClass = Class.forName("android.media.projection.IMediaProjection") + val mockImpl = org.mockito.Mockito.mock(imClass) + val mpConstructor = MediaProjection::class.java.getConstructor(Context::class.java, imClass) + val realMp = mpConstructor.newInstance(context, mockImpl) as MediaProjection + + val mockMediaProjection = mock() + whenever(mockMediaProjection.getActualMediaProjection()).doReturn(realMp) + + val mockMediaProjectionManager = mock() + whenever(mockMediaProjectionManager.getMediaProjection(any(), any())) + .doReturn(mockMediaProjection) + + val controller = Robolectric.buildService(ScreenShareService::class.java, intent) + val service = controller.get() + service.mediaProjectionManagerWrapper = mockMediaProjectionManager + + controller.create().startCommand(0, 1) + + assertThat(service.getMediaProjection()).isEqualTo(mockMediaProjection) + verify(mockMediaProjection).getActualMediaProjection() + } +} diff --git a/android_env/apps/java/com/google/androidenv/core/StreamedSignal.kt b/android_env/apps/java/com/google/androidenv/core/StreamedSignal.kt new file mode 100644 index 00000000..3ca9bb2d --- /dev/null +++ b/android_env/apps/java/com/google/androidenv/core/StreamedSignal.kt @@ -0,0 +1,31 @@ +// Copyright 2026 DeepMind Technologies Limited. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package com.google.androidenv.core + +/** + * Wraps a signal with its own metadata, acknowledging it is optional and has independent frequency. + */ +data class StreamedSignal( + // The raw value. Null if unavailable, restricted, or permission denied. + val value: T?, + // The exact epoch timestamp when this specific value was captured. + val timestampUs: Long, + // Monotonically increasing ID to easily track updates. + val sequenceId: Long, +) { + companion object { + fun empty(): StreamedSignal = StreamedSignal(null, 0, 0) + } +} diff --git a/android_env/apps/java/com/google/androidenv/core/res/xml/InputService.xml b/android_env/apps/java/com/google/androidenv/core/res/xml/InputService.xml new file mode 100644 index 00000000..7a610d0d --- /dev/null +++ b/android_env/apps/java/com/google/androidenv/core/res/xml/InputService.xml @@ -0,0 +1,22 @@ + + + + + diff --git a/android_env/components/BUILD.bazel b/android_env/components/BUILD.bazel index e7c80617..a966c487 100644 --- a/android_env/components/BUILD.bazel +++ b/android_env/components/BUILD.bazel @@ -386,3 +386,124 @@ py_test( "@pip//numpy", ], ) + +py_library( + name = "telemetry", + srcs = ["telemetry.py"], + deps = [ + "//android_env/proto:device_env_service_py_pb2", + "@pip//absl_py", + ], +) + +py_test( + name = "telemetry_test", + srcs = ["telemetry_test.py"], + legacy_create_init = False, + deps = [ + ":telemetry", + "//android_env/proto:device_env_service_py_pb2", + "@pip//absl_py", + ], +) + +py_library( + name = "device_connection", + srcs = ["device_connection.py"], + deps = [ + ":config_classes", + "//android_env/proto:device_env_service_py_pb2", + "@pip//absl_py", + "@pip//numpy", + "@pip//opencv_python_headless", + "@pip//websocket_client", + ], +) + +py_test( + name = "device_connection_test", + srcs = ["device_connection_test.py"], + legacy_create_init = False, + deps = [ + ":config_classes", + ":device_connection", + "//android_env/proto:device_env_service_py_pb2", + "@pip//absl_py", + "@pip//numpy", + "@pip//opencv_python_headless", + "@pip//websocket_client", + ], +) + +py_library( + name = "android_device", + srcs = ["android_device.py"], + deps = [ + ":adb_call_parser", + ":device_connection", + ":errors", + "//android_env/proto:adb_py_pb2", + "//android_env/proto:device_env_service_py_pb2", + "@pip//absl_py", + "@pip//numpy", + "@pip//opencv_python_headless", + ], +) + +py_test( + name = "android_device_test", + srcs = ["android_device_test.py"], + legacy_create_init = False, + deps = [ + ":adb_call_parser", + ":android_device", + ":device_connection", + ":errors", + "//android_env/proto:adb_py_pb2", + "//android_env/proto:device_env_service_py_pb2", + "@pip//absl_py", + "@pip//numpy", + "@pip//opencv_python_headless", + ], +) + +py_library( + name = "android_device_loader", + srcs = ["android_device_loader.py"], + data = ["device_prep_steps.textproto"], + deps = [ + ":adb_call_parser", + ":adb_controller", + ":android_device", + ":config_classes", + ":device_connection", + ":errors", + ":setup_step_interpreter", + "//android_env/components/simulators/emulator:emulator_launcher", + "//android_env/components/simulators/emulator:tcp_ports", + "//android_env/proto:adb_py_pb2", + "//android_env/proto:task_py_pb2", + "@pip//absl_py", + "@pip//portpicker", + "@pip//protobuf", + ], +) + +py_test( + name = "android_device_loader_test", + srcs = ["android_device_loader_test.py"], + legacy_create_init = False, + deps = [ + ":adb_call_parser", + ":adb_controller", + ":android_device", + ":android_device_loader", + ":config_classes", + ":device_connection", + ":errors", + ":setup_step_interpreter", + "//android_env/components/simulators/emulator:emulator_launcher", + "//android_env/components/simulators/emulator:tcp_ports", + "@pip//absl_py", + ], +) diff --git a/android_env/components/action_type.py b/android_env/components/action_type.py index ceeea974..51358d61 100644 --- a/android_env/components/action_type.py +++ b/android_env/components/action_type.py @@ -50,3 +50,8 @@ class ActionType(enum.IntEnum): KEYDOWN = 3 KEYUP = 4 KEYPRESS = 5 + TEXT = 6 + SET_SELECTION = 7 + CUT = 8 + COPY = 9 + PASTE = 10 diff --git a/android_env/components/android_device.py b/android_env/components/android_device.py new file mode 100644 index 00000000..4eb7cc97 --- /dev/null +++ b/android_env/components/android_device.py @@ -0,0 +1,150 @@ +# coding=utf-8 +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""First-class ergonomic handle to an Android device or emulator.""" + +from collections.abc import Callable, Sequence +import time +from typing import Any, Self + +from absl import logging +from android_env.components import adb_call_parser as adb_parser +from android_env.components import device_connection +from android_env.components import errors +from android_env.proto import adb_pb2 +from android_env.proto import device_env_service_pb2 +import cv2 +import numpy as np + + +class AndroidDevice: + """Ergonomic handle to an Android device running the Core on-device server.""" + + def __init__( + self, + connection: device_connection.DeviceConnection, + adb_call_parser: adb_parser.AdbCallParser | None = None, + on_close_callbacks: Sequence[Callable[[], None]] | None = None, + on_error_callbacks: ( + Sequence[Callable[[Any, Any, Any], None]] | None + ) = None, + ): + """Initializes the AndroidDevice handle. + + Args: + connection: Low-level WebSocket transport connection to the device. + adb_call_parser: Optional AdbCallParser instance for ADB calls. + on_close_callbacks: Callbacks executed when close() is called. + on_error_callbacks: Callbacks executed when exiting context manager with + an error. + """ + self._connection = connection + self._adb_call_parser = adb_call_parser + self._on_close_callbacks = list(on_close_callbacks or []) + self._on_error_callbacks = list(on_error_callbacks or []) + self._clock_offset_ms = 0 + self._sync_clock() + + def execute_adb_call(self, call: adb_pb2.AdbRequest) -> adb_pb2.AdbResponse: + """Executes an AdbRequest using the device's AdbCallParser.""" + if self._adb_call_parser is None: + raise errors.AdbControllerError('ADB is not configured for this device.') + return self._adb_call_parser.parse(call) + + def get_state(self) -> device_env_service_pb2.DeviceState: + """Returns the latest device state (telemetry signals, decoded screenshot, etc.).""" + state = self._connection.get_device_state() + self._decode_screenshot_signals(state) + return state + + def send_action(self, action: device_env_service_pb2.Action) -> None: + """Sends an action (touch, key event, etc.) to the device.""" + req = device_env_service_pb2.InjectActionRequest(action=action) + msg = device_env_service_pb2.ClientMessage(inject_action=req) + self._connection.send_message(msg) + + def update_subscriptions( + self, active_signals: set[device_env_service_pb2.DeviceSignal] + ) -> None: + """Updates active telemetry signal subscriptions on the device.""" + req = device_env_service_pb2.UpdateSubscriptionsRequest() + req.active_signals.extend(active_signals) + msg = device_env_service_pb2.ClientMessage(update_subscriptions=req) + self._connection.send_message(msg) + + def close(self) -> None: + """Closes the connection and runs all cleanup callbacks.""" + logging.info('Closing AndroidDevice...') + self._connection.close() + for cb in self._on_close_callbacks: + cb() + + def __enter__(self) -> Self: + return self + + def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: + if exc_type is not None: + for cb in self._on_error_callbacks: + cb(exc_type, exc_value, traceback) + self.close() + + def _decode_screenshot_signals( + self, state: device_env_service_pb2.DeviceState + ) -> None: + """Decodes compressed screenshot bytes into in-place decoded pixels.""" + for signal in state.signals: + if ( + signal.type == device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT + and signal.HasField('screenshot') + and signal.screenshot.encoded_bytes + ): + img_bytes = signal.screenshot.encoded_bytes + img_arr = np.frombuffer(img_bytes, dtype=np.uint8) + frame = cv2.imdecode(img_arr, cv2.IMREAD_COLOR) + if frame is not None: + rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + h, w, _ = rgb_frame.shape + signal.screenshot.decoded.raw_pixels = rgb_frame.tobytes() + signal.screenshot.decoded.width = w + signal.screenshot.decoded.height = h + break + + def _sync_clock(self) -> None: + """Synchronizes device wall clock offset.""" + t0 = time.time_ns() // 1_000_000 + try: + state = self._connection.get_device_state() + except (RuntimeError, ValueError) as e: + logging.warning('Clock sync skipped: %s', e) + return + + t1 = time.time_ns() // 1_000_000 + rtt = t1 - t0 + settings = self._get_device_settings(state) + device_time_str = settings.get('device_time_ms') + if device_time_str: + t_device = int(device_time_str) + self._clock_offset_ms = t_device - (t0 + rtt // 2) + logging.info( + 'Clock sync: RTT=%dms, offset=%dms', rtt, self._clock_offset_ms + ) + + def _get_device_settings( + self, state: device_env_service_pb2.DeviceState + ) -> dict[str, str]: + for signal in state.signals: + if signal.type == device_env_service_pb2.DEVICE_SIGNAL_DEVICE_SETTINGS: + return dict(signal.device_settings.entries) + return {} diff --git a/android_env/components/android_device_loader.py b/android_env/components/android_device_loader.py new file mode 100644 index 00000000..0c857a13 --- /dev/null +++ b/android_env/components/android_device_loader.py @@ -0,0 +1,227 @@ +# coding=utf-8 +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Loader function for creating an AndroidDevice instance.""" + +import os +import time +from typing import Any + +from absl import logging +from android_env.components import adb_call_parser +from android_env.components import adb_controller +from android_env.components import android_device +from android_env.components import config_classes +from android_env.components import device_connection +from android_env.components import errors +from android_env.components import setup_step_interpreter +from android_env.components.simulators.emulator import emulator_launcher +from android_env.components.simulators.emulator import tcp_ports +from android_env.proto import adb_pb2 +from android_env.proto import task_pb2 +from google.protobuf import text_format +import portpicker + + +def load_android_device( + config: config_classes.AndroidDeviceConfig, +) -> android_device.AndroidDevice: + """Loads an AndroidDevice handle configured according to `config`.""" + on_close_callbacks = [] + on_error_callbacks = [] + + if config.emulator_launcher is not None: + if config.adb_controller is None: + config.adb_controller = config_classes.AdbControllerConfig() + + launcher_config = config.emulator_launcher + tcp_ports.pick_emulator_ports(launcher_config) + + if not config.adb_controller.device_name: + config.adb_controller.device_name = 'emulator-%d' % ( + launcher_config.adb_port - 1 + ) + + launcher = emulator_launcher.EmulatorLauncher( + config=config.emulator_launcher, + adb_controller_config=config.adb_controller, + ) + serial = config.adb_controller.device_name + adb = adb_controller.AdbController(config.adb_controller) + adb.init_server() + + launcher.launch_emulator_process() + try: + _wait_for_boot(adb) + except (errors.AdbControllerError, RuntimeError): + logging.exception('Failed to boot emulator (or aborted).') + launcher.close() + raise + + on_close_callbacks.append(launcher.close) + + def dump_logcat_cb(exc_type: Any, exc_value: Any, traceback: Any) -> None: + del exc_type, exc_value, traceback + logging.error( + 'Error occurred, dumping logcat to /tmp/androidenv_logcat.txt' + ) + try: + logcat = adb.execute_command(['logcat', '-d', '-t', '50000']) + with open('/tmp/androidenv_logcat.txt', 'wb') as f: + f.write(logcat) + except (errors.AdbControllerError, OSError): + logging.exception('Failed to dump logcat.') + + on_error_callbacks.append(dump_logcat_cb) + + local_port = portpicker.pick_unused_port() + config.adb_path = config.adb_controller.adb_path + config.serial = serial + config.device_connection.websocket_port = local_port + config.auto_setup = True + + try: + adb = adb_controller.AdbController( + config_classes.AdbControllerConfig( + adb_path=config.adb_path, + device_name=config.serial, + ) + ) + + parser = adb_call_parser.AdbCallParser(adb) + if config.auto_setup and config.device_connection.host in ( + 'localhost', + '127.0.0.1', + ): + _prep_device(config, parser) + + conn = _connect_to_device(config, adb) + return android_device.AndroidDevice( + connection=conn, + adb_call_parser=parser, + on_close_callbacks=on_close_callbacks, + on_error_callbacks=on_error_callbacks, + ) + except (errors.AdbControllerError, RuntimeError, OSError, ValueError): + logging.exception('Failed to load android device.') + for cb in on_close_callbacks: + cb() + raise + + +def _wait_for_boot( + adb: adb_controller.AdbController, timeout_sec: float = 300.0 +) -> None: + """Waits for the emulator device to finish booting.""" + deadline = time.time() + timeout_sec + while time.time() < deadline: + try: + output = adb.execute_command(['shell', 'getprop', 'sys.boot_completed']) + if output.strip() == b'1': + return + except errors.AdbControllerError: + pass + time.sleep(1.0) + raise RuntimeError('Timed out waiting for emulator boot.') + + +def _get_core_apk_from_runfiles() -> str: + """Returns the path to the android_env_server APK in runfiles.""" + return '' + + +def _get_prep_steps_path() -> str: + """Returns the path to the default `device_prep_steps.textproto`.""" + return os.path.join( + os.path.dirname(os.path.abspath(__file__)), 'device_prep_steps.textproto' + ) + + +def _prep_device( + config: config_classes.AndroidDeviceConfig, + parser: adb_call_parser.AdbCallParser, +) -> None: + """Executes default setup steps on the device.""" + textproto_path = _get_prep_steps_path() + with open(textproto_path, 'r') as f: + text_content = f.read() + + task = task_pb2.Task() + text_format.Parse(text_content, task) + steps = list(task.setup_steps) + + apk_path = _get_core_apk_from_runfiles() + if apk_path: + steps.insert( + 1, + task_pb2.SetupStep( + adb_request=adb_pb2.AdbRequest( + install_apk=adb_pb2.AdbRequest.InstallApk( + filesystem=adb_pb2.AdbRequest.InstallApk.Filesystem( + path=apk_path + ) + ) + ) + ), + ) + + forward_ws_step = task_pb2.SetupStep( + adb_request=adb_pb2.AdbRequest( + generic=adb_pb2.AdbRequest.GenericRequest( + args=[ + 'forward', + f'tcp:{config.device_connection.websocket_port}', + 'tcp:8000', + ] + ) + ) + ) + steps.insert(-2, forward_ws_step) + + interpreter = setup_step_interpreter.SetupStepInterpreter(parser) + interpreter.interpret(steps) + + +def _connect_to_device( + config: config_classes.AndroidDeviceConfig, + adb: adb_controller.AdbController, + max_retries: int = 3, +) -> device_connection.DeviceConnection: + """Establishes WebSocket connection to device with retries.""" + for attempt in range(max_retries): + try: + return device_connection.DeviceConnection(config.device_connection) + except (errors.AdbControllerError, RuntimeError, OSError) as e: + logging.warning( + 'Attempt %d/%d failed to connect: %s', attempt + 1, max_retries, e + ) + if attempt == max_retries - 1: + raise + if config.auto_setup and config.device_connection.host in ( + 'localhost', + '127.0.0.1', + ): + try: + adb.execute_command([ + 'shell', + 'am', + 'startservice', + '-a', + 'com.google.androidenv.core.START', + 'com.google.androidenv.core/.AndroidEnvService', + ]) + except errors.AdbControllerError: + pass + raise RuntimeError('Failed to connect to on-device server.') diff --git a/android_env/components/android_device_loader_test.py b/android_env/components/android_device_loader_test.py new file mode 100644 index 00000000..94fe7d3e --- /dev/null +++ b/android_env/components/android_device_loader_test.py @@ -0,0 +1,347 @@ +# coding=utf-8 +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for android_device_loader.""" + +import os +import time +from unittest import mock + +from absl.testing import absltest +from android_env.components import adb_call_parser +from android_env.components import adb_controller +from android_env.components import android_device +from android_env.components import android_device_loader +from android_env.components import config_classes +from android_env.components import device_connection +from android_env.components import errors +from android_env.components import setup_step_interpreter +from android_env.components.simulators.emulator import emulator_launcher +from android_env.components.simulators.emulator import tcp_ports + + +class AndroidDeviceLoaderTest(absltest.TestCase): + + def setUp(self): + super().setUp() + self._mock_conn_init = mock.patch.object( + device_connection, 'DeviceConnection', autospec=True + ).start() + self._mock_conn = self._mock_conn_init.return_value + + self._mock_launcher_cls = mock.patch.object( + emulator_launcher, 'EmulatorLauncher', autospec=True + ).start() + self._mock_launcher = self._mock_launcher_cls.return_value + + self._mock_adb_cls = mock.patch.object( + adb_controller, 'AdbController', autospec=True + ).start() + self._mock_adb = self._mock_adb_cls.return_value + + self._mock_device_cls = mock.patch.object( + android_device, 'AndroidDevice', autospec=True + ).start() + self._mock_device = self._mock_device_cls.return_value + + self._mock_pick_ports = mock.patch.object( + tcp_ports, 'pick_emulator_ports' + ).start() + + def tearDown(self): + mock.patch.stopall() + super().tearDown() + + @mock.patch.object(android_device_loader, '_prep_device', autospec=True) + def test_load_physical_device_success(self, mock_prep): + config = config_classes.AndroidDeviceConfig( + serial='device-1', + auto_setup=False, + ) + device = android_device_loader.load_android_device(config) + + self.assertEqual(device, self._mock_device) + self._mock_launcher_cls.assert_not_called() + mock_prep.assert_not_called() + + @mock.patch.object(android_device_loader, '_wait_for_boot', autospec=True) + @mock.patch.object(android_device_loader, '_prep_device', autospec=True) + def test_load_emulator_launches_process(self, mock_prep, mock_wait_boot): + def fake_pick_ports(launcher_config): + launcher_config.adb_port = 5555 + + self._mock_pick_ports.side_effect = fake_pick_ports + + config = config_classes.AndroidDeviceConfig( + emulator_launcher=config_classes.EmulatorLauncherConfig( + avd_name='test_avd', + emulator_console_port=5554, + ), + adb_controller=config_classes.AdbControllerConfig(device_name=''), + ) + device = android_device_loader.load_android_device(config) + + self.assertEqual(device, self._mock_device) + self._mock_pick_ports.assert_called_once_with(config.emulator_launcher) + self.assertIsNotNone(config.adb_controller) + self.assertEqual(config.adb_controller.device_name, 'emulator-5554') + self._mock_adb.init_server.assert_called_once() + self._mock_launcher_cls.assert_called_once() + self._mock_launcher.launch_emulator_process.assert_called_once() + mock_wait_boot.assert_called_once() + mock_prep.assert_called_once() + + on_error_cbs = self._mock_device_cls.call_args[1]['on_error_callbacks'] + self.assertLen(on_error_cbs, 1) + + with mock.patch('builtins.open', mock.mock_open()): + on_error_cbs[0](None, None, None) + self._mock_adb.execute_command.assert_called_with( + ['logcat', '-d', '-t', '50000'] + ) + + @mock.patch.object(android_device_loader, '_wait_for_boot', autospec=True) + def test_emulator_launch_failure_closes_launcher(self, mock_wait_boot): + mock_wait_boot.side_effect = RuntimeError('Boot timeout') + config = config_classes.AndroidDeviceConfig( + emulator_launcher=config_classes.EmulatorLauncherConfig( + avd_name='test_avd' + ) + ) + with self.assertRaises(RuntimeError): + android_device_loader.load_android_device(config) + + self._mock_launcher.close.assert_called_once() + + @mock.patch.object( + android_device_loader, '_get_prep_steps_path', autospec=True + ) + @mock.patch.object( + android_device_loader, '_get_core_apk_from_runfiles', autospec=True + ) + @mock.patch.object( + setup_step_interpreter, 'SetupStepInterpreter', autospec=True + ) + def test_prep_device_configures_steps_and_interprets( + self, mock_interpreter_cls, mock_get_apk, mock_get_prep_steps_path + ): + mock_get_prep_steps_path.return_value = ( + '/path/to/device_prep_steps.textproto' + ) + mock_get_apk.return_value = '/path/to/server.apk' + + sample_textproto = ( + 'setup_steps { adb_request { generic { args: "root" } } }\n' + 'setup_steps { adb_request { generic { args: "shell" } } }\n' + ) + + mock_parser = mock.create_autospec( + adb_call_parser.AdbCallParser, instance=True + ) + config = config_classes.AndroidDeviceConfig() + config.device_connection.websocket_port = 12345 + + with mock.patch( + 'builtins.open', mock.mock_open(read_data=sample_textproto) + ): + android_device_loader._prep_device(config, mock_parser) + + mock_interpreter_cls.assert_called_once_with(mock_parser) + interpreter_instance = mock_interpreter_cls.return_value + interpreter_instance.interpret.assert_called_once() + + steps = interpreter_instance.interpret.call_args[0][0] + root_step = [ + s + for s in steps + if s.adb_request.HasField('generic') + and s.adb_request.generic.args + and s.adb_request.generic.args[0] == 'root' + ] + self.assertLen(root_step, 1) + + install_step = [s for s in steps if s.adb_request.HasField('install_apk')][ + 0 + ] + self.assertEqual( + install_step.adb_request.install_apk.filesystem.path, + '/path/to/server.apk', + ) + forward_step = [ + s + for s in steps + if s.adb_request.HasField('generic') + and s.adb_request.generic.args + and s.adb_request.generic.args[0] == 'forward' + ][0] + self.assertIn('tcp:12345', forward_step.adb_request.generic.args) + + def test_connect_to_device_retries_and_starts_service(self): + self._mock_conn_init.side_effect = [ + RuntimeError('Connection refused'), + self._mock_conn, + ] + config = config_classes.AndroidDeviceConfig( + auto_setup=True, + ) + conn = android_device_loader._connect_to_device(config, self._mock_adb) + self.assertEqual(conn, self._mock_conn) + self.assertEqual(self._mock_conn_init.call_count, 2) + self._mock_adb.execute_command.assert_called_once_with([ + 'shell', + 'am', + 'startservice', + '-a', + 'com.google.androidenv.core.START', + 'com.google.androidenv.core/.AndroidEnvService', + ]) + + def test_connect_to_device_fails_after_max_retries(self): + self._mock_conn_init.side_effect = RuntimeError('Connection failed') + config = config_classes.AndroidDeviceConfig(auto_setup=False) + with self.assertRaisesRegex(RuntimeError, 'Connection failed'): + android_device_loader._connect_to_device( + config, self._mock_adb, max_retries=2 + ) + self.assertEqual(self._mock_conn_init.call_count, 2) + + def test_wait_for_boot_success(self): + self._mock_adb.execute_command.return_value = b'1\n' + android_device_loader._wait_for_boot(self._mock_adb, timeout_sec=5.0) + self._mock_adb.execute_command.assert_called_with( + ['shell', 'getprop', 'sys.boot_completed'] + ) + + def test_wait_for_boot_handles_adb_error_then_succeeds(self): + self._mock_adb.execute_command.side_effect = [ + errors.AdbControllerError('Boot check failed'), + b'1\n', + ] + android_device_loader._wait_for_boot(self._mock_adb, timeout_sec=5.0) + self.assertEqual(self._mock_adb.execute_command.call_count, 2) + + def test_wait_for_boot_succeeds_just_before_deadline(self): + self._mock_adb.execute_command.side_effect = [b'0\n', b'1\n'] + with mock.patch.object(time, 'time', side_effect=[100.0, 100.0, 104.5]): + android_device_loader._wait_for_boot(self._mock_adb, timeout_sec=5.0) + self.assertEqual(self._mock_adb.execute_command.call_count, 2) + + def test_wait_for_boot_timeout(self): + self._mock_adb.execute_command.side_effect = [b'0\n', b'1\n'] + with mock.patch.object(time, 'time', side_effect=[100.0, 100.0, 105.5]): + with self.assertRaisesRegex( + RuntimeError, 'Timed out waiting for emulator boot' + ): + android_device_loader._wait_for_boot(self._mock_adb, timeout_sec=5.0) + + @mock.patch.object( + android_device_loader, '_get_prep_steps_path', autospec=True + ) + @mock.patch.object( + android_device_loader, '_get_core_apk_from_runfiles', autospec=True + ) + @mock.patch.object( + setup_step_interpreter, 'SetupStepInterpreter', autospec=True + ) + def test_prep_device_without_apk( + self, mock_interpreter_cls, mock_get_apk, mock_get_prep_steps_path + ): + mock_get_prep_steps_path.return_value = ( + '/path/to/device_prep_steps.textproto' + ) + mock_get_apk.return_value = '' + + sample_textproto = ( + 'setup_steps { adb_request { generic { args: "root" } } }\n' + ) + mock_parser = mock.create_autospec( + adb_call_parser.AdbCallParser, instance=True + ) + config = config_classes.AndroidDeviceConfig() + + with mock.patch( + 'builtins.open', mock.mock_open(read_data=sample_textproto) + ): + android_device_loader._prep_device(config, mock_parser) + + interpreter_instance = mock_interpreter_cls.return_value + steps = interpreter_instance.interpret.call_args[0][0] + install_steps = [s for s in steps if s.adb_request.HasField('install_apk')] + self.assertEmpty(install_steps) + + def test_get_prep_steps_path_exists(self): + self.assertTrue( + os.path.exists(android_device_loader._get_prep_steps_path()) + ) + + def test_dump_logcat_cb_handles_adb_error(self): + self._mock_adb.execute_command.side_effect = errors.AdbControllerError( + 'Logcat failed' + ) + config = config_classes.AndroidDeviceConfig( + emulator_launcher=config_classes.EmulatorLauncherConfig( + avd_name='test_avd', + emulator_console_port=5554, + ), + adb_controller=config_classes.AdbControllerConfig(device_name=''), + ) + with mock.patch.object(android_device_loader, '_wait_for_boot'): + with mock.patch.object(android_device_loader, '_prep_device'): + android_device_loader.load_android_device(config) + + on_error_cbs = self._mock_device_cls.call_args[1]['on_error_callbacks'] + self.assertLen(on_error_cbs, 1) + on_error_cbs[0](None, None, None) + + def test_load_android_device_error_runs_on_close_callbacks(self): + self._mock_conn_init.side_effect = errors.AdbControllerError('Conn failed') + config = config_classes.AndroidDeviceConfig( + emulator_launcher=config_classes.EmulatorLauncherConfig( + avd_name='test_avd', + emulator_console_port=5554, + ), + adb_controller=config_classes.AdbControllerConfig(device_name=''), + ) + with mock.patch.object(android_device_loader, '_wait_for_boot'): + with mock.patch.object(android_device_loader, '_prep_device'): + with self.assertRaises(errors.AdbControllerError): + android_device_loader.load_android_device(config) + + self._mock_launcher.close.assert_called_once() + + def test_connect_to_device_handles_adb_start_service_error(self): + self._mock_conn_init.side_effect = [ + RuntimeError('Connection refused'), + self._mock_conn, + ] + self._mock_adb.execute_command.side_effect = errors.AdbControllerError( + 'Start service error' + ) + config = config_classes.AndroidDeviceConfig(auto_setup=True) + conn = android_device_loader._connect_to_device(config, self._mock_adb) + self.assertEqual(conn, self._mock_conn) + + def test_connect_to_device_raises_runtime_error_if_max_retries_zero(self): + config = config_classes.AndroidDeviceConfig(auto_setup=False) + with self.assertRaisesRegex( + RuntimeError, 'Failed to connect to on-device server' + ): + android_device_loader._connect_to_device( + config, self._mock_adb, max_retries=0 + ) + + +if __name__ == '__main__': + absltest.main() diff --git a/android_env/components/android_device_test.py b/android_env/components/android_device_test.py new file mode 100644 index 00000000..c06ed89d --- /dev/null +++ b/android_env/components/android_device_test.py @@ -0,0 +1,217 @@ +# coding=utf-8 +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for android_device.AndroidDevice.""" + +import time +from unittest import mock + +from absl.testing import absltest +from android_env.components import adb_call_parser +from android_env.components import android_device +from android_env.components import device_connection +from android_env.components import errors +from android_env.proto import adb_pb2 +from android_env.proto import device_env_service_pb2 +import cv2 +import numpy as np + + +class AndroidDeviceTest(absltest.TestCase): + + def setUp(self): + super().setUp() + self._mock_connection = mock.create_autospec( + device_connection.DeviceConnection, instance=True + ) + self._mock_connection.get_device_state.return_value = ( + device_env_service_pb2.DeviceState() + ) + + def test_get_state_decodes_screenshot(self): + img = np.zeros((50, 50, 3), dtype=np.uint8) + img[5:15, 5:15] = [0, 255, 0] + success, encoded_img = cv2.imencode('.png', img) + self.assertTrue(success) + + state = device_env_service_pb2.DeviceState() + sig = state.signals.add() + sig.type = device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT + sig.screenshot.encoded_bytes = encoded_img.tobytes() + + self._mock_connection.get_device_state.return_value = state + device = android_device.AndroidDevice(connection=self._mock_connection) + + dev_state = device.get_state() + screenshot_sig = dev_state.signals[0] + decoded = screenshot_sig.screenshot.decoded + pixels = np.frombuffer(decoded.raw_pixels, dtype=np.uint8).reshape( + (decoded.height, decoded.width, 3) + ) + expected_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + np.testing.assert_array_equal(pixels, expected_rgb) + + def test_decode_screenshot_invalid_image(self): + state = device_env_service_pb2.DeviceState() + sig = state.signals.add() + sig.type = device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT + sig.screenshot.encoded_bytes = b'invalid_image_bytes' + + self._mock_connection.get_device_state.return_value = state + device = android_device.AndroidDevice(connection=self._mock_connection) + + dev_state = device.get_state() + screenshot_sig = dev_state.signals[0] + self.assertFalse(screenshot_sig.screenshot.HasField('decoded')) + + @mock.patch.object(time, 'time_ns', autospec=True) + def test_clock_sync_calculates_offset(self, mock_time): + def make_time_func(t1_ns): + times = [1_000_000_000, t1_ns] + return lambda: times.pop(0) if times else 1_030_000_000 + + for t1_ns, expected_offset in [ + (1_020_000_000, 3990), # rtt = 20ms -> rtt//2 = 10ms -> offset = 3990ms + (1_021_000_000, 3990), # rtt = 21ms -> rtt//2 = 10ms -> offset = 3990ms + (1_022_000_000, 3989), # rtt = 22ms -> rtt//2 = 11ms -> offset = 3989ms + ]: + with self.subTest(t1_ns=t1_ns): + mock_time.side_effect = make_time_func(t1_ns) + state = device_env_service_pb2.DeviceState() + sig = state.signals.add() + sig.type = device_env_service_pb2.DEVICE_SIGNAL_DEVICE_SETTINGS + sig.device_settings.entries['device_time_ms'] = '5000' + + self._mock_connection.get_device_state.return_value = state + device = android_device.AndroidDevice(connection=self._mock_connection) + + self.assertEqual(device._clock_offset_ms, expected_offset) + + def test_clock_sync_handles_connection_error(self): + self._mock_connection.get_device_state.side_effect = RuntimeError( + 'Connection error' + ) + device = android_device.AndroidDevice(connection=self._mock_connection) + self.assertEqual(device._clock_offset_ms, 0) + + def test_clock_sync_handles_value_error(self): + self._mock_connection.get_device_state.side_effect = ValueError( + 'Value error' + ) + device = android_device.AndroidDevice(connection=self._mock_connection) + self.assertEqual(device._clock_offset_ms, 0) + + def test_clock_sync_missing_device_time_setting(self): + state = device_env_service_pb2.DeviceState() + sig = state.signals.add() + sig.type = device_env_service_pb2.DEVICE_SIGNAL_DEVICE_SETTINGS + # entries lacks device_time_ms + self._mock_connection.get_device_state.return_value = state + device = android_device.AndroidDevice(connection=self._mock_connection) + self.assertEqual(device._clock_offset_ms, 0) + + def test_send_action_sends_client_message(self): + device = android_device.AndroidDevice(connection=self._mock_connection) + self._mock_connection.send_message.reset_mock() + + action = device_env_service_pb2.Action() + action.action_type = device_env_service_pb2.ACTION_TYPE_TOUCH_DOWN + device.send_action(action) + + self._mock_connection.send_message.assert_called_once() + sent_msg = self._mock_connection.send_message.call_args[0][0] + self.assertTrue(sent_msg.HasField('inject_action')) + self.assertEqual(sent_msg.inject_action.action, action) + + def test_update_subscriptions_sends_client_message(self): + device = android_device.AndroidDevice(connection=self._mock_connection) + self._mock_connection.send_message.reset_mock() + + signals = {device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT} + device.update_subscriptions(signals) + + self._mock_connection.send_message.assert_called_once() + sent_msg = self._mock_connection.send_message.call_args[0][0] + self.assertTrue(sent_msg.HasField('update_subscriptions')) + self.assertCountEqual(sent_msg.update_subscriptions.active_signals, signals) + + def test_close_runs_callbacks(self): + cb1 = mock.MagicMock() + cb2 = mock.MagicMock() + device = android_device.AndroidDevice( + connection=self._mock_connection, + on_close_callbacks=[cb1, cb2], + ) + device.close() + + self._mock_connection.close.assert_called_once() + cb1.assert_called_once() + cb2.assert_called_once() + + def test_context_manager_normal_exit(self): + cb = mock.MagicMock() + with android_device.AndroidDevice( + connection=self._mock_connection, + on_close_callbacks=[cb], + ): + pass + + self._mock_connection.close.assert_called_once() + cb.assert_called_once() + + def test_context_manager_error_exit(self): + on_close = mock.MagicMock() + on_error = mock.MagicMock() + try: + with android_device.AndroidDevice( + connection=self._mock_connection, + on_close_callbacks=[on_close], + on_error_callbacks=[on_error], + ): + raise ValueError('Test error') + except ValueError: + pass + + on_error.assert_called_once() + on_close.assert_called_once() + self._mock_connection.close.assert_called_once() + + def test_execute_adb_call_with_adb_call_parser(self): + mock_parser = mock.create_autospec( + adb_call_parser.AdbCallParser, instance=True + ) + device = android_device.AndroidDevice( + connection=self._mock_connection, + adb_call_parser=mock_parser, + ) + + req = adb_pb2.AdbRequest( + generic=adb_pb2.AdbRequest.GenericRequest(args=['shell', 'ls']) + ) + mock_parser.parse.return_value = adb_pb2.AdbResponse() + res = device.execute_adb_call(req) + mock_parser.parse.assert_called_once_with(req) + self.assertIsInstance(res, adb_pb2.AdbResponse) + + def test_execute_adb_call_without_adb_call_parser_raises(self): + device = android_device.AndroidDevice(connection=self._mock_connection) + + req = adb_pb2.AdbRequest() + with self.assertRaises(errors.AdbControllerError): + device.execute_adb_call(req) + + +if __name__ == '__main__': + absltest.main() diff --git a/android_env/components/config_classes.py b/android_env/components/config_classes.py index a1792dc1..eef86ba4 100644 --- a/android_env/components/config_classes.py +++ b/android_env/components/config_classes.py @@ -172,6 +172,39 @@ class FakeSimulatorConfig(SimulatorConfig): screen_dimensions: tuple[int, int] = (0, 0) +@dataclasses.dataclass +class DeviceConnectionConfig: + """Configuration for network connection to the on-device AndroidEnv server. + + Attributes: + host: Hostname or IP address where the AndroidEnv server runs (e.g. + 'localhost' when port-forwarded via ADB, or remote device IP). + websocket_port: Port for the WebSocket control channel served by the device. + The Python client connects to `ws://:`. + Payloads: Binary Protobuf `ClientMessage` and `ServerMessage` for action + injection, subscription updates, video frames, and device + state/metadata. + """ + + host: str = 'localhost' + websocket_port: int = 8000 + + +@dataclasses.dataclass +class AndroidDeviceConfig(SimulatorConfig): + """Config class for AndroidDevice (on-device Core server).""" + + device_connection: DeviceConnectionConfig = dataclasses.field( + default_factory=DeviceConnectionConfig + ) + adb_path: str = 'adb' + serial: str = '' + auto_setup: bool = True + # Optional launcher config. If set, loader will launch emulator. + emulator_launcher: EmulatorLauncherConfig | None = None + adb_controller: AdbControllerConfig | None = None + + @dataclasses.dataclass class TaskManagerConfig: """Config class for TaskManager.""" diff --git a/android_env/components/device_connection.py b/android_env/components/device_connection.py new file mode 100644 index 00000000..39b12599 --- /dev/null +++ b/android_env/components/device_connection.py @@ -0,0 +1,117 @@ +# coding=utf-8 +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Manages low-level WebSocket connection to the on-device server.""" + +import socket +import threading + +from absl import logging +from android_env.components import config_classes +from android_env.proto import device_env_service_pb2 +import websocket + + +class DeviceConnection: + """Manages low-level WebSocket connection to the on-device AndroidEnv server.""" + + def __init__(self, config: config_classes.DeviceConnectionConfig): + self._config = config + self._url = f'ws://{self._config.host}:{self._config.websocket_port}' + self._ws = None + self._state_lock = threading.Lock() + self._state_condition = threading.Condition(self._state_lock) + self._send_lock = threading.Lock() + self._latest_state: device_env_service_pb2.DeviceState | None = None + self._latest_error: str | None = None + self._running = False + self._reader_thread: threading.Thread | None = None + self.connect() + + def connect(self) -> None: + """Establishes connection to the on-device WebSocket server.""" + try: + self._ws = websocket.create_connection(self._url) + except (websocket.WebSocketException, socket.error) as e: + raise RuntimeError( + f'Failed to connect to WebSocket at {self._url}: {e}' + ) from e + + self._running = True + self._reader_thread = threading.Thread( + target=self._receive_loop, daemon=True + ) + self._reader_thread.start() + + def _receive_loop(self) -> None: + """Background loop to continuously read incoming WebSocket messages.""" + while self._running and self._ws: + try: + resp_bytes = self._ws.recv() + if resp_bytes is None or not self._running: + continue + resp = device_env_service_pb2.ServerMessage() + resp.ParseFromString(resp_bytes) + + with self._state_condition: + if resp.HasField('error'): + self._latest_error = f'Server error: {resp.error.message}' + self._latest_state = None + elif resp.HasField('state'): + self._latest_state = resp.state + self._latest_error = None + else: + self._latest_error = 'Server response missing state' + self._latest_state = None + self._state_condition.notify_all() + except (websocket.WebSocketException, socket.error, AttributeError): + logging.warning('WebSocket connection closed or error encountered.') + break + except (ValueError, TypeError): + logging.exception('Error decoding ServerMessage.') + + def close(self) -> None: + """Closes the WebSocket connection.""" + self._running = False + if self._ws: + try: + self._ws.close() + except (websocket.WebSocketException, socket.error): + pass + self._ws = None + + def send_message(self, msg: device_env_service_pb2.ClientMessage) -> None: + """Sends a binary ClientMessage to the on-device server.""" + if not self._ws: + raise RuntimeError('Not connected') + with self._send_lock: + self._ws.send_binary(msg.SerializeToString()) + + def get_device_state(self) -> device_env_service_pb2.DeviceState: + """Returns the latest raw device state received from the server.""" + if not self._ws: + raise RuntimeError('Not connected') + + with self._state_condition: + if self._latest_state is None and self._latest_error is None: + req = device_env_service_pb2.GetStateRequest() + msg = device_env_service_pb2.ClientMessage(get_state=req) + self.send_message(msg) + self._state_condition.wait(timeout=5.0) + if self._latest_error: + raise RuntimeError(self._latest_error) + if self._latest_state is None: + raise RuntimeError('Timed out waiting for initial device state') + return self._latest_state diff --git a/android_env/components/device_connection_test.py b/android_env/components/device_connection_test.py new file mode 100644 index 00000000..de55b20d --- /dev/null +++ b/android_env/components/device_connection_test.py @@ -0,0 +1,192 @@ +# coding=utf-8 +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for device_connection, verifying pure WebSocket transport.""" + +from unittest import mock + +from absl.testing import absltest +from android_env.components import config_classes +from android_env.components import device_connection +from android_env.proto import device_env_service_pb2 +import websocket + + +class DeviceConnectionTest(absltest.TestCase): + + def setUp(self): + super().setUp() + self._mock_ws_create = mock.patch.object( + websocket, 'create_connection', autospec=True + ).start() + self._mock_ws = self._mock_ws_create.return_value + self._set_mock_state_response() + + def tearDown(self): + mock.patch.stopall() + super().tearDown() + + def _set_mock_state_response(self): + msg = device_env_service_pb2.ServerMessage() + sig = msg.state.signals.add() + sig.type = device_env_service_pb2.DEVICE_SIGNAL_ACTIVE_PACKAGE + sig.active_package = 'com.example.app' + self._mock_ws.recv.return_value = msg.SerializeToString() + + def test_init_success(self): + config = config_classes.DeviceConnectionConfig( + host='localhost', websocket_port=8000 + ) + conn = device_connection.DeviceConnection(config=config) + + self._mock_ws_create.assert_called_once_with('ws://localhost:8000') + conn.close() + + def test_init_failure_raises_error(self): + self._mock_ws_create.side_effect = websocket.WebSocketException( + 'WS Connection failed' + ) + + config = config_classes.DeviceConnectionConfig( + host='localhost', websocket_port=8000 + ) + with self.assertRaises(RuntimeError) as context: + device_connection.DeviceConnection(config=config) + + self.assertIn('Failed to connect to WebSocket', str(context.exception)) + + def test_send_message(self): + config = config_classes.DeviceConnectionConfig( + host='localhost', websocket_port=8000 + ) + conn = device_connection.DeviceConnection(config=config) + self._mock_ws.send_binary.reset_mock() + + msg = device_env_service_pb2.ClientMessage() + msg.get_state.SetInParent() + conn.send_message(msg) + + self._mock_ws.send_binary.assert_called_once_with(msg.SerializeToString()) + conn.close() + + def test_get_device_state(self): + config = config_classes.DeviceConnectionConfig( + host='localhost', websocket_port=8000 + ) + conn = device_connection.DeviceConnection(config=config) + state = conn.get_device_state() + + self.assertLen(state.signals, 1) + self.assertEqual( + state.signals[0].type, + device_env_service_pb2.DEVICE_SIGNAL_ACTIVE_PACKAGE, + ) + conn.close() + + def test_get_device_state_sends_request_when_state_is_none(self): + self._mock_ws.recv.side_effect = websocket.WebSocketException('No state') + + config = config_classes.DeviceConnectionConfig( + host='localhost', websocket_port=8000 + ) + conn = device_connection.DeviceConnection(config=config) + self._mock_ws.send_binary.reset_mock() + + def mock_send_binary(data): + sent_msg = device_env_service_pb2.ClientMessage() + sent_msg.ParseFromString(data) + self.assertTrue(sent_msg.HasField('get_state')) + state = device_env_service_pb2.DeviceState() + sig = state.signals.add() + sig.type = device_env_service_pb2.DEVICE_SIGNAL_ACTIVE_PACKAGE + sig.active_package = 'com.test.app' + conn._latest_state = state + + self._mock_ws.send_binary.side_effect = mock_send_binary + + state = conn.get_device_state() + + self.assertEqual(state.signals[0].active_package, 'com.test.app') + self._mock_ws.send_binary.assert_called_once() + conn.close() + + def test_get_device_state_raises_server_error(self): + msg = device_env_service_pb2.ServerMessage() + msg.error.message = 'Internal server error' + self._mock_ws.recv.return_value = msg.SerializeToString() + + config = config_classes.DeviceConnectionConfig( + host='localhost', websocket_port=8000 + ) + conn = device_connection.DeviceConnection(config=config) + + with self.assertRaisesRegex( + RuntimeError, 'Server error: Internal server error' + ): + conn.get_device_state() + + conn.close() + + def test_get_device_state_raises_when_missing_state_or_error(self): + msg = device_env_service_pb2.ServerMessage() + self._mock_ws.recv.return_value = msg.SerializeToString() + + config = config_classes.DeviceConnectionConfig( + host='localhost', websocket_port=8000 + ) + conn = device_connection.DeviceConnection(config=config) + + with self.assertRaisesRegex(RuntimeError, 'Server response missing state'): + conn.get_device_state() + + conn.close() + + def test_get_device_state_timeout(self): + self._mock_ws.recv.side_effect = websocket.WebSocketException('No state') + + config = config_classes.DeviceConnectionConfig( + host='localhost', websocket_port=8000 + ) + conn = device_connection.DeviceConnection(config=config) + + def mock_wait(*args, **kwargs): + del args, kwargs + + with mock.patch.object( + conn._state_condition, 'wait', side_effect=mock_wait + ): + with self.assertRaisesRegex( + RuntimeError, 'Timed out waiting for initial device state' + ): + conn.get_device_state() + + conn.close() + + def test_not_connected_raises_error(self): + config = config_classes.DeviceConnectionConfig( + host='localhost', websocket_port=8000 + ) + conn = device_connection.DeviceConnection(config=config) + conn.close() + + with self.assertRaisesRegex(RuntimeError, 'Not connected'): + conn.get_device_state() + + with self.assertRaisesRegex(RuntimeError, 'Not connected'): + conn.send_message(device_env_service_pb2.ClientMessage()) + + +if __name__ == '__main__': + absltest.main() diff --git a/android_env/components/device_prep_steps.textproto b/android_env/components/device_prep_steps.textproto new file mode 100644 index 00000000..9fe6df67 --- /dev/null +++ b/android_env/components/device_prep_steps.textproto @@ -0,0 +1,104 @@ +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# proto-file: android_env/proto/task.proto +# proto-message: android_env.Task + +# 1. Verify device state +setup_steps { + adb_request { + generic { + args: "get-state" + } + } +} + +# 2. Grant Media Projection permission +setup_steps { + adb_request { + generic { + args: "shell" + args: "cmd" + args: "appops" + args: "set" + args: "com.google.androidenv.core" + args: "PROJECT_MEDIA" + args: "allow" + } + } +} + +# 3. Enable Accessibility Service +setup_steps { + adb_request { + settings { + name_space: SECURE + put { + key: "enabled_accessibility_services" + value: "com.google.androidenv.core/com.google.androidenv.core.InputService" + } + } + } +} +setup_steps { + adb_request { + settings { + name_space: SECURE + put { + key: "accessibility_enabled" + value: "1" + } + } + } +} + +# 4. Enable show touches +setup_steps { + adb_request { + settings { + name_space: SYSTEM + put { + key: "show_touches" + value: "1" + } + } + } +} + +# 5. Clear logcat and start MainActivity SERVER +setup_steps { + adb_request { + generic { + args: "logcat" + args: "-c" + } + } +} +setup_steps { + adb_request { + start_activity { + full_activity: "com.google.androidenv.core/.MainActivity" + extra_args: "--es" + extra_args: "COMMAND" + extra_args: "START_SERVER" + } + } + success_condition { + check_dumpsys { + service: "accessibility" + expected_string: "Input Service" + timeout_sec: 30.0 + } + } +} diff --git a/android_env/components/task_manager.py b/android_env/components/task_manager.py index 4b873a37..2200a62c 100644 --- a/android_env/components/task_manager.py +++ b/android_env/components/task_manager.py @@ -91,6 +91,7 @@ def __init__( 'extra': collections.defaultdict(list), 'episode_end': False, } + self._latest_active_package: Any = None logging.info('Task config: %s', self._task) @@ -207,6 +208,7 @@ def rl_step(self, observation: dict[str, Any]) -> dm_env.TimeStep: ), 'LogcatThread is None. `start()` must be called before `rl_step()`.' self._logcat_thread.line_ready().wait() with self._lock: + self._latest_active_package = observation.get('active_package') reward = self._get_current_reward() extras = self._get_current_extras() transition_fn = self._determine_transition_fn() @@ -230,6 +232,19 @@ def _get_current_extras(self) -> dict[str, Any]: self._latest_values['extra'] = collections.defaultdict(list) return extras + def _check_user_exited(self) -> bool: + """Returns True if the user exited the expected app screen.""" + active_package = self._latest_active_package + if active_package is not None: + expected_package = self._task.expected_app_screen.activity.split('/')[0] + return bool( + expected_package + and active_package + and active_package != expected_package + ) + assert self._dumpsys_thread is not None + return self._dumpsys_thread.check_user_exited() + def _determine_transition_fn(self) -> Callable[..., dm_env.TimeStep]: """Determines the type of RL transition will be used.""" @@ -239,7 +254,7 @@ def _determine_transition_fn(self) -> Callable[..., dm_env.TimeStep]: ) # Check if user existed the task - if self._dumpsys_thread.check_user_exited(): + if self._check_user_exited(): self._increment_bad_state() self._stats['reset_count_user_exited'] += 1 logging.warning('User exited the task. Truncating the episode.') diff --git a/android_env/components/telemetry.py b/android_env/components/telemetry.py new file mode 100644 index 00000000..5b0e76e6 --- /dev/null +++ b/android_env/components/telemetry.py @@ -0,0 +1,169 @@ +# coding=utf-8 +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Structured telemetry data classes for AndroidEnv. + +This module defines structured telemetry data classes such as +`BatteryPowerMetrics`. Metrics support both attribute access +(e.g., `metrics.level_percent`, `metrics.power_mw`) and dictionary subscript +indexing (e.g., `metrics['level_percent']`). +""" + +from collections.abc import Mapping +import dataclasses +from typing import Any + +from absl import logging +from android_env.proto import device_env_service_pb2 + + +@dataclasses.dataclass(frozen=True) +class BatteryPowerMetrics: + """Structured telemetry metrics for Android device battery and power state. + + Supports attribute access (`metrics.level_percent`, `metrics.power_mw`). + + Example: + >>> metrics = telemetry.BatteryPowerMetrics.from_proto(proto) + >>> print(metrics.level_percent) + 85 + >>> print(metrics.power_mw) + 1890.0 + >>> print(metrics.voltage_mv) + 4200 + + Attributes: + level_percent: Battery charge level percentage [0-100]. + current_now_ua: Instantaneous current draw in microamperes (uA). Positive + for discharge, negative for charging. + voltage_mv: Battery voltage in millivolts (mV). + power_mw: Instantaneous power in milliwatts (mW). + temperature_celsius: Battery temperature in degrees Celsius (°C). + status: Battery status string (e.g., 'CHARGING', 'DISCHARGING', 'FULL'). + plugged_source: Power source ('AC', 'USB', 'WIRELESS', 'UNPLUGGED'). + charge_counter_uah: Remaining charge counter in microampere-hours (uAh). + thermal_status: Thermal status integer code (0 = NONE to 6 = SHUTDOWN). + energy_counter_nwh: Total remaining battery energy in nanowatt-hours (nWh). + """ + + level_percent: int + current_now_ua: int + voltage_mv: int + power_mw: float + temperature_celsius: float + status: str + plugged_source: str + charge_counter_uah: int + thermal_status: int + energy_counter_nwh: int + + @classmethod + def from_proto( + cls, proto: device_env_service_pb2.BatteryPowerMetrics + ) -> 'BatteryPowerMetrics': + """Creates a BatteryPowerMetrics instance from a protobuf message. + + Example: + >>> metrics = BatteryPowerMetrics.from_proto(proto) + >>> print(metrics.level_percent) + 85 + + Args: + proto: Protobuf message containing battery power metrics. + + Returns: + An immutable BatteryPowerMetrics instance. + """ + return cls( + level_percent=proto.level_percent, + current_now_ua=proto.current_now_ua, + voltage_mv=proto.voltage_mv, + power_mw=proto.power_mw, + temperature_celsius=proto.temperature_celsius, + status=proto.status, + plugged_source=proto.plugged_source, + charge_counter_uah=proto.charge_counter_uah, + thermal_status=proto.thermal_status, + energy_counter_nwh=proto.energy_counter_nwh, + ) + + +def _extract_from_obj_or_mapping(bp: Any) -> BatteryPowerMetrics | None: + """Extracts BatteryPowerMetrics from a proto, dataclass, mock, or dict object. + + Args: + bp: Source object containing battery power metric attributes or keys. + + Returns: + Extracted BatteryPowerMetrics instance, or None if extraction fails. + """ + if bp is None: + return None + if isinstance(bp, BatteryPowerMetrics): + return bp + if isinstance(bp, device_env_service_pb2.BatteryPowerMetrics): + return BatteryPowerMetrics.from_proto(bp) + + def _get_val(obj: Any, name: str, default: Any) -> Any: + val = getattr(obj, name, None) + if val is None and isinstance(obj, (dict, Mapping)) and name in obj: + val = obj[name] + if hasattr(val, 'item'): + val = val.item() + if val is not None and isinstance(val, (int, float, str, bool)): + return val + return default + + level_percent = _get_val(bp, 'level_percent', None) + if level_percent is not None: + return BatteryPowerMetrics( + level_percent=int(level_percent), + current_now_ua=int(_get_val(bp, 'current_now_ua', 0)), + voltage_mv=int(_get_val(bp, 'voltage_mv', 4000)), + power_mw=float(_get_val(bp, 'power_mw', 0.0)), + temperature_celsius=float(_get_val(bp, 'temperature_celsius', 25.0)), + status=str(_get_val(bp, 'status', 'UNKNOWN')), + plugged_source=str(_get_val(bp, 'plugged_source', 'UNPLUGGED')), + charge_counter_uah=int(_get_val(bp, 'charge_counter_uah', 0)), + thermal_status=int(_get_val(bp, 'thermal_status', 0)), + energy_counter_nwh=int(_get_val(bp, 'energy_counter_nwh', 0)), + ) + return None + + +def fetch_battery_power_metrics(env: Any) -> BatteryPowerMetrics | None: + """Fetches battery metrics from task_extras or returns None if missing. + + Args: + env: Android environment instance. + + Returns: + BatteryPowerMetrics instance if present in task_extras, else None. + """ + try: + if hasattr(env, 'task_extras'): + extras = env.task_extras(latest_only=True) + if extras and extras.get('battery_power') is not None: + bp = extras['battery_power'] + extracted = _extract_from_obj_or_mapping(bp) + if extracted is not None: + return extracted + except (AttributeError, KeyError, TypeError, RuntimeError) as e: + logging.vlog(1, 'Could not extract battery_power from task_extras: %s', e) + + return None + + +get_battery_metrics = fetch_battery_power_metrics diff --git a/android_env/components/telemetry_test.py b/android_env/components/telemetry_test.py new file mode 100644 index 00000000..37df6ec2 --- /dev/null +++ b/android_env/components/telemetry_test.py @@ -0,0 +1,182 @@ +# coding=utf-8 +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for structured telemetry data classes.""" + +import dataclasses +from unittest import mock +from absl.testing import absltest +from android_env.components import telemetry +from android_env.proto import device_env_service_pb2 + + +class BatteryPowerMetricsTest(absltest.TestCase): + + def _make_sample_proto(self) -> device_env_service_pb2.BatteryPowerMetrics: + proto = device_env_service_pb2.BatteryPowerMetrics() + proto.level_percent = 85 + proto.current_now_ua = 450000 + proto.voltage_mv = 4200 + proto.power_mw = 1890.0 + proto.temperature_celsius = 28.5 + proto.status = 'DISCHARGING' + proto.plugged_source = 'UNPLUGGED' + proto.charge_counter_uah = 3000000 + proto.thermal_status = 0 + proto.energy_counter_nwh = 15000000000 + return proto + + def test_from_proto_and_attribute_access(self): + proto = self._make_sample_proto() + metrics = telemetry.BatteryPowerMetrics.from_proto(proto) + + self.assertEqual(metrics.level_percent, 85) + self.assertEqual(metrics.current_now_ua, 450000) + self.assertEqual(metrics.voltage_mv, 4200) + self.assertAlmostEqual(metrics.power_mw, 1890.0) + self.assertAlmostEqual(metrics.temperature_celsius, 28.5) + self.assertEqual(metrics.status, 'DISCHARGING') + self.assertEqual(metrics.plugged_source, 'UNPLUGGED') + self.assertEqual(metrics.charge_counter_uah, 3000000) + self.assertEqual(metrics.thermal_status, 0) + self.assertEqual(metrics.energy_counter_nwh, 15000000000) + + def test_immutability_and_replace(self): + proto = self._make_sample_proto() + metrics = telemetry.BatteryPowerMetrics.from_proto(proto) + + with self.assertRaises((dataclasses.FrozenInstanceError, AttributeError)): + setattr(metrics, 'level_percent', 50) + + replaced = dataclasses.replace(metrics, level_percent=90) + self.assertEqual(replaced.level_percent, 90) + self.assertEqual(metrics.level_percent, 85) + self.assertNotEqual(metrics, replaced) + + def test_extract_from_obj_or_mapping(self): + proto = self._make_sample_proto() + metrics = telemetry.BatteryPowerMetrics.from_proto(proto) + + # Test None input + self.assertIsNone(telemetry._extract_from_obj_or_mapping(None)) + + # Test returning existing BatteryPowerMetrics instance + self.assertIs(telemetry._extract_from_obj_or_mapping(metrics), metrics) + + # Test extracting from proto message and verify from_proto call + with mock.patch.object( + telemetry.BatteryPowerMetrics, + 'from_proto', + wraps=telemetry.BatteryPowerMetrics.from_proto, + ) as mock_from_proto: + extracted_proto = telemetry._extract_from_obj_or_mapping(proto) + mock_from_proto.assert_called_once_with(proto) + self.assertEqual(extracted_proto, metrics) + + # Test extracting from dict + as_dict = { + 'level_percent': 85, + 'current_now_ua': 450000, + 'voltage_mv': 4200, + 'power_mw': 1890.0, + 'temperature_celsius': 28.5, + 'status': 'DISCHARGING', + 'plugged_source': 'UNPLUGGED', + 'charge_counter_uah': 3000000, + 'thermal_status': 0, + 'energy_counter_nwh': 15000000000, + } + extracted_dict = telemetry._extract_from_obj_or_mapping(as_dict) + self.assertEqual(extracted_dict, metrics) + + # Test extracting from object with .item() scalar wrappers (e.g. numpy scalars) + class FakeScalar: + + def __init__(self, val): + self._val = val + + def item(self): + return self._val + + class FakeNumpyObj: + + def __init__(self): + self.level_percent = FakeScalar(85) + self.current_now_ua = FakeScalar(450000) + self.voltage_mv = FakeScalar(4200) + self.power_mw = FakeScalar(1890.0) + self.temperature_celsius = FakeScalar(28.5) + self.status = 'DISCHARGING' + self.plugged_source = 'UNPLUGGED' + self.charge_counter_uah = FakeScalar(3000000) + self.thermal_status = FakeScalar(0) + self.energy_counter_nwh = FakeScalar(15000000000) + + extracted_numpy = telemetry._extract_from_obj_or_mapping(FakeNumpyObj()) + self.assertEqual(extracted_numpy, metrics) + + # Test invalid object (missing level_percent) returns None + self.assertIsNone( + telemetry._extract_from_obj_or_mapping({'other_key': 123}) + ) + + # Test object with invalid non-primitive attribute types (e.g. list/dict) returns None + class FakeInvalidTypesObj: + level_percent = [85, 90] + status = {'invalid': 'type'} + + self.assertIsNone( + telemetry._extract_from_obj_or_mapping(FakeInvalidTypesObj()) + ) + self.assertIsNone( + telemetry._extract_from_obj_or_mapping({'level_percent': [85, 90]}) + ) + + # Verify attribute access takes precedence over dictionary subscripting + class HybridObj(dict): + level_percent = 95 + status = 'CHARGING' + + hybrid = HybridObj({'level_percent': 10, 'status': 'DISCHARGING'}) + extracted_hybrid = telemetry._extract_from_obj_or_mapping(hybrid) + self.assertIsNotNone(extracted_hybrid) + self.assertEqual(extracted_hybrid.level_percent, 95) + self.assertEqual(extracted_hybrid.status, 'CHARGING') + + def test_fetch_battery_power_metrics_task_extras(self): + mock_env = mock.Mock() + mock_battery = self._make_sample_proto() + mock_env.task_extras.return_value = {'battery_power': mock_battery} + + metrics = telemetry.fetch_battery_power_metrics(mock_env) + self.assertIsNotNone(metrics) + self.assertEqual(metrics.level_percent, 85) + + def test_fetch_battery_power_metrics_fallback_and_exception(self): + mock_env = mock.Mock() + mock_env.task_extras.return_value = {} + self.assertIsNone(telemetry.fetch_battery_power_metrics(mock_env)) + + mock_env_none = mock.Mock() + mock_env_none.task_extras.return_value = {'battery_power': None} + self.assertIsNone(telemetry.fetch_battery_power_metrics(mock_env_none)) + + mock_env_err = mock.Mock() + mock_env_err.task_extras.side_effect = RuntimeError('SDK error') + self.assertIsNone(telemetry.fetch_battery_power_metrics(mock_env_err)) + + +if __name__ == '__main__': + absltest.main() diff --git a/android_env/mcp/BUILD.bazel b/android_env/mcp/BUILD.bazel new file mode 100644 index 00000000..02d44d1d --- /dev/null +++ b/android_env/mcp/BUILD.bazel @@ -0,0 +1,54 @@ +load("@rules_python//python:defs.bzl", "py_binary", "py_library", "py_test") + +package(default_visibility = ["//visibility:public"]) + +py_library( + name = "errors", + srcs = ["errors.py"], +) + +py_library( + name = "android_mcp_tools", + srcs = ["android_mcp_tools.py"], + deps = [ + ":errors", + "//android_env/components:android_device", + "//android_env/proto:device_env_service_py_pb2", + "//android_env/proto/a11y:android_accessibility_node_info_py_pb2", + "@pip//mcp", + "@pip//numpy", + "@pip//pillow", + "@pip//protobuf", + ], +) + +py_test( + name = "android_mcp_tools_test", + srcs = ["android_mcp_tools_test.py"], + legacy_create_init = False, + deps = [ + ":android_mcp_tools", + ":errors", + "//android_env/components:android_device", + "//android_env/proto:device_env_service_py_pb2", + "//android_env/proto/a11y:android_accessibility_forest_py_pb2", + "//android_env/proto/a11y:android_accessibility_node_info_py_pb2", + "@pip//absl_py", + "@pip//mcp", + "@pip//numpy", + "@pip//pillow", + "@pip//protobuf", + ], +) + +py_binary( + name = "android_mcp_server", + srcs = ["android_mcp_server.py"], + deps = [ + ":android_mcp_tools", + "//android_env/components:android_device_loader", + "//android_env/components:config_classes", + "@pip//absl_py", + "@pip//mcp", + ], +) diff --git a/android_env/mcp/android_mcp_server.py b/android_env/mcp/android_mcp_server.py new file mode 100644 index 00000000..4b96cf0a --- /dev/null +++ b/android_env/mcp/android_mcp_server.py @@ -0,0 +1,64 @@ +# coding=utf-8 +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""FastMCP server binary for AndroidEnv device interaction. + +Provides a stdio FastMCP server entrypoint for AndroidEnv device management +and interaction. + +Usage: + python -m android_env.mcp.android_mcp_server --serial=emulator-5554 +""" + +from collections.abc import Sequence + +from absl import app +from absl import flags +from android_env.components import android_device_loader +from android_env.components import config_classes +from android_env.mcp import android_mcp_tools +from mcp.server import fastmcp + +_SERIAL = flags.DEFINE_string( + 'serial', None, 'ADB serial number of target device.' +) +_WEBSOCKET_PORT = flags.DEFINE_integer( + 'websocket_port', 8000, 'On-device WebSocket port.' +) +_HOST = flags.DEFINE_string('host', 'localhost', 'On-device WebSocket host.') + + +def main(argv: Sequence[str]) -> None: + if len(argv) > 1: + raise app.UsageError('Too many command-line arguments.') + + config = config_classes.AndroidDeviceConfig( + serial=_SERIAL.value, # pyrefly: ignore[bad-argument-type] + auto_setup=True, + device_connection=config_classes.DeviceConnectionConfig( + host=_HOST.value, + websocket_port=_WEBSOCKET_PORT.value, + ), + ) + + device = android_device_loader.load_android_device(config) + mcp_server = fastmcp.FastMCP(name='android_mcp_server') + tools_handler = android_mcp_tools.AndroidMcpTools(device) + tools_handler.register_tools(mcp_server) + mcp_server.run(transport='stdio') + + +if __name__ == '__main__': + app.run(main) diff --git a/android_env/mcp/android_mcp_tools.py b/android_env/mcp/android_mcp_tools.py new file mode 100644 index 00000000..a7a1d862 --- /dev/null +++ b/android_env/mcp/android_mcp_tools.py @@ -0,0 +1,419 @@ +# coding=utf-8 +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""FastMCP tool implementations for direct AndroidDevice connection. + +Exposes Android device inspection and interaction capabilities to LLM agents +by receiving a connected AndroidDevice instance, bypassing the RL interface +(dm_env). +""" + +import base64 +from collections.abc import Mapping +import io +import json +from typing import Any + +from android_env.components import android_device as android_device_mod +from android_env.mcp import errors +from android_env.proto import device_env_service_pb2 +from android_env.proto.a11y import android_accessibility_node_info_pb2 +from google.protobuf import json_format +from mcp.server import fastmcp +import numpy as np +from PIL import Image + +_DEFAULT_SIGNALS = frozenset({ + device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT, + device_env_service_pb2.DEVICE_SIGNAL_ACCESSIBILITY_TREE, + device_env_service_pb2.DEVICE_SIGNAL_ACTIVE_PACKAGE, + device_env_service_pb2.DEVICE_SIGNAL_DEVICE_SETTINGS, +}) + + +def _extract_screenshot( + state: device_env_service_pb2.DeviceState, +) -> np.ndarray: + """Extracts decoded RGB screenshot pixels from DeviceState signals. + + Args: + state: A `DeviceState` protobuf containing device signals. + + Returns: + A NumPy array of shape `(height, width, 3)` with uint8 dtype containing + decoded RGB pixel values. + + Raises: + errors.DeviceConnectionError: If no decoded screenshot signal is found in + the provided device state. + """ + for signal in state.signals: + if ( + signal.type == device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT + and signal.HasField('screenshot') + and signal.screenshot.HasField('decoded') + ): + decoded = signal.screenshot.decoded + return np.frombuffer(decoded.raw_pixels, dtype=np.uint8).reshape( + (decoded.height, decoded.width, 3) + ) + raise errors.DeviceConnectionError( + 'No decoded screenshot signal found in device state.' + ) + + +def _render_node( + nodes_by_id: Mapping[ + int, android_accessibility_node_info_pb2.AndroidAccessibilityNodeInfo + ], + node_id: int, + indent: int, + lines: list[str], + visited: set[int] | None = None, +) -> None: + """Recursively renders an accessibility node and its children into text lines. + + Args: + nodes_by_id: Mapping from unique node ID to `AndroidAccessibilityNodeInfo` + protobuf messages. + node_id: The ID of the current accessibility node to render. + indent: The current indentation level in the tree display. + lines: Output list of strings where rendered line entries are appended. + visited: Set of already visited node IDs used to prevent infinite recursion + in cyclic graphs. + """ + if visited is None: + visited = set() + if node_id in visited: + return + visited.add(node_id) + + node = nodes_by_id.get(node_id) + if node is None: + return + + prefix = ' ' * indent + text = node.text or node.content_description or '' + bounds = node.bounds_in_screen + bounds_str = f'{{{bounds.left},{bounds.top},{bounds.right},{bounds.bottom}}}' + + flags = [] + if node.is_clickable: + flags.append('clickable') + if node.is_scrollable: + flags.append('scrollable') + if node.is_checked: + flags.append('checked') + if node.is_focused: + flags.append('focused') + if node.is_enabled: + flags.append('enabled') + if node.is_editable: + flags.append('editable') + + flag_str = f' [{" ".join(flags)}]' if flags else '' + lines.append(f'{prefix}[{node.class_name}] "{text}" {bounds_str}{flag_str}') + + for child_id in node.child_ids: + _render_node(nodes_by_id, child_id, indent + 1, lines, visited) + + +def _format_accessibility_tree( + state: device_env_service_pb2.DeviceState, + output_format: str = 'tree', +) -> str: + """Extracts and formats the accessibility tree from DeviceState signals. + + Args: + state: A `DeviceState` protobuf message containing device signals. + output_format: Desired string format of the accessibility tree output; + either `'tree'` for human-readable tree lines or `'json'` for serialized + JSON. Defaults to `'tree'`. + + Returns: + Formatted string representation of the accessibility tree. + + Raises: + ValueError: If `output_format` is not `'tree'` or `'json'`. + errors.DeviceConnectionError: If no accessibility tree signal is present in + the device state. + """ + if output_format not in ('tree', 'json'): + raise ValueError( + f"Invalid output_format {output_format!r}. Must be 'tree' or 'json'." + ) + + forest = None + for signal in state.signals: + if ( + signal.type == device_env_service_pb2.DEVICE_SIGNAL_ACCESSIBILITY_TREE + and signal.HasField('accessibility_forest') + ): + forest = signal.accessibility_forest + break + + if forest is None: + raise errors.DeviceConnectionError( + 'No accessibility tree signal found in device state.' + ) + + if output_format == 'json': + return json_format.MessageToJson(forest) + + lines = [] + for window in forest.windows: + lines.append(f'[Window] "{window.title}" id={window.id}') + if window.HasField('tree'): + nodes_by_id = {n.unique_id: n for n in window.tree.nodes} + root_ids = [n.unique_id for n in window.tree.nodes if n.depth == 0] + visited = set() + for root_id in root_ids: + _render_node( + nodes_by_id, root_id, indent=1, lines=lines, visited=visited + ) + return '\n'.join(lines) + + +class AndroidMcpTools: + """FastMCP tools foundation class and device state manager via AndroidDevice.""" + + def __init__( + self, device: android_device_mod.AndroidDevice | None = None + ) -> None: + """Initializes AndroidMcpTools with an optional AndroidDevice. + + Args: + device: Optional `AndroidDevice` instance used for querying device state. + """ + self._device = device + if self._device is not None: + self._device.update_subscriptions(set(_DEFAULT_SIGNALS)) + + def get_screenshot(self) -> str: + """Captures and returns a base64-encoded PNG screenshot of the device. + + Returns: + Base64-encoded PNG image string. + + Raises: + errors.DeviceConnectionError: If not connected to any device or if no + decoded screenshot signal is present in the device state. + """ + if self._device is None: + raise errors.DeviceConnectionError('Not connected to any device.') + + state = self._device.get_state() + pixels = _extract_screenshot(state) + + img = Image.fromarray(pixels, mode='RGB') + buffer = io.BytesIO() + img.save(buffer, format='PNG') + return base64.b64encode(buffer.getvalue()).decode('ascii') + + def get_ui_elements(self, output_format: str = 'tree') -> str: + """Returns accessibility hierarchy string of the device screen. + + Args: + output_format: Desired format of the returned hierarchy string; either + `'tree'` for human-readable tree lines or `'json'` for serialized JSON. + Defaults to `'tree'`. + + Returns: + Formatted accessibility hierarchy string. + + Raises: + errors.DeviceConnectionError: If not connected to any device or if no + accessibility tree signal is present in the device state. + ValueError: If `output_format` is invalid. + """ + if self._device is None: + raise errors.DeviceConnectionError('Not connected to any device.') + + state = self._device.get_state() + return _format_accessibility_tree(state, output_format) + + def get_device_info(self) -> str: + """Returns device settings and metadata. + + Returns: + JSON string containing device orientation, settings entries, and active + package metadata. + + Raises: + errors.DeviceConnectionError: If not connected to any device. + """ + if self._device is None: + raise errors.DeviceConnectionError('Not connected to any device.') + + state = self._device.get_state() + info: dict[str, Any] = {'orientation': state.orientation} + + for signal in state.signals: + if ( + signal.type == device_env_service_pb2.DEVICE_SIGNAL_DEVICE_SETTINGS + and signal.HasField('device_settings') + ): + info['settings'] = dict(signal.device_settings.entries) + elif ( + signal.type == device_env_service_pb2.DEVICE_SIGNAL_ACTIVE_PACKAGE + and signal.HasField('active_package') + ): + info['active_package'] = signal.active_package + + return json.dumps(info, indent=2) + + def tap(self, x: float, y: float) -> str: + """Taps at the given normalized screen coordinates. + + Args: + x: Horizontal position, normalized [0.0, 1.0] (left to right). + y: Vertical position, normalized [0.0, 1.0] (top to bottom). + + Returns: + Confirmation message. + + Raises: + errors.DeviceConnectionError: If not connected. + ValueError: If coordinates are out of bounds [0.0, 1.0]. + """ + if self._device is None: + raise errors.DeviceConnectionError('Not connected to any device.') + + if not (0.0 <= x <= 1.0 and 0.0 <= y <= 1.0): + raise ValueError( + f'Coordinates ({x}, {y}) out of bounds. Must be in range [0.0, 1.0].' + ) + + point = device_env_service_pb2.Point(x=x, y=y) + self._device.send_action( + device_env_service_pb2.Action( + action_type=device_env_service_pb2.ACTION_TYPE_TOUCH_DOWN, + touch_position=point, + ) + ) + self._device.send_action( + device_env_service_pb2.Action( + action_type=device_env_service_pb2.ACTION_TYPE_TOUCH_UP, + touch_position=point, + ) + ) + return f'Tapped at ({x}, {y})' + + def swipe( + self, + start_x: float, + start_y: float, + end_x: float, + end_y: float, + steps: int = 10, + ) -> str: + """Swipes from start to end coordinates. + + Args: + start_x: Starting horizontal position, normalized [0.0, 1.0] (left to + right). + start_y: Starting vertical position, normalized [0.0, 1.0] (top to + bottom). + end_x: Ending horizontal position, normalized [0.0, 1.0] (left to right). + end_y: Ending vertical position, normalized [0.0, 1.0] (top to bottom). + steps: Number of steps to take from start to end. + + Returns: + Confirmation message. + + Raises: + errors.DeviceConnectionError: If not connected. + ValueError: If coordinates are out of bounds [0.0, 1.0] or steps is less + than 1. + """ + if self._device is None: + raise errors.DeviceConnectionError('Not connected to any device.') + + if not ( + 0.0 <= start_x <= 1.0 + and 0.0 <= start_y <= 1.0 + and 0.0 <= end_x <= 1.0 + and 0.0 <= end_y <= 1.0 + ): + raise ValueError( + 'Coordinates must be in range [0.0, 1.0]. Got' + f' start=({start_x}, {start_y}), end=({end_x}, {end_y}).' + ) + if steps < 1: + raise ValueError(f'steps must be at least 1, got {steps}.') + + self._device.send_action( + device_env_service_pb2.Action( + action_type=device_env_service_pb2.ACTION_TYPE_TOUCH_DOWN, + touch_position=device_env_service_pb2.Point(x=start_x, y=start_y), + ) + ) + + for i in range(1, steps + 1): + t = i / steps + x = start_x + (end_x - start_x) * t + y = start_y + (end_y - start_y) * t + self._device.send_action( + device_env_service_pb2.Action( + action_type=device_env_service_pb2.ACTION_TYPE_TOUCH_MOVE, + touch_position=device_env_service_pb2.Point(x=x, y=y), + ) + ) + + self._device.send_action( + device_env_service_pb2.Action( + action_type=device_env_service_pb2.ACTION_TYPE_TOUCH_UP, + touch_position=device_env_service_pb2.Point(x=end_x, y=end_y), + ) + ) + + return f'Swiped from ({start_x}, {start_y}) to ({end_x}, {end_y})' + + def press_key(self, keycode: int) -> str: + """Presses the given Android keycode. + + Args: + keycode: The Android keycode to press. + + Returns: + Confirmation message. + + Raises: + errors.DeviceConnectionError: If not connected. + """ + if self._device is None: + raise errors.DeviceConnectionError('Not connected to any device.') + + self._device.send_action( + device_env_service_pb2.Action( + action_type=device_env_service_pb2.ACTION_TYPE_KEY_EVENT, + keycode=keycode, + ) + ) + return f'Pressed key {keycode}' + + def register_tools(self, mcp_server: fastmcp.FastMCP) -> None: + """Registers all tool methods on the given FastMCP server. + + Args: + mcp_server: FastMCP server instance where tool handlers will be + registered. + """ + mcp_server.tool()(self.get_screenshot) + mcp_server.tool()(self.get_ui_elements) + mcp_server.tool()(self.get_device_info) + mcp_server.tool()(self.tap) + mcp_server.tool()(self.swipe) + mcp_server.tool()(self.press_key) diff --git a/android_env/mcp/android_mcp_tools_test.py b/android_env/mcp/android_mcp_tools_test.py new file mode 100644 index 00000000..8ff44943 --- /dev/null +++ b/android_env/mcp/android_mcp_tools_test.py @@ -0,0 +1,660 @@ +# coding=utf-8 +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for FastMCP Android device interaction tools in android_mcp_tools.""" + +import base64 +import io +import json +from unittest import mock + +from absl.testing import absltest +from android_env.components import android_device as android_device_mod +from android_env.mcp import android_mcp_tools +from android_env.mcp import errors +from android_env.proto import device_env_service_pb2 +from android_env.proto.a11y import android_accessibility_forest_pb2 +from android_env.proto.a11y import android_accessibility_node_info_pb2 +from google.protobuf import json_format +from mcp.server import fastmcp +import numpy as np +from PIL import Image + + +class AndroidMcpToolsTest(absltest.TestCase): + + def test_init_updates_subscriptions(self): + mock_device = mock.create_autospec( + android_device_mod.AndroidDevice, instance=True + ) + tools = android_mcp_tools.AndroidMcpTools(mock_device) + self.assertIs(tools._device, mock_device) + mock_device.update_subscriptions.assert_called_once_with( + set(android_mcp_tools._DEFAULT_SIGNALS) + ) + + def test_init_with_none_device(self): + tools = android_mcp_tools.AndroidMcpTools(device=None) + self.assertIsNone(tools._device) + + +class GetScreenshotTest(absltest.TestCase): + + def test_get_screenshot_not_connected(self): + tools = android_mcp_tools.AndroidMcpTools() + with self.assertRaises(errors.DeviceConnectionError): + tools.get_screenshot() + + def test_get_screenshot_missing_signal(self): + tools = android_mcp_tools.AndroidMcpTools() + mock_device = mock.create_autospec( + android_device_mod.AndroidDevice, instance=True + ) + tools._device = mock_device + mock_device.get_state.return_value = device_env_service_pb2.DeviceState() + with self.assertRaises(errors.DeviceConnectionError): + tools.get_screenshot() + + def test_get_screenshot_signal_without_decoded(self): + tools = android_mcp_tools.AndroidMcpTools() + mock_device = mock.create_autospec( + android_device_mod.AndroidDevice, instance=True + ) + tools._device = mock_device + state = device_env_service_pb2.DeviceState() + sig = state.signals.add() + sig.type = device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT + mock_device.get_state.return_value = state + with self.assertRaises(errors.DeviceConnectionError): + tools.get_screenshot() + + def test_get_screenshot_incomplete_signal(self): + mock_device = mock.create_autospec( + android_device_mod.AndroidDevice, instance=True + ) + tools = android_mcp_tools.AndroidMcpTools(mock_device) + state = device_env_service_pb2.DeviceState() + sig1 = state.signals.add() + sig1.type = device_env_service_pb2.DEVICE_SIGNAL_ACTIVE_PACKAGE + sig2 = state.signals.add() + sig2.type = device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT + sig3 = state.signals.add() + sig3.type = device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT + sig3.screenshot.encoded_bytes = b'123' + mock_device.get_state.return_value = state + with self.assertRaises(errors.DeviceConnectionError): + tools.get_screenshot() + + def test_get_screenshot_success(self): + mock_device = mock.create_autospec( + android_device_mod.AndroidDevice, instance=True + ) + tools = android_mcp_tools.AndroidMcpTools(mock_device) + + height, width = 100, 100 + raw_pixels = np.zeros((height, width, 3), dtype=np.uint8).tobytes() + + state = device_env_service_pb2.DeviceState() + sig_other = state.signals.add() + sig_other.type = device_env_service_pb2.DEVICE_SIGNAL_ACTIVE_PACKAGE + + sig = state.signals.add() + sig.type = device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT + sig.screenshot.decoded.height = height + sig.screenshot.decoded.width = width + sig.screenshot.decoded.raw_pixels = raw_pixels + + mock_device.get_state.return_value = state + + result = tools.get_screenshot() + decoded_data = base64.b64decode(result.encode('ascii')) + img = Image.open(io.BytesIO(decoded_data)) + self.assertEqual(img.size, (width, height)) + # Check exact pixels (since PNG is lossless) + expected_pixels = np.frombuffer(raw_pixels, dtype=np.uint8).reshape( + (height, width, 3) + ) + np.testing.assert_array_equal(np.array(img), expected_pixels) + + +class GetUiElementsTest(absltest.TestCase): + + def setUp(self): + super().setUp() + self._forest = android_accessibility_forest_pb2.AndroidAccessibilityForest() + window = self._forest.windows.add() + window.id = 1 + window.title = 'MainActivity' + + node = window.tree.nodes.add() + node.unique_id = 0 + node.class_name = 'android.widget.TextView' + node.text = 'Hello World' + node.depth = 0 + node.bounds_in_screen.left = 0 + node.bounds_in_screen.top = 0 + node.bounds_in_screen.right = 100 + node.bounds_in_screen.bottom = 50 + node.is_clickable = True + + def test_get_ui_elements_not_connected(self): + mock_device = mock.create_autospec( + android_device_mod.AndroidDevice, instance=True + ) + tools = android_mcp_tools.AndroidMcpTools(mock_device) + tools._device = None + with self.assertRaises(errors.DeviceConnectionError): + tools.get_ui_elements() + + def test_get_ui_elements_missing_signal(self): + mock_device = mock.create_autospec( + android_device_mod.AndroidDevice, instance=True + ) + tools = android_mcp_tools.AndroidMcpTools(mock_device) + state = device_env_service_pb2.DeviceState() + sig = state.signals.add() + sig.type = device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT + mock_device.get_state.return_value = state + with self.assertRaises(errors.DeviceConnectionError): + tools.get_ui_elements() + + def test_get_ui_elements_invalid_format(self): + mock_device = mock.create_autospec( + android_device_mod.AndroidDevice, instance=True + ) + tools = android_mcp_tools.AndroidMcpTools(mock_device) + with self.assertRaises(ValueError): + tools.get_ui_elements(output_format='invalid') + + def test_get_ui_elements_tree_format(self): + mock_device = mock.create_autospec( + android_device_mod.AndroidDevice, instance=True + ) + tools = android_mcp_tools.AndroidMcpTools(mock_device) + + state = device_env_service_pb2.DeviceState() + sig = state.signals.add() + sig.type = device_env_service_pb2.DEVICE_SIGNAL_ACCESSIBILITY_TREE + sig.accessibility_forest.CopyFrom(self._forest) + mock_device.get_state.return_value = state + + result = tools.get_ui_elements(output_format='tree') + self.assertIn('[Window] "MainActivity" id=1', result) + self.assertIn( + ' [android.widget.TextView] "Hello World" {0,0,100,50} [clickable]', + result, + ) + + def test_get_ui_elements_json_format(self): + mock_device = mock.create_autospec( + android_device_mod.AndroidDevice, instance=True + ) + tools = android_mcp_tools.AndroidMcpTools(mock_device) + + state = device_env_service_pb2.DeviceState() + sig = state.signals.add() + sig.type = device_env_service_pb2.DEVICE_SIGNAL_ACCESSIBILITY_TREE + sig.accessibility_forest.CopyFrom(self._forest) + mock_device.get_state.return_value = state + + result = tools.get_ui_elements(output_format='json') + expected_json = json_format.MessageToJson(self._forest) + self.assertEqual(json.loads(result), json.loads(expected_json)) + + def test_get_ui_elements_with_cycle(self): + mock_device = mock.create_autospec( + android_device_mod.AndroidDevice, instance=True + ) + tools = android_mcp_tools.AndroidMcpTools(mock_device) + + forest = android_accessibility_forest_pb2.AndroidAccessibilityForest() + window = forest.windows.add() + window.id = 1 + window.title = 'MainActivity' + + node = window.tree.nodes.add() + node.unique_id = 0 + node.class_name = 'android.widget.TextView' + node.text = 'Cyclic Node' + node.depth = 0 + node.bounds_in_screen.left = 0 + node.bounds_in_screen.top = 0 + node.bounds_in_screen.right = 100 + node.bounds_in_screen.bottom = 50 + node.child_ids.append(0) + + state = device_env_service_pb2.DeviceState() + sig = state.signals.add() + sig.type = device_env_service_pb2.DEVICE_SIGNAL_ACCESSIBILITY_TREE + sig.accessibility_forest.CopyFrom(forest) + mock_device.get_state.return_value = state + + result = tools.get_ui_elements(output_format='tree') + self.assertIn('[Window] "MainActivity" id=1', result) + self.assertIn( + ' [android.widget.TextView] "Cyclic Node" {0,0,100,50}', result + ) + + def test_get_ui_elements_content_description_fallback(self): + mock_device = mock.create_autospec( + android_device_mod.AndroidDevice, instance=True + ) + tools = android_mcp_tools.AndroidMcpTools(mock_device) + + forest = android_accessibility_forest_pb2.AndroidAccessibilityForest() + window = forest.windows.add() + window.id = 1 + window.title = 'MainActivity' + + node = window.tree.nodes.add() + node.unique_id = 0 + node.class_name = 'android.widget.ImageView' + node.text = '' + node.content_description = 'Content Description Text' + node.depth = 0 + node.bounds_in_screen.left = 0 + node.bounds_in_screen.top = 0 + node.bounds_in_screen.right = 100 + node.bounds_in_screen.bottom = 50 + + state = device_env_service_pb2.DeviceState() + sig = state.signals.add() + sig.type = device_env_service_pb2.DEVICE_SIGNAL_ACCESSIBILITY_TREE + sig.accessibility_forest.CopyFrom(forest) + mock_device.get_state.return_value = state + + result = tools.get_ui_elements(output_format='tree') + self.assertIn( + ' [android.widget.ImageView] "Content Description Text" {0,0,100,50}', + result, + ) + + def test_get_ui_elements_no_text_or_content_description(self): + mock_device = mock.create_autospec( + android_device_mod.AndroidDevice, instance=True + ) + tools = android_mcp_tools.AndroidMcpTools(mock_device) + + forest = android_accessibility_forest_pb2.AndroidAccessibilityForest() + window = forest.windows.add() + window.id = 1 + window.title = 'MainActivity' + + node = window.tree.nodes.add() + node.unique_id = 0 + node.class_name = 'android.view.View' + node.text = '' + node.content_description = '' + node.depth = 0 + node.bounds_in_screen.left = 0 + node.bounds_in_screen.top = 0 + node.bounds_in_screen.right = 100 + node.bounds_in_screen.bottom = 50 + + state = device_env_service_pb2.DeviceState() + sig = state.signals.add() + sig.type = device_env_service_pb2.DEVICE_SIGNAL_ACCESSIBILITY_TREE + sig.accessibility_forest.CopyFrom(forest) + mock_device.get_state.return_value = state + + result = tools.get_ui_elements(output_format='tree') + self.assertIn(' [android.view.View] "" {0,0,100,50}', result) + + def test_get_ui_elements_with_all_flags(self): + mock_device = mock.create_autospec( + android_device_mod.AndroidDevice, instance=True + ) + tools = android_mcp_tools.AndroidMcpTools(mock_device) + + forest = android_accessibility_forest_pb2.AndroidAccessibilityForest() + window = forest.windows.add() + window.id = 1 + window.title = 'MainActivity' + + node = window.tree.nodes.add() + node.unique_id = 0 + node.class_name = 'android.widget.EditText' + node.text = 'Editable Text' + node.depth = 0 + node.bounds_in_screen.left = 0 + node.bounds_in_screen.top = 0 + node.bounds_in_screen.right = 100 + node.bounds_in_screen.bottom = 50 + node.is_clickable = True + node.is_scrollable = True + node.is_checked = True + node.is_focused = True + node.is_enabled = True + node.is_editable = True + + state = device_env_service_pb2.DeviceState() + sig = state.signals.add() + sig.type = device_env_service_pb2.DEVICE_SIGNAL_ACCESSIBILITY_TREE + sig.accessibility_forest.CopyFrom(forest) + mock_device.get_state.return_value = state + + result = tools.get_ui_elements(output_format='tree') + self.assertIn( + ' [android.widget.EditText] "Editable Text" {0,0,100,50}' + ' [clickable scrollable checked focused enabled editable]', + result, + ) + + def test_get_ui_elements_missing_child_node(self): + mock_device = mock.create_autospec( + android_device_mod.AndroidDevice, instance=True + ) + tools = android_mcp_tools.AndroidMcpTools(mock_device) + + forest = android_accessibility_forest_pb2.AndroidAccessibilityForest() + window = forest.windows.add() + window.id = 1 + window.title = 'MainActivity' + + node = window.tree.nodes.add() + node.unique_id = 0 + node.class_name = 'android.widget.LinearLayout' + node.text = 'Parent Node' + node.depth = 0 + node.bounds_in_screen.left = 0 + node.bounds_in_screen.top = 0 + node.bounds_in_screen.right = 100 + node.bounds_in_screen.bottom = 50 + node.child_ids.append(999) + + state = device_env_service_pb2.DeviceState() + sig = state.signals.add() + sig.type = device_env_service_pb2.DEVICE_SIGNAL_ACCESSIBILITY_TREE + sig.accessibility_forest.CopyFrom(forest) + mock_device.get_state.return_value = state + + result = tools.get_ui_elements(output_format='tree') + self.assertIn('[Window] "MainActivity" id=1', result) + self.assertIn( + ' [android.widget.LinearLayout] "Parent Node" {0,0,100,50}', result + ) + + def test_get_ui_elements_cycle_detection(self): + mock_device = mock.create_autospec( + android_device_mod.AndroidDevice, instance=True + ) + tools = android_mcp_tools.AndroidMcpTools(mock_device) + + forest = android_accessibility_forest_pb2.AndroidAccessibilityForest() + window = forest.windows.add() + window.id = 1 + window.title = 'MainActivity' + + parent = window.tree.nodes.add() + parent.unique_id = 0 + parent.class_name = 'android.widget.LinearLayout' + parent.text = 'Parent Node' + parent.depth = 0 + parent.bounds_in_screen.left = 0 + parent.bounds_in_screen.top = 0 + parent.bounds_in_screen.right = 100 + parent.bounds_in_screen.bottom = 100 + parent.child_ids.append(1) + + child = window.tree.nodes.add() + child.unique_id = 1 + child.class_name = 'android.widget.TextView' + child.text = 'Child Node' + child.depth = 1 + child.bounds_in_screen.left = 0 + child.bounds_in_screen.top = 0 + child.bounds_in_screen.right = 100 + child.bounds_in_screen.bottom = 50 + child.child_ids.append(0) + + state = device_env_service_pb2.DeviceState() + sig = state.signals.add() + sig.type = device_env_service_pb2.DEVICE_SIGNAL_ACCESSIBILITY_TREE + sig.accessibility_forest.CopyFrom(forest) + mock_device.get_state.return_value = state + + result = tools.get_ui_elements(output_format='tree') + self.assertIn('[Window] "MainActivity" id=1', result) + self.assertIn( + ' [android.widget.LinearLayout] "Parent Node" {0,0,100,100}', result + ) + self.assertIn( + ' [android.widget.TextView] "Child Node" {0,0,100,50}', result + ) + + def test_get_ui_elements_window_without_tree(self): + mock_device = mock.create_autospec( + android_device_mod.AndroidDevice, instance=True + ) + tools = android_mcp_tools.AndroidMcpTools(mock_device) + + forest = android_accessibility_forest_pb2.AndroidAccessibilityForest() + window = forest.windows.add() + window.id = 2 + window.title = 'EmptyWindow' + + state = device_env_service_pb2.DeviceState() + sig = state.signals.add() + sig.type = device_env_service_pb2.DEVICE_SIGNAL_ACCESSIBILITY_TREE + sig.accessibility_forest.CopyFrom(forest) + mock_device.get_state.return_value = state + + result = tools.get_ui_elements(output_format='tree') + self.assertEqual(result, '[Window] "EmptyWindow" id=2') + + def test_render_node_visited_default(self): + node = android_accessibility_node_info_pb2.AndroidAccessibilityNodeInfo() + node.unique_id = 1 + node.class_name = 'android.view.View' + node.text = 'Direct Render' + node.depth = 0 + node.bounds_in_screen.left = 0 + node.bounds_in_screen.top = 0 + node.bounds_in_screen.right = 10 + node.bounds_in_screen.bottom = 10 + + nodes_by_id = {1: node} + lines = [] + android_mcp_tools._render_node(nodes_by_id, 1, indent=0, lines=lines) + self.assertEqual(lines, ['[android.view.View] "Direct Render" {0,0,10,10}']) + + +class GetDeviceInfoTest(absltest.TestCase): + + def test_get_device_info_not_connected(self): + mock_device = mock.create_autospec( + android_device_mod.AndroidDevice, instance=True + ) + tools = android_mcp_tools.AndroidMcpTools(mock_device) + tools._device = None + with self.assertRaises(errors.DeviceConnectionError): + tools.get_device_info() + + def test_get_device_info_success(self): + mock_device = mock.create_autospec( + android_device_mod.AndroidDevice, instance=True + ) + tools = android_mcp_tools.AndroidMcpTools(mock_device) + + state = device_env_service_pb2.DeviceState() + state.orientation = 1 + + sig_pkg = state.signals.add() + sig_pkg.type = device_env_service_pb2.DEVICE_SIGNAL_ACTIVE_PACKAGE + sig_pkg.active_package = 'com.example.app' + + sig_settings = state.signals.add() + sig_settings.type = device_env_service_pb2.DEVICE_SIGNAL_DEVICE_SETTINGS + sig_settings.device_settings.entries['system_theme'] = 'dark' + + mock_device.get_state.return_value = state + + result = tools.get_device_info() + data = json.loads(result) + self.assertEqual(data['orientation'], 1) + self.assertEqual(data['active_package'], 'com.example.app') + self.assertEqual(data['settings']['system_theme'], 'dark') + + def test_get_device_info_signals_without_payloads(self): + mock_device = mock.create_autospec( + android_device_mod.AndroidDevice, instance=True + ) + tools = android_mcp_tools.AndroidMcpTools(mock_device) + + state = device_env_service_pb2.DeviceState() + state.orientation = 0 + + sig_settings = state.signals.add() + sig_settings.type = device_env_service_pb2.DEVICE_SIGNAL_DEVICE_SETTINGS + + sig_pkg = state.signals.add() + sig_pkg.type = device_env_service_pb2.DEVICE_SIGNAL_ACTIVE_PACKAGE + + sig_other = state.signals.add() + sig_other.type = device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT + + mock_device.get_state.return_value = state + + result = tools.get_device_info() + data = json.loads(result) + self.assertEqual(data, {'orientation': 0}) + + +class TapTest(absltest.TestCase): + + def test_tap_not_connected(self): + mock_device = mock.create_autospec( + android_device_mod.AndroidDevice, instance=True + ) + tools = android_mcp_tools.AndroidMcpTools(mock_device) + tools._device = None + with self.assertRaises(errors.DeviceConnectionError): + tools.tap(0.5, 0.5) + + def test_tap_out_of_bounds(self): + mock_device = mock.create_autospec( + android_device_mod.AndroidDevice, instance=True + ) + tools = android_mcp_tools.AndroidMcpTools(mock_device) + with self.assertRaises(ValueError): + tools.tap(-0.1, 0.5) + with self.assertRaises(ValueError): + tools.tap(0.5, 1.1) + + def test_tap_success(self): + mock_device = mock.create_autospec( + android_device_mod.AndroidDevice, instance=True + ) + tools = android_mcp_tools.AndroidMcpTools(mock_device) + + result = tools.tap(0.25, 0.75) + self.assertEqual(result, 'Tapped at (0.25, 0.75)') + + self.assertEqual(mock_device.send_action.call_count, 2) + down_action = mock_device.send_action.call_args_list[0][0][0] + up_action = mock_device.send_action.call_args_list[1][0][0] + + self.assertEqual( + down_action.action_type, device_env_service_pb2.ACTION_TYPE_TOUCH_DOWN + ) + self.assertEqual(down_action.touch_position.x, 0.25) + self.assertEqual(down_action.touch_position.y, 0.75) + + self.assertEqual( + up_action.action_type, device_env_service_pb2.ACTION_TYPE_TOUCH_UP + ) + self.assertEqual(up_action.touch_position.x, 0.25) + self.assertEqual(up_action.touch_position.y, 0.75) + + +class SwipeTest(absltest.TestCase): + + def test_swipe_not_connected(self): + mock_device = mock.create_autospec( + android_device_mod.AndroidDevice, instance=True + ) + tools = android_mcp_tools.AndroidMcpTools(mock_device) + tools._device = None + with self.assertRaises(errors.DeviceConnectionError): + tools.swipe(0.1, 0.1, 0.9, 0.9) + + def test_swipe_invalid_steps(self): + mock_device = mock.create_autospec( + android_device_mod.AndroidDevice, instance=True + ) + tools = android_mcp_tools.AndroidMcpTools(mock_device) + with self.assertRaises(ValueError): + tools.swipe(0.1, 0.1, 0.9, 0.9, steps=0) + + def test_swipe_out_of_bounds(self): + mock_device = mock.create_autospec( + android_device_mod.AndroidDevice, instance=True + ) + tools = android_mcp_tools.AndroidMcpTools(mock_device) + with self.assertRaises(ValueError): + tools.swipe(-0.1, 0.1, 0.9, 0.9) + + def test_swipe_success(self): + mock_device = mock.create_autospec( + android_device_mod.AndroidDevice, instance=True + ) + tools = android_mcp_tools.AndroidMcpTools(mock_device) + + result = tools.swipe(0.0, 0.0, 1.0, 1.0, steps=2) + self.assertEqual(result, 'Swiped from (0.0, 0.0) to (1.0, 1.0)') + self.assertEqual(mock_device.send_action.call_count, 4) + + +class PressKeyTest(absltest.TestCase): + + def test_press_key_not_connected(self): + mock_device = mock.create_autospec( + android_device_mod.AndroidDevice, instance=True + ) + tools = android_mcp_tools.AndroidMcpTools(mock_device) + tools._device = None + with self.assertRaises(errors.DeviceConnectionError): + tools.press_key(3) + + def test_press_key_success(self): + mock_device = mock.create_autospec( + android_device_mod.AndroidDevice, instance=True + ) + tools = android_mcp_tools.AndroidMcpTools(mock_device) + + result = tools.press_key(3) + self.assertEqual(result, 'Pressed key 3') + mock_device.send_action.assert_called_once() + action = mock_device.send_action.call_args[0][0] + self.assertEqual( + action.action_type, device_env_service_pb2.ACTION_TYPE_KEY_EVENT + ) + self.assertEqual(action.keycode, 3) + + +class RegisterToolsTest(absltest.TestCase): + + def test_register_tools_count(self): + mock_device = mock.create_autospec( + android_device_mod.AndroidDevice, instance=True + ) + tools = android_mcp_tools.AndroidMcpTools(mock_device) + server = mock.MagicMock(spec=fastmcp.FastMCP) + tools.register_tools(server) + self.assertEqual(server.tool.call_count, 6) + + +if __name__ == '__main__': + absltest.main() diff --git a/android_env/mcp/errors.py b/android_env/mcp/errors.py new file mode 100644 index 00000000..f90a1673 --- /dev/null +++ b/android_env/mcp/errors.py @@ -0,0 +1,28 @@ +# coding=utf-8 +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Exception hierarchy for AndroidEnv MCP server operations.""" + + +class AndroidEnvMcpError(Exception): + """Base exception for Android MCP operations.""" + + +class AdbError(AndroidEnvMcpError): + """Raised when ADB commands fail or execution errors occur.""" + + +class DeviceConnectionError(AndroidEnvMcpError): + """Raised when connecting or disconnecting a device fails.""" diff --git a/android_env/playground/BUILD.bazel b/android_env/playground/BUILD.bazel new file mode 100644 index 00000000..4642c29f --- /dev/null +++ b/android_env/playground/BUILD.bazel @@ -0,0 +1,222 @@ +load("@rules_python//python:defs.bzl", "py_binary", "py_library", "py_test") + +package(default_visibility = ["//visibility:public"]) + +py_library( + name = "actions", + srcs = ["actions.py"], + deps = [ + "//android_env/components:action_type", + "//android_env/proto:device_env_service_py_pb2", + "@pip//numpy", + ], +) + +py_library( + name = "coordinates", + srcs = ["coordinates.py"], + deps = [ + "@pip//numpy", + ], +) + +py_library( + name = "frame_processing", + srcs = ["frame_processing.py"], + deps = [ + ":coordinates", + "@pip//numpy", + "@pip//pillow", + ], +) + +py_library( + name = "state", + srcs = ["state.py"], +) + +py_library( + name = "websocket", + srcs = ["websocket.py"], + deps = [ + ":actions", + ":state", + "@pip//absl_py", + ], +) + +py_library( + name = "log_handler", + srcs = ["log_handler.py"], + deps = [ + ":state", + ":websocket", + ], +) + +py_library( + name = "server", + srcs = ["server.py"], + data = [ + "//android_env/playground/frontend", + ], + deps = [ + ":state", + ":websocket", + "@pip//absl_py", + ], +) + +py_library( + name = "device_runner", + srcs = ["device_runner.py"], + deps = [ + ":actions", + ":coordinates", + ":frame_processing", + ":state", + ":websocket", + "//android_env/components:android_device", + "//android_env/components:android_device_loader", + "//android_env/components:config_classes", + "//android_env/components:errors", + "//android_env/components:setup_step_interpreter", + "//android_env/proto:adb_py_pb2", + "//android_env/proto:device_env_service_py_pb2", + "//android_env/proto:task_py_pb2", + "@pip//absl_py", + "@pip//numpy", + ], +) + +py_library( + name = "playground_lib", + srcs = ["playground.py"], + deps = [ + ":device_runner", + ":log_handler", + ":server", + ":state", + ":websocket", + "@pip//absl_py", + "@pip//ml_collections", + ], +) + +py_binary( + name = "playground", + srcs = ["playground.py"], + deps = [ + ":playground_lib", + ], +) + +py_test( + name = "coordinates_test", + srcs = ["coordinates_test.py"], + legacy_create_init = False, + deps = [ + ":coordinates", + "@pip//absl_py", + "@pip//numpy", + ], +) + +py_test( + name = "actions_test", + srcs = ["actions_test.py"], + legacy_create_init = False, + deps = [ + ":actions", + "//android_env/components:action_type", + "@pip//absl_py", + ], +) + +py_test( + name = "frame_processing_test", + srcs = ["frame_processing_test.py"], + legacy_create_init = False, + deps = [ + ":frame_processing", + "@pip//absl_py", + "@pip//numpy", + "@pip//pillow", + ], +) + +py_test( + name = "log_handler_test", + srcs = ["log_handler_test.py"], + legacy_create_init = False, + deps = [ + ":log_handler", + ":state", + "@pip//absl_py", + ], +) + +py_test( + name = "state_test", + srcs = ["state_test.py"], + legacy_create_init = False, + deps = [ + ":state", + "@pip//absl_py", + ], +) + +py_test( + name = "websocket_test", + srcs = ["websocket_test.py"], + legacy_create_init = False, + deps = [ + ":state", + ":websocket", + "@pip//absl_py", + ], +) + +py_test( + name = "server_test", + srcs = ["server_test.py"], + legacy_create_init = False, + deps = [ + ":server", + ":state", + ":websocket", + "@pip//absl_py", + ], +) + +py_test( + name = "device_runner_test", + srcs = ["device_runner_test.py"], + legacy_create_init = False, + deps = [ + ":actions", + ":device_runner", + ":state", + "//android_env/components:android_device", + "//android_env/components:config_classes", + "//android_env/components:errors", + "//android_env/proto:adb_py_pb2", + "//android_env/proto:device_env_service_py_pb2", + "//android_env/proto:task_py_pb2", + "@pip//absl_py", + "@pip//numpy", + ], +) + +py_test( + name = "playground_test", + srcs = ["playground_test.py"], + legacy_create_init = False, + deps = [ + ":playground_lib", + ":state", + "//android_env/components:android_device", + "@pip//absl_py", + "@pip//numpy", + ], +) diff --git a/android_env/playground/README.md b/android_env/playground/README.md new file mode 100644 index 00000000..81aa7a74 --- /dev/null +++ b/android_env/playground/README.md @@ -0,0 +1,59 @@ +# AndroidEnv Playground + +Interactive playground for AndroidEnv, allowing human interaction via a web +browser. It streams the emulator screen and forwards your clicks, touches, and +keyboard inputs over WebSockets, avoiding X11 and Pygame display issues. + +## Quick Start (Open Source) + +### Prerequisites + +1. **Android SDK and Tools**: + Ensure `adb` and `emulator` are installed and available in `$PATH` + (or located under `$ANDROID_HOME` or `~/Android/Sdk`). +2. **Android Virtual Device (AVD)**: + Create an AVD using Android Studio or the `avdmanager` command-line tool. + For AndroidWorld evaluation, ensure your AVD is named `androidworld_avd` + (or customize via the config parameter). + +### Option 1: Running with Bazel + +Run the Playground binary directly using Bazel: + +```bash +bazel run //android_env/playground:playground -- \ + --config=android_env/playground/configs/emulator_config.py \ + --port=8080 \ + --alsologtostderr +``` + +You can pass a custom AVD name as a configuration argument: + +```bash +bazel run //android_env/playground:playground -- \ + --config=android_env/playground/configs/emulator_config.py:my_custom_avd \ + --port=8080 +``` + +### Option 2: Running with Python + +After installing AndroidEnv from source or via pip: + +```bash +python -m android_env.playground.playground \ + --config=android_env/playground/configs/emulator_config.py \ + --port=8080 +``` + +### Accessing the Interface + +Once the server logs `Serving on port 8080...`, open your browser to: + +[http://localhost:8080](http://localhost:8080) + +If you are running on a remote server or cloud workstation, forward the port +via SSH: + +```bash +ssh -L 8080:localhost:8080 +``` diff --git a/android_env/playground/actions.py b/android_env/playground/actions.py new file mode 100644 index 00000000..c9dd38e0 --- /dev/null +++ b/android_env/playground/actions.py @@ -0,0 +1,137 @@ +# coding=utf-8 +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Action mapping functions for AndroidEnv Playground.""" + +import dataclasses +from android_env.components import action_type +from android_env.proto import device_env_service_pb2 +import numpy as np + + +@dataclasses.dataclass(frozen=True) +class ClientAction: + """Represents an action received from the playground client. + + Attributes: + action_type: The type of action (e.g. 'touch_down', 'touch_move', + 'touch_up', 'text', 'keydown'). + x: Normalized X coordinate (0.0 to 1.0). + y: Normalized Y coordinate (0.0 to 1.0). + client_timestamp_ms: Timestamp in milliseconds when the action was generated + on the client. + recv_timestamp_ms: Timestamp in milliseconds when the action was received by + the server. + text: Optional text to type for 'text' action_type. + keycode: Optional integer Android keycode for 'keydown' action_type. + selection: Optional tuple (start, end) for 'set_selection' action_type. + """ + + action_type: str + x: float + y: float + client_timestamp_ms: float + recv_timestamp_ms: float + text: str | None = None + keycode: int | None = None + selection: tuple[int, int] | None = None + + +def coalesce_actions(actions: list[ClientAction]) -> list[ClientAction]: + """Coalesces consecutive touch_move actions in a list, preserving timestamps.""" + if not actions: + return [] + compressed = [] + for act in actions: + if not compressed: + compressed.append(act) + continue + + prev_act = compressed[-1] + if prev_act.action_type == 'touch_move' and act.action_type == 'touch_move': + compressed[-1] = act + else: + compressed.append(act) + return compressed + + +def map_action_type(act_str: str) -> action_type.ActionType: + """Maps client action string to AndroidEnv ActionType.""" + match act_str: + case 'touch_down' | 'touch_move': + return action_type.ActionType.TOUCH + case 'touch_up': + return action_type.ActionType.LIFT + case 'text': + return action_type.ActionType.TEXT + case 'keydown': + return action_type.ActionType.KEYDOWN + case 'set_selection': + return action_type.ActionType.SET_SELECTION + case 'cut': + return action_type.ActionType.CUT + case 'copy': + return action_type.ActionType.COPY + case 'paste': + return action_type.ActionType.PASTE + case _: + return action_type.ActionType.REPEAT + + +def map_client_action_to_proto( + act: ClientAction, pos: np.ndarray | None = None +) -> device_env_service_pb2.Action: + """Maps a ClientAction directly to an AndroidEnv protobuf Action.""" + proto_action = device_env_service_pb2.Action() + match act.action_type: + case 'touch_down': + proto_action.action_type = device_env_service_pb2.ACTION_TYPE_TOUCH_DOWN + if pos is not None: + proto_action.touch_position.x = float(pos[0]) + proto_action.touch_position.y = float(pos[1]) + case 'touch_move': + proto_action.action_type = device_env_service_pb2.ACTION_TYPE_TOUCH_MOVE + if pos is not None: + proto_action.touch_position.x = float(pos[0]) + proto_action.touch_position.y = float(pos[1]) + case 'touch_up': + proto_action.action_type = device_env_service_pb2.ACTION_TYPE_TOUCH_UP + if pos is not None: + proto_action.touch_position.x = float(pos[0]) + proto_action.touch_position.y = float(pos[1]) + case 'keydown': + proto_action.action_type = device_env_service_pb2.ACTION_TYPE_KEY_EVENT + if act.keycode is not None: + proto_action.keycode = int(act.keycode) + case 'text': + proto_action.action_type = device_env_service_pb2.ACTION_TYPE_SET_TEXT + if act.text is not None: + proto_action.text = str(act.text) + case 'set_selection': + proto_action.action_type = ( + device_env_service_pb2.ACTION_TYPE_SET_SELECTION + ) + if act.selection is not None: + proto_action.selection_start = int(act.selection[0]) + proto_action.selection_end = int(act.selection[1]) + case 'cut': + proto_action.action_type = device_env_service_pb2.ACTION_TYPE_CUT + case 'copy': + proto_action.action_type = device_env_service_pb2.ACTION_TYPE_COPY + case 'paste': + proto_action.action_type = device_env_service_pb2.ACTION_TYPE_PASTE + case _: + proto_action.action_type = device_env_service_pb2.ACTION_TYPE_UNSPECIFIED + return proto_action diff --git a/android_env/playground/actions_test.py b/android_env/playground/actions_test.py new file mode 100644 index 00000000..3cd0bae8 --- /dev/null +++ b/android_env/playground/actions_test.py @@ -0,0 +1,212 @@ +# coding=utf-8 +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for action mapping and coalescing in AndroidEnv Playground.""" + +from absl.testing import absltest +from absl.testing import parameterized +from android_env.components import action_type +from android_env.playground import actions as actions_lib +from android_env.proto import device_env_service_pb2 +import numpy as np + + +class ActionsTest(parameterized.TestCase): + + def test_coalesce_actions_empty(self): + self.assertEqual(actions_lib.coalesce_actions([]), []) + + def test_coalesce_actions_single(self): + action = actions_lib.ClientAction( + action_type='touch_down', + x=0.1, + y=0.2, + client_timestamp_ms=1000.0, + recv_timestamp_ms=1001.0, + ) + self.assertEqual(actions_lib.coalesce_actions([action]), [action]) + + def test_coalesce_actions_compress_moves(self): + actions = [ + actions_lib.ClientAction( + action_type='touch_down', + x=0.1, + y=0.2, + client_timestamp_ms=1000.0, + recv_timestamp_ms=1001.0, + ), + actions_lib.ClientAction( + action_type='touch_move', + x=0.15, + y=0.25, + client_timestamp_ms=1016.0, + recv_timestamp_ms=1017.0, + ), + actions_lib.ClientAction( + action_type='touch_move', + x=0.2, + y=0.3, + client_timestamp_ms=1032.0, + recv_timestamp_ms=1033.0, + ), + actions_lib.ClientAction( + action_type='touch_up', + x=0.2, + y=0.3, + client_timestamp_ms=1048.0, + recv_timestamp_ms=1049.0, + ), + ] + expected = [ + actions[0], + actions[2], + actions[3], + ] + self.assertEqual(actions_lib.coalesce_actions(actions), expected) + + @parameterized.parameters( + ('touch_down', action_type.ActionType.TOUCH), + ('touch_move', action_type.ActionType.TOUCH), + ('touch_up', action_type.ActionType.LIFT), + ('text', action_type.ActionType.TEXT), + ('keydown', action_type.ActionType.KEYDOWN), + ('set_selection', action_type.ActionType.SET_SELECTION), + ('cut', action_type.ActionType.CUT), + ('copy', action_type.ActionType.COPY), + ('paste', action_type.ActionType.PASTE), + ('random', action_type.ActionType.REPEAT), + ) + def test_map_action_type(self, act_str, expected): + self.assertEqual(actions_lib.map_action_type(act_str), expected) + + @parameterized.parameters( + ( + 'touch_down', + [0.1, 0.2], + None, + None, + None, + device_env_service_pb2.ACTION_TYPE_TOUCH_DOWN, + ), + ( + 'touch_move', + [0.3, 0.4], + None, + None, + None, + device_env_service_pb2.ACTION_TYPE_TOUCH_MOVE, + ), + ( + 'touch_up', + [0.5, 0.6], + None, + None, + None, + device_env_service_pb2.ACTION_TYPE_TOUCH_UP, + ), + ( + 'keydown', + None, + 66, + None, + None, + device_env_service_pb2.ACTION_TYPE_KEY_EVENT, + ), + ( + 'text', + None, + None, + 'hello', + None, + device_env_service_pb2.ACTION_TYPE_SET_TEXT, + ), + ( + 'set_selection', + None, + None, + None, + (1, 4), + device_env_service_pb2.ACTION_TYPE_SET_SELECTION, + ), + ( + 'cut', + None, + None, + None, + None, + device_env_service_pb2.ACTION_TYPE_CUT, + ), + ( + 'copy', + None, + None, + None, + None, + device_env_service_pb2.ACTION_TYPE_COPY, + ), + ( + 'paste', + None, + None, + None, + None, + device_env_service_pb2.ACTION_TYPE_PASTE, + ), + ( + 'unsupported_unknown', + None, + None, + None, + None, + device_env_service_pb2.ACTION_TYPE_UNSPECIFIED, + ), + ) + def test_map_client_action_to_proto( + self, + act_type, + pos_list, + keycode, + text, + selection, + expected_proto_type, + ): + client_action = actions_lib.ClientAction( + action_type=act_type, + x=0.1, + y=0.2, + client_timestamp_ms=100.0, + recv_timestamp_ms=102.0, + keycode=keycode, + text=text, + selection=selection, + ) + pos = np.array(pos_list, dtype=np.float32) if pos_list else None + proto = actions_lib.map_client_action_to_proto(client_action, pos) + + self.assertEqual(proto.action_type, expected_proto_type) + if pos_list: + self.assertAlmostEqual(proto.touch_position.x, pos_list[0]) + self.assertAlmostEqual(proto.touch_position.y, pos_list[1]) + if keycode is not None: + self.assertEqual(proto.keycode, keycode) + if text is not None: + self.assertEqual(proto.text, text) + if selection is not None: + self.assertEqual(proto.selection_start, selection[0]) + self.assertEqual(proto.selection_end, selection[1]) + + +if __name__ == '__main__': + absltest.main() diff --git a/android_env/playground/configs/BUILD.bazel b/android_env/playground/configs/BUILD.bazel new file mode 100644 index 00000000..099c557a --- /dev/null +++ b/android_env/playground/configs/BUILD.bazel @@ -0,0 +1,22 @@ +load("@rules_python//python:defs.bzl", "py_library", "py_test") + +package(default_visibility = ["//visibility:public"]) + +py_library( + name = "emulator_config", + srcs = ["emulator_config.py"], + deps = [ + "//android_env/components:config_classes", + ], +) + +py_test( + name = "emulator_config_test", + srcs = ["emulator_config_test.py"], + legacy_create_init = False, + deps = [ + ":emulator_config", + "//android_env/components:config_classes", + "@pip//absl_py", + ], +) diff --git a/android_env/playground/configs/emulator_config.py b/android_env/playground/configs/emulator_config.py new file mode 100644 index 00000000..d2d0d7d7 --- /dev/null +++ b/android_env/playground/configs/emulator_config.py @@ -0,0 +1,53 @@ +# coding=utf-8 +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Sample open-source emulator configuration for AndroidEnv Playground.""" + +import os +from android_env.components import config_classes + + +def get_config( + avd_name: str = 'androidworld_avd', +) -> config_classes.AndroidDeviceConfig: + """Returns an AndroidDeviceConfig for open-source emulators and AndroidWorld. + + Args: + avd_name: Name of the local Android Virtual Device (AVD). Defaults to + 'androidworld_avd', matching the standard AndroidWorld benchmark AVD. Can + be customized via --config=path/to/emulator_config.py:your_avd_name. + + Returns: + An AndroidDeviceConfig instance configured for local open-source execution. + """ + android_home = os.environ.get( + 'ANDROID_HOME', os.path.expanduser('~/Android/Sdk') + ) + android_avd_home = os.environ.get( + 'ANDROID_AVD_HOME', os.path.expanduser('~/.android/avd') + ) + + return config_classes.AndroidDeviceConfig( + emulator_launcher=config_classes.EmulatorLauncherConfig( + avd_name=avd_name, + android_avd_home=android_avd_home, + android_sdk_root=android_home, + emulator_path=os.path.join(android_home, 'emulator', 'emulator'), + run_headless=True, + ), + adb_controller=config_classes.AdbControllerConfig( + adb_path=os.path.join(android_home, 'platform-tools', 'adb'), + ), + ) diff --git a/android_env/playground/configs/emulator_config_test.py b/android_env/playground/configs/emulator_config_test.py new file mode 100644 index 00000000..9edc40a1 --- /dev/null +++ b/android_env/playground/configs/emulator_config_test.py @@ -0,0 +1,65 @@ +# coding=utf-8 +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests verifying open-source emulator configuration and environment variables.""" + +from unittest import mock + +from absl.testing import absltest +from android_env.components import config_classes +from android_env.playground.configs import emulator_config + + +class EmulatorConfigTest(absltest.TestCase): + + def test_get_config_default(self): + config = emulator_config.get_config() + self.assertIsInstance(config, config_classes.AndroidDeviceConfig) + self.assertIsNotNone(config.emulator_launcher) + self.assertEqual(config.emulator_launcher.avd_name, 'androidworld_avd') + self.assertTrue(config.emulator_launcher.run_headless) + + def test_get_config_custom_avd(self): + config = emulator_config.get_config(avd_name='custom_test_avd') + self.assertIsInstance(config, config_classes.AndroidDeviceConfig) + self.assertIsNotNone(config.emulator_launcher) + self.assertEqual(config.emulator_launcher.avd_name, 'custom_test_avd') + + def test_get_config_environment_variables(self): + with mock.patch.dict( + emulator_config.os.environ, + { + 'ANDROID_HOME': '/custom/android/home', + 'ANDROID_AVD_HOME': '/custom/avd/home', + }, + ): + config = emulator_config.get_config() + self.assertIsInstance(config, config_classes.AndroidDeviceConfig) + self.assertIsNotNone(config.emulator_launcher) + self.assertIsNotNone(config.adb_controller) + launcher = config.emulator_launcher + self.assertEqual(launcher.android_sdk_root, '/custom/android/home') + self.assertEqual(launcher.android_avd_home, '/custom/avd/home') + self.assertEqual( + launcher.emulator_path, '/custom/android/home/emulator/emulator' + ) + self.assertEqual( + config.adb_controller.adb_path, + '/custom/android/home/platform-tools/adb', + ) + + +if __name__ == '__main__': + absltest.main() diff --git a/android_env/playground/coordinates.py b/android_env/playground/coordinates.py new file mode 100644 index 00000000..2329bcfa --- /dev/null +++ b/android_env/playground/coordinates.py @@ -0,0 +1,71 @@ +# coding=utf-8 +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Coordinate mapping functions for AndroidEnv Playground.""" + +import numpy as np + + +def get_crop_metrics(raw_shape: tuple[int, int]) -> tuple[int, int, int, int]: + """Returns (crop_y_start, crop_y_end, crop_x_start, crop_x_end) for landscape.""" + h_v, w_v = raw_shape[:2] + target_ar = 16.0 / 9.0 + current_ar = w_v / h_v + + if current_ar > target_ar: + # Too wide, crop width + crop_w = int(h_v * target_ar) + crop_x_start = (w_v - crop_w) // 2 + crop_x_end = crop_x_start + crop_w + return 0, h_v, crop_x_start, crop_x_end + else: + # Too tall (or matches), crop height + crop_h = int(w_v / target_ar) + crop_y_start = (h_v - crop_h) // 2 + crop_y_end = crop_y_start + crop_h + return crop_y_start, crop_y_end, 0, w_v + + +def map_client_coordinates( + x: float, + y: float, + orientation: int, + raw_shape: tuple[int, int], +) -> np.ndarray: + """Maps client-side normalized coordinates to raw emulator coordinates.""" + pos = np.array([x, y], dtype=np.float32) + + # Adjust coordinates based on cropping in landscape mode. + if orientation in {1, 3}: + h_rot, w_rot = raw_shape[1], raw_shape[0] + crop_y_start, crop_y_end, crop_x_start, crop_x_end = get_crop_metrics( + (h_rot, w_rot) + ) + if crop_x_start > 0 or crop_x_end < w_rot: + crop_width = crop_x_end - crop_x_start + pos[0] = (crop_x_start + pos[0] * crop_width) / w_rot + if crop_y_start > 0 or crop_y_end < h_rot: + crop_height = crop_y_end - crop_y_start + pos[1] = (crop_y_start + pos[1] * crop_height) / h_rot + + # Coordinate mapping based on orientation + if orientation == 1: # LANDSCAPE_90 + pos = np.array([1.0 - pos[1], pos[0]], dtype=np.float32) + elif orientation == 3: # LANDSCAPE_270 + pos = np.array([pos[1], 1.0 - pos[0]], dtype=np.float32) + elif orientation == 2: # PORTRAIT_180 + pos = np.array([1.0 - pos[0], 1.0 - pos[1]], dtype=np.float32) + + return pos diff --git a/android_env/playground/coordinates_test.py b/android_env/playground/coordinates_test.py new file mode 100644 index 00000000..15a6590c --- /dev/null +++ b/android_env/playground/coordinates_test.py @@ -0,0 +1,56 @@ +# coding=utf-8 +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for coordinate mapping and cropping in AndroidEnv Playground.""" + +from absl.testing import absltest +from absl.testing import parameterized +from android_env.playground import coordinates +import numpy as np + + +class CoordinatesTest(parameterized.TestCase): + + @parameterized.parameters( + ((1080, 2209), (0, 1080, 144, 2064)), + ((720, 1280), (0, 720, 0, 1280)), + ((768, 1024), (96, 672, 0, 1024)), + ) + def test_get_crop_metrics(self, raw_shape, expected): + self.assertEqual(coordinates.get_crop_metrics(raw_shape), expected) + + def test_map_client_coordinates_portrait(self): + raw_shape = (1280, 720) + pos = coordinates.map_client_coordinates(0.1, 0.2, 0, raw_shape) + np.testing.assert_allclose(pos, [0.1, 0.2]) + + def test_map_client_coordinates_portrait_180(self): + raw_shape = (1280, 720) + pos = coordinates.map_client_coordinates(0.1, 0.2, 2, raw_shape) + np.testing.assert_allclose(pos, [0.9, 0.8]) + + def test_map_client_coordinates_landscape_90_wide(self): + raw_shape = (2209, 1080) + pos = coordinates.map_client_coordinates(0.5, 0.5, 1, raw_shape) + np.testing.assert_allclose(pos, [0.5, 1104.0 / 2209.0], rtol=1e-5) + + def test_map_client_coordinates_landscape_270_tall(self): + raw_shape = (1024, 768) + pos = coordinates.map_client_coordinates(0.5, 0.5, 3, raw_shape) + np.testing.assert_allclose(pos, [0.5, 0.5]) + + +if __name__ == '__main__': + absltest.main() diff --git a/android_env/playground/device_runner.py b/android_env/playground/device_runner.py new file mode 100644 index 00000000..e1553cba --- /dev/null +++ b/android_env/playground/device_runner.py @@ -0,0 +1,506 @@ +# coding=utf-8 +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Device runner for AndroidEnv Playground.""" + +import os +import queue +import shlex +import tempfile +import time +from typing import Any +from absl import logging +from android_env.components import android_device +from android_env.components import android_device_loader +from android_env.components import config_classes +from android_env.components import errors +from android_env.components import setup_step_interpreter + +from google.protobuf import text_format +from android_env.playground import actions as actions_lib +from android_env.playground import coordinates +from android_env.playground import frame_processing +from android_env.playground import state as state_lib +from android_env.playground import websocket +from android_env.proto import adb_pb2 +from android_env.proto import device_env_service_pb2 +from android_env.proto import task_pb2 +import numpy as np + + +class DeviceSetupError(errors.AndroidEnvError, RuntimeError): + """Raised when setting up or resetting the device fails.""" + + +def _update_orientation(curr_orient: int, orientation_container: list[int]): + """Updates orientation in container if it changed, logging the change.""" + if curr_orient != orientation_container[0]: + logging.info( + 'Orientation changed: %d -> %d', orientation_container[0], curr_orient + ) + orientation_container[0] = curr_orient + + +def _extract_screenshot( + device_state: device_env_service_pb2.DeviceState, +) -> np.ndarray | None: + """Extracts RGB pixels from decoded screenshot signal in DeviceState.""" + for signal in device_state.signals: + if ( + signal.type == device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT + and signal.HasField('screenshot') + and signal.screenshot.HasField('decoded') + and signal.screenshot.decoded.raw_pixels + ): + decoded = signal.screenshot.decoded + arr = np.frombuffer(decoded.raw_pixels, dtype=np.uint8) + return arr.reshape((decoded.height, decoded.width, 3)) + return None + + +def _update_frame( + state: state_lib.PlaygroundState, + pixels: np.ndarray, + orientation: int, +): + """Processes the frame, updates state, and sends it to the client.""" + jpeg_bytes = frame_processing.process_frame(pixels, orientation) + with state.frame_lock: + state.current_frame = jpeg_bytes + websocket.try_send_frame(state, jpeg_bytes) + + +def _get_actions( + state: state_lib.PlaygroundState, frame_rate_secs: float +) -> list[actions_lib.ClientAction]: + """Reads and coalesces actions from the state queue.""" + actions_to_take = [] + try: + actions_to_take.append( + state.action_queue.get(timeout=frame_rate_secs / 2.0) + ) + while True: + actions_to_take.append(state.action_queue.get_nowait()) + except queue.Empty: + pass + + raw_queue_size = len(actions_to_take) + actions_to_take = actions_lib.coalesce_actions(actions_to_take) + if raw_queue_size > 1: + logging.vlog( + 1, + 'Coalesced %d actions into %d', + raw_queue_size, + len(actions_to_take), + ) + return actions_to_take + + +def _send_client_actions( + device: android_device.AndroidDevice, + actions: list[actions_lib.ClientAction], + state: state_lib.PlaygroundState, + orientation_container: list[int], + last_pos: np.ndarray, + raw_shape: tuple[int, int], +) -> tuple[np.ndarray, np.ndarray | None]: + """Sends client actions to the device and updates timing stats.""" + device_state = None + for i, act_data in enumerate(actions): + pos = coordinates.map_client_coordinates( + act_data.x, act_data.y, orientation_container[0], raw_shape + ) + last_pos = pos + + now_ms = time.time() * 1000 + transit_ms = act_data.recv_timestamp_ms - act_data.client_timestamp_ms + queue_ms = now_ms - act_data.recv_timestamp_ms + total_pre_step_ms = now_ms - act_data.client_timestamp_ms + logging.vlog( + 1, + 'Action %d/%d [%s]: Transit: %.2f ms, Queue: %.2f ms, Total' + ' pre-step: %.2f ms', + i + 1, + len(actions), + act_data.action_type, + transit_ms, + queue_ms, + total_pre_step_ms, + ) + state.add_stat('action_transit_ms', transit_ms) + state.add_stat('action_queue_ms', queue_ms) + state.add_stat('action_total_pre_step_ms', total_pre_step_ms) + + proto_action = actions_lib.map_client_action_to_proto(act_data, pos) + + start_step = time.time() + device.send_action(proto_action) + step_duration_ms = (time.time() - start_step) * 1000 + logging.vlog(1, 'device.send_action took: %.2f ms', step_duration_ms) + state.add_stat('device_send_action_ms', step_duration_ms) + state.add_stat('env_step_ms', step_duration_ms) + + device_state = device.get_state() + _update_orientation(device_state.orientation, orientation_container) + + if device_state is None: + device_state = device.get_state() + _update_orientation(device_state.orientation, orientation_container) + + pixels = _extract_screenshot(device_state) + if pixels is not None: + _update_frame(state, pixels, orientation_container[0]) + + return last_pos, pixels + + +def _update_battery_stats( + device_state: device_env_service_pb2.DeviceState, + state: state_lib.PlaygroundState, +) -> None: + """Fetches battery power metrics from device_state and updates state if present.""" + for signal in device_state.signals: + if ( + signal.type == device_env_service_pb2.DEVICE_SIGNAL_BATTERY_POWER + and signal.HasField('battery_power') + ): + state.set_battery_power_stats(signal.battery_power) + return + + +def _load_task_proto(task_id: str, tmp_dir: str) -> task_pb2.Task: + """Resolves a task_id or file path into a Task protobuf.""" + expanded_path = os.path.expanduser(task_id) + del tmp_dir + task = task_pb2.Task() + with open(expanded_path, 'r') as proto_file: + text_format.Parse(proto_file.read(), task) + return task + + +def _get_interpreter_for_device( + device: android_device.AndroidDevice, +) -> setup_step_interpreter.SetupStepInterpreter: + """Creates a SetupStepInterpreter configured for the given AndroidDevice.""" + parser = getattr(device, '_adb_call_parser', None) + if parser is None: + + class _AdbParserAdapter: + + def __init__(self, dev: android_device.AndroidDevice): + self._dev = dev + + def parse(self, call: adb_pb2.AdbRequest) -> adb_pb2.AdbResponse: + return self._dev.execute_adb_call(call) + + parser = _AdbParserAdapter(device) + return setup_step_interpreter.SetupStepInterpreter(parser) # pyrefly: ignore[bad-argument-type] + + +def _handle_task_reload( + device: android_device.AndroidDevice, + state: state_lib.PlaygroundState, + orientation_container: list[int], + current_task: task_pb2.Task | None, +) -> tuple[task_pb2.Task | None, np.ndarray | None]: + """Checks state for a pending task reload request and executes setup steps.""" + task_id = state.get_reload() + if not task_id: + return current_task, None + + state.clear_reload() + websocket.try_send_status(state, f'Loading task: {task_id}...') + with tempfile.TemporaryDirectory(prefix='android_env_task_') as tmp_dir: + try: + task = _load_task_proto(task_id, tmp_dir) + except (errors.AndroidEnvError, RuntimeError, ValueError, OSError) as e: + logging.exception('Failed to load task %s: %s', task_id, e) + websocket.try_send_status(state, f'ERROR loading task {task_id}: {e}') + return current_task, None + + try: + interpreter = _get_interpreter_for_device(device) + interpreter.interpret(task.setup_steps) + except (errors.AndroidEnvError, RuntimeError, ValueError, OSError) as e: + logging.exception( + 'Failed to execute task setup steps for %s: %s', task_id, e + ) + websocket.try_send_status(state, f'ERROR setting up task {task_id}: {e}') + return current_task, None + + pixels = None + try: + device_state = device.get_state() + _update_orientation(device_state.orientation, orientation_container) + pixels = _extract_screenshot(device_state) + if pixels is not None: + _update_frame(state, pixels, orientation_container[0]) + except (errors.AndroidEnvError, RuntimeError, ValueError, OSError) as e: + logging.exception( + 'Failed to get device state after loading task %s: %s', task_id, e + ) + + websocket.try_send_status(state, f'Ready! Task: {task_id}') + return task, pixels + + +def _handle_episode_reset( + device: android_device.AndroidDevice, + state: state_lib.PlaygroundState, + orientation_container: list[int], + current_task: task_pb2.Task | None, +) -> np.ndarray | None: + """Checks state for a pending episode reset request and executes reset steps.""" + if not state.get_reset(): + return None + + state.clear_reset() + websocket.try_send_status(state, 'Resetting episode...') + try: + if current_task is not None and current_task.reset_steps: + interpreter = _get_interpreter_for_device(device) + interpreter.interpret(current_task.reset_steps) + except (errors.AndroidEnvError, RuntimeError, ValueError, OSError) as e: + logging.exception('Failed to reset episode: %s', e) + websocket.try_send_status(state, f'ERROR resetting episode: {e}') + return None + + pixels = None + try: + device_state = device.get_state() + _update_orientation(device_state.orientation, orientation_container) + pixels = _extract_screenshot(device_state) + if pixels is not None: + _update_frame(state, pixels, orientation_container[0]) + except (errors.AndroidEnvError, RuntimeError, ValueError, OSError) as e: + logging.exception('Failed to get device state after reset: %s', e) + websocket.try_send_status(state, f'ERROR resetting episode: {e}') + return None + + websocket.try_send_status(state, 'Ready! Episode reset.') + return pixels + + +def _execute_single_adb_command( + device: android_device.AndroidDevice, + state: state_lib.PlaygroundState, + cmd: str, +) -> None: + """Parses and executes a single ADB command against device and sends response.""" + try: + args = shlex.split(cmd) + except ValueError as e: + websocket.try_send_adb_response( + state, + command=cmd, + status='FAILED_PRECONDITION', + output='', + error=f'Invalid command syntax: {e}', + ) + return + + if args[:1] == ['adb']: + args = args[1:] + + if not args: + websocket.try_send_adb_response( + state, + command=cmd, + status='FAILED_PRECONDITION', + output='', + error='Empty ADB command.', + ) + return + + start_ms = time.time() * 1000 + request = adb_pb2.AdbRequest( + generic=adb_pb2.AdbRequest.GenericRequest(args=args) + ) + try: + response = device.execute_adb_call(request) + except (errors.AndroidEnvError, RuntimeError, ValueError, OSError) as e: + duration_ms = time.time() * 1000 - start_ms + websocket.try_send_adb_response( + state, + command=cmd, + status='INTERNAL_ERROR', + output='', + error=str(e), + duration_ms=duration_ms, + ) + return + + duration_ms = time.time() * 1000 - start_ms + try: + status_name = adb_pb2.AdbResponse.Status.Name(response.status) + except ValueError: + status_name = str(response.status) + output_text = response.generic.output.decode('utf-8', errors='replace') + websocket.try_send_adb_response( + state, + command=cmd, + status=status_name, + output=output_text, + error=response.error_message, + duration_ms=duration_ms, + ) + + +def _handle_adb_commands( + device: android_device.AndroidDevice, state: state_lib.PlaygroundState +) -> None: + """Drains and executes all queued ADB commands.""" + while True: + try: + cmd = state.adb_queue.get_nowait() + except queue.Empty: + break + _execute_single_adb_command(device, state, cmd) + + +def run_device( + device: android_device.AndroidDevice, + initial_orientation: int, + initial_pixels: np.ndarray, + state: state_lib.PlaygroundState, +) -> None: + """Runs the device loop, handling actions and updating frames.""" + frame_rate_secs = 1.0 / 30.0 + prev_frame_time = time.time() + last_pos = np.array([0.0, 0.0], dtype=np.float32) + + orientation_container = [initial_orientation] + last_stats_send_time = time.time() + current_pixels = initial_pixels + current_task: task_pb2.Task | None = None + + # Send initial frame + _update_frame(state, current_pixels, orientation_container[0]) + + while state.running: + current_task, new_pixels = _handle_task_reload( + device, state, orientation_container, current_task + ) + if new_pixels is not None: + current_pixels = new_pixels + + new_pixels = _handle_episode_reset( + device, state, orientation_container, current_task + ) + if new_pixels is not None: + current_pixels = new_pixels + + _handle_adb_commands(device, state) + + actions = _get_actions(state, frame_rate_secs) + + if actions: + last_pos, new_pixels = _send_client_actions( + device, + actions, + state, + orientation_container, + last_pos, + current_pixels.shape[:2], + ) + if new_pixels is not None: + current_pixels = new_pixels + else: + # Idle: poll device state for screen changes without injecting dummy + # actions. + try: + device_state = device.get_state() + _update_orientation(device_state.orientation, orientation_container) + pixels = _extract_screenshot(device_state) + if pixels is not None: + current_pixels = pixels + _update_frame(state, current_pixels, orientation_container[0]) + except (errors.AndroidEnvError, RuntimeError, ValueError, OSError) as e: + logging.vlog(1, 'Failed to fetch device state during idle poll: %s', e) + + # Limit frame rate + now = time.time() + elapsed = now - prev_frame_time + if elapsed < frame_rate_secs: + time.sleep(frame_rate_secs - elapsed) + prev_frame_time = time.time() + + # Periodically send accumulated stats (every 500ms) + now = time.time() + if now - last_stats_send_time > 0.5: + try: + device_state = device.get_state() + _update_battery_stats(device_state, state) + except (errors.AndroidEnvError, RuntimeError, ValueError, OSError) as e: + logging.vlog(1, 'Failed to fetch device state for stats: %s', e) + stats = state.flush_stats() + if stats: + websocket.try_send_stats(state, stats) + last_stats_send_time = now + + +def setup_device( + state: state_lib.PlaygroundState, + config: Any, +) -> tuple[android_device.AndroidDevice, int, np.ndarray]: + """Loads and connects to AndroidDevice, returning the device, orientation and first frame pixels.""" + + def set_status(text: str): + logging.info('Status: %s', text) + websocket.try_send_status(state, text) + + set_status('Loading configuration...') + if not config: + logging.error('No configuration provided.') + set_status('ERROR: No config') + raise DeviceSetupError('No configuration provided.') + + # Extract device config if wrapped in an AndroidEnvConfig. + if isinstance(config, config_classes.AndroidEnvConfig): + device_config = config.simulator + else: + device_config = config + + if not isinstance(device_config, config_classes.AndroidDeviceConfig): + raise DeviceSetupError( + f'Expected AndroidDeviceConfig, got {type(device_config).__name__}' + ) + + set_status('Launching Device / Connecting (this may take 1-2 mins)...') + try: + device = android_device_loader.load_android_device(device_config) + except (errors.AndroidEnvError, RuntimeError, ValueError, OSError) as e: + logging.exception('Error during device setup') + set_status(f'ERROR: {e}') + raise DeviceSetupError(f'Error during device setup: {e}') from e + + set_status('Fetching initial device state...') + try: + device_state = device.get_state() + orientation = device_state.orientation + pixels = _extract_screenshot(device_state) + if pixels is None: + raise errors.ReadObservationError( + 'Initial device state is missing decoded screenshot.' + ) + except (errors.AndroidEnvError, RuntimeError, ValueError, OSError) as e: + logging.exception('Error after loading device, closing device.') + device.close() + set_status(f'ERROR: {e}') + raise DeviceSetupError(f'Error fetching initial device state: {e}') from e + + _update_battery_stats(device_state, state) + set_status(f'Ready! Orientation: {orientation}') + return device, orientation, pixels diff --git a/android_env/playground/device_runner_test.py b/android_env/playground/device_runner_test.py new file mode 100644 index 00000000..218ba823 --- /dev/null +++ b/android_env/playground/device_runner_test.py @@ -0,0 +1,828 @@ +# coding=utf-8 +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for playground device_runner.""" + +import typing +from unittest import mock + +from absl.testing import absltest +from android_env.components import android_device +from android_env.components import config_classes +from android_env.components import errors +from android_env.playground import actions as actions_lib +from android_env.playground import device_runner +from android_env.playground import state as state_lib +from android_env.proto import adb_pb2 +from android_env.proto import device_env_service_pb2 +from android_env.proto import task_pb2 +import numpy as np + + +class DeviceRunnerTest(absltest.TestCase): + + def test_update_orientation_no_change(self): + orientation_container = [0] + device_runner._update_orientation(0, orientation_container) + self.assertEqual(orientation_container[0], 0) + + def test_update_orientation_changed(self): + orientation_container = [0] + with mock.patch.object(device_runner.logging, 'info') as mock_info: + device_runner._update_orientation(1, orientation_container) + self.assertEqual(orientation_container[0], 1) + mock_info.assert_called_once() + + def test_extract_screenshot_valid(self): + state = device_env_service_pb2.DeviceState() + signal = state.signals.add() + signal.type = device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT + img_data = np.zeros((10, 20, 3), dtype=np.uint8) + signal.screenshot.decoded.raw_pixels = img_data.tobytes() + signal.screenshot.decoded.width = 20 + signal.screenshot.decoded.height = 10 + + pixels = device_runner._extract_screenshot(state) + self.assertIsNotNone(pixels) + self.assertEqual(pixels.shape, (10, 20, 3)) + + def test_extract_screenshot_missing(self): + state = device_env_service_pb2.DeviceState() + pixels = device_runner._extract_screenshot(state) + self.assertIsNone(pixels) + + @mock.patch.object( + device_runner.frame_processing, 'process_frame', autospec=True + ) + @mock.patch.object(device_runner.websocket, 'try_send_frame', autospec=True) + def test_update_frame(self, mock_send_frame, mock_process_frame): + state = state_lib.PlaygroundState() + pixels = np.zeros((10, 10, 3), dtype=np.uint8) + mock_process_frame.return_value = b'jpeg_bytes' + + device_runner._update_frame(state, pixels, 0) + + mock_process_frame.assert_called_once_with(pixels, 0) + self.assertEqual(state.current_frame, b'jpeg_bytes') + mock_send_frame.assert_called_once_with(state, b'jpeg_bytes') + + def test_get_actions_empty(self): + state = state_lib.PlaygroundState() + actions = device_runner._get_actions(state, frame_rate_secs=0.1) + self.assertEqual(actions, []) + + def test_get_actions_single(self): + state = state_lib.PlaygroundState() + act = actions_lib.ClientAction( + action_type='touch_down', + x=0.5, + y=0.5, + client_timestamp_ms=100.0, + recv_timestamp_ms=105.0, + ) + state.action_queue.put(act) + actions = device_runner._get_actions(state, frame_rate_secs=0.1) + self.assertEqual(actions, [act]) + + def test_get_actions_multiple_coalesced(self): + state = state_lib.PlaygroundState() + act1 = actions_lib.ClientAction( + action_type='touch_move', + x=0.5, + y=0.5, + client_timestamp_ms=100.0, + recv_timestamp_ms=105.0, + ) + act2 = actions_lib.ClientAction( + action_type='touch_move', + x=0.6, + y=0.6, + client_timestamp_ms=110.0, + recv_timestamp_ms=115.0, + ) + state.action_queue.put(act1) + state.action_queue.put(act2) + actions = device_runner._get_actions(state, frame_rate_secs=0.1) + self.assertLen(actions, 1) + self.assertEqual(actions[0].x, 0.6) + + @mock.patch.object(device_runner, '_update_frame', autospec=True) + @mock.patch.object( + device_runner.coordinates, 'map_client_coordinates', autospec=True + ) + @mock.patch.object( + device_runner.actions_lib, 'map_client_action_to_proto', autospec=True + ) + def test_send_client_actions( + self, mock_map_proto, mock_map_coords, mock_update_frame + ): + mock_dev = mock.create_autospec(android_device.AndroidDevice, instance=True) + state = state_lib.PlaygroundState() + orientation_container = [0] + last_pos = np.array([0.0, 0.0]) + + dev_state = device_env_service_pb2.DeviceState(orientation=0) + sig = dev_state.signals.add() + sig.type = device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT + img = np.zeros((10, 10, 3), dtype=np.uint8) + sig.screenshot.decoded.raw_pixels = img.tobytes() + sig.screenshot.decoded.width = 10 + sig.screenshot.decoded.height = 10 + mock_dev.get_state.return_value = dev_state + + proto_action = device_env_service_pb2.Action( + action_type=device_env_service_pb2.ACTION_TYPE_TOUCH_DOWN + ) + mock_map_proto.return_value = proto_action + mock_map_coords.return_value = np.array([5.0, 5.0]) + + act_data = actions_lib.ClientAction( + action_type='touch_down', + x=0.5, + y=0.5, + client_timestamp_ms=100.0, + recv_timestamp_ms=105.0, + ) + + res_last_pos, res_pixels = device_runner._send_client_actions( + mock_dev, + [act_data], + state, + orientation_container, + last_pos, + (10, 10), + ) + + mock_dev.send_action.assert_called_once_with(proto_action) + np.testing.assert_array_equal(res_last_pos, np.array([5.0, 5.0])) + self.assertIsNotNone(res_pixels) + mock_update_frame.assert_called_once() + + @mock.patch.object(device_runner, '_update_frame', autospec=True) + @mock.patch.object( + device_runner.coordinates, 'map_client_coordinates', autospec=True + ) + @mock.patch.object( + device_runner.actions_lib, 'map_client_action_to_proto', autospec=True + ) + def test_send_client_actions_updates_orientation_per_action( + self, mock_map_proto, mock_map_coords, mock_update_frame + ): + mock_dev = mock.create_autospec(android_device.AndroidDevice, instance=True) + state = state_lib.PlaygroundState() + orientation_container = [0] + last_pos = np.array([0.0, 0.0]) + + dev_state_0 = device_env_service_pb2.DeviceState(orientation=1) + dev_state_1 = device_env_service_pb2.DeviceState(orientation=1) + mock_dev.get_state.side_effect = [dev_state_0, dev_state_1] + + mock_map_proto.return_value = device_env_service_pb2.Action() + mock_map_coords.side_effect = [np.array([1.0, 1.0]), np.array([2.0, 2.0])] + + act_1 = actions_lib.ClientAction( + action_type='touch_down', + x=0.1, + y=0.1, + client_timestamp_ms=100.0, + recv_timestamp_ms=105.0, + ) + act_2 = actions_lib.ClientAction( + action_type='touch_move', + x=0.2, + y=0.2, + client_timestamp_ms=110.0, + recv_timestamp_ms=115.0, + ) + + device_runner._send_client_actions( + mock_dev, + [act_1, act_2], + state, + orientation_container, + last_pos, + (10, 10), + ) + + self.assertEqual(mock_dev.send_action.call_count, 2) + self.assertEqual(mock_dev.get_state.call_count, 2) + self.assertEqual(orientation_container[0], 1) + self.assertEqual(mock_map_coords.call_args_list[0][0][2], 0) + self.assertEqual(mock_map_coords.call_args_list[1][0][2], 1) + mock_update_frame.assert_not_called() + + def test_update_battery_stats(self): + state = state_lib.PlaygroundState() + dev_state = device_env_service_pb2.DeviceState() + sig = dev_state.signals.add() + sig.type = device_env_service_pb2.DEVICE_SIGNAL_BATTERY_POWER + sig.battery_power.level_percent = 85 + sig.battery_power.voltage_mv = 4100 + + device_runner._update_battery_stats(dev_state, state) + + battery_stats = state.get_battery_power_stats() + self.assertIsNotNone(battery_stats) + self.assertEqual(battery_stats.get('level_percent'), 85) + self.assertEqual(battery_stats.get('voltage_mv'), 4100) + + def test_update_battery_stats_missing(self): + state = state_lib.PlaygroundState() + dev_state = device_env_service_pb2.DeviceState() + device_runner._update_battery_stats(dev_state, state) + self.assertIsNone(state.get_battery_power_stats()) + + def test_load_task_proto_from_textproto_file(self): + task_file = self.create_tempfile( + 'task.textproto', content='id: "my_task"\nname: "My Task"\n' + ) + task = device_runner._load_task_proto( + task_file.full_path, self.create_tempdir().full_path + ) + self.assertEqual(task.id, 'my_task') + self.assertEqual(task.name, 'My Task') + + def test_get_interpreter_for_device_fallback(self): + mock_dev = mock.create_autospec(android_device.AndroidDevice, instance=True) + mock_dev.execute_adb_call.return_value = adb_pb2.AdbResponse( + status=adb_pb2.AdbResponse.Status.OK + ) + interpreter = device_runner._get_interpreter_for_device(mock_dev) + self.assertIsNotNone(interpreter) + req = adb_pb2.AdbRequest() + res = interpreter._adb_call_parser.parse(req) + self.assertEqual(res.status, adb_pb2.AdbResponse.Status.OK) + mock_dev.execute_adb_call.assert_called_once_with(req) + + def test_handle_task_reload_no_request(self): + mock_dev = mock.create_autospec(android_device.AndroidDevice, instance=True) + state = state_lib.PlaygroundState() + current_task = task_pb2.Task() + orientation_container = [0] + + task, pixels = device_runner._handle_task_reload( + mock_dev, state, orientation_container, current_task + ) + self.assertEqual(task, current_task) + self.assertIsNone(pixels) + + @mock.patch.object(device_runner.websocket, 'try_send_status', autospec=True) + @mock.patch.object(device_runner, '_load_task_proto', autospec=True) + @mock.patch.object( + device_runner, '_get_interpreter_for_device', autospec=True + ) + def test_handle_task_reload_success( + self, mock_get_interp, mock_load_proto, mock_send_status + ): + mock_dev = mock.create_autospec(android_device.AndroidDevice, instance=True) + state = state_lib.PlaygroundState() + state.set_reload('sample_task') + + loaded_task = task_pb2.Task() + mock_load_proto.return_value = loaded_task + + mock_interp = mock.MagicMock() + mock_get_interp.return_value = mock_interp + + dev_state = device_env_service_pb2.DeviceState(orientation=1) + sig = dev_state.signals.add() + sig.type = device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT + img = np.zeros((10, 10, 3), dtype=np.uint8) + sig.screenshot.decoded.raw_pixels = img.tobytes() + sig.screenshot.decoded.width = 10 + sig.screenshot.decoded.height = 10 + mock_dev.get_state.return_value = dev_state + + orientation_container = [0] + task, pixels = device_runner._handle_task_reload( + mock_dev, state, orientation_container, None + ) + + self.assertEqual(task, loaded_task) + self.assertIsNotNone(pixels) + self.assertEqual(orientation_container[0], 1) + mock_interp.interpret.assert_called_once_with(loaded_task.setup_steps) + mock_send_status.assert_any_call(state, 'Ready! Task: sample_task') + + @mock.patch.object(device_runner.websocket, 'try_send_status', autospec=True) + @mock.patch.object(device_runner, '_load_task_proto', autospec=True) + def test_handle_task_reload_load_error( + self, mock_load_proto, mock_send_status + ): + mock_dev = mock.create_autospec(android_device.AndroidDevice, instance=True) + state = state_lib.PlaygroundState() + state.set_reload('bad_task') + mock_load_proto.side_effect = RuntimeError('Failed to load') + + orientation_container = [0] + task, pixels = device_runner._handle_task_reload( + mock_dev, state, orientation_container, None + ) + + self.assertIsNone(task) + self.assertIsNone(pixels) + mock_send_status.assert_any_call( + state, 'ERROR loading task bad_task: Failed to load' + ) + + def test_handle_episode_reset_no_request(self): + mock_dev = mock.create_autospec(android_device.AndroidDevice, instance=True) + state = state_lib.PlaygroundState() + orientation_container = [0] + + pixels = device_runner._handle_episode_reset( + mock_dev, state, orientation_container, None + ) + self.assertIsNone(pixels) + + @mock.patch.object(device_runner.websocket, 'try_send_status', autospec=True) + @mock.patch.object( + device_runner, '_get_interpreter_for_device', autospec=True + ) + def test_handle_episode_reset_success( + self, mock_get_interp, mock_send_status + ): + mock_dev = mock.create_autospec(android_device.AndroidDevice, instance=True) + state = state_lib.PlaygroundState() + state.set_reset() + + current_task = task_pb2.Task() + current_task.reset_steps.add() + + mock_interp = mock.MagicMock() + mock_get_interp.return_value = mock_interp + + dev_state = device_env_service_pb2.DeviceState(orientation=0) + sig = dev_state.signals.add() + sig.type = device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT + img = np.zeros((10, 10, 3), dtype=np.uint8) + sig.screenshot.decoded.raw_pixels = img.tobytes() + sig.screenshot.decoded.width = 10 + sig.screenshot.decoded.height = 10 + mock_dev.get_state.return_value = dev_state + + orientation_container = [0] + pixels = device_runner._handle_episode_reset( + mock_dev, state, orientation_container, current_task + ) + + self.assertIsNotNone(pixels) + mock_interp.interpret.assert_called_once_with(current_task.reset_steps) + mock_send_status.assert_any_call(state, 'Ready! Episode reset.') + + @mock.patch.object( + device_runner.websocket, 'try_send_adb_response', autospec=True + ) + def test_execute_single_adb_command_syntax_error(self, mock_send_adb): + mock_dev = mock.create_autospec(android_device.AndroidDevice, instance=True) + state = state_lib.PlaygroundState() + device_runner._execute_single_adb_command( + mock_dev, state, 'shell "unclosed' + ) + mock_send_adb.assert_called_once() + self.assertEqual( + mock_send_adb.call_args[1]['status'], 'FAILED_PRECONDITION' + ) + + @mock.patch.object( + device_runner.websocket, 'try_send_adb_response', autospec=True + ) + def test_execute_single_adb_command_empty(self, mock_send_adb): + mock_dev = mock.create_autospec(android_device.AndroidDevice, instance=True) + state = state_lib.PlaygroundState() + device_runner._execute_single_adb_command(mock_dev, state, 'adb') + device_runner._execute_single_adb_command(mock_dev, state, ' ') + self.assertEqual(mock_send_adb.call_count, 2) + self.assertEqual( + mock_send_adb.call_args_list[0][1]['status'], 'FAILED_PRECONDITION' + ) + self.assertEqual( + mock_send_adb.call_args_list[1][1]['status'], 'FAILED_PRECONDITION' + ) + + @mock.patch.object( + device_runner.websocket, 'try_send_adb_response', autospec=True + ) + def test_execute_single_adb_command_execution_error(self, mock_send_adb): + mock_dev = mock.create_autospec(android_device.AndroidDevice, instance=True) + mock_dev.execute_adb_call.side_effect = RuntimeError('ADB connection lost') + state = state_lib.PlaygroundState() + with mock.patch.object( + device_runner.time, 'time', autospec=True, side_effect=[10.0, 10.15] + ): + device_runner._execute_single_adb_command( + mock_dev, state, 'adb shell getprop' + ) + mock_send_adb.assert_called_once() + self.assertEqual(mock_send_adb.call_args[1]['status'], 'INTERNAL_ERROR') + self.assertEqual(mock_send_adb.call_args[1]['duration_ms'], 150.0) + + @mock.patch.object( + device_runner.websocket, 'try_send_adb_response', autospec=True + ) + def test_execute_single_adb_command_success(self, mock_send_adb): + mock_dev = mock.create_autospec(android_device.AndroidDevice, instance=True) + resp = adb_pb2.AdbResponse( + status=adb_pb2.AdbResponse.Status.OK, + generic=adb_pb2.AdbResponse.GenericResponse(output=b'google/pixel\n'), + ) + mock_dev.execute_adb_call.return_value = resp + state = state_lib.PlaygroundState() + with mock.patch.object( + device_runner.time, 'time', autospec=True, side_effect=[20.0, 20.05] + ): + device_runner._execute_single_adb_command( + mock_dev, state, 'shell getprop ro.product.name' + ) + mock_send_adb.assert_called_once() + self.assertEqual(mock_send_adb.call_args[1]['status'], 'OK') + self.assertEqual(mock_send_adb.call_args[1]['output'], 'google/pixel\n') + self.assertEqual(mock_send_adb.call_args[1]['duration_ms'], 50.0) + + @mock.patch.object( + device_runner.websocket, 'try_send_adb_response', autospec=True + ) + def test_execute_single_adb_command_unrecognized_status(self, mock_send_adb): + mock_dev = mock.create_autospec(android_device.AndroidDevice, instance=True) + resp = adb_pb2.AdbResponse( + status=typing.cast(typing.Any, 999), + generic=adb_pb2.AdbResponse.GenericResponse(output=b''), + ) + mock_dev.execute_adb_call.return_value = resp + state = state_lib.PlaygroundState() + with mock.patch.object( + device_runner.time, 'time', autospec=True, side_effect=[30.0, 30.01] + ): + device_runner._execute_single_adb_command(mock_dev, state, 'adb devices') + mock_send_adb.assert_called_once() + self.assertEqual(mock_send_adb.call_args[1]['status'], '999') + self.assertEqual(mock_send_adb.call_args[1]['duration_ms'], 10.0) + + @mock.patch.object( + device_runner, '_execute_single_adb_command', autospec=True + ) + def test_handle_adb_commands(self, mock_execute_cmd): + mock_dev = mock.create_autospec(android_device.AndroidDevice, instance=True) + state = state_lib.PlaygroundState() + state.adb_queue.put('adb devices') + state.adb_queue.put('adb shell') + device_runner._handle_adb_commands(mock_dev, state) + self.assertEqual(mock_execute_cmd.call_count, 2) + + def test_setup_device_no_config_raises(self): + state = state_lib.PlaygroundState() + with self.assertRaises(device_runner.DeviceSetupError): + device_runner.setup_device(state, None) + + @mock.patch.object( + device_runner.android_device_loader, + 'load_android_device', + autospec=True, + ) + def test_setup_device_load_error_raises(self, mock_loader): + mock_loader.side_effect = errors.AdbControllerError('ADB failed') + state = state_lib.PlaygroundState() + config = config_classes.AndroidDeviceConfig() + with self.assertRaises(device_runner.DeviceSetupError): + device_runner.setup_device(state, config) + + @mock.patch.object( + device_runner.android_device_loader, + 'load_android_device', + autospec=True, + ) + def test_setup_device_missing_screenshot_raises(self, mock_loader): + mock_dev = mock.create_autospec(android_device.AndroidDevice, instance=True) + mock_dev.get_state.return_value = device_env_service_pb2.DeviceState() + mock_loader.return_value = mock_dev + state = state_lib.PlaygroundState() + config = config_classes.AndroidDeviceConfig() + with self.assertRaises(device_runner.DeviceSetupError): + device_runner.setup_device(state, config) + mock_dev.close.assert_called_once() + + @mock.patch.object( + device_runner.android_device_loader, + 'load_android_device', + autospec=True, + ) + def test_setup_device_success(self, mock_loader): + mock_dev = mock.create_autospec(android_device.AndroidDevice, instance=True) + dev_state = device_env_service_pb2.DeviceState(orientation=1) + sig = dev_state.signals.add() + sig.type = device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT + img = np.zeros((10, 10, 3), dtype=np.uint8) + sig.screenshot.decoded.raw_pixels = img.tobytes() + sig.screenshot.decoded.width = 10 + sig.screenshot.decoded.height = 10 + mock_dev.get_state.return_value = dev_state + mock_loader.return_value = mock_dev + + state = state_lib.PlaygroundState() + config = config_classes.AndroidEnvConfig( + simulator=config_classes.AndroidDeviceConfig() + ) + + dev, orientation, pixels = device_runner.setup_device(state, config) + self.assertEqual(dev, mock_dev) + self.assertEqual(orientation, 1) + self.assertEqual(pixels.shape, (10, 10, 3)) + + def test_setup_device_invalid_config_type(self): + state = state_lib.PlaygroundState() + with self.assertRaises(device_runner.DeviceSetupError): + device_runner.setup_device(state, 'invalid_config_type') + + @mock.patch.object(device_runner, '_update_frame', autospec=True) + @mock.patch.object(device_runner, '_send_client_actions', autospec=True) + @mock.patch.object(device_runner, '_get_actions', autospec=True) + def test_run_device_loop_with_actions( + self, mock_get_actions, mock_send_actions, unused_mock_frame + ): + mock_dev = mock.create_autospec(android_device.AndroidDevice, instance=True) + state = state_lib.PlaygroundState() + act = actions_lib.ClientAction( + action_type='touch_down', + x=0.5, + y=0.5, + client_timestamp_ms=100.0, + recv_timestamp_ms=105.0, + ) + mock_get_actions.return_value = [act] + pixels = np.zeros((10, 10, 3), dtype=np.uint8) + mock_send_actions.return_value = (np.array([5.0, 5.0]), pixels) + + def _stop_running(*unused_args, **unused_kwargs): + state.running = False + return np.array([5.0, 5.0]), pixels + + mock_send_actions.side_effect = _stop_running + + device_runner.run_device(mock_dev, 0, pixels, state) + + mock_send_actions.assert_called_once() + + @mock.patch.object(device_runner, '_update_frame', autospec=True) + @mock.patch.object(device_runner, '_get_actions', autospec=True) + def test_run_device_loop_idle(self, mock_get_actions, unused_mock_frame): + mock_dev = mock.create_autospec(android_device.AndroidDevice, instance=True) + dev_state = device_env_service_pb2.DeviceState(orientation=0) + sig = dev_state.signals.add() + sig.type = device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT + img = np.zeros((10, 10, 3), dtype=np.uint8) + sig.screenshot.decoded.raw_pixels = img.tobytes() + sig.screenshot.decoded.width = 10 + sig.screenshot.decoded.height = 10 + mock_dev.get_state.return_value = dev_state + + state = state_lib.PlaygroundState() + mock_get_actions.return_value = [] + + def _stop_running_on_get_state(): + state.running = False + return dev_state + + mock_dev.get_state.side_effect = _stop_running_on_get_state + pixels = np.zeros((10, 10, 3), dtype=np.uint8) + + device_runner.run_device(mock_dev, 0, pixels, state) + + mock_dev.get_state.assert_called_once() + + @mock.patch.object(device_runner.websocket, 'try_send_stats', autospec=True) + @mock.patch.object(device_runner, '_update_frame', autospec=True) + @mock.patch.object(device_runner, '_get_actions', autospec=True) + def test_run_device_loop_stats_flush( + self, mock_get_actions, unused_mock_frame, mock_send_stats + ): + mock_dev = mock.create_autospec(android_device.AndroidDevice, instance=True) + dev_state = device_env_service_pb2.DeviceState(orientation=0) + sig = dev_state.signals.add() + sig.type = device_env_service_pb2.DEVICE_SIGNAL_BATTERY_POWER + sig.battery_power.level_percent = 90 + mock_dev.get_state.return_value = dev_state + + state = state_lib.PlaygroundState() + state.add_stat('test_stat', 42.0) + mock_get_actions.return_value = [] + + iteration = [0] + + def _advance_and_stop(): + iteration[0] += 1 + if iteration[0] >= 2: + state.running = False + return dev_state + + mock_dev.get_state.side_effect = _advance_and_stop + + pixels = np.zeros((10, 10, 3), dtype=np.uint8) + with mock.patch.object( + device_runner.time, + 'time', + side_effect=[0.0, 0.0, 0.1, 1.0, 1.0, 1.0, 1.0], + ): + device_runner.run_device(mock_dev, 0, pixels, state) + + mock_send_stats.assert_called_once() + + @mock.patch.object(device_runner.websocket, 'try_send_status', autospec=True) + @mock.patch.object(device_runner, '_load_task_proto', autospec=True) + @mock.patch.object( + device_runner, '_get_interpreter_for_device', autospec=True + ) + def test_handle_task_reload_interpret_error( + self, mock_get_interp, mock_load_proto, mock_send_status + ): + mock_dev = mock.create_autospec(android_device.AndroidDevice, instance=True) + state = state_lib.PlaygroundState() + state.set_reload('bad_task') + task = task_pb2.Task() + task.setup_steps.add() + mock_load_proto.return_value = task + mock_interp = mock.MagicMock() + mock_interp.interpret.side_effect = RuntimeError('Setup step failed') + mock_get_interp.return_value = mock_interp + + orientation_container = [0] + res_task, pixels = device_runner._handle_task_reload( + mock_dev, state, orientation_container, None + ) + self.assertIsNone(res_task) + self.assertIsNone(pixels) + mock_send_status.assert_any_call( + state, 'ERROR setting up task bad_task: Setup step failed' + ) + + @mock.patch.object(device_runner.websocket, 'try_send_status', autospec=True) + @mock.patch.object(device_runner, '_load_task_proto', autospec=True) + @mock.patch.object( + device_runner, '_get_interpreter_for_device', autospec=True + ) + def test_handle_task_reload_get_state_error( + self, mock_get_interp, mock_load_proto, mock_send_status + ): + mock_dev = mock.create_autospec(android_device.AndroidDevice, instance=True) + state = state_lib.PlaygroundState() + state.set_reload('task_id') + task = task_pb2.Task() + mock_load_proto.return_value = task + mock_get_interp.return_value = mock.MagicMock() + mock_dev.get_state.side_effect = RuntimeError('State error') + + orientation_container = [0] + res_task, pixels = device_runner._handle_task_reload( + mock_dev, state, orientation_container, None + ) + self.assertEqual(res_task, task) + self.assertIsNone(pixels) + mock_send_status.assert_any_call(state, 'Ready! Task: task_id') + + @mock.patch.object(device_runner.websocket, 'try_send_status', autospec=True) + @mock.patch.object( + device_runner, '_get_interpreter_for_device', autospec=True + ) + def test_handle_episode_reset_interpret_error( + self, mock_get_interp, mock_send_status + ): + mock_dev = mock.create_autospec(android_device.AndroidDevice, instance=True) + state = state_lib.PlaygroundState() + state.set_reset() + task = task_pb2.Task() + task.reset_steps.add() + mock_interp = mock.MagicMock() + mock_interp.interpret.side_effect = RuntimeError('Reset step failed') + mock_get_interp.return_value = mock_interp + + orientation_container = [0] + pixels = device_runner._handle_episode_reset( + mock_dev, state, orientation_container, task + ) + self.assertIsNone(pixels) + mock_send_status.assert_any_call( + state, 'ERROR resetting episode: Reset step failed' + ) + + @mock.patch.object(device_runner.websocket, 'try_send_status', autospec=True) + @mock.patch.object( + device_runner, '_get_interpreter_for_device', autospec=True + ) + def test_handle_episode_reset_get_state_error( + self, mock_get_interp, mock_send_status + ): + mock_dev = mock.create_autospec(android_device.AndroidDevice, instance=True) + state = state_lib.PlaygroundState() + state.set_reset() + mock_get_interp.return_value = mock.MagicMock() + mock_dev.get_state.side_effect = RuntimeError('State error') + + orientation_container = [0] + pixels = device_runner._handle_episode_reset( + mock_dev, state, orientation_container, None + ) + self.assertIsNone(pixels) + mock_send_status.assert_any_call( + state, 'ERROR resetting episode: State error' + ) + + @mock.patch.object(device_runner, '_handle_task_reload', autospec=True) + @mock.patch.object(device_runner, '_handle_episode_reset', autospec=True) + @mock.patch.object(device_runner, '_update_frame', autospec=True) + @mock.patch.object(device_runner, '_get_actions', autospec=True) + def test_run_device_loop_updates_pixels_on_reload_and_reset( + self, mock_get_actions, unused_mock_frame, mock_reset, mock_reload + ): + mock_dev = mock.create_autospec(android_device.AndroidDevice, instance=True) + state = state_lib.PlaygroundState() + reload_pixels = np.ones((10, 10, 3), dtype=np.uint8) + reset_pixels = np.full((10, 10, 3), 2, dtype=np.uint8) + + def _side_effect_reload(*unused_args, **unused_kwargs): + return task_pb2.Task(), reload_pixels + + def _side_effect_reset(*unused_args, **unused_kwargs): + state.running = False + return reset_pixels + + mock_reload.side_effect = _side_effect_reload + mock_reset.side_effect = _side_effect_reset + mock_get_actions.return_value = [] + initial_pixels = np.zeros((10, 10, 3), dtype=np.uint8) + + device_runner.run_device(mock_dev, 0, initial_pixels, state) + mock_reload.assert_called_once() + mock_reset.assert_called_once() + + @mock.patch.object(device_runner, '_update_frame', autospec=True) + @mock.patch.object(device_runner, '_get_actions', autospec=True) + def test_run_device_loop_idle_error( + self, mock_get_actions, unused_mock_frame + ): + mock_dev = mock.create_autospec(android_device.AndroidDevice, instance=True) + state = state_lib.PlaygroundState() + mock_get_actions.return_value = [] + + def _stop_running(*unused_args, **unused_kwargs): + state.running = False + raise RuntimeError('Idle get_state failed') + + mock_dev.get_state.side_effect = _stop_running + pixels = np.zeros((10, 10, 3), dtype=np.uint8) + + device_runner.run_device(mock_dev, 0, pixels, state) + mock_dev.get_state.assert_called_once() + + @mock.patch.object(device_runner.websocket, 'try_send_stats', autospec=True) + @mock.patch.object(device_runner, '_update_frame', autospec=True) + @mock.patch.object(device_runner, '_get_actions', autospec=True) + def test_run_device_loop_stats_error( + self, mock_get_actions, unused_mock_frame, unused_mock_send_stats + ): + mock_dev = mock.create_autospec(android_device.AndroidDevice, instance=True) + state = state_lib.PlaygroundState() + mock_get_actions.return_value = [] + + dev_state_idle = device_env_service_pb2.DeviceState(orientation=0) + sig = dev_state_idle.signals.add() + sig.type = device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT + img = np.zeros((10, 10, 3), dtype=np.uint8) + sig.screenshot.decoded.raw_pixels = img.tobytes() + sig.screenshot.decoded.width = 10 + sig.screenshot.decoded.height = 10 + + call_count = [0] + + def _mock_get_state(): + call_count[0] += 1 + if call_count[0] == 1: + return dev_state_idle + state.running = False + raise RuntimeError('Stats get_state error') + + mock_dev.get_state.side_effect = _mock_get_state + + pixels = np.zeros((10, 10, 3), dtype=np.uint8) + with mock.patch.object( + device_runner.time, + 'time', + side_effect=[0.0, 0.0, 0.1, 1.0, 1.0, 1.0, 1.0], + ): + device_runner.run_device(mock_dev, 0, pixels, state) + + self.assertGreaterEqual(call_count[0], 2) + + +if __name__ == '__main__': + absltest.main() diff --git a/android_env/playground/frame_processing.py b/android_env/playground/frame_processing.py new file mode 100644 index 00000000..53660d4f --- /dev/null +++ b/android_env/playground/frame_processing.py @@ -0,0 +1,48 @@ +# coding=utf-8 +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Frame processing functions for AndroidEnv Playground.""" + +import io +from android_env.playground import coordinates +import numpy as np +from PIL import Image + + +def process_frame( + pixels: np.ndarray, + orientation: int, +) -> bytes: + """Processes raw emulator pixels (rotates, crops) to JPEG bytes.""" + # Rotate pixels if needed based on orientation + if orientation == 1: # LANDSCAPE_90 + pixels = np.rot90(pixels, k=1) + elif orientation == 3: # LANDSCAPE_270 + pixels = np.rot90(pixels, k=3) + elif orientation == 2: # PORTRAIT_180 + pixels = np.rot90(pixels, k=2) + + # If landscape, crop to 16:9 aspect ratio + if orientation in {1, 3}: + h_v, w_v = pixels.shape[:2] + crop_y_start, crop_y_end, crop_x_start, crop_x_end = ( + coordinates.get_crop_metrics((h_v, w_v)) + ) + pixels = pixels[crop_y_start:crop_y_end, crop_x_start:crop_x_end] + + img = Image.fromarray(pixels) + output = io.BytesIO() + img.save(output, format='JPEG') + return output.getvalue() diff --git a/android_env/playground/frame_processing_test.py b/android_env/playground/frame_processing_test.py new file mode 100644 index 00000000..f1ffdd92 --- /dev/null +++ b/android_env/playground/frame_processing_test.py @@ -0,0 +1,70 @@ +# coding=utf-8 +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for frame processing and image transformation in AndroidEnv Playground.""" + +import io +from absl.testing import absltest +from android_env.playground import frame_processing +import numpy as np +from PIL import Image + + +class FrameProcessingTest(absltest.TestCase): + + def test_process_frame_portrait(self): + pixels = np.zeros((1280, 720, 3), dtype=np.uint8) + pixels[:, :, 0] = 255 # Red + + jpeg_bytes = frame_processing.process_frame(pixels, 0) + img = Image.open(io.BytesIO(jpeg_bytes)) + self.assertEqual(img.size, (720, 1280)) + data = np.array(img) + np.testing.assert_allclose(data[10, 10], [255, 0, 0], atol=5) + + def test_process_frame_landscape_crop_and_rotate(self): + pixels = np.zeros((2209, 1080, 3), dtype=np.uint8) + pixels[0:144, :, 2] = 255 # Blue + pixels[144:2064, :, 1] = 255 # Green + pixels[2064:2209, :, 2] = 255 # Blue + + jpeg_bytes = frame_processing.process_frame(pixels, 1) + img = Image.open(io.BytesIO(jpeg_bytes)) + self.assertEqual(img.size, (1920, 1080)) + + data = np.array(img) + np.testing.assert_allclose(data[10, 10], [0, 255, 0], atol=5) + np.testing.assert_allclose(data[180, 320], [0, 255, 0], atol=5) + self.assertTrue(np.all(data[:, :, 2] < 50)) + + def test_process_frame_portrait_180(self): + pixels = np.zeros((1280, 720, 3), dtype=np.uint8) + pixels[:, :, 1] = 255 # Green + + jpeg_bytes = frame_processing.process_frame(pixels, 2) + img = Image.open(io.BytesIO(jpeg_bytes)) + self.assertEqual(img.size, (720, 1280)) + + def test_process_frame_landscape_270(self): + pixels = np.zeros((2209, 1080, 3), dtype=np.uint8) + pixels[144:2064, :, 0] = 255 # Red + + jpeg_bytes = frame_processing.process_frame(pixels, 3) + img = Image.open(io.BytesIO(jpeg_bytes)) + self.assertEqual(img.size, (1920, 1080)) + + +if __name__ == '__main__': + absltest.main() diff --git a/android_env/playground/frontend/BUILD.bazel b/android_env/playground/frontend/BUILD.bazel new file mode 100644 index 00000000..93bbdcb5 --- /dev/null +++ b/android_env/playground/frontend/BUILD.bazel @@ -0,0 +1,34 @@ +load("@aspect_rules_ts//ts:defs.bzl", "ts_project") +load("@npm//:defs.bzl", "npm_link_all_packages") + +package(default_visibility = ["//android_env/playground:__subpackages__"]) + +npm_link_all_packages(name = "node_modules") + +ts_project( + name = "frontend_ts", + srcs = [ + "adb_console.ts", + "controls.ts", + "logs_viewer.ts", + "main.ts", + "socket.ts", + "stats.ts", + "task_loader.ts", + "theme.ts", + ], + declaration = True, + source_map = True, + transpiler = "tsc", + tsconfig = ":tsconfig.json", + deps = [":node_modules/safevalues"], +) + +filegroup( + name = "frontend", + srcs = [ + "index.html", + ":frontend_ts", + ], + visibility = ["//visibility:public"], +) diff --git a/android_env/playground/frontend/adb_console.ts b/android_env/playground/frontend/adb_console.ts new file mode 100644 index 00000000..6434e629 --- /dev/null +++ b/android_env/playground/frontend/adb_console.ts @@ -0,0 +1,230 @@ +import {AdbResponseData, SocketHandler} from './socket'; + +/** + * Interface representing an entry in the ADB console history. + */ +export interface AdbConsoleEntry { + command: string; + status: string; + output: string; + error?: string; + duration_ms?: number; + timestamp: string; +} + +const PRESET_COMMANDS: ReadonlyArray<{label: string; cmd: string}> = [ + {label: 'Current Activity', cmd: 'shell dumpsys activity activities'}, + {label: 'List Packages', cmd: 'shell pm list packages'}, + {label: 'Battery State', cmd: 'shell dumpsys battery'}, + {label: 'System Props', cmd: 'shell getprop'}, + {label: 'Screen Size', cmd: 'shell wm size'}, + {label: 'Devices', cmd: 'devices -l'}, +]; + +/** + * Custom Web Component for ADB shell execution and history in the Playground UI. + */ +export class AdbConsoleElement extends HTMLElement { + private socket: SocketHandler | null = null; + private outputDiv: HTMLElement | null = null; + private inputEl: HTMLInputElement | null = null; + private runBtn: HTMLButtonElement | null = null; + private readonly history: AdbConsoleEntry[] = []; + private readonly commandHistory: string[] = []; + private historyIndex = -1; + + constructor() { + super(); + } + + /** @export */ + connectedCallback(): void { + if (!this.outputDiv) { + this.classList.add('adb-console-root'); + + const presetChips = document.createElement('div'); + presetChips.className = 'preset-chips'; + for (const preset of PRESET_COMMANDS) { + const btn = document.createElement('button'); + btn.className = 'chip'; + btn.setAttribute('data-adb', preset.cmd); + btn.textContent = preset.label; + presetChips.appendChild(btn); + } + this.appendChild(presetChips); + + const output = document.createElement('div'); + output.id = 'adb-output'; + output.className = 'adb-output'; + this.appendChild(output); + this.outputDiv = output; + + const inputRow = document.createElement('div'); + inputRow.className = 'adb-input-row'; + + const prompt = document.createElement('span'); + prompt.className = 'adb-prompt'; + prompt.textContent = '$ adb'; + inputRow.appendChild(prompt); + + const input = document.createElement('input'); + input.type = 'text'; + input.id = 'adb-input'; + input.placeholder = + 'shell pm list packages -3 (Press Enter to run)'; + inputRow.appendChild(input); + this.inputEl = input; + + const runBtn = document.createElement('button'); + runBtn.id = 'adb-run-btn'; + runBtn.className = 'btn btn-primary'; + runBtn.textContent = 'Run'; + inputRow.appendChild(runBtn); + this.runBtn = runBtn; + + this.appendChild(inputRow); + this.initEventListeners(); + } + } + + /** @export */ + setSocket(socket: SocketHandler): void { + this.socket = socket; + } + + private initEventListeners(): void { + this.runBtn?.addEventListener('click', () => { + this.runInput(); + }); + + this.inputEl?.addEventListener('keydown', (e: KeyboardEvent) => { + if (e.key === 'Enter') { + this.runInput(); + } else if (e.key === 'ArrowUp') { + this.navigateHistory(-1); + e.preventDefault(); + } else if (e.key === 'ArrowDown') { + this.navigateHistory(1); + e.preventDefault(); + } + }); + + this.querySelectorAll('.chip[data-adb]').forEach((chip) => { + chip.addEventListener('click', () => { + const cmd = chip.getAttribute('data-adb'); + if (cmd) { + this.executeCommand(cmd); + } + }); + }); + } + + /** @export */ + executeCommand(command: string): void { + const clean = command.trim(); + if (!clean) return; + this.commandHistory.push(clean); + this.historyIndex = this.commandHistory.length; + this.socket?.sendAdbCommand(clean); + } + + /** @export */ + handleAdbResponse(resp: AdbResponseData): void { + const entry: AdbConsoleEntry = { + command: resp.command, + status: resp.status, + output: resp.output, + error: resp.error, + duration_ms: resp.duration_ms, + timestamp: new Date().toLocaleTimeString(), + }; + this.history.push(entry); + this.renderEntry(entry); + } + + private renderEntry(entry: AdbConsoleEntry): void { + if (!this.outputDiv) return; + const div = document.createElement('div'); + div.className = 'adb-entry'; + const statusColor = + entry.status === 'OK' ? 'var(--success)' : 'var(--danger)'; + const duration = + entry.duration_ms !== undefined + ? ` (${entry.duration_ms.toFixed(1)}ms)` + : ''; + + const cmdHeader = document.createElement('div'); + cmdHeader.className = 'adb-cmd-line'; + + const cmdSpan = document.createElement('span'); + cmdSpan.textContent = `$ adb ${entry.command}`; + + const statusSpan = document.createElement('span'); + statusSpan.style.color = statusColor; + statusSpan.style.fontSize = '11px'; + statusSpan.textContent = `${entry.status}${duration}`; + + cmdHeader.appendChild(cmdSpan); + cmdHeader.appendChild(statusSpan); + + const body = document.createElement('div'); + body.className = 'adb-body'; + body.textContent = entry.output || entry.error || '(no output)'; + + div.appendChild(cmdHeader); + div.appendChild(body); + this.outputDiv.appendChild(div); + this.outputDiv.scrollTop = this.outputDiv.scrollHeight; + } + + private runInput(): void { + if (!this.inputEl) return; + const val = this.inputEl.value.trim(); + if (val) { + this.executeCommand(val); + this.inputEl.value = ''; + } + } + + private navigateHistory(direction: number): void { + if (this.commandHistory.length === 0 || !this.inputEl) return; + this.historyIndex += direction; + if (this.historyIndex < 0) this.historyIndex = 0; + if (this.historyIndex > this.commandHistory.length) { + this.historyIndex = this.commandHistory.length; + } + if (this.historyIndex === this.commandHistory.length) { + this.inputEl.value = ''; + } else { + this.inputEl.value = this.commandHistory[this.historyIndex]; + } + } + + /** @export */ + getHistory(): AdbConsoleEntry[] { + return [...this.history]; + } + + /** @export */ + clear(): void { + this.history.length = 0; + if (this.outputDiv) { + this.outputDiv.replaceChildren(); + } + } +} + +if (!customElements.get('adb-console')) { + customElements.define('adb-console', AdbConsoleElement); +} + +/** + * Helper factory function to instantiate and upgrade an AdbConsoleElement. + */ +export function createAdbConsole(socket?: SocketHandler): AdbConsoleElement { + const el = document.createElement('adb-console') as AdbConsoleElement; + if (socket) { + el.setSocket(socket); + } + return el; +} diff --git a/android_env/playground/frontend/adb_console_test.ts b/android_env/playground/frontend/adb_console_test.ts new file mode 100644 index 00000000..36697192 --- /dev/null +++ b/android_env/playground/frontend/adb_console_test.ts @@ -0,0 +1,105 @@ +import 'jasmine'; +import {createAdbConsole, AdbConsoleElement} from './adb_console'; +import {SocketHandler} from './socket'; + +describe('AdbConsoleElement', () => { + let socketMock: jasmine.SpyObj; + let element: AdbConsoleElement; + + beforeEach(() => { + socketMock = jasmine.createSpyObj('SocketHandler', [ + 'sendAdbCommand', + ]); + + element = createAdbConsole(socketMock); + document.body.appendChild(element); + }); + + afterEach(() => { + document.body.innerHTML = ''; + }); + + it('renders adb console with chips, output area, and input controls', () => { + expect(element.querySelector('.preset-chips')).toBeTruthy(); + expect(element.querySelector('#adb-output')).toBeTruthy(); + expect(element.querySelector('#adb-input')).toBeTruthy(); + expect(element.querySelector('#adb-run-btn')).toBeTruthy(); + }); + + it('sends adb command on Run button click and Enter key', () => { + const input = element.querySelector('#adb-input') as HTMLInputElement; + const runBtn = element.querySelector('#adb-run-btn') as HTMLButtonElement; + + input.value = 'shell wm size'; + runBtn.click(); + expect(socketMock.sendAdbCommand).toHaveBeenCalledWith('shell wm size'); + expect(input.value).toBe(''); + + socketMock.sendAdbCommand.calls.reset(); + input.value = 'shell getprop'; + input.dispatchEvent(new KeyboardEvent('keydown', {key: 'Enter'})); + expect(socketMock.sendAdbCommand).toHaveBeenCalledWith('shell getprop'); + expect(input.value).toBe(''); + }); + + it('sends adb command when chip is clicked', () => { + const chip = element.querySelector( + '.chip[data-adb="devices -l"]' + ) as HTMLElement; + expect(chip).toBeTruthy(); + chip.click(); + expect(socketMock.sendAdbCommand).toHaveBeenCalledWith('devices -l'); + }); + + it('renders adb response and clears history', () => { + element.handleAdbResponse({ + command: 'shell wm size', + status: 'OK', + output: 'Physical size: 1080x1920', + duration_ms: 5.4, + }); + + const outputDiv = element.querySelector('#adb-output') as HTMLElement; + expect(outputDiv.children.length).toBe(1); + expect(outputDiv.textContent).toContain('shell wm size'); + expect(outputDiv.textContent).toContain('Physical size: 1080x1920'); + expect(element.getHistory().length).toBe(1); + + element.clear(); + expect(outputDiv.children.length).toBe(0); + expect(element.getHistory().length).toBe(0); + }); + + it('navigates command history with ArrowUp and ArrowDown', () => { + const input = element.querySelector('#adb-input') as HTMLInputElement; + + element.executeCommand('cmd1'); + element.executeCommand('cmd2'); + + input.dispatchEvent( + new KeyboardEvent('keydown', {key: 'ArrowUp', cancelable: true}) + ); + expect(input.value).toBe('cmd2'); + + input.dispatchEvent( + new KeyboardEvent('keydown', {key: 'ArrowUp', cancelable: true}) + ); + expect(input.value).toBe('cmd1'); + + input.dispatchEvent( + new KeyboardEvent('keydown', {key: 'ArrowDown', cancelable: true}) + ); + expect(input.value).toBe('cmd2'); + + input.dispatchEvent( + new KeyboardEvent('keydown', {key: 'ArrowDown', cancelable: true}) + ); + expect(input.value).toBe(''); + + // Beyond history boundary + input.dispatchEvent( + new KeyboardEvent('keydown', {key: 'ArrowDown', cancelable: true}) + ); + expect(input.value).toBe(''); + }); +}); diff --git a/android_env/playground/frontend/controls.ts b/android_env/playground/frontend/controls.ts new file mode 100644 index 00000000..f71dd8b2 --- /dev/null +++ b/android_env/playground/frontend/controls.ts @@ -0,0 +1,304 @@ +import {setAnchorHref} from 'safevalues/dom'; +import {SocketHandler} from './socket'; + +interface Point { + x: number; + y: number; +} + +const KEY_CODES: Readonly> = { + 'ArrowUp': 19, + 'ArrowDown': 20, + 'ArrowLeft': 21, + 'ArrowRight': 22, + 'Enter': 66, + 'Backspace': 67, + 'Delete': 112, + 'Tab': 61, + 'Escape': 4, + 'Home': 122, + 'End': 123, + 'PageUp': 92, + 'PageDown': 93, +}; + +/** + * Custom Web Component for handling mouse, touch, and keyboard controls for AndroidEnv. + */ +export class ScreenControlsElement extends HTMLElement { + private screenImg: HTMLImageElement | null = null; + private socket: SocketHandler | null = null; + private isDown = false; + private lastMoveTime = 0; + private readonly throttleMs = 16; // ~60 Hz + + private readonly boundMouseMove: (e: MouseEvent) => void; + private readonly boundMouseUp: (e: MouseEvent) => void; + private readonly boundTouchMove: (e: TouchEvent) => void; + private readonly boundTouchEnd: (e: TouchEvent) => void; + private readonly boundKeyDown: (e: KeyboardEvent) => void; + + constructor() { + super(); + this.boundMouseMove = this.onMouseMove.bind(this); + this.boundMouseUp = this.onMouseUp.bind(this); + this.boundTouchMove = this.onTouchMove.bind(this); + this.boundTouchEnd = this.onTouchEnd.bind(this); + this.boundKeyDown = this.onKeyDown.bind(this); + } + + /** @export */ + connectedCallback(): void { + if (!this.screenImg) { + const container = document.createElement('div'); + container.id = 'container'; + const img = document.createElement('img'); + img.id = 'screen'; + container.appendChild(img); + this.appendChild(container); + this.screenImg = img; + this.initEventListeners(); + } + + window.addEventListener('mouseup', this.boundMouseUp); + window.addEventListener('touchend', this.boundTouchEnd); + window.addEventListener('keydown', this.boundKeyDown); + } + + /** @export */ + disconnectedCallback(): void { + window.removeEventListener('mouseup', this.boundMouseUp); + window.removeEventListener('touchend', this.boundTouchEnd); + window.removeEventListener('keydown', this.boundKeyDown); + } + + /** @export */ + setSocket(socket: SocketHandler): void { + this.socket = socket; + } + + /** @export */ + getScreenImage(): HTMLImageElement | null { + return this.screenImg; + } + + private getCoords(clientX: number, clientY: number): Point { + const rect = this.screenImg!.getBoundingClientRect(); + const x = (clientX - rect.left) / rect.width; + const y = (clientY - rect.top) / rect.height; + return { + x: Math.max(0, Math.min(1, x)), + y: Math.max(0, Math.min(1, y)), + }; + } + + private initEventListeners(): void { + this.screenImg!.addEventListener('mousedown', (e: MouseEvent) => { + this.isDown = true; + const {x, y} = this.getCoords(e.clientX, e.clientY); + this.socket?.sendAction('touch_down', x, y); + }); + + this.screenImg!.addEventListener('mousemove', this.boundMouseMove); + this.screenImg!.addEventListener('dragstart', (e: Event) => + e.preventDefault() + ); + + this.screenImg!.addEventListener('touchstart', (e: TouchEvent) => { + this.isDown = true; + if (e.touches.length > 0) { + const touch = e.touches[0]; + const {x, y} = this.getCoords(touch.clientX, touch.clientY); + this.socket?.sendAction('touch_down', x, y); + } + e.preventDefault(); + }); + + this.screenImg!.addEventListener('touchmove', this.boundTouchMove); + } + + private onMouseMove(e: MouseEvent): void { + if (!this.isDown) return; + const now = Date.now(); + if (now - this.lastMoveTime < this.throttleMs) return; + this.lastMoveTime = now; + const {x, y} = this.getCoords(e.clientX, e.clientY); + this.socket?.sendAction('touch_move', x, y); + } + + private onMouseUp(e: MouseEvent): void { + if (!this.isDown) return; + this.isDown = false; + const {x, y} = this.getCoords(e.clientX, e.clientY); + this.socket?.sendAction('touch_up', x, y); + } + + private onTouchMove(e: TouchEvent): void { + if (!this.isDown) return; + const now = Date.now(); + if (now - this.lastMoveTime < this.throttleMs) return; + this.lastMoveTime = now; + if (e.touches.length > 0) { + const touch = e.touches[0]; + const {x, y} = this.getCoords(touch.clientX, touch.clientY); + this.socket?.sendAction('touch_move', x, y); + } + e.preventDefault(); + } + + private onTouchEnd(e: TouchEvent): void { + if (!this.isDown) return; + this.isDown = false; + if (e.changedTouches.length > 0) { + const touch = e.changedTouches[0]; + const {x, y} = this.getCoords(touch.clientX, touch.clientY); + this.socket?.sendAction('touch_up', x, y); + } + } + + private onKeyDown(e: KeyboardEvent): void { + const activeEl = document.activeElement; + if ( + activeEl && + (activeEl.tagName === 'INPUT' || activeEl.tagName === 'TEXTAREA') + ) { + return; + } + + if (e.ctrlKey || e.metaKey) { + const keyLower = e.key.toLowerCase(); + if (keyLower === 'c') { + this.socket?.sendClipboard('copy'); + e.preventDefault(); + } else if (keyLower === 'x') { + this.socket?.sendClipboard('cut'); + e.preventDefault(); + } else if (keyLower === 'v') { + this.socket?.sendClipboard('paste'); + e.preventDefault(); + } else if (keyLower === 'a') { + this.socket?.sendSelection(0, -1); + e.preventDefault(); + } + return; + } + + if (e.altKey) { + return; + } + + const keycode = KEY_CODES[e.key]; + if (keycode !== undefined) { + this.socket?.sendKeydown(keycode); + e.preventDefault(); + } else if (e.key.length === 1) { + this.socket?.sendText(e.key); + e.preventDefault(); + } + } +} + +if (!customElements.get('screen-controls')) { + customElements.define('screen-controls', ScreenControlsElement); +} + +/** + * Helper factory function to instantiate and upgrade a ScreenControlsElement. + */ +export function createScreenControls( + socket?: SocketHandler +): ScreenControlsElement { + const el = document.createElement('screen-controls') as ScreenControlsElement; + if (socket) { + el.setSocket(socket); + } + return el; +} + +/** + * Captures screenshot from the screen element and triggers download. + */ +export function captureScreenshot( + screenImg: HTMLImageElement | null, + filename = 'android_env_screenshot.jpg' +): boolean { + if (!screenImg || !screenImg.src) return false; + try { + const canvas = document.createElement('canvas'); + canvas.width = screenImg.naturalWidth || 1080; + canvas.height = screenImg.naturalHeight || 1920; + const ctx = canvas.getContext('2d'); + if (!ctx) return false; + ctx.drawImage(screenImg, 0, 0); + const dataUrl = canvas.toDataURL('image/jpeg', 0.95); + const a = document.createElement('a'); + // `setAnchorHref` blocks `javascript:` URLs. + setAnchorHref(a, dataUrl); + a.download = filename; + document.body.appendChild(a); + a.click(); + setTimeout(() => a.remove(), 100); + return true; + } catch (err) { + console.error('Failed to capture screenshot:', err); + return false; + } +} + +/** + * Initializes hardware button rail, episode reset, text input, and screenshot buttons. + */ +export function initHardwareControls( + socket: SocketHandler, + screenControls?: ScreenControlsElement | null +): void { + // Bind hardware buttons + document + .querySelectorAll('.hw-btn[data-key]') + .forEach((btn) => { + btn.addEventListener('click', () => { + const key = Number(btn.getAttribute('data-key')); + if (!isNaN(key)) { + socket.sendKeydown(key); + } + }); + }); + + // Bind screenshot button + const shotBtn = document.getElementById('shot-btn'); + if (shotBtn) { + shotBtn.addEventListener('click', () => { + const screenImg = + screenControls?.getScreenImage() || + (document.getElementById('screen') as HTMLImageElement | null); + captureScreenshot(screenImg); + }); + } + + // Bind episode reset button + const resetBtn = document.getElementById('reset-btn'); + if (resetBtn) { + resetBtn.addEventListener('click', () => { + socket.sendReset(); + }); + } + + // Bind send text input & button + const sendTextBtn = document.getElementById('send-text-btn'); + const sendTextInput = document.getElementById( + 'send-text-input' + ) as HTMLInputElement | null; + const doSendText = () => { + if (sendTextInput && sendTextInput.value) { + socket.sendText(sendTextInput.value); + sendTextInput.value = ''; + } + }; + sendTextBtn?.addEventListener('click', doSendText); + sendTextInput?.addEventListener('keydown', (e: KeyboardEvent) => { + if (e.key === 'Enter') { + doSendText(); + } + }); +} + diff --git a/android_env/playground/frontend/controls_test.ts b/android_env/playground/frontend/controls_test.ts new file mode 100644 index 00000000..687458ac --- /dev/null +++ b/android_env/playground/frontend/controls_test.ts @@ -0,0 +1,453 @@ +import 'jasmine'; +import { + captureScreenshot, + createScreenControls, + initHardwareControls, + ScreenControlsElement, +} from './controls'; +import {SocketHandler} from './socket'; + +describe('ScreenControlsElement', () => { + let screenImg: HTMLImageElement; + let socketMock: jasmine.SpyObj; + let element: ScreenControlsElement; + + beforeEach(() => { + socketMock = jasmine.createSpyObj('SocketHandler', [ + 'sendAction', + 'sendText', + 'sendKeydown', + 'sendClipboard', + 'sendSelection', + ]); + + element = createScreenControls(socketMock); + element.id = 'screen-controls'; + document.body.appendChild(element); + + screenImg = element.querySelector('#screen') as HTMLImageElement; + + // Mock getBoundingClientRect for deterministic coordinates calculation + spyOn(screenImg, 'getBoundingClientRect').and.returnValue({ + left: 100, + top: 100, + right: 300, + bottom: 500, + width: 200, + height: 400, + x: 100, + y: 100, + toJSON: () => {}, + }); + }); + + afterEach(() => { + document.body.innerHTML = ''; + }); + + it('initializes custom element successfully', () => { + expect(element).toBeTruthy(); + expect(element.tagName.toLowerCase()).toBe('screen-controls'); + expect(screenImg).toBeTruthy(); + expect(element.getScreenImage()).toBe(screenImg); + expect(element.innerHTML).toBe('
'); + + // Calling connectedCallback again is idempotent + element.connectedCallback(); + expect(element.children.length).toBe(1); + }); + + it('removes window listeners on disconnectedCallback', () => { + screenImg.dispatchEvent( + new MouseEvent('mousedown', {clientX: 150, clientY: 200}) + ); + socketMock.sendAction.calls.reset(); + + element.remove(); + + window.dispatchEvent( + new MouseEvent('mouseup', {clientX: 150, clientY: 200}) + ); + const endEvent: Event = new CustomEvent('touchend', {cancelable: true}); + Object.defineProperty(endEvent, 'changedTouches', { + value: [{clientX: 150, clientY: 200}], + }); + window.dispatchEvent(endEvent); + window.dispatchEvent(new KeyboardEvent('keydown', {key: 'a'})); + + expect(socketMock.sendAction).not.toHaveBeenCalled(); + expect(socketMock.sendText).not.toHaveBeenCalled(); + }); + + it('sends touch_down action on mousedown', () => { + const event = new MouseEvent('mousedown', { + clientX: 200, // (200-100)/200 = 0.5 + clientY: 300, // (300-100)/400 = 0.5 + }); + screenImg.dispatchEvent(event); + + expect(socketMock.sendAction).toHaveBeenCalledWith('touch_down', 0.5, 0.5); + }); + + it('sends touch_move action when mouse move occurs while pressed', () => { + // First press down + screenImg.dispatchEvent( + new MouseEvent('mousedown', {clientX: 150, clientY: 200}) + ); + socketMock.sendAction.calls.reset(); + + // Now move + const event = new MouseEvent('mousemove', { + clientX: 200, // (200-100)/200 = 0.5 + clientY: 200, // (200-100)/400 = 0.25 + }); + screenImg.dispatchEvent(event); + + expect(socketMock.sendAction).toHaveBeenCalledWith('touch_move', 0.5, 0.25); + }); + + it('sends touch_up action on mouseup', () => { + screenImg.dispatchEvent( + new MouseEvent('mousedown', {clientX: 150, clientY: 200}) + ); + socketMock.sendAction.calls.reset(); + + const event = new MouseEvent('mouseup', { + clientX: 150, + clientY: 200, + }); + window.dispatchEvent(event); + + expect(socketMock.sendAction).toHaveBeenCalledWith('touch_up', 0.25, 0.25); + }); + + it('prevents default on dragstart', () => { + const event = new Event('dragstart', {cancelable: true}); + screenImg.dispatchEvent(event); + expect(event.defaultPrevented).toBeTrue(); + }); + + it('handles touchstart, touchmove, and touchend events', () => { + const touchObj = {clientX: 200, clientY: 300}; + + const startEvent: Event = new CustomEvent('touchstart', { + cancelable: true, + }); + Object.defineProperty(startEvent, 'touches', {value: [touchObj]}); + screenImg.dispatchEvent(startEvent); + + expect(startEvent.defaultPrevented).toBeTrue(); + expect(socketMock.sendAction).toHaveBeenCalledWith('touch_down', 0.5, 0.5); + socketMock.sendAction.calls.reset(); + + const moveTouchObj = {clientX: 200, clientY: 200}; + const moveEvent: Event = new CustomEvent('touchmove', {cancelable: true}); + Object.defineProperty(moveEvent, 'touches', {value: [moveTouchObj]}); + screenImg.dispatchEvent(moveEvent); + + expect(moveEvent.defaultPrevented).toBeTrue(); + expect(socketMock.sendAction).toHaveBeenCalledWith('touch_move', 0.5, 0.25); + socketMock.sendAction.calls.reset(); + + const endTouchObj = {clientX: 150, clientY: 200}; + const endEvent: Event = new CustomEvent('touchend', {cancelable: true}); + Object.defineProperty(endEvent, 'changedTouches', {value: [endTouchObj]}); + window.dispatchEvent(endEvent); + + expect(socketMock.sendAction).toHaveBeenCalledWith('touch_up', 0.25, 0.25); + }); + + it('ignores mousemove, mouseup, touchmove, and touchend when not down', () => { + screenImg.dispatchEvent( + new MouseEvent('mousemove', {clientX: 200, clientY: 200}) + ); + window.dispatchEvent( + new MouseEvent('mouseup', {clientX: 200, clientY: 200}) + ); + const moveEvent: Event = new CustomEvent('touchmove', {cancelable: true}); + Object.defineProperty(moveEvent, 'touches', { + value: [{clientX: 200, clientY: 200}], + }); + screenImg.dispatchEvent(moveEvent); + const endEvent: Event = new CustomEvent('touchend', {cancelable: true}); + Object.defineProperty(endEvent, 'changedTouches', { + value: [{clientX: 200, clientY: 200}], + }); + window.dispatchEvent(endEvent); + expect(socketMock.sendAction).not.toHaveBeenCalled(); + }); + + it('throttles mousemove events within throttle window', () => { + screenImg.dispatchEvent( + new MouseEvent('mousedown', {clientX: 150, clientY: 200}) + ); + socketMock.sendAction.calls.reset(); + + screenImg.dispatchEvent( + new MouseEvent('mousemove', {clientX: 200, clientY: 200}) + ); + expect(socketMock.sendAction).toHaveBeenCalledTimes(1); + + // Immediate second move should be throttled out + screenImg.dispatchEvent( + new MouseEvent('mousemove', {clientX: 220, clientY: 220}) + ); + expect(socketMock.sendAction).toHaveBeenCalledTimes(1); + }); + + it('throttles touchmove events within throttle window', () => { + const startEvent: Event = new CustomEvent('touchstart', { + cancelable: true, + }); + Object.defineProperty(startEvent, 'touches', { + value: [{clientX: 150, clientY: 200}], + }); + screenImg.dispatchEvent(startEvent); + socketMock.sendAction.calls.reset(); + + const moveEvent1: Event = new CustomEvent('touchmove', { + cancelable: true, + }); + Object.defineProperty(moveEvent1, 'touches', { + value: [{clientX: 200, clientY: 200}], + }); + screenImg.dispatchEvent(moveEvent1); + expect(socketMock.sendAction).toHaveBeenCalledTimes(1); + + const moveEvent2: Event = new CustomEvent('touchmove', { + cancelable: true, + }); + Object.defineProperty(moveEvent2, 'touches', { + value: [{clientX: 220, clientY: 220}], + }); + screenImg.dispatchEvent(moveEvent2); + expect(socketMock.sendAction).toHaveBeenCalledTimes(1); + }); + + it('handles empty touches or changedTouches list gracefully', () => { + const startEvent: Event = new CustomEvent('touchstart', { + cancelable: true, + }); + Object.defineProperty(startEvent, 'touches', {value: []}); + screenImg.dispatchEvent(startEvent); + + const moveEvent: Event = new CustomEvent('touchmove', {cancelable: true}); + Object.defineProperty(moveEvent, 'touches', {value: []}); + screenImg.dispatchEvent(moveEvent); + + const endEvent: Event = new CustomEvent('touchend', {cancelable: true}); + Object.defineProperty(endEvent, 'changedTouches', {value: []}); + window.dispatchEvent(endEvent); + + expect(socketMock.sendAction).not.toHaveBeenCalled(); + }); + + it('clamps coordinates to [0, 1] range', () => { + // Negative coordinates (outside top-left) + screenImg.dispatchEvent( + new MouseEvent('mousedown', {clientX: 50, clientY: 50}) + ); + expect(socketMock.sendAction).toHaveBeenCalledWith('touch_down', 0, 0); + + // Over-boundary coordinates (outside bottom-right) + screenImg.dispatchEvent( + new MouseEvent('mousedown', {clientX: 500, clientY: 600}) + ); + expect(socketMock.sendAction).toHaveBeenCalledWith('touch_down', 1, 1); + }); + + it('sends text on printable keydown event and prevents default', () => { + const event = new KeyboardEvent('keydown', {key: 'a', cancelable: true}); + window.dispatchEvent(event); + expect(socketMock.sendText).toHaveBeenCalledWith('a'); + expect(event.defaultPrevented).toBeTrue(); + }); + + it('sends keydown on Enter and Backspace keydown events', () => { + const enterEvent = new KeyboardEvent('keydown', { + key: 'Enter', + cancelable: true, + }); + window.dispatchEvent(enterEvent); + expect(socketMock.sendKeydown).toHaveBeenCalledWith(66); + expect(enterEvent.defaultPrevented).toBeTrue(); + + const backspaceEvent = new KeyboardEvent('keydown', { + key: 'Backspace', + cancelable: true, + }); + window.dispatchEvent(backspaceEvent); + expect(socketMock.sendKeydown).toHaveBeenCalledWith(67); + expect(backspaceEvent.defaultPrevented).toBeTrue(); + }); + + it('sends clipboard commands on Ctrl/Meta+C, X, V, A and prevents default', () => { + const copyEvent = new KeyboardEvent('keydown', { + key: 'C', + metaKey: true, + cancelable: true, + }); + window.dispatchEvent(copyEvent); + expect(socketMock.sendClipboard).toHaveBeenCalledWith('copy'); + expect(copyEvent.defaultPrevented).toBeTrue(); + + const cutEvent = new KeyboardEvent('keydown', { + key: 'x', + ctrlKey: true, + cancelable: true, + }); + window.dispatchEvent(cutEvent); + expect(socketMock.sendClipboard).toHaveBeenCalledWith('cut'); + expect(cutEvent.defaultPrevented).toBeTrue(); + + const pasteEvent = new KeyboardEvent('keydown', { + key: 'v', + ctrlKey: true, + cancelable: true, + }); + window.dispatchEvent(pasteEvent); + expect(socketMock.sendClipboard).toHaveBeenCalledWith('paste'); + expect(pasteEvent.defaultPrevented).toBeTrue(); + + const selectEvent = new KeyboardEvent('keydown', { + key: 'a', + ctrlKey: true, + cancelable: true, + }); + window.dispatchEvent(selectEvent); + expect(socketMock.sendSelection).toHaveBeenCalledWith(0, -1); + expect(selectEvent.defaultPrevented).toBeTrue(); + }); + + it('ignores keydown when focused on input or textarea element', () => { + const input = document.createElement('input'); + document.body.appendChild(input); + input.focus(); + + window.dispatchEvent(new KeyboardEvent('keydown', {key: 'a'})); + expect(socketMock.sendText).not.toHaveBeenCalled(); + + const textarea = document.createElement('textarea'); + document.body.appendChild(textarea); + textarea.focus(); + + window.dispatchEvent(new KeyboardEvent('keydown', {key: 'b'})); + expect(socketMock.sendText).not.toHaveBeenCalled(); + }); + + it('ignores keydown when altKey is pressed', () => { + window.dispatchEvent( + new KeyboardEvent('keydown', {key: 'a', altKey: true}) + ); + expect(socketMock.sendText).not.toHaveBeenCalled(); + expect(socketMock.sendKeydown).not.toHaveBeenCalled(); + }); + + it('sends DPAD keydown on ArrowLeft, ArrowRight, ArrowUp, and ArrowDown', () => { + const arrowMap: Array<[string, number]> = [ + ['ArrowLeft', 21], + ['ArrowRight', 22], + ['ArrowUp', 19], + ['ArrowDown', 20], + ]; + for (const [key, code] of arrowMap) { + const event = new KeyboardEvent('keydown', {key, cancelable: true}); + window.dispatchEvent(event); + expect(socketMock.sendKeydown).toHaveBeenCalledWith(code); + expect(event.defaultPrevented).toBeTrue(); + } + }); + + it('sends keydown on Delete, Tab, Escape, Home, End, PageUp, and PageDown', () => { + const keyMap: Array<[string, number]> = [ + ['Delete', 112], + ['Tab', 61], + ['Escape', 4], + ['Home', 122], + ['End', 123], + ['PageUp', 92], + ['PageDown', 93], + ]; + for (const [key, code] of keyMap) { + const event = new KeyboardEvent('keydown', {key, cancelable: true}); + window.dispatchEvent(event); + expect(socketMock.sendKeydown).toHaveBeenCalledWith(code); + expect(event.defaultPrevented).toBeTrue(); + } + }); + + it('captures screenshot cleanly', () => { + expect(captureScreenshot(null)).toBeFalse(); + + const img = document.createElement('img'); + expect(captureScreenshot(img)).toBeFalse(); + + img.src = 'data:image/png;base64,iVBORw0KGgo='; + const result = captureScreenshot(img); + expect(result).toBeTrue(); + }); + + it('downloads the screenshot via an anchor with a JPEG data URL', () => { + const clickSpy = spyOn(HTMLAnchorElement.prototype, 'click'); + const img = document.createElement('img'); + img.src = 'data:image/png;base64,iVBORw0KGgo='; + + expect(captureScreenshot(img, 'shot.jpg')).toBeTrue(); + + expect(clickSpy).toHaveBeenCalledTimes(1); + const anchor = clickSpy.calls.mostRecent().object as HTMLAnchorElement; + expect(anchor.href).toMatch(/^data:image\/jpeg/); + expect(anchor.download).toBe('shot.jpg'); + }); + + it('binds hardware controls, shot, reset, and send text buttons', () => { + document.body.innerHTML = ` + + + + + + + `; + + const socketSpy = jasmine.createSpyObj('SocketHandler', [ + 'sendKeydown', + 'sendReset', + 'sendText', + ]); + + initHardwareControls(socketSpy); + + const hwBtn = document.querySelector('.hw-btn') as HTMLElement; + hwBtn.click(); + expect(socketSpy.sendKeydown).toHaveBeenCalledWith(26); + + const resetBtn = document.getElementById('reset-btn') as HTMLElement; + resetBtn.click(); + expect(socketSpy.sendReset).toHaveBeenCalled(); + + const shotBtn = document.getElementById('shot-btn') as HTMLElement; + shotBtn.click(); + + const input = document.getElementById('send-text-input') as HTMLInputElement; + input.value = 'typed text'; + input.dispatchEvent(new KeyboardEvent('keydown', {key: 'Enter'})); + expect(socketSpy.sendText).toHaveBeenCalledWith('typed text'); + }); + + it('returns false when captureScreenshot encounters an error', () => { + const img = document.createElement('img'); + img.src = 'data:image/png;base64,iVBORw0KGgo='; + const origCreate = document.createElement.bind(document); + spyOn(document, 'createElement').and.callFake((tagName: string) => { + if (tagName === 'canvas') { + throw new Error('Canvas not supported'); + } + return origCreate(tagName); + }); + expect(captureScreenshot(img)).toBeFalse(); + }); +}); + + + diff --git a/android_env/playground/frontend/index.html b/android_env/playground/frontend/index.html new file mode 100644 index 00000000..4de70ab6 --- /dev/null +++ b/android_env/playground/frontend/index.html @@ -0,0 +1,650 @@ + + + + + + + AndroidEnv Playground + + + + +
+
+
+

AndroidEnv Playground

+ Connecting... + pkg: launcher +
+
+ + + + + +
+
+ +
+ + + + +
+
+
+ +
+
+ + + +
+ + + +
+ +
+
+
+
+ + + + diff --git a/android_env/playground/frontend/logs_viewer.ts b/android_env/playground/frontend/logs_viewer.ts new file mode 100644 index 00000000..8f2cb36f --- /dev/null +++ b/android_env/playground/frontend/logs_viewer.ts @@ -0,0 +1,225 @@ + +/** + * Interface representing a log entry received from the server. + */ +export interface LogData { + message: string; + level: string; + timestamp: number; +} + +/** + * Custom Web Component for rendering and managing console log messages in the Playground UI. + */ +export class LogViewerElement extends HTMLElement { + private panel: HTMLElement | null = null; + private contentDiv: HTMLElement | null = null; + private toggleBtn: HTMLElement | null = null; + private clearBtn: HTMLElement | null = null; + private levelFilter: HTMLSelectElement | null = null; + private autoscrollCheck: HTMLInputElement | null = null; + private readonly logs: LogData[] = []; + private readonly maxLogs: number = 10000; + + constructor() { + super(); + } + + /** @export */ + connectedCallback(): void { + if (!this.panel) { + this.panel = document.createElement('div'); + this.panel.id = 'logs-panel'; + + const header = document.createElement('div'); + header.className = 'logs-header'; + + const title = document.createElement('span'); + title.textContent = 'Console Logs'; + + const controls = document.createElement('div'); + controls.className = 'logs-controls'; + + const levelLabel = document.createElement('label'); + levelLabel.setAttribute('for', 'logs-level-filter'); + levelLabel.textContent = 'Level: '; + + this.levelFilter = document.createElement('select'); + this.levelFilter.id = 'logs-level-filter'; + [ + {val: 'ALL', label: 'All Levels'}, + {val: 'INFO', label: 'INFO+'}, + {val: 'WARNING', label: 'WARNING+'}, + {val: 'ERROR', label: 'ERROR'}, + ].forEach((item) => { + const opt = document.createElement('option'); + opt.value = item.val; + opt.textContent = item.label; + this.levelFilter!.appendChild(opt); + }); + + const autoScrollLabel = document.createElement('label'); + this.autoscrollCheck = document.createElement('input'); + this.autoscrollCheck.type = 'checkbox'; + this.autoscrollCheck.id = 'logs-autoscroll'; + this.autoscrollCheck.checked = true; + autoScrollLabel.appendChild(this.autoscrollCheck); + autoScrollLabel.appendChild(document.createTextNode(' Auto-scroll')); + + this.clearBtn = document.createElement('button'); + this.clearBtn.id = 'logs-clear-btn'; + this.clearBtn.textContent = 'Clear'; + + this.toggleBtn = document.createElement('button'); + this.toggleBtn.id = 'logs-toggle-btn'; + this.toggleBtn.textContent = 'Hide'; + + controls.appendChild(levelLabel); + controls.appendChild(this.levelFilter); + controls.appendChild(autoScrollLabel); + controls.appendChild(this.clearBtn); + controls.appendChild(this.toggleBtn); + + header.appendChild(title); + header.appendChild(controls); + + this.contentDiv = document.createElement('div'); + this.contentDiv.id = 'logs-content'; + + this.panel.appendChild(header); + this.panel.appendChild(this.contentDiv); + + this.appendChild(this.panel); + this.initEventListeners(); + } + } + + private initEventListeners(): void { + if (this.toggleBtn) { + this.toggleBtn.addEventListener('click', () => { + this.hidePanel(); + }); + } + + if (this.clearBtn) { + this.clearBtn.addEventListener('click', () => { + this.clearLogs(); + }); + } + + if (this.levelFilter) { + this.levelFilter.addEventListener('change', () => { + this.renderAll(); + }); + } + } + + /** @export */ + showPanel(): void { + if (this.panel) { + this.panel.style.display = 'flex'; + this.scrollToBottom(); + } + const openBtn = document.getElementById('logs-open-btn'); + if (openBtn) { + openBtn.style.display = 'none'; + } + } + + /** @export */ + hidePanel(): void { + if (this.panel) { + this.panel.style.display = 'none'; + } + const openBtn = document.getElementById('logs-open-btn'); + if (openBtn) { + openBtn.style.display = 'inline-block'; + } + } + + /** @export */ + addLog(log: LogData): void { + this.logs.push(log); + if (this.logs.length > this.maxLogs) { + this.logs.shift(); + } + if (this.shouldDisplayLog(log)) { + this.appendLogElement(log); + } + } + + /** @export */ + addLogs(logs: LogData[]): void { + for (const log of logs) { + this.logs.push(log); + if (this.logs.length > this.maxLogs) { + this.logs.shift(); + } + } + this.renderAll(); + } + + /** @export */ + clearLogs(): void { + this.logs.length = 0; + if (this.contentDiv) { + this.contentDiv.replaceChildren(); + } + } + + private shouldDisplayLog(log: LogData): boolean { + if (!this.levelFilter) return true; + const filter = this.levelFilter.value; + if (filter === 'ALL') return true; + const levels = ['DEBUG', 'INFO', 'WARNING', 'ERROR', 'FATAL', 'CRITICAL']; + const minIndex = levels.indexOf(filter); + const logIndex = levels.indexOf(log.level.toUpperCase()); + if (minIndex === -1 || logIndex === -1) return true; + return logIndex >= minIndex; + } + + private renderAll(): void { + if (!this.contentDiv) return; + this.contentDiv.replaceChildren(); + for (const log of this.logs) { + if (this.shouldDisplayLog(log)) { + this.appendLogElement(log); + } + } + this.scrollToBottom(); + } + + private appendLogElement(log: LogData): void { + if (!this.contentDiv) return; + const line = document.createElement('div'); + line.className = `log-line log-${log.level.toLowerCase()}`; + line.textContent = log.message; + this.contentDiv.appendChild(line); + this.scrollToBottom(); + } + + private scrollToBottom(): void { + if ( + this.contentDiv && + (!this.autoscrollCheck || this.autoscrollCheck.checked) + ) { + this.contentDiv.scrollTop = this.contentDiv.scrollHeight; + } + } + + /** @export */ + getLogCount(): number { + return this.logs.length; + } +} + +if (!customElements.get('log-viewer')) { + customElements.define('log-viewer', LogViewerElement); +} + +/** + * Helper factory function to instantiate and upgrade a LogViewerElement. + */ +export function createLogViewer(): LogViewerElement { + return document.createElement('log-viewer') as LogViewerElement; +} diff --git a/android_env/playground/frontend/logs_viewer_test.ts b/android_env/playground/frontend/logs_viewer_test.ts new file mode 100644 index 00000000..5df74f28 --- /dev/null +++ b/android_env/playground/frontend/logs_viewer_test.ts @@ -0,0 +1,125 @@ +import 'jasmine'; +import {createLogViewer, LogData, LogViewerElement} from './logs_viewer'; + +describe('LogViewerElement', () => { + let viewer: LogViewerElement; + + beforeEach(() => { + document.body.innerHTML = ` + + `; + viewer = createLogViewer(); + viewer.id = 'logs-viewer'; + document.body.appendChild(viewer); + }); + + afterEach(() => { + document.body.innerHTML = ''; + }); + + it('initializes custom element and renders internal DOM on connect', () => { + expect(viewer).toBeTruthy(); + expect(viewer.tagName.toLowerCase()).toBe('log-viewer'); + expect(viewer.getLogCount()).toBe(0); + const content = viewer.querySelector('#logs-content'); + expect(content).not.toBeNull(); + }); + + it('adds and renders a single log line', () => { + const log: LogData = { + message: 'Test log line', + level: 'INFO', + timestamp: 1000, + }; + viewer.addLog(log); + + expect(viewer.getLogCount()).toBe(1); + const content = viewer.querySelector('#logs-content') as HTMLElement; + expect(content.children.length).toBe(1); + expect(content.children[0].textContent).toBe('Test log line'); + expect(content.children[0].className).toContain('log-info'); + }); + + it('adds multiple log lines', () => { + const logs: LogData[] = [ + {message: 'Log 1', level: 'INFO', timestamp: 1000}, + {message: 'Log 2', level: 'WARNING', timestamp: 1001}, + {message: 'Log 3', level: 'ERROR', timestamp: 1002}, + ]; + viewer.addLogs(logs); + + expect(viewer.getLogCount()).toBe(3); + const content = viewer.querySelector('#logs-content') as HTMLElement; + expect(content.children.length).toBe(3); + expect(content.children[1].className).toContain('log-warning'); + expect(content.children[2].className).toContain('log-error'); + }); + + it('clears logs', () => { + viewer.addLog({message: 'Log 1', level: 'INFO', timestamp: 1000}); + expect(viewer.getLogCount()).toBe(1); + + const clearBtn = viewer.querySelector( + '#logs-clear-btn' + ) as HTMLButtonElement; + clearBtn.click(); + + expect(viewer.getLogCount()).toBe(0); + const content = viewer.querySelector('#logs-content') as HTMLElement; + expect(content.children.length).toBe(0); + }); + + it('filters logs by level', () => { + const logs: LogData[] = [ + {message: 'Info log', level: 'INFO', timestamp: 1000}, + {message: 'Warn log', level: 'WARNING', timestamp: 1001}, + {message: 'Error log', level: 'ERROR', timestamp: 1002}, + ]; + viewer.addLogs(logs); + + const content = viewer.querySelector('#logs-content') as HTMLElement; + expect(content.children.length).toBe(3); + + const levelFilter = viewer.querySelector( + '#logs-level-filter' + ) as HTMLSelectElement; + levelFilter.value = 'WARNING'; + levelFilter.dispatchEvent(new Event('change')); + + expect(content.children.length).toBe(2); + expect(content.children[0].textContent).toBe('Warn log'); + expect(content.children[1].textContent).toBe('Error log'); + }); + + it('toggles panel visibility', () => { + const toggleBtn = viewer.querySelector( + '#logs-toggle-btn' + ) as HTMLButtonElement; + const panel = viewer.querySelector('#logs-panel') as HTMLElement; + const openBtn = document.getElementById( + 'logs-open-btn' + ) as HTMLButtonElement; + + toggleBtn.click(); + expect(panel.style.display).toBe('none'); + expect(openBtn.style.display).toBe('inline-block'); + + viewer.showPanel(); + expect(panel.style.display).toBe('flex'); + expect(openBtn.style.display).toBe('none'); + }); + + it('evicts oldest logs when maxLogs threshold is exceeded', () => { + (viewer as unknown as {maxLogs: number}).maxLogs = 2; + viewer.addLog({message: 'Log 1', level: 'INFO', timestamp: 1000}); + viewer.addLog({message: 'Log 2', level: 'INFO', timestamp: 1001}); + viewer.addLog({message: 'Log 3', level: 'INFO', timestamp: 1002}); + expect(viewer.getLogCount()).toBe(2); + + viewer.addLogs([ + {message: 'Log 4', level: 'INFO', timestamp: 1003}, + {message: 'Log 5', level: 'INFO', timestamp: 1004}, + ]); + expect(viewer.getLogCount()).toBe(2); + }); +}); diff --git a/android_env/playground/frontend/main.ts b/android_env/playground/frontend/main.ts new file mode 100644 index 00000000..1b43db0c --- /dev/null +++ b/android_env/playground/frontend/main.ts @@ -0,0 +1,142 @@ +import './adb_console'; +import './controls'; +import './logs_viewer'; +import './stats'; +import './task_loader'; +import './theme'; +import {AdbConsoleElement} from './adb_console'; +import {initHardwareControls, ScreenControlsElement} from './controls'; +import {LogViewerElement} from './logs_viewer'; +import {SocketHandler} from './socket'; +import {StatsViewerElement} from './stats'; +import {TaskLoaderElement} from './task_loader'; +import {setupThemeListeners} from './theme'; + +/** + * Interface representing all initialized playground components. + */ +export interface PlaygroundComponents { + controlsHandler: ScreenControlsElement | null; + logViewer: LogViewerElement | null; + statsViewer: StatsViewerElement | null; + taskLoaderHandler: TaskLoaderElement | null; + adbConsole: AdbConsoleElement | null; + socketHandler: SocketHandler; +} + +/** + * Initializes the playground frontend components and event handlers. + */ +export function initPlayground(): PlaygroundComponents { + const statsViewer = document.getElementById( + 'stats-viewer' + ) as StatsViewerElement | null; + const logViewer = document.getElementById( + 'logs-viewer' + ) as LogViewerElement | null; + const taskLoader = document.getElementById( + 'task-loader' + ) as TaskLoaderElement | null; + const screenControls = document.getElementById( + 'screen-controls' + ) as ScreenControlsElement | null; + const adbConsole = document.getElementById( + 'adb-console' + ) as AdbConsoleElement | null; + + const openLogsBtn = document.getElementById('logs-open-btn'); + if (openLogsBtn && logViewer) { + openLogsBtn.addEventListener('click', () => { + logViewer.showPanel(); + openLogsBtn.style.display = 'none'; + }); + } + + const openStatsBtn = document.getElementById('stats-toggle'); + if (openStatsBtn && statsViewer) { + openStatsBtn.addEventListener('click', () => { + statsViewer.togglePanel(); + }); + } + + // Bind tab switching + document.querySelectorAll('.tab-btn').forEach((btn) => { + btn.addEventListener('click', () => { + document + .querySelectorAll('.tab-btn') + .forEach((b) => b.classList.remove('active')); + document + .querySelectorAll('.tab-content') + .forEach((c) => c.classList.remove('active')); + btn.classList.add('active'); + const targetId = btn.getAttribute('data-tab'); + if (targetId) { + document.getElementById(targetId)?.classList.add('active'); + } + }); + }); + + const socketHandler = new SocketHandler('screen', 'status', { + onStatsReceived: (stats) => statsViewer?.handleStats(stats), + onLogsReceived: (logs) => logViewer?.addLogs(logs), + onAdbResponseReceived: (resp) => adbConsole?.handleAdbResponse(resp), + onStatusChange: (status) => { + const connDot = document.getElementById('conn-dot'); + if (connDot) { + connDot.style.background = status.startsWith('Connected') + ? 'var(--success)' + : 'var(--warning)'; + } + }, + }); + + screenControls?.setSocket(socketHandler); + taskLoader?.setSocket(socketHandler); + adbConsole?.setSocket(socketHandler); + + initHardwareControls(socketHandler, screenControls); + setupThemeListeners(); + + return { + controlsHandler: screenControls, + logViewer, + statsViewer, + taskLoaderHandler: taskLoader, + adbConsole, + socketHandler, + }; +} + +/** + * Auto-initializes the playground when the DOM is ready. + */ +export function setupAutoInit( + targetDoc: Document | null = typeof document !== 'undefined' ? document : null +): boolean { + if (!targetDoc) { + return false; + } + if (targetDoc.readyState === 'loading') { + targetDoc.addEventListener( + 'DOMContentLoaded', + () => { + if ( + targetDoc.getElementById('screen') || + targetDoc.getElementById('screen-controls') + ) { + initPlayground(); + } + }, + {once: true} + ); + } else if ( + targetDoc.getElementById('screen') || + targetDoc.getElementById('screen-controls') + ) { + initPlayground(); + } + return true; +} + +// Unconditionally execute setupAutoInit on module load +setupAutoInit(); diff --git a/android_env/playground/frontend/main_test.ts b/android_env/playground/frontend/main_test.ts new file mode 100644 index 00000000..fd3d62ff --- /dev/null +++ b/android_env/playground/frontend/main_test.ts @@ -0,0 +1,280 @@ +import 'jasmine'; +import {initPlayground, setupAutoInit} from './main'; + +describe('main', () => { + let result: ReturnType | null = null; + + beforeEach(() => { + document.body.innerHTML = ` +
Connecting...
+
+ + + +
Logs Content
+
ADB Content
+ + + + + + + + `; + + spyOn(window, 'WebSocket').and.returnValue({ + readyState: 1, + send: jasmine.createSpy('send'), + binaryType: '', + onopen: null, + onclose: null, + onmessage: null, + } as unknown as WebSocket); + }); + + afterEach(() => { + document.body.innerHTML = ''; + result = null; + }); + + it('initializes playground components successfully and handles socket callbacks and bindings', () => { + const holder: { + cb: ((event: {data: string}) => void) | null; + sendSpy: jasmine.Spy; + } = {cb: null, sendSpy: jasmine.createSpy('send')}; + + const ws = { + readyState: 1, + send: holder.sendSpy, + binaryType: '', + onopen: null, + onclose: null, + onmessage: null, + }; + Object.defineProperty(ws, 'onmessage', { + set: (fn: (event: {data: string}) => void) => { + holder.cb = fn; + }, + get: () => holder.cb, + }); + // `returnValue` (unlike an arrow `callFake`) supports `new WebSocket(...)`. + (window.WebSocket as unknown as jasmine.Spy).and.returnValue( + ws as unknown as WebSocket, + ); + + result = initPlayground(); + expect(result.controlsHandler).toBeTruthy(); + expect(result.taskLoaderHandler).toBeTruthy(); + expect(result.statsViewer).toBeTruthy(); + expect(result.logViewer).toBeTruthy(); + expect(result.adbConsole).toBeTruthy(); + expect(result.socketHandler).toBeTruthy(); + + if (holder.cb) { + holder.cb({ + data: JSON.stringify({type: 'stats', stats: {ack_latency: 12.34}}), + }); + const statAck = document.getElementById('stat-ack') as HTMLElement; + expect(statAck.innerText).toBe('12.34'); + + holder.cb({ + data: JSON.stringify({ + type: 'logs', + logs: [{message: 'hello', level: 'INFO', timestamp: 1}], + }), + }); + expect(result.logViewer!.getLogCount()).toBe(1); + + // Verify adb response callback + spyOn(result.adbConsole!, 'handleAdbResponse'); + holder.cb({ + data: JSON.stringify({ + type: 'adb_response', + command: 'shell pm list packages', + output: 'package:com.google.android.gms', + status: 'OK', + }), + }); + expect(result.adbConsole!.handleAdbResponse).toHaveBeenCalledWith( + jasmine.objectContaining({command: 'shell pm list packages'}) + ); + + // Verify status change updates conn-dot + const connDot = document.getElementById('conn-dot') as HTMLElement; + holder.cb({ + data: JSON.stringify({type: 'status', message: 'Connected to emulator'}), + }); + expect(connDot.style.background).toBe('var(--success)'); + + holder.cb({ + data: JSON.stringify({type: 'status', message: 'Disconnected from emulator'}), + }); + expect(connDot.style.background).toBe('var(--warning)'); + } + + // Verify screenControls.setSocket(socketHandler) was bound + const screenImg = document.getElementById('screen') as HTMLImageElement; + screenImg.dispatchEvent( + new MouseEvent('mousedown', {clientX: 10, clientY: 10}) + ); + expect(holder.sendSpy).toHaveBeenCalled(); + holder.sendSpy.calls.reset(); + + // Verify taskLoader.setSocket(socketHandler) was bound + const taskInput = document.getElementById('task-name') as HTMLInputElement; + const loadBtn = document.getElementById( + 'load-task-btn' + ) as HTMLButtonElement; + taskInput.value = 'classic_2048'; + loadBtn.click(); + expect(holder.sendSpy).toHaveBeenCalledWith('load_task,classic_2048'); + }); + + it('handles initPlayground gracefully when optional elements are absent', () => { + document.body.innerHTML = ''; + const emptyResult = initPlayground(); + expect(emptyResult.controlsHandler).toBeNull(); + expect(emptyResult.taskLoaderHandler).toBeNull(); + expect(emptyResult.statsViewer).toBeNull(); + expect(emptyResult.logViewer).toBeNull(); + expect(emptyResult.adbConsole).toBeNull(); + + // Verify callbacks with absent elements do not throw + const callbacks = (emptyResult.socketHandler as unknown as { + callbacks: { + onStatsReceived: (s: unknown) => void; + onLogsReceived: (l: unknown) => void; + onAdbResponseReceived: (r: unknown) => void; + onStatusChange: (msg: string) => void; + }; + }).callbacks; + expect(() => { + callbacks.onStatsReceived({}); + callbacks.onLogsReceived([]); + callbacks.onAdbResponseReceived({command: 'test', status: 'OK', output: ''}); + callbacks.onStatusChange('Connected'); + }).not.toThrow(); + }); + + it('switches tabs on tab button click', () => { + result = initPlayground(); + const tabLogs = document.getElementById('tab-logs') as HTMLElement; + const tabAdb = document.getElementById('tab-adb') as HTMLElement; + const tabBtns = document.querySelectorAll('.tab-btn'); + const logsBtn = tabBtns[0] as HTMLButtonElement; + const adbBtn = tabBtns[1] as HTMLButtonElement; + const noTargetBtn = document.getElementById( + 'btn-no-target' + ) as HTMLButtonElement; + + expect(tabLogs.classList.contains('active')).toBeTrue(); + expect(tabAdb.classList.contains('active')).toBeFalse(); + + adbBtn.click(); + expect(adbBtn.classList.contains('active')).toBeTrue(); + expect(logsBtn.classList.contains('active')).toBeFalse(); + expect(tabLogs.classList.contains('active')).toBeFalse(); + expect(tabAdb.classList.contains('active')).toBeTrue(); + + // Clicking a button without a target attribute should not throw and still update active button + noTargetBtn.click(); + expect(noTargetBtn.classList.contains('active')).toBeTrue(); + expect(adbBtn.classList.contains('active')).toBeFalse(); + }); + + it('opens log viewer panel when logs-open-btn is clicked', () => { + result = initPlayground(); + const openLogsBtn = document.getElementById( + 'logs-open-btn' + ) as HTMLButtonElement; + spyOn(result.logViewer!, 'showPanel'); + expect(openLogsBtn.style.display).toBe('inline-block'); + openLogsBtn.click(); + expect(result.logViewer!.showPanel).toHaveBeenCalled(); + expect(openLogsBtn.style.display).toBe('none'); + }); + + it('toggles stats viewer panel when stats-toggle is clicked', () => { + result = initPlayground(); + const openStatsBtn = document.getElementById( + 'stats-toggle' + ) as HTMLButtonElement; + const statsPanel = document.getElementById( + 'stats-panel' + ) as HTMLDialogElement; + spyOn(statsPanel, 'show'); + spyOn(statsPanel, 'close'); + openStatsBtn.click(); + expect(statsPanel.show).toHaveBeenCalledTimes(1); + openStatsBtn.click(); + expect(statsPanel.close).toHaveBeenCalledTimes(1); + }); + + it('handles setupAutoInit when document.readyState is loading', () => { + Object.defineProperty(document, 'readyState', { + value: 'loading', + configurable: true, + }); + + const wsSpy = window.WebSocket as unknown as jasmine.Spy; + wsSpy.calls.reset(); + + setupAutoInit(); + + expect(wsSpy).not.toHaveBeenCalled(); + + document.dispatchEvent(new Event('DOMContentLoaded')); + + expect(wsSpy).toHaveBeenCalled(); + }); + + it('does not initialize on DOMContentLoaded if screen element is missing', () => { + document.body.innerHTML = ''; + Object.defineProperty(document, 'readyState', { + value: 'loading', + configurable: true, + }); + + const wsSpy = window.WebSocket as unknown as jasmine.Spy; + wsSpy.calls.reset(); + + setupAutoInit(); + document.dispatchEvent(new Event('DOMContentLoaded')); + + expect(wsSpy).not.toHaveBeenCalled(); + }); + + it('handles setupAutoInit when document.readyState is complete', () => { + Object.defineProperty(document, 'readyState', { + value: 'complete', + configurable: true, + }); + + const wsSpy = window.WebSocket as unknown as jasmine.Spy; + wsSpy.calls.reset(); + + setupAutoInit(); + + expect(wsSpy).toHaveBeenCalled(); + }); + + it('handles setupAutoInit when screen element does not exist and readyState is complete', () => { + document.body.innerHTML = ''; + Object.defineProperty(document, 'readyState', { + value: 'complete', + configurable: true, + }); + const wsSpy = window.WebSocket as unknown as jasmine.Spy; + wsSpy.calls.reset(); + + setupAutoInit(); + + expect(wsSpy).not.toHaveBeenCalled(); + }); + + it('returns false when document is null or undefined', () => { + expect(setupAutoInit(null)).toBeFalse(); + }); +}); + + diff --git a/android_env/playground/frontend/package.json b/android_env/playground/frontend/package.json new file mode 100644 index 00000000..30957b98 --- /dev/null +++ b/android_env/playground/frontend/package.json @@ -0,0 +1,10 @@ +{ + "name": "android-env-playground-frontend", + "private": true, + "dependencies": { + "safevalues": "1.2.0" + }, + "pnpm": { + "onlyBuiltDependencies": [] + } +} diff --git a/android_env/playground/frontend/pnpm-lock.yaml b/android_env/playground/frontend/pnpm-lock.yaml new file mode 100644 index 00000000..563f8e9a --- /dev/null +++ b/android_env/playground/frontend/pnpm-lock.yaml @@ -0,0 +1,22 @@ +lockfileVersion: '9.0' + +settings: + autoInstallPeers: true + excludeLinksFromLockfile: false + +importers: + + .: + dependencies: + safevalues: + specifier: 1.2.0 + version: 1.2.0 + +packages: + + safevalues@1.2.0: + resolution: {integrity: sha512-zIsuhjYvJCjfsfjoim2ab6gLKFYAnTiDSJGh0cC3T44L/4kNLL90hBG2BzrXPrHA3f8Ms8FSJ1mljKH5dVR1cw==} + +snapshots: + + safevalues@1.2.0: {} diff --git a/android_env/playground/frontend/socket.ts b/android_env/playground/frontend/socket.ts new file mode 100644 index 00000000..639d28cb --- /dev/null +++ b/android_env/playground/frontend/socket.ts @@ -0,0 +1,213 @@ +import {objectUrlFromSafeSource} from 'safevalues/dom'; +import {LogData} from './logs_viewer'; +import {StatsData} from './stats'; + +/** + * Interface representing an ADB command execution response. + */ +export interface AdbResponseData { + command: string; + status: string; + output: string; + error?: string; + duration_ms?: number; +} + +/** + * Interface for socket event callbacks. + */ +export interface SocketCallbacks { + onStatusChange?: (status: string) => void; + onStatsReceived?: (stats: StatsData) => void; + onLogsReceived?: (logs: LogData[]) => void; + onAdbResponseReceived?: (response: AdbResponseData) => void; +} + +/** + * Generates the WebSocket URL based on protocol and host. + */ +export function getWebSocketUrl(protocol: string, host: string): string { + const wsProto = protocol === 'https:' ? 'wss' : 'ws'; + return `${wsProto}://${host}/ws`; +} + +/** + * Handles WebSocket communications with the AndroidEnv backend. + */ +export class SocketHandler { + private readonly ws: WebSocket; + private readonly screenImg: HTMLImageElement | null; + private readonly statusDiv: HTMLElement | null; + private readonly callbacks: SocketCallbacks; + + constructor( + screenImgId = 'screen', + statusDivId = 'status', + callbacks: SocketCallbacks = {} + ) { + this.screenImg = document.getElementById( + screenImgId + ) as HTMLImageElement | null; + this.statusDiv = document.getElementById(statusDivId) as HTMLElement | null; + this.callbacks = callbacks; + + this.ws = new WebSocket( + getWebSocketUrl(window.location.protocol, window.location.host) + ); + this.ws.binaryType = 'blob'; + + this.initWebSocket(); + } + + private initWebSocket(): void { + this.ws.onopen = () => { + console.log('WS Connected'); + this.updateStatus('Connected. Waiting for emulator...'); + }; + + this.ws.onclose = () => { + console.log('WS Closed'); + this.updateStatus('Disconnected'); + }; + + this.ws.onmessage = (event: MessageEvent) => { + if (event.data instanceof Blob) { + this.handleBlobFrame(event.data); + } else if (typeof event.data === 'string') { + this.handleTextMessage(event.data); + } + }; + } + + private updateStatus(msg: string): void { + if (this.statusDiv) { + this.statusDiv.innerText = msg; + } + if (this.callbacks.onStatusChange) { + this.callbacks.onStatusChange(msg); + } + } + + private handleBlobFrame(blob: Blob): void { + const screenImg = this.screenImg; + if (!screenImg) { + return; + } + const safeBlob = blob.type ? blob : new Blob([blob], {type: 'image/jpeg'}); + const url = objectUrlFromSafeSource(safeBlob); + const oldUrl = screenImg.src; + + const tempImg = new Image(); + tempImg.src = url; + + const onFinish = () => { + screenImg.src = url; + if (oldUrl && oldUrl.startsWith('blob:')) { + URL.revokeObjectURL(oldUrl); + } + requestAnimationFrame(() => { + requestAnimationFrame(() => { + this.sendAck(); + }); + }); + }; + + tempImg.decode().then(onFinish).catch((err) => { + console.error('Decode failed:', err); + onFinish(); + }); + } + + private handleTextMessage(data: string): void { + try { + const msg = JSON.parse(data) as { + type?: string; + message?: string; + stats?: StatsData; + logs?: LogData[]; + command?: string; + status?: string; + output?: string; + error?: string; + duration_ms?: number; + }; + if (msg.type === 'status' && msg.message !== undefined) { + this.updateStatus(msg.message); + } else if (msg.type === 'stats' && msg.stats && this.callbacks.onStatsReceived) { + this.callbacks.onStatsReceived(msg.stats); + } else if (msg.type === 'logs' && msg.logs && this.callbacks.onLogsReceived) { + this.callbacks.onLogsReceived(msg.logs); + } else if ( + msg.type === 'adb_response' && + msg.command !== undefined && + this.callbacks.onAdbResponseReceived + ) { + this.callbacks.onAdbResponseReceived({ + command: msg.command, + status: msg.status ?? 'OK', + output: msg.output ?? '', + error: msg.error, + duration_ms: msg.duration_ms, + }); + } + } catch (err) { + console.error('Failed to parse message:', err); + } + } + + + sendAck(): void { + if (this.ws.readyState === WebSocket.OPEN) { + this.ws.send('ack'); + } + } + + sendAction(type: string, x: number, y: number): void { + if (this.ws.readyState === WebSocket.OPEN) { + this.ws.send(`${type},${x},${y},${Date.now()}`); + } + } + + sendText(text: string): void { + if (this.ws.readyState === WebSocket.OPEN) { + this.ws.send(`text,${text},${Date.now()}`); + } + } + + sendKeydown(keycode: number): void { + if (this.ws.readyState === WebSocket.OPEN) { + this.ws.send(`keydown,${keycode},${Date.now()}`); + } + } + + sendSelection(start: number, end: number): void { + if (this.ws.readyState === WebSocket.OPEN) { + this.ws.send(`set_selection,${start},${end},${Date.now()}`); + } + } + + sendClipboard(action: 'cut' | 'copy' | 'paste'): void { + if (this.ws.readyState === WebSocket.OPEN) { + this.ws.send(`${action},${Date.now()}`); + } + } + + loadTask(taskId: string): void { + if (this.ws.readyState === WebSocket.OPEN) { + console.log('Sending load_task request:', taskId); + this.ws.send(`load_task,${taskId}`); + } + } + + sendAdbCommand(command: string): void { + if (this.ws.readyState === WebSocket.OPEN) { + this.ws.send(`adb,${command}`); + } + } + + sendReset(): void { + if (this.ws.readyState === WebSocket.OPEN) { + this.ws.send('reset'); + } + } +} diff --git a/android_env/playground/frontend/socket_test.ts b/android_env/playground/frontend/socket_test.ts new file mode 100644 index 00000000..57496edf --- /dev/null +++ b/android_env/playground/frontend/socket_test.ts @@ -0,0 +1,329 @@ +import 'jasmine'; +import {getWebSocketUrl, SocketHandler} from './socket'; + +interface FakeWebSocket { + readyState: number; + send: jasmine.Spy; + binaryType: string; + onopen: (() => void) | null; + onclose: (() => void) | null; + onmessage: ((event: {data: unknown}) => void) | null; +} + +describe('SocketHandler', () => { + let screenImg: HTMLImageElement; + let statusDiv: HTMLElement; + let fakeWs: FakeWebSocket; + + beforeEach(() => { + document.body.innerHTML = ` + +
Connecting...
+ `; + screenImg = document.getElementById('screen') as HTMLImageElement; + statusDiv = document.getElementById('status') as HTMLElement; + + fakeWs = { + readyState: 1, // WebSocket.OPEN + send: jasmine.createSpy('send'), + binaryType: '', + onopen: null, + onclose: null, + onmessage: null, + }; + + spyOn(window, 'WebSocket').and.returnValue(fakeWs as unknown as WebSocket); + }); + + afterEach(() => { + document.body.innerHTML = ''; + }); + + it('connects WebSocket and updates status on open', () => { + expect(screenImg).toBeTruthy(); + const handler = new SocketHandler('screen', 'status'); + expect(handler).toBeTruthy(); + + expect(fakeWs.binaryType).toBe('blob'); + if (fakeWs.onopen) { + fakeWs.onopen(); + } + expect(statusDiv.innerText).toBe('Connected. Waiting for emulator...'); + }); + + it('generates correct WebSocket URL for http and https protocols', () => { + expect(getWebSocketUrl('http:', 'localhost:8080')).toBe('ws://localhost:8080/ws'); + expect(getWebSocketUrl('https:', 'localhost:8080')).toBe('wss://localhost:8080/ws'); + }); + + + it('updates status on close', () => { + const handler = new SocketHandler('screen', 'status'); + expect(handler).toBeTruthy(); + + if (fakeWs.onclose) { + fakeWs.onclose(); + } + expect(statusDiv.innerText).toBe('Disconnected'); + }); + + it('sends action string over WebSocket', () => { + const handler = new SocketHandler('screen', 'status'); + handler.sendAction('touch_down', 0.5, 0.25); + + expect(fakeWs.send).toHaveBeenCalled(); + const sentData: string = fakeWs.send.calls.mostRecent().args[0]; + expect(sentData.startsWith('touch_down,0.5,0.25,')).toBeTrue(); + }); + + it('sends load_task string over WebSocket', () => { + const handler = new SocketHandler('screen', 'status'); + handler.loadTask('classic_2048'); + + expect(fakeWs.send).toHaveBeenCalledWith('load_task,classic_2048'); + }); + + it('sends text string over WebSocket', () => { + const handler = new SocketHandler('screen', 'status'); + handler.sendText('hello'); + + expect(fakeWs.send).toHaveBeenCalled(); + const sentData: string = fakeWs.send.calls.mostRecent().args[0]; + expect(sentData.startsWith('text,hello,')).toBeTrue(); + }); + + it('sends keydown string over WebSocket', () => { + const handler = new SocketHandler('screen', 'status'); + handler.sendKeydown(66); + + expect(fakeWs.send).toHaveBeenCalled(); + const sentData: string = fakeWs.send.calls.mostRecent().args[0]; + expect(sentData.startsWith('keydown,66,')).toBeTrue(); + }); + + it('sends set_selection string over WebSocket', () => { + const handler = new SocketHandler('screen', 'status'); + handler.sendSelection(2, 5); + + expect(fakeWs.send).toHaveBeenCalled(); + const sentData: string = fakeWs.send.calls.mostRecent().args[0]; + expect(sentData.startsWith('set_selection,2,5,')).toBeTrue(); + }); + + it('sends clipboard string over WebSocket', () => { + const handler = new SocketHandler('screen', 'status'); + handler.sendClipboard('copy'); + + expect(fakeWs.send).toHaveBeenCalled(); + const sentData: string = fakeWs.send.calls.mostRecent().args[0]; + expect(sentData.startsWith('copy,')).toBeTrue(); + }); + + it('handles status JSON text messages', () => { + let statusReceived = ''; + const handler = new SocketHandler('screen', 'status', { + onStatusChange: (status) => { + statusReceived = status; + }, + }); + expect(handler).toBeTruthy(); + + if (fakeWs.onmessage) { + fakeWs.onmessage({ + data: JSON.stringify({type: 'status', message: 'Emulator Running'}), + }); + } + + expect(statusDiv.innerText).toBe('Emulator Running'); + expect(statusReceived).toBe('Emulator Running'); + }); + + it('handles stats JSON text messages', () => { + let statsReceived: unknown = null; + const handler = new SocketHandler('screen', 'status', { + onStatsReceived: (stats) => { + statsReceived = stats; + }, + }); + expect(handler).toBeTruthy(); + + if (fakeWs.onmessage) { + fakeWs.onmessage({ + data: JSON.stringify({type: 'stats', stats: {ack_latency: 15.0}}), + }); + } + + expect(statsReceived).toEqual({ack_latency: 15.0}); + }); + + it('handles binary Blob frame message and cleans up old blob URL', async () => { + screenImg.src = 'blob:http://localhost/old-blob'; + const handler = new SocketHandler('screen', 'status'); + expect(handler).toBeTruthy(); + + spyOn(URL, 'createObjectURL').and.returnValue('blob:http://localhost/new-blob'); + const revokeSpy = spyOn(URL, 'revokeObjectURL'); + + let srcSetBeforeDecode = false; + spyOn(Image.prototype, 'decode').and.callFake(function (this: HTMLImageElement) { + if (this.src && this.src.includes('blob:http://localhost/new-blob')) { + srcSetBeforeDecode = true; + } + return Promise.resolve(); + }); + + // Mock requestAnimationFrame to execute immediately + spyOn(window, 'requestAnimationFrame').and.callFake((cb: FrameRequestCallback) => { + cb(0); + return 0; + }); + + const fakeBlob = new Blob(['fake image data']); // No type specified (empty type) + + if (fakeWs.onmessage) { + fakeWs.onmessage({data: fakeBlob}); + } + + // Wait for promise resolution + await Promise.resolve(); + + expect(URL.createObjectURL).toHaveBeenCalled(); + expect(srcSetBeforeDecode).toBeTrue(); + expect(screenImg.src).toContain('blob:http://localhost/new-blob'); + expect(revokeSpy).toHaveBeenCalledWith('blob:http://localhost/old-blob'); + expect(fakeWs.send).toHaveBeenCalledWith('ack'); + }); + + it('does not revoke URL when oldUrl is not a blob URL', async () => { + screenImg.src = 'http://localhost/static/screen.png'; + const handler = new SocketHandler('screen', 'status'); + expect(handler).toBeTruthy(); + + spyOn(URL, 'createObjectURL').and.returnValue('blob:http://localhost/new-blob'); + const revokeSpy = spyOn(URL, 'revokeObjectURL'); + spyOn(Image.prototype, 'decode').and.returnValue(Promise.resolve()); + + if (fakeWs.onmessage) { + fakeWs.onmessage({data: new Blob(['data'], {type: 'image/png'})}); + } + + await Promise.resolve(); + + expect(revokeSpy).not.toHaveBeenCalled(); + expect(screenImg.src).toContain('blob:http://localhost/new-blob'); + }); + + it('handles image decode failure gracefully in handleBlobFrame', async () => { + const handler = new SocketHandler('screen', 'status'); + expect(handler).toBeTruthy(); + + spyOn(URL, 'createObjectURL').and.returnValue('blob:http://localhost/new-blob'); + spyOn(Image.prototype, 'decode').and.rejectWith(new Error('Decode error')); + spyOn(console, 'error'); + + const fakeBlob = new Blob(['corrupt image data'], {type: 'image/png'}); + + if (fakeWs.onmessage) { + fakeWs.onmessage({data: fakeBlob}); + } + + await new Promise((resolve) => setTimeout(resolve, 10)); + + expect(console.error).toHaveBeenCalledWith('Decode failed:', jasmine.any(Error)); + expect(screenImg.src).toContain('blob:http://localhost/new-blob'); + }); + + + + it('handles invalid JSON text messages gracefully', () => { + const handler = new SocketHandler('screen', 'status'); + expect(handler).toBeTruthy(); + spyOn(console, 'error'); + + if (fakeWs.onmessage) { + fakeWs.onmessage({data: 'invalid json {'}); + } + + expect(console.error).toHaveBeenCalledWith('Failed to parse message:', jasmine.any(Error)); + }); + + it('does not send actions, ack, or loadTask when socket is not OPEN', () => { + fakeWs.readyState = 3; // CLOSED + const handler = new SocketHandler('screen', 'status'); + + handler.sendAck(); + handler.sendAction('touch_down', 0.5, 0.5); + handler.sendText('hello'); + handler.sendKeydown(66); + handler.sendSelection(2, 5); + handler.sendClipboard('copy'); + handler.loadTask('task1'); + handler.sendAdbCommand('shell wm size'); + handler.sendReset(); + + expect(fakeWs.send).not.toHaveBeenCalled(); + }); + + it('handles missing screenImg and statusDiv elements gracefully', () => { + let statusReceived = ''; + const handler = new SocketHandler('missing-screen', 'missing-status', { + onStatusChange: (status) => { + statusReceived = status; + }, + }); + expect(handler).toBeTruthy(); + + if (fakeWs.onopen) { + fakeWs.onopen(); + } + expect(statusReceived).toBe('Connected. Waiting for emulator...'); + + const createUrlSpy = spyOn(URL, 'createObjectURL'); + if (fakeWs.onmessage) { + fakeWs.onmessage({data: new Blob(['data'], {type: 'image/jpeg'})}); + } + expect(createUrlSpy).not.toHaveBeenCalled(); + }); + + it('sends adb and reset commands over WebSocket', () => { + const handler = new SocketHandler('screen', 'status'); + handler.sendAdbCommand('shell wm size'); + expect(fakeWs.send).toHaveBeenCalledWith('adb,shell wm size'); + + handler.sendReset(); + expect(fakeWs.send).toHaveBeenCalledWith('reset'); + }); + + it('handles adb_response JSON text messages', () => { + let adbReceived: unknown = null; + const handler = new SocketHandler('screen', 'status', { + onAdbResponseReceived: (resp) => { + adbReceived = resp; + }, + }); + expect(handler).toBeTruthy(); + + if (fakeWs.onmessage) { + fakeWs.onmessage({ + data: JSON.stringify({ + type: 'adb_response', + command: 'shell wm size', + status: 'OK', + output: 'Physical size: 1080x1920', + duration_ms: 4.2, + }), + }); + } + + expect(adbReceived).toEqual({ + command: 'shell wm size', + status: 'OK', + output: 'Physical size: 1080x1920', + error: undefined, + duration_ms: 4.2, + }); + }); +}); + + + diff --git a/android_env/playground/frontend/stats.ts b/android_env/playground/frontend/stats.ts new file mode 100644 index 00000000..ada90947 --- /dev/null +++ b/android_env/playground/frontend/stats.ts @@ -0,0 +1,533 @@ + +/** + * Interface representing battery and power telemetry statistics. + */ +export interface BatteryPowerStatsData { + level_percent?: number; + current_now_ua?: number; + voltage_mv?: number; + power_mw?: number; + temperature_celsius?: number; + status?: string; + plugged_source?: string; + charge_counter_uah?: number; + thermal_status?: number; + energy_counter_nwh?: number; +} + +/** + * Interface representing RL episode telemetry metrics. + */ +export interface EpisodeStatsData { + step?: number; + reward?: number; + cumulative_reward?: number; + step_type?: string; + active_package?: string; +} + +/** + * Interface representing performance stats from the server. + */ +export interface StatsData { + ack_latency?: number; + env_step_ms?: number; + action_transit_ms?: number; + action_queue_ms?: number; + battery_power?: BatteryPowerStatsData; + episode?: EpisodeStatsData; +} + +/** + * Custom Web Component for rendering and managing "Stats for Nerds" in the Playground UI. + */ +export class StatsViewerElement extends HTMLElement { + private statsPanel: HTMLDialogElement | null = null; + private statsClose: HTMLButtonElement | null = null; + private statAck: HTMLElement | null = null; + private statStep: HTMLElement | null = null; + private statTransit: HTMLElement | null = null; + private statQueue: HTMLElement | null = null; + private statsCanvas: HTMLCanvasElement | null = null; + private ctx: CanvasRenderingContext2D | null = null; + private statBatteryLevel: HTMLElement | null = null; + private statBatteryStatus: HTMLElement | null = null; + private statBatterySource: HTMLElement | null = null; + private statBatteryVoltage: HTMLElement | null = null; + private statBatteryCurrent: HTMLElement | null = null; + private statBatteryPower: HTMLElement | null = null; + private statBatteryTemp: HTMLElement | null = null; + private statBatteryCounter: HTMLElement | null = null; + private statBatteryEnergy: HTMLElement | null = null; + private batteryCanvas: HTMLCanvasElement | null = null; + private batteryCtx: CanvasRenderingContext2D | null = null; + + private showStats = false; + private readonly maxDataPoints = 100; + private readonly statsHistory: { + ack: number[]; + step: number[]; + transit: number[]; + charge: number[]; + current: number[]; + } = { + ack: [], + step: [], + transit: [], + charge: [], + current: [], + }; + + constructor() { + super(); + } + + /** @export */ + connectedCallback(): void { + if (!this.statsPanel) { + const panel = document.createElement('dialog'); + panel.id = 'stats-panel'; + + const header = document.createElement('div'); + header.className = 'stats-header'; + + const title = document.createElement('h3'); + title.textContent = 'Stats for Nerds'; + + this.statsClose = document.createElement('button'); + this.statsClose.id = 'stats-close'; + this.statsClose.textContent = '\u00d7'; + + header.appendChild(title); + header.appendChild(this.statsClose); + panel.appendChild(header); + + const addStatLine = ( + label: string, + id: string, + unit: string, + tooltip: string + ): HTMLElement => { + const line = document.createElement('div'); + line.className = 'stat-line'; + const labelSpan = document.createElement('span'); + labelSpan.title = tooltip; + labelSpan.textContent = label + ':'; + const valContainer = document.createElement('span'); + const valSpan = document.createElement('span'); + valSpan.id = id; + valSpan.textContent = '-'; + valContainer.appendChild(valSpan); + if (unit) { + valContainer.appendChild(document.createTextNode(' ' + unit)); + } + line.appendChild(labelSpan); + line.appendChild(valContainer); + panel.appendChild(line); + return valSpan; + }; + + this.statAck = addStatLine( + 'ACK Latency', + 'stat-ack', + 'ms', + 'Round-trip time for frames. Measures network latency + client rendering time.' + ); + this.statStep = addStatLine( + 'Env Step', + 'stat-step', + 'ms', + 'Time taken by the emulator to process the action and return the next frame.' + ); + this.statTransit = addStatLine( + 'Action Transit', + 'stat-transit', + 'ms', + 'Time taken for the action to travel from the browser client to the server.' + ); + this.statQueue = addStatLine( + 'Action Queue', + 'stat-queue', + 'ms', + 'Time the action spent waiting in the server queue before being processed.' + ); + + this.statsCanvas = document.createElement('canvas'); + this.statsCanvas.id = 'stats-chart'; + this.statsCanvas.width = 400; + this.statsCanvas.height = 150; + this.statsCanvas.style.cssText = 'border: 1px solid #555; margin-top: 10px;'; + panel.appendChild(this.statsCanvas); + + const hr = document.createElement('hr'); + hr.style.cssText = 'border-color: #444; margin: 8px 0;'; + panel.appendChild(hr); + + const batteryHeader = document.createElement('div'); + batteryHeader.style.cssText = + 'font-weight: bold; margin-bottom: 4px; color: #8ab4f8;'; + batteryHeader.textContent = 'Phone / Device Metrics'; + panel.appendChild(batteryHeader); + + this.statBatteryLevel = addStatLine( + 'Battery Level', + 'stat-battery-level', + '%', + 'Battery charge level percentage [0-100]' + ); + this.statBatteryStatus = addStatLine( + 'Status', + 'stat-battery-status', + '', + 'Battery charging status (e.g. CHARGING, DISCHARGING)' + ); + this.statBatterySource = addStatLine( + 'Source', + 'stat-battery-source', + '', + 'Power source (AC, USB, WIRELESS, UNPLUGGED)' + ); + this.statBatteryVoltage = addStatLine( + 'Voltage', + 'stat-battery-voltage', + 'mV', + 'Battery voltage in millivolts' + ); + this.statBatteryCurrent = addStatLine( + 'Current', + 'stat-battery-current', + '\u03bcA', + 'Instantaneous current in microamperes' + ); + this.statBatteryPower = addStatLine( + 'Power', + 'stat-battery-power', + 'mW', + 'Instantaneous power in milliwatts' + ); + this.statBatteryTemp = addStatLine( + 'Temp', + 'stat-battery-temp', + '\u00b0C', + 'Battery temperature in degrees Celsius' + ); + this.statBatteryCounter = addStatLine( + 'Charge Counter', + 'stat-battery-counter', + '\u03bcAh', + 'Remaining charge counter in microampere-hours' + ); + this.statBatteryEnergy = addStatLine( + 'Energy', + 'stat-battery-energy', + 'nWh', + 'Total remaining energy in nanowatt-hours' + ); + + this.batteryCanvas = document.createElement('canvas'); + this.batteryCanvas.id = 'stats-battery-chart'; + this.batteryCanvas.width = 400; + this.batteryCanvas.height = 150; + this.batteryCanvas.style.cssText = + 'border: 1px solid #555; margin-top: 10px;'; + panel.appendChild(this.batteryCanvas); + + this.appendChild(panel); + this.statsPanel = panel; + this.ctx = this.statsCanvas.getContext('2d'); + this.batteryCtx = this.batteryCanvas.getContext('2d'); + this.initEventListeners(); + } + } + + private initEventListeners(): void { + this.statsClose!.addEventListener('click', () => { + this.hidePanel(); + }); + + this.statsPanel!.addEventListener('close', () => { + this.showStats = false; + }); + } + + /** @export */ + showPanel(): void { + this.showStats = true; + this.statsPanel!.show(); + this.drawChart(); + this.drawBatteryChart(); + } + + /** @export */ + hidePanel(): void { + this.showStats = false; + this.statsPanel!.close(); + } + + /** @export */ + togglePanel(): void { + if (this.showStats) { + this.hidePanel(); + } else { + this.showPanel(); + } + } + + /** @export */ + handleStats(stats: StatsData): void { + if (stats.ack_latency !== undefined) { + this.statAck!.innerText = stats.ack_latency.toFixed(2); + this.pushStat(this.statsHistory.ack, stats.ack_latency); + } + if (stats.env_step_ms !== undefined) { + this.statStep!.innerText = stats.env_step_ms.toFixed(2); + this.pushStat(this.statsHistory.step, stats.env_step_ms); + } + if (stats.action_transit_ms !== undefined) { + this.statTransit!.innerText = stats.action_transit_ms.toFixed(2); + this.pushStat(this.statsHistory.transit, stats.action_transit_ms); + } + if (stats.action_queue_ms !== undefined) { + this.statQueue!.innerText = stats.action_queue_ms.toFixed(2); + } + if (stats.battery_power) { + const b = stats.battery_power; + if (b.level_percent !== undefined) { + this.statBatteryLevel!.innerText = b.level_percent.toString(); + this.pushStat(this.statsHistory.charge, b.level_percent); + } + if (b.status !== undefined) { + this.statBatteryStatus!.innerText = b.status; + } + if (b.plugged_source !== undefined) { + this.statBatterySource!.innerText = b.plugged_source; + } + if (b.voltage_mv !== undefined) { + this.statBatteryVoltage!.innerText = b.voltage_mv.toString(); + } + if (b.current_now_ua !== undefined) { + this.statBatteryCurrent!.innerText = b.current_now_ua.toString(); + this.pushStat(this.statsHistory.current, b.current_now_ua); + } + if (b.power_mw !== undefined) { + this.statBatteryPower!.innerText = b.power_mw.toFixed(2); + } + if (b.temperature_celsius !== undefined) { + this.statBatteryTemp!.innerText = b.temperature_celsius.toFixed(1); + } + if (b.charge_counter_uah !== undefined) { + this.statBatteryCounter!.innerText = b.charge_counter_uah.toString(); + } + if (b.energy_counter_nwh !== undefined) { + this.statBatteryEnergy!.innerText = b.energy_counter_nwh.toString(); + } + } + + if (stats.episode) { + const ep = stats.episode; + const kpiStep = document.getElementById('kpi-step'); + if (kpiStep && ep.step !== undefined) { + kpiStep.textContent = `${ep.step} (${ep.step_type ?? 'MID'})`; + } + const kpiReward = document.getElementById('kpi-reward'); + if (kpiReward && ep.cumulative_reward !== undefined) { + const r = ep.reward ?? 0; + kpiReward.textContent = `${ep.cumulative_reward.toFixed(1)} (+${r.toFixed(1)})`; + } + const pkgBadge = document.getElementById('pkg-badge'); + if (pkgBadge && ep.active_package) { + const shortName = + ep.active_package.split('.').pop() || ep.active_package; + pkgBadge.textContent = `pkg: ${shortName}`; + } + } + + const kpiLatency = document.getElementById('kpi-latency'); + if ( + kpiLatency && + (stats.ack_latency !== undefined || stats.env_step_ms !== undefined) + ) { + kpiLatency.textContent = `${(stats.ack_latency ?? 0).toFixed(1)} / ${(stats.env_step_ms ?? 0).toFixed(1)} ms`; + } + + const kpiBattery = document.getElementById('kpi-battery'); + if ( + kpiBattery && + stats.battery_power && + stats.battery_power.level_percent !== undefined + ) { + kpiBattery.textContent = `${stats.battery_power.level_percent}% · ${(stats.battery_power.power_mw ?? 0).toFixed(0)}mW`; + } + + if (this.showStats) { + this.drawChart(); + this.drawBatteryChart(); + } + } + + private pushStat(arr: number[], val: number): void { + arr.push(val); + if (arr.length > this.maxDataPoints) { + arr.shift(); + } + } + + private drawChart(): void { + const width = this.statsCanvas!.width; + const height = this.statsCanvas!.height; + const ctx = this.ctx!; + ctx.clearRect(0, 0, width, height); + + let maxVal = 100; + for (const key of ['ack', 'step', 'transit'] as const) { + if (this.statsHistory[key].length > 0) { + maxVal = Math.max(maxVal, ...this.statsHistory[key]); + } + } + maxVal = maxVal * 1.2; + + ctx.strokeStyle = '#444'; + ctx.lineWidth = 1; + for (let i = 1; i < 4; i++) { + const y = height * (i / 4); + ctx.beginPath(); + ctx.moveTo(0, y); + ctx.lineTo(width, y); + ctx.stroke(); + + ctx.fillStyle = '#aaa'; + ctx.font = '10px sans-serif'; + ctx.fillText((maxVal * (1 - i / 4)).toFixed(0) + ' ms', 5, y - 2); + } + + this.drawPlot(this.statsHistory.ack, '#ff8a80', maxVal); + this.drawPlot(this.statsHistory.step, '#b9f6ca', maxVal); + this.drawPlot(this.statsHistory.transit, '#82b1ff', maxVal); + + ctx.fillStyle = '#ff8a80'; + ctx.fillRect(10, 10, 10, 10); + ctx.fillStyle = '#e8eaed'; + ctx.fillText('ACK', 25, 18); + + ctx.fillStyle = '#b9f6ca'; + ctx.fillRect(70, 10, 10, 10); + ctx.fillStyle = '#e8eaed'; + ctx.fillText('Step', 85, 18); + + ctx.fillStyle = '#82b1ff'; + ctx.fillRect(130, 10, 10, 10); + ctx.fillStyle = '#e8eaed'; + ctx.fillText('Transit', 145, 18); + } + + private drawPlot(data: number[], color: string, maxVal: number): void { + this.drawSeries(this.ctx!, this.statsCanvas!, data, color, 0, maxVal); + } + + private drawSeries( + ctx: CanvasRenderingContext2D, + canvas: HTMLCanvasElement, + data: number[], + color: string, + minVal: number, + maxVal: number + ): void { + if (data.length < 2) return; + const width = canvas.width; + const height = canvas.height; + const range = maxVal - minVal || 1; + ctx.strokeStyle = color; + ctx.lineWidth = 2; + ctx.beginPath(); + const stepX = width / (this.maxDataPoints - 1); + const startX = width - (data.length - 1) * stepX; + + const getY = (val: number) => { + const fraction = (val - minVal) / range; + const clamped = Math.max(0, Math.min(1, fraction)); + return height - clamped * height; + }; + + ctx.moveTo(startX, getY(data[0])); + for (let i = 1; i < data.length; i++) { + const x = startX + i * stepX; + ctx.lineTo(x, getY(data[i])); + } + ctx.stroke(); + } + + private drawBatteryChart(): void { + const width = this.batteryCanvas!.width; + const height = this.batteryCanvas!.height; + const batteryCtx = this.batteryCtx!; + batteryCtx.clearRect(0, 0, width, height); + + let minCurrent = 0; + let maxCurrent = 100; + if (this.statsHistory.current.length > 0) { + minCurrent = Math.min(0, ...this.statsHistory.current); + maxCurrent = Math.max(100, ...this.statsHistory.current); + } + maxCurrent = maxCurrent * 1.2; + + batteryCtx.strokeStyle = '#444'; + batteryCtx.lineWidth = 1; + for (let i = 1; i < 4; i++) { + const y = height * (i / 4); + batteryCtx.beginPath(); + batteryCtx.moveTo(0, y); + batteryCtx.lineTo(width, y); + batteryCtx.stroke(); + + batteryCtx.fillStyle = '#aaa'; + batteryCtx.font = '10px sans-serif'; + batteryCtx.fillText(`${(100 * (1 - i / 4)).toFixed(0)}%`, 5, y - 2); + + const currentVal = `${( + minCurrent + + (maxCurrent - minCurrent) * (1 - i / 4) + ).toFixed(0)} \u03bcA`; + const textWidth = batteryCtx.measureText(currentVal).width; + batteryCtx.fillText(currentVal, width - textWidth - 5, y - 2); + } + + this.drawSeries( + batteryCtx, + this.batteryCanvas!, + this.statsHistory.charge, + '#00e676', + 0, + 100 + ); + this.drawSeries( + batteryCtx, + this.batteryCanvas!, + this.statsHistory.current, + '#ffd54f', + minCurrent, + maxCurrent + ); + + batteryCtx.fillStyle = '#00e676'; + batteryCtx.fillRect(10, 10, 10, 10); + batteryCtx.fillStyle = '#e8eaed'; + batteryCtx.font = '10px sans-serif'; + batteryCtx.fillText('Charge (%)', 25, 18); + + batteryCtx.fillStyle = '#ffd54f'; + batteryCtx.fillRect(100, 10, 10, 10); + batteryCtx.fillStyle = '#e8eaed'; + batteryCtx.fillText('Current (\u03bcA)', 115, 18); + } +} + +if (!customElements.get('stats-viewer')) { + customElements.define('stats-viewer', StatsViewerElement); +} + +/** + * Helper factory function to instantiate and upgrade a StatsViewerElement. + */ +export function createStatsViewer(): StatsViewerElement { + return document.createElement('stats-viewer') as StatsViewerElement; +} diff --git a/android_env/playground/frontend/stats_test.ts b/android_env/playground/frontend/stats_test.ts new file mode 100644 index 00000000..772b4cd6 --- /dev/null +++ b/android_env/playground/frontend/stats_test.ts @@ -0,0 +1,353 @@ +import 'jasmine'; +import {createStatsViewer, StatsViewerElement} from './stats'; + +interface InternalStatsViewer { + batteryCtx: CanvasRenderingContext2D | null; + batteryCanvas: HTMLCanvasElement | null; + drawBatteryChart: () => void; + statsHistory: { + ack: number[]; + step: number[]; + transit: number[]; + charge: number[]; + current: number[]; + }; +} + +describe('StatsViewerElement', () => { + let statsViewer: StatsViewerElement; + let statsToggle: HTMLButtonElement; + let statsPanel: HTMLDialogElement; + let statsClose: HTMLButtonElement; + let statAck: HTMLElement; + let statStep: HTMLElement; + let statTransit: HTMLElement; + let statQueue: HTMLElement; + let statsCanvas: HTMLCanvasElement; + let batteryCanvas: HTMLCanvasElement; + + beforeEach(() => { + document.body.innerHTML = ` + + `; + statsViewer = createStatsViewer(); + statsViewer.id = 'stats-viewer'; + document.body.appendChild(statsViewer); + + statsToggle = document.getElementById('stats-toggle') as HTMLButtonElement; + statsPanel = statsViewer.querySelector('#stats-panel') as HTMLDialogElement; + statsClose = statsViewer.querySelector('#stats-close') as HTMLButtonElement; + statAck = statsViewer.querySelector('#stat-ack') as HTMLElement; + statStep = statsViewer.querySelector('#stat-step') as HTMLElement; + statTransit = statsViewer.querySelector('#stat-transit') as HTMLElement; + statQueue = statsViewer.querySelector('#stat-queue') as HTMLElement; + statsCanvas = statsViewer.querySelector( + '#stats-chart' + ) as HTMLCanvasElement; + batteryCanvas = statsViewer.querySelector( + '#stats-battery-chart' + ) as HTMLCanvasElement; + + statsPanel.show = jasmine.createSpy('show'); + statsPanel.close = jasmine.createSpy('close'); + + statsToggle.addEventListener('click', () => { + statsViewer.showPanel(); + }); + }); + + afterEach(() => { + document.body.innerHTML = ''; + }); + + it('initializes custom element, canvas and elements with exact DOM structure', () => { + expect(statsViewer).toBeTruthy(); + expect(statsViewer.tagName.toLowerCase()).toBe('stats-viewer'); + expect(statsCanvas).toBeTruthy(); + expect(statsPanel).toBeTruthy(); + + const expectedHtml = + '

Stats for Nerds

' + + '
ACK Latency:- ms
' + + '
Env Step:- ms
' + + '
Action Transit:- ms
' + + '
Action Queue:- ms
' + + '' + + '
' + + '
Phone / Device Metrics
' + + '
Battery Level:- %
' + + '
Status:-
' + + '
Source:-
' + + '
Voltage:- mV
' + + '
Current:- \u03bcA
' + + '
Power:- mW
' + + '
Temp:- \u00b0C
' + + '
Charge Counter:- \u03bcAh
' + + '
Energy:- nWh
' + + ''; + expect(statsPanel.innerHTML).toBe(expectedHtml); + + // Calling connectedCallback again is idempotent + statsViewer.connectedCallback(); + expect(statsViewer.children.length).toBe(1); + }); + + it('updates stat text content when receiving stats data', () => { + statsViewer.handleStats({ + ack_latency: 12.34, + env_step_ms: 5.67, + action_transit_ms: 8.91, + action_queue_ms: 1.23, + }); + + expect(statAck.innerText).toBe('12.34'); + expect(statStep.innerText).toBe('5.67'); + expect(statTransit.innerText).toBe('8.91'); + expect(statQueue.innerText).toBe('1.23'); + }); + + it('opens panel and renders chart on new stats', () => { + const ctx = statsCanvas.getContext('2d')!; + spyOn(ctx, 'clearRect').and.callThrough(); + spyOn(ctx, 'fillText').and.callThrough(); + spyOn(ctx, 'fillRect').and.callThrough(); + spyOn(ctx, 'stroke').and.callThrough(); + + statsToggle.click(); + expect(statsPanel.show).toHaveBeenCalled(); + + for (let i = 0; i < 105; i++) { + statsViewer.handleStats({ + ack_latency: 10 + i, + env_step_ms: 20 + i, + action_transit_ms: 30 + i, + action_queue_ms: 5, + }); + } + + expect(statAck.innerText).toBe('114.00'); + const internalViewer = statsViewer as unknown as InternalStatsViewer; + expect(internalViewer.statsHistory.ack.length).toBe(100); + expect(ctx.clearRect).toHaveBeenCalledWith(0, 0, 400, 150); + expect(ctx.fillText).toHaveBeenCalledWith('ACK', 25, 18); + expect(ctx.fillText).toHaveBeenCalledWith('Step', 85, 18); + expect(ctx.fillText).toHaveBeenCalledWith('Transit', 145, 18); + expect(ctx.fillRect).toHaveBeenCalledWith(10, 10, 10, 10); + expect(ctx.fillRect).toHaveBeenCalledWith(70, 10, 10, 10); + expect(ctx.fillRect).toHaveBeenCalledWith(130, 10, 10, 10); + expect(ctx.stroke).toHaveBeenCalled(); + }); + + it('closes panel when close button is clicked or close event triggers and stops drawing', () => { + const ctx = statsCanvas.getContext('2d')!; + const clearSpy = spyOn(ctx, 'clearRect').and.callThrough(); + + statsViewer.showPanel(); + clearSpy.calls.reset(); + + statsClose.click(); + expect(statsPanel.close).toHaveBeenCalled(); + statsViewer.handleStats({ack_latency: 10}); + expect(clearSpy).not.toHaveBeenCalled(); + + statsViewer.showPanel(); + clearSpy.calls.reset(); + + statsPanel.dispatchEvent(new Event('close')); + statsViewer.handleStats({ack_latency: 20}); + expect(clearSpy).not.toHaveBeenCalled(); + }); + + it('handles empty or partial stats objects and single data point chart rendering', () => { + statsToggle.click(); + + // Partial stats object + statsViewer.handleStats({ + ack_latency: 10.0, + }); + + expect(statAck.innerText).toBe('10.00'); + expect(statStep.innerText).toBe('-'); + + // Empty stats object + statsViewer.handleStats({}); + }); + + it('updates phone / device metrics when receiving battery stats data', () => { + const internalViewer = statsViewer as unknown as InternalStatsViewer; + expect(internalViewer.batteryCtx).not.toBeNull(); + + statsViewer.handleStats({ + battery_power: { + level_percent: 85, + status: 'CHARGING', + plugged_source: 'AC', + voltage_mv: 4200, + current_now_ua: 500000, + power_mw: 2100.0, + temperature_celsius: 28.5, + charge_counter_uah: 3000000, + energy_counter_nwh: 11100000000, + }, + }); + + expect( + (statsViewer.querySelector('#stat-battery-level') as HTMLElement).innerText + ).toBe('85'); + expect( + (statsViewer.querySelector('#stat-battery-status') as HTMLElement).innerText + ).toBe('CHARGING'); + expect( + (statsViewer.querySelector('#stat-battery-source') as HTMLElement).innerText + ).toBe('AC'); + expect( + (statsViewer.querySelector('#stat-battery-voltage') as HTMLElement).innerText + ).toBe('4200'); + expect( + (statsViewer.querySelector('#stat-battery-current') as HTMLElement).innerText + ).toBe('500000'); + expect( + (statsViewer.querySelector('#stat-battery-power') as HTMLElement).innerText + ).toBe('2100.00'); + expect( + (statsViewer.querySelector('#stat-battery-temp') as HTMLElement).innerText + ).toBe('28.5'); + expect( + (statsViewer.querySelector('#stat-battery-counter') as HTMLElement).innerText + ).toBe('3000000'); + expect( + (statsViewer.querySelector('#stat-battery-energy') as HTMLElement).innerText + ).toBe('11100000000'); + }); + + it('handles zero current_now_ua correctly', () => { + statsViewer.handleStats({ + battery_power: { + current_now_ua: 0, + }, + }); + + expect( + (statsViewer.querySelector('#stat-battery-current') as HTMLElement).innerText + ).toBe('0'); + const internalViewer = statsViewer as unknown as InternalStatsViewer; + expect(internalViewer.statsHistory.current).toContain(0); + }); + + it('positions phone / device metrics below stats-chart in DOM order', () => { + const chartElem = statsViewer.querySelector('#stats-chart') as HTMLElement; + const batteryElem = statsViewer.querySelector( + '#stat-battery-level' + ) as HTMLElement; + const isAfter = + (chartElem.compareDocumentPosition(batteryElem) & + Node.DOCUMENT_POSITION_FOLLOWING) !== + 0; + expect(isAfter).toBeTrue(); + }); + + it('renders battery charge and current line diagram on stats-battery-chart', () => { + const internalViewer = statsViewer as unknown as InternalStatsViewer; + const batteryCtx = batteryCanvas.getContext('2d')!; + spyOn(batteryCtx, 'clearRect').and.callThrough(); + spyOn(batteryCtx, 'fillText').and.callThrough(); + spyOn(batteryCtx, 'fillRect').and.callThrough(); + spyOn(batteryCtx, 'stroke').and.callThrough(); + const drawBatterySpy = spyOn( + internalViewer, + 'drawBatteryChart' + ).and.callThrough(); + statsToggle.click(); + expect(statsPanel.show).toHaveBeenCalled(); + expect(drawBatterySpy).toHaveBeenCalled(); + + for (let i = 0; i < 5; i++) { + statsViewer.handleStats({ + battery_power: { + level_percent: 80 + i, + current_now_ua: i % 2 === 0 ? 500000 : -200000, + }, + }); + } + + expect( + (statsViewer.querySelector('#stat-battery-level') as HTMLElement).innerText + ).toBe('84'); + expect( + (statsViewer.querySelector('#stat-battery-current') as HTMLElement).innerText + ).toBe('500000'); + expect(batteryCanvas).toBeTruthy(); + expect(internalViewer.statsHistory.charge).toEqual([ + 80, 81, 82, 83, 84, + ]); + expect(batteryCtx.clearRect).toHaveBeenCalledWith(0, 0, 400, 150); + expect(batteryCtx.fillText).toHaveBeenCalledWith('Charge (%)', 25, 18); + expect(batteryCtx.fillText).toHaveBeenCalledWith( + 'Current (\u03bcA)', + 115, + 18 + ); + expect(batteryCtx.fillRect).toHaveBeenCalledWith(10, 10, 10, 10); + expect(batteryCtx.fillRect).toHaveBeenCalledWith(100, 10, 10, 10); + expect(batteryCtx.stroke).toHaveBeenCalled(); + }); + + it('renders battery chart correctly when current values are uniform', () => { + statsToggle.click(); + expect(statsPanel.show).toHaveBeenCalled(); + + for (let i = 0; i < 3; i++) { + statsViewer.handleStats({ + battery_power: { + level_percent: 50, + current_now_ua: 300000, + }, + }); + } + + expect( + (statsViewer.querySelector('#stat-battery-level') as HTMLElement).innerText + ).toBe('50'); + expect( + (statsViewer.querySelector('#stat-battery-current') as HTMLElement).innerText + ).toBe('300000'); + const internalViewer = statsViewer as unknown as InternalStatsViewer; + expect(internalViewer.statsHistory.current).toEqual([ + 300000, 300000, 300000, + ]); + }); + + it('updates episode and KPI badge cards when present in DOM', () => { + document.body.innerHTML += ` +
+
+
+
+
+ `; + + statsViewer.handleStats({ + ack_latency: 15.5, + env_step_ms: 22.3, + episode: { + step: 42, + step_type: 'MID', + reward: 1.0, + cumulative_reward: 10.5, + active_package: 'com.android.settings', + }, + battery_power: { + level_percent: 85, + power_mw: 1200, + }, + }); + + expect(document.getElementById('kpi-step')!.textContent).toBe('42 (MID)'); + expect(document.getElementById('kpi-reward')!.textContent).toBe('10.5 (+1.0)'); + expect(document.getElementById('pkg-badge')!.textContent).toBe('pkg: settings'); + expect(document.getElementById('kpi-latency')!.textContent).toBe('15.5 / 22.3 ms'); + expect(document.getElementById('kpi-battery')!.textContent).toBe('85% · 1200mW'); + }); +}); + + diff --git a/android_env/playground/frontend/task_loader.ts b/android_env/playground/frontend/task_loader.ts new file mode 100644 index 00000000..83706368 --- /dev/null +++ b/android_env/playground/frontend/task_loader.ts @@ -0,0 +1,91 @@ +import {SocketHandler} from './socket'; + +/** + * Custom Web Component for loading tasks in AndroidEnv Playground. + */ +export class TaskLoaderElement extends HTMLElement { + private taskInput: HTMLInputElement | null = null; + private loadTaskBtn: HTMLButtonElement | null = null; + private socket: SocketHandler | null = null; + + constructor() { + super(); + } + + /** @export */ + connectedCallback(): void { + if (!this.taskInput) { + const container = document.createElement('div'); + container.style.cssText = 'display: flex; gap: 10px; align-items: center;'; + + const input = document.createElement('input'); + input.type = 'text'; + input.id = 'task-name'; + input.placeholder = 'Task ID (e.g. classic_2048)'; + input.style.cssText = + 'background: var(--bg-input, #12151b); color: var(--text-main, #e6edf3); border: 1px solid var(--border, #2a2f3a); padding: 6px 10px; border-radius: 6px; font-size: 12px; font-family: var(--font-sans); outline: none;'; + + const button = document.createElement('button'); + button.id = 'load-task-btn'; + button.className = 'btn btn-primary'; + button.textContent = 'Load Task'; + button.style.cssText = + 'padding: 6px 11px; border-radius: 6px; font-size: 12px; font-weight: 600; cursor: pointer;'; + + container.appendChild(input); + container.appendChild(button); + this.appendChild(container); + + this.taskInput = input; + this.loadTaskBtn = button; + this.initEventListeners(); + } + } + + /** @export */ + setSocket(socket: SocketHandler): void { + this.socket = socket; + } + + private initEventListeners(): void { + this.loadTaskBtn!.addEventListener('click', () => { + this.triggerLoadTask(); + }); + + this.taskInput!.addEventListener('keydown', (e: KeyboardEvent) => { + if (e.key === 'Enter') { + this.triggerLoadTask(); + } + }); + } + + /** @export */ + triggerLoadTask(): void { + const taskId = this.taskInput!.value.trim(); + if (taskId) { + this.socket?.loadTask(taskId); + this.dispatchEvent( + new CustomEvent('load-task', { + bubbles: true, + composed: true, + detail: {taskId}, + }) + ); + } + } +} + +if (!customElements.get('task-loader')) { + customElements.define('task-loader', TaskLoaderElement); +} + +/** + * Helper factory function to instantiate and upgrade a TaskLoaderElement. + */ +export function createTaskLoader(socket?: SocketHandler): TaskLoaderElement { + const el = document.createElement('task-loader') as TaskLoaderElement; + if (socket) { + el.setSocket(socket); + } + return el; +} diff --git a/android_env/playground/frontend/task_loader_test.ts b/android_env/playground/frontend/task_loader_test.ts new file mode 100644 index 00000000..15743af3 --- /dev/null +++ b/android_env/playground/frontend/task_loader_test.ts @@ -0,0 +1,93 @@ +import 'jasmine'; +import {SocketHandler} from './socket'; +import {createTaskLoader, TaskLoaderElement} from './task_loader'; + +describe('TaskLoaderElement', () => { + let taskInput: HTMLInputElement; + let loadTaskBtn: HTMLButtonElement; + let socketMock: jasmine.SpyObj; + let element: TaskLoaderElement; + + beforeEach(() => { + socketMock = jasmine.createSpyObj('SocketHandler', [ + 'loadTask', + ]); + + element = createTaskLoader(socketMock); + element.id = 'task-loader'; + document.body.appendChild(element); + + taskInput = element.querySelector('#task-name') as HTMLInputElement; + loadTaskBtn = element.querySelector('#load-task-btn') as HTMLButtonElement; + }); + + afterEach(() => { + document.body.innerHTML = ''; + }); + + it('initializes custom element successfully with expected DOM and styles', () => { + expect(element).toBeTruthy(); + expect(element.tagName.toLowerCase()).toBe('task-loader'); + expect(taskInput).toBeTruthy(); + expect(taskInput.type).toBe('text'); + expect(taskInput.placeholder).toBe('Task ID (e.g. classic_2048)'); + expect(taskInput.style.padding).toBe('6px 10px'); + expect(taskInput.style.background).toContain('var(--bg-input'); + expect(loadTaskBtn).toBeTruthy(); + expect(loadTaskBtn.className).toBe('btn btn-primary'); + expect(loadTaskBtn.textContent).toBe('Load Task'); + expect(loadTaskBtn.style.fontWeight).toBe('600'); + const container = element.firstElementChild as HTMLElement; + expect(container.style.display).toBe('flex'); + + // Calling connectedCallback again is idempotent + element.connectedCallback(); + expect(element.children.length).toBe(1); + }); + + it('triggers socket.loadTask on button click with trimmed text', () => { + taskInput.value = ' classic_2048 '; + loadTaskBtn.click(); + + expect(socketMock.loadTask).toHaveBeenCalledWith('classic_2048'); + }); + + it('triggers socket.loadTask on Enter keydown in input', () => { + taskInput.value = 'snake'; + const event = new KeyboardEvent('keydown', {key: 'Enter'}); + taskInput.dispatchEvent(event); + + expect(socketMock.loadTask).toHaveBeenCalledWith('snake'); + }); + + it('does not trigger socket.loadTask on non-Enter keydown in input', () => { + taskInput.value = 'snake'; + const event = new KeyboardEvent('keydown', {key: 'Escape'}); + taskInput.dispatchEvent(event); + + expect(socketMock.loadTask).not.toHaveBeenCalled(); + }); + + it('does not trigger socket.loadTask if input is empty', () => { + taskInput.value = ' '; + loadTaskBtn.click(); + + expect(socketMock.loadTask).not.toHaveBeenCalled(); + }); + + it('dispatches load-task event when task is loaded', () => { + let capturedEvent: CustomEvent<{taskId: string}> | null = null; + element.addEventListener('load-task', (e: Event) => { + capturedEvent = e as CustomEvent<{taskId: string}>; + }); + + taskInput.value = 'classic_2048'; + loadTaskBtn.click(); + + expect(capturedEvent).not.toBeNull(); + expect(capturedEvent!.detail.taskId).toBe('classic_2048'); + expect(capturedEvent!.bubbles).toBeTrue(); + expect(capturedEvent!.composed).toBeTrue(); + }); +}); + diff --git a/android_env/playground/frontend/theme.ts b/android_env/playground/frontend/theme.ts new file mode 100644 index 00000000..093736b8 --- /dev/null +++ b/android_env/playground/frontend/theme.ts @@ -0,0 +1,101 @@ +/** + * Strongly-typed theme management for AndroidEnv Playground. + */ + +/** + * Theme mode supported by the Playground. + */ +export type Theme = 'light' | 'dark'; + +/** + * Key used to persist theme choice in localStorage. + */ +export const THEME_STORAGE_KEY = 'aenv_theme'; + +/** + * Returns the current theme from localStorage, defaulting to 'light'. + */ +export function getTheme(): Theme { + try { + const saved = window.localStorage.getItem(THEME_STORAGE_KEY); + return saved === 'dark' ? 'dark' : 'light'; + } catch (_) { + return 'light'; + } +} + +/** + * Applies the given theme to the document and updates theme toggle buttons. + */ +export function applyTheme(theme: Theme): Theme { + const validTheme: Theme = theme === 'dark' ? 'dark' : 'light'; + document.documentElement.setAttribute('data-theme', validTheme); + if (document.body) { + document.body.setAttribute('data-theme', validTheme); + } + try { + window.localStorage.setItem(THEME_STORAGE_KEY, validTheme); + } catch (_) {} + + document.querySelectorAll('.theme-toggle-btn').forEach((btn) => { + btn.textContent = validTheme === 'light' ? '🌙 Dark' : '☀️ Light'; + btn.title = `Switch to ${validTheme === 'light' ? 'Dark' : 'Light'} theme (Ctrl+Shift+D)`; + btn.setAttribute( + 'aria-label', + `Switch to ${validTheme === 'light' ? 'Dark' : 'Light'} theme` + ); + }); + + // NOMUTANTS(SBR): Tested in frontend_test_chrome-linux (theme_test.ts). + window.dispatchEvent( + new CustomEvent('aenv-theme-changed', {detail: {theme: validTheme}}) + ); + return validTheme; +} + +/** + * Toggles between light and dark theme. + */ +export function toggleTheme(): Theme { + const current = getTheme(); + const next: Theme = current === 'light' ? 'dark' : 'light'; + return applyTheme(next); +} + +/** + * Handles global keyboard shortcuts for theme toggling (Ctrl+Shift+D / Cmd+Shift+D). + */ +export function handleThemeKeydown(e: KeyboardEvent): void { + if ( + (e.ctrlKey || e.metaKey) && + e.shiftKey && + (e.key === 'D' || e.key === 'd' || e.key === 'T' || e.key === 't') + ) { + e.preventDefault(); + toggleTheme(); + } +} + +/** + * Click handler for theme toggle buttons. + */ +export function handleThemeToggleClick(): void { + toggleTheme(); +} + +/** + * Sets up click and keyboard listeners for theme toggling. + */ +export function setupThemeListeners(): void { + document.querySelectorAll('.theme-toggle-btn').forEach((btn) => { + btn.removeEventListener('click', handleThemeToggleClick); + btn.addEventListener('click', handleThemeToggleClick); + }); + + window.removeEventListener('keydown', handleThemeKeydown); + window.addEventListener('keydown', handleThemeKeydown); + + // NOMUTANTS(SBR): Tested in frontend_test_chrome-linux (theme_test.ts). + applyTheme(getTheme()); +} + diff --git a/android_env/playground/frontend/theme_test.ts b/android_env/playground/frontend/theme_test.ts new file mode 100644 index 00000000..1f4cd61a --- /dev/null +++ b/android_env/playground/frontend/theme_test.ts @@ -0,0 +1,207 @@ +import 'jasmine'; +import { + applyTheme, + getTheme, + handleThemeKeydown, + handleThemeToggleClick, + setupThemeListeners, + THEME_STORAGE_KEY, + toggleTheme, +} from './theme'; + +describe('theme', () => { + beforeEach(() => { + window.localStorage.clear(); + document.documentElement.removeAttribute('data-theme'); + document.body.innerHTML = ` + + `; + }); + + afterEach(() => { + window.localStorage.clear(); + document.documentElement.removeAttribute('data-theme'); + document.body.innerHTML = ''; + }); + + it('defaults to light theme when nothing is in localStorage', () => { + expect(getTheme()).toBe('light'); + }); + + it('retrieves saved dark theme from localStorage', () => { + window.localStorage.setItem(THEME_STORAGE_KEY, 'dark'); + expect(getTheme()).toBe('dark'); + }); + + it('falls back to light theme if localStorage throws', () => { + spyOn(window.localStorage, 'getItem').and.throwError('SecurityError'); + expect(getTheme()).toBe('light'); + }); + + it('applies theme and updates DOM attributes and toggle buttons', () => { + let receivedTheme: string | null = null; + const listener = (e: Event) => { + receivedTheme = (e as CustomEvent<{theme: string}>).detail.theme; + }; + window.addEventListener('aenv-theme-changed', listener); + + try { + const applied = applyTheme('dark'); + expect(applied).toBe('dark'); + expect(receivedTheme as string | null).toBe('dark'); + expect(document.documentElement.getAttribute('data-theme')).toBe('dark'); + expect(document.body.getAttribute('data-theme')).toBe('dark'); + expect(window.localStorage.getItem(THEME_STORAGE_KEY)).toBe('dark'); + + const btn = document.querySelector('.theme-toggle-btn') as HTMLElement; + expect(btn.textContent).toBe('☀️ Light'); + expect(btn.title).toBe('Switch to Light theme (Ctrl+Shift+D)'); + expect(btn.getAttribute('aria-label')).toBe('Switch to Light theme'); + + const appliedLight = applyTheme('light'); + expect(appliedLight).toBe('light'); + expect(receivedTheme as string | null).toBe('light'); + expect(document.documentElement.getAttribute('data-theme')).toBe('light'); + expect(btn.textContent).toBe('🌙 Dark'); + expect(btn.title).toBe('Switch to Dark theme (Ctrl+Shift+D)'); + expect(btn.getAttribute('aria-label')).toBe('Switch to Dark theme'); + } finally { + window.removeEventListener('aenv-theme-changed', listener); + } + }); + + it('does not throw when localStorage.setItem throws during applyTheme', () => { + spyOn(window.localStorage, 'setItem').and.throwError('QuotaExceeded'); + expect(() => applyTheme('dark')).not.toThrow(); + expect(document.documentElement.getAttribute('data-theme')).toBe('dark'); + }); + + it('toggles theme between light and dark', () => { + applyTheme('light'); + expect(getTheme()).toBe('light'); + + const next = toggleTheme(); + expect(next).toBe('dark'); + expect(getTheme()).toBe('dark'); + + const nextAgain = toggleTheme(); + expect(nextAgain).toBe('light'); + expect(getTheme()).toBe('light'); + }); + + it('handles keyboard shortcuts in handleThemeKeydown', () => { + applyTheme('light'); + + // Ctrl+Shift+D triggers toggle to dark + const eventCtrlD = new KeyboardEvent('keydown', { + key: 'D', + ctrlKey: true, + shiftKey: true, + cancelable: true, + }); + handleThemeKeydown(eventCtrlD); + expect(eventCtrlD.defaultPrevented).toBeTrue(); + expect(getTheme()).toBe('dark'); + + // Meta+Shift+T triggers toggle to light + const eventMetaT = new KeyboardEvent('keydown', { + key: 'T', + metaKey: true, + shiftKey: true, + cancelable: true, + }); + handleThemeKeydown(eventMetaT); + expect(eventMetaT.defaultPrevented).toBeTrue(); + expect(getTheme()).toBe('light'); + + // Lowercase 'd' with Ctrl+Shift + const eventLowerD = new KeyboardEvent('keydown', { + key: 'd', + ctrlKey: true, + shiftKey: true, + cancelable: true, + }); + handleThemeKeydown(eventLowerD); + expect(eventLowerD.defaultPrevented).toBeTrue(); + expect(getTheme()).toBe('dark'); + + // Negative case: key D without Ctrl/Meta does not toggle + const eventNoMod = new KeyboardEvent('keydown', { + key: 'D', + ctrlKey: false, + metaKey: false, + shiftKey: true, + cancelable: true, + }); + handleThemeKeydown(eventNoMod); + expect(eventNoMod.defaultPrevented).toBeFalse(); + expect(getTheme()).toBe('dark'); + + // Negative case: Ctrl+D without Shift does not toggle + const eventNoShift = new KeyboardEvent('keydown', { + key: 'D', + ctrlKey: true, + shiftKey: false, + cancelable: true, + }); + handleThemeKeydown(eventNoShift); + expect(eventNoShift.defaultPrevented).toBeFalse(); + expect(getTheme()).toBe('dark'); + + // Negative case: Unrelated key (Ctrl+Shift+X) does not toggle + const eventOtherKey = new KeyboardEvent('keydown', { + key: 'X', + ctrlKey: true, + shiftKey: true, + cancelable: true, + }); + handleThemeKeydown(eventOtherKey); + expect(eventOtherKey.defaultPrevented).toBeFalse(); + expect(getTheme()).toBe('dark'); + }); + + it('handles click events via handleThemeToggleClick', () => { + applyTheme('light'); + handleThemeToggleClick(); + expect(getTheme()).toBe('dark'); + handleThemeToggleClick(); + expect(getTheme()).toBe('light'); + }); + + it('binds click and window keydown listeners in setupThemeListeners idempotently', () => { + // Calling setupThemeListeners twice should not attach duplicate listeners. + setupThemeListeners(); + setupThemeListeners(); + const btn = document.querySelector('.theme-toggle-btn') as HTMLElement; + + // A single click should toggle once to dark (not toggle twice back to light) + btn.click(); + expect(getTheme()).toBe('dark'); + + const keyEvent = new KeyboardEvent('keydown', { + key: 'D', + ctrlKey: true, + shiftKey: true, + cancelable: true, + }); + window.dispatchEvent(keyEvent); + expect(keyEvent.defaultPrevented).toBeTrue(); + expect(getTheme()).toBe('light'); + + spyOn(window, 'removeEventListener').and.callThrough(); + setupThemeListeners(); + expect(window.removeEventListener).toHaveBeenCalledWith( + 'keydown', + handleThemeKeydown + ); + }); + + it('applies saved theme during setupThemeListeners', () => { + window.localStorage.setItem(THEME_STORAGE_KEY, 'dark'); + expect(document.documentElement.getAttribute('data-theme')).toBeNull(); + + setupThemeListeners(); + + expect(document.documentElement.getAttribute('data-theme')).toBe('dark'); + }); +}); diff --git a/android_env/playground/frontend/tsconfig.json b/android_env/playground/frontend/tsconfig.json new file mode 100644 index 00000000..63743617 --- /dev/null +++ b/android_env/playground/frontend/tsconfig.json @@ -0,0 +1,18 @@ +{ + "compilerOptions": { + "target": "ES2022", + "module": "ES2022", + "moduleResolution": "node", + "strict": true, + "declaration": true, + "sourceMap": true, + "lib": ["ES2022", "DOM"], + "skipLibCheck": true + }, + "include": [ + "**/*.ts" + ], + "exclude": [ + "**/*_test.ts" + ] +} diff --git a/android_env/playground/log_handler.py b/android_env/playground/log_handler.py new file mode 100644 index 00000000..d02fc51c --- /dev/null +++ b/android_env/playground/log_handler.py @@ -0,0 +1,66 @@ +# coding=utf-8 +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Logging handler for AndroidEnv Playground.""" + +import logging +import threading + +from android_env.playground import state as state_lib +from android_env.playground import websocket + + +class _LogState(threading.local): + active: bool = False + + +class PlaygroundLogHandler(logging.Handler): + """Custom logging handler to store logs in PlaygroundState and stream via WS.""" + + def __init__(self, state: state_lib.PlaygroundState): + super().__init__() + self.state = state + self._is_logging = _LogState() + + def emit(self, record: logging.LogRecord): + if self._is_logging.active: + return + # Avoid feedback loops and lock contention from internal networking logs. + if getattr(record, 'filename', '') in ('websocket.py', 'server.py'): + return + self._is_logging.active = True + try: + try: + msg = self.format(record) + except ( + TypeError, + ValueError, + AttributeError, + KeyError, + RuntimeError, + IndexError, + ): + self.handleError(record) + return + + log_entry = { + 'message': str(msg), + 'level': str(record.levelname), + 'timestamp': record.created, + } + self.state.add_log(log_entry) + websocket.try_send_logs(self.state, [log_entry]) + finally: + self._is_logging.active = False diff --git a/android_env/playground/log_handler_test.py b/android_env/playground/log_handler_test.py new file mode 100644 index 00000000..068df2fa --- /dev/null +++ b/android_env/playground/log_handler_test.py @@ -0,0 +1,100 @@ +# coding=utf-8 +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the custom PlaygroundLogHandler log streaming component.""" + +import logging +import socket +from unittest import mock + +from absl.testing import absltest +from android_env.playground import log_handler +from android_env.playground import state as state_lib + + +class LogHandlerTest(absltest.TestCase): + + def test_playground_log_handler(self): + state = state_lib.PlaygroundState() + mock_sock = mock.create_autospec(socket.socket, instance=True) + state.active_ws = mock_sock + + handler = log_handler.PlaygroundLogHandler(state) + logger = logging.getLogger('test_logger') + logger.addHandler(handler) + logger.setLevel(logging.INFO) + + logger.info('Test log entry') + + logs = state.get_logs() + self.assertLen(logs, 1) + self.assertEqual(logs[0]['message'], 'Test log entry') + self.assertEqual(logs[0]['level'], 'INFO') + + mock_sock.sendall.assert_called_once() + logger.removeHandler(handler) + + def test_playground_log_handler_prevents_reentrant_emit(self): + state = state_lib.PlaygroundState() + handler = log_handler.PlaygroundLogHandler(state) + record = logging.makeLogRecord({'msg': 'First log', 'levelname': 'INFO'}) + + def recursive_send_logs(unused_state, unused_logs): + handler.emit(record) + + with mock.patch.object( + log_handler.websocket, 'try_send_logs', side_effect=recursive_send_logs + ): + handler.emit(record) + + self.assertLen(state.get_logs(), 1) + + def test_playground_log_handler_format_error(self): + state = state_lib.PlaygroundState() + handler = log_handler.PlaygroundLogHandler(state) + bad_record = logging.makeLogRecord( + {'msg': 'Bad format %d', 'args': ('not_an_int',), 'levelname': 'INFO'} + ) + with mock.patch.object(handler, 'handleError', autospec=True) as mock_err: + handler.emit(bad_record) + mock_err.assert_called_once_with(bad_record) + self.assertEmpty(state.get_logs()) + + good_record = logging.makeLogRecord( + {'msg': 'Recovered', 'levelname': 'INFO'} + ) + handler.emit(good_record) + self.assertLen(state.get_logs(), 1) + + def test_playground_log_handler_skips_internal_websocket_logs(self): + state = state_lib.PlaygroundState() + mock_sock = mock.create_autospec(socket.socket, instance=True) + state.active_ws = mock_sock + + handler = log_handler.PlaygroundLogHandler(state) + ws_record = logging.makeLogRecord({ + 'msg': 'ACK timeout, resetting waiting_for_ack', + 'levelname': 'WARNING', + 'pathname': '/path/to/websocket.py', + }) + ws_record.filename = 'websocket.py' + handler.emit(ws_record) + + mock_sock.sendall.assert_not_called() + self.assertEmpty(state.get_logs()) + + +if __name__ == '__main__': + absltest.main() diff --git a/android_env/playground/playground.py b/android_env/playground/playground.py new file mode 100644 index 00000000..4e3d2350 --- /dev/null +++ b/android_env/playground/playground.py @@ -0,0 +1,139 @@ +# coding=utf-8 +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Playground for AndroidEnv, allowing human interaction via a web browser.""" + +import logging as std_logging +import os +import threading + +from absl import app +from absl import flags +from absl import logging +from android_env.playground import device_runner +from android_env.playground import log_handler +from android_env.playground import server +from android_env.playground import state as state_lib +from android_env.playground import websocket +from ml_collections.config_flags import config_flags + +_PORT = flags.DEFINE_integer('port', 8080, 'Port to serve the web interface.') +_CONFIG = config_flags.DEFINE_config_file( + 'config', + None, + 'File path to the AndroidEnv config.', + lock_config=False, +) + + +def _run_device_thread(state: state_lib.PlaygroundState, config): + """Orchestrates device setup, running, and cleanup in background.""" + try: + device, orientation, initial_pixels = device_runner.setup_device( + state, config + ) + except device_runner.DeviceSetupError as e: + logging.exception('Error in device setup: %s', e) + return + + if not state.running: + logging.warning( + 'Server is not running; closing device without starting run loop.' + ) + device.close() + return + + try: + device_runner.run_device(device, orientation, initial_pixels, state) + except (RuntimeError, ValueError, KeyError, IndexError, OSError) as e: + logging.exception('Error in device execution: %s', e) + websocket.try_send_status(state, f'ERROR: {e}') + finally: + logging.info('Closing device...') + device.close() + + +def _serve_http(state: state_lib.PlaygroundState, static_dir: str, port: int): + """Starts the HTTP/WebSocket server and blocks until interrupted.""" + server_address = ('', port) + server.PlaygroundServer.allow_reuse_address = True + server.PlaygroundServer.daemon_threads = True + with server.PlaygroundServer( + server_address, + server.StreamHandler, + state, + os.path.expanduser(static_dir), + ) as httpd: + logging.info('Serving static files from %s...', static_dir) + logging.info( + 'Serving on port %s (http://localhost:%s/)...', + port, + port, + ) + try: + httpd.serve_forever() + except KeyboardInterrupt: + logging.info('KeyboardInterrupt caught, setting running=False') + state.running = False + finally: + logging.info('Shutting down server...') + + +def main(_): + static_dir = server.get_default_static_dir() + state = state_lib.PlaygroundState() + + formatter = std_logging.Formatter( + '%(asctime)s [%(levelname)s] %(name)s: %(message)s' + ) + custom_log_handler = log_handler.PlaygroundLogHandler(state) + custom_log_handler.setFormatter(formatter) + stderr_log_handler = std_logging.StreamHandler() + stderr_log_handler.setLevel(std_logging.INFO) + stderr_log_handler.setFormatter(formatter) + + absl_logger = logging.get_absl_logger() + absl_logger.addHandler(custom_log_handler) + absl_logger.addHandler(stderr_log_handler) + + # Start setup in background (must be daemon to avoid hanging on server + # bind failure). + setup_thread = threading.Thread( + target=_run_device_thread, + args=( + state, + _CONFIG.value, + ), + daemon=True, + ) + setup_thread.start() + try: + _serve_http(state, static_dir, _PORT.value) + except OSError as e: + logging.exception('Server failed to start on port %s: %s', _PORT.value, e) + finally: + logging.info('Shutting down agent...') + state.running = False + # Wait for setup thread to exit, but not forever (in case it is blocked + # in boot). + setup_thread.join(timeout=5.0) + absl_logger.removeHandler(stderr_log_handler) + absl_logger.removeHandler(custom_log_handler) + stderr_log_handler.close() + custom_log_handler.close() + + +if __name__ == '__main__': + app.run(main) diff --git a/android_env/playground/playground_test.py b/android_env/playground/playground_test.py new file mode 100644 index 00000000..7a282632 --- /dev/null +++ b/android_env/playground/playground_test.py @@ -0,0 +1,167 @@ +# coding=utf-8 +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for playground entrypoint and thread orchestration.""" + +import runpy +import sys +from unittest import mock + +from absl.testing import absltest +from android_env.components import android_device +from android_env.playground import playground +from android_env.playground import state as state_lib +import numpy as np + + +class PlaygroundTest(absltest.TestCase): + + @mock.patch.object(playground.device_runner, 'run_device', autospec=True) + @mock.patch.object(playground.device_runner, 'setup_device', autospec=True) + def test_run_device_thread_success_closes_device( + self, mock_setup_device, mock_run_device + ): + state = state_lib.PlaygroundState() + mock_dev = mock.create_autospec(android_device.AndroidDevice, instance=True) + pixels = np.zeros((10, 10, 3), dtype=np.uint8) + mock_setup_device.return_value = (mock_dev, 0, pixels) + + playground._run_device_thread(state, 'dummy_config') + + mock_setup_device.assert_called_once_with(state, 'dummy_config') + mock_run_device.assert_called_once_with(mock_dev, 0, pixels, state) + mock_dev.close.assert_called_once() + + @mock.patch.object(playground.device_runner, 'run_device', autospec=True) + @mock.patch.object(playground.device_runner, 'setup_device', autospec=True) + def test_run_device_thread_setup_error_returns_early( + self, mock_setup_device, mock_run_device + ): + state = state_lib.PlaygroundState() + mock_setup_device.side_effect = playground.device_runner.DeviceSetupError( + 'bad orientation' + ) + + playground._run_device_thread(state, 'dummy_config') + + mock_run_device.assert_not_called() + + @mock.patch.object(playground.websocket, 'try_send_status', autospec=True) + @mock.patch.object(playground.device_runner, 'run_device', autospec=True) + @mock.patch.object(playground.device_runner, 'setup_device', autospec=True) + def test_run_device_thread_run_error_sends_status_and_closes_device( + self, mock_setup_device, mock_run_device, mock_send_status + ): + state = state_lib.PlaygroundState() + mock_dev = mock.create_autospec(android_device.AndroidDevice, instance=True) + pixels = np.zeros((10, 10, 3), dtype=np.uint8) + mock_setup_device.return_value = (mock_dev, 0, pixels) + mock_run_device.side_effect = RuntimeError('step failed') + + playground._run_device_thread(state, 'dummy_config') + + mock_send_status.assert_called_once_with(state, 'ERROR: step failed') + mock_dev.close.assert_called_once() + + @mock.patch.object(playground.device_runner, 'run_device', autospec=True) + @mock.patch.object(playground.device_runner, 'setup_device', autospec=True) + def test_run_device_thread_server_not_running_closes_device_without_running( + self, mock_setup_device, mock_run_device + ): + state = state_lib.PlaygroundState() + state.running = False + mock_dev = mock.create_autospec(android_device.AndroidDevice, instance=True) + pixels = np.zeros((10, 10, 3), dtype=np.uint8) + mock_setup_device.return_value = (mock_dev, 0, pixels) + + playground._run_device_thread(state, 'dummy_config') + + mock_setup_device.assert_called_once_with(state, 'dummy_config') + mock_run_device.assert_not_called() + mock_dev.close.assert_called_once() + + @mock.patch.object(playground.server, 'PlaygroundServer', autospec=True) + def test_serve_http_handles_keyboard_interrupt(self, mock_server_cls): + state = state_lib.PlaygroundState() + mock_httpd = mock_server_cls.return_value.__enter__.return_value + mock_httpd.serve_forever.side_effect = KeyboardInterrupt() + + playground._serve_http(state, '/fake/static', 8080) + + self.assertTrue(mock_server_cls.allow_reuse_address) + self.assertTrue(mock_server_cls.daemon_threads) + mock_server_cls.assert_called_once_with( + ('', 8080), + playground.server.StreamHandler, + state, + '/fake/static', + ) + self.assertFalse(state.running) + mock_httpd.serve_forever.assert_called_once() + + @mock.patch.object(playground, '_serve_http', autospec=True) + @mock.patch.object( + playground.server, + 'get_default_static_dir', + autospec=True, + return_value='/fake/static', + ) + @mock.patch.object(playground.threading, 'Thread', autospec=True) + def test_main_starts_thread_and_serves_http( + self, mock_thread_cls, unused_mock_static_dir, mock_serve_http + ): + mock_serve_http.side_effect = OSError('bind failed') + absl_logger = playground.logging.get_absl_logger() + initial_handlers = list(absl_logger.handlers) + + playground.main(['playground']) + + self.assertEqual(list(absl_logger.handlers), initial_handlers) + mock_thread_cls.assert_called_once() + self.assertTrue(mock_thread_cls.call_args.kwargs.get('daemon')) + mock_thread_cls.return_value.start.assert_called_once() + mock_serve_http.assert_called_once_with(mock.ANY, '/fake/static', 8080) + mock_thread_cls.return_value.join.assert_called_once_with(timeout=5.0) + + def test_main_execution(self): + original_setitem = playground.flags.FlagValues.__setitem__ + + def _safe_setitem( + fv_self: playground.flags.FlagValues, + name: str, + flag: playground.flags.Flag, + ) -> None: + try: + original_setitem(fv_self, name, flag) + except playground.flags.DuplicateFlagError: + pass + + with mock.patch.object( + playground.flags.FlagValues, '__setitem__', _safe_setitem + ): + with mock.patch.object( + playground.app, 'run', autospec=True + ) as mock_app_run: + with mock.patch.object(sys, 'argv', ['playground.py']): + runpy.run_module( + 'android_env.playground.playground', + run_name='__main__', + ) + mock_app_run.assert_called_once() + self.assertEqual(mock_app_run.call_args[0][0].__name__, 'main') + + +if __name__ == '__main__': + absltest.main() diff --git a/android_env/playground/server.py b/android_env/playground/server.py new file mode 100644 index 00000000..6748a9c6 --- /dev/null +++ b/android_env/playground/server.py @@ -0,0 +1,101 @@ +# coding=utf-8 +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""HTTP and WebSocket server for AndroidEnv Playground.""" + +import http.server +import os +import socketserver +import urllib.parse + +from absl import logging +from android_env.playground import state as state_lib +from android_env.playground import websocket as websocket_lib + +_FRONTEND_INDEX_RESOURCE = 'android_env/playground/frontend/index.html' + + +def _get_resource_filename(path: str) -> str: + """Resolves a resource filename from the runfiles, if available.""" + return '' + + +def get_default_static_dir() -> str: + """Resolves the default frontend static files directory from runfiles.""" + resource_path = _get_resource_filename(_FRONTEND_INDEX_RESOURCE) + if resource_path and os.path.exists(resource_path): + return os.path.dirname(resource_path) + + for env_var in ('RUNFILES_DIR', 'PYTHON_RUNFILES'): + runfiles_dir = os.environ.get(env_var) + if runfiles_dir: + for rel_path in ( + 'android_env/playground/frontend', + ): + candidate = os.path.join(runfiles_dir, rel_path) + if os.path.exists(os.path.join(candidate, 'index.html')): + return candidate + + # Fallback to package-relative directory for standard pip/wheel installs. + pkg_candidate = os.path.join(os.path.dirname(__file__), 'frontend') + if os.path.exists(os.path.join(pkg_candidate, 'index.html')): + return pkg_candidate + + raise FileNotFoundError( + 'Could not resolve static files directory from runfiles or package.' + ) + + +class PlaygroundServer(socketserver.ThreadingTCPServer): + + """Threading TCP Server that holds playground state.""" + + def __init__( + self, + server_address, + RequestHandlerClass, + playground_state: state_lib.PlaygroundState, + static_dir: str, + ): + super().__init__(server_address, RequestHandlerClass) + self.playground_state = playground_state + self.static_dir = static_dir + + +class StreamHandler(http.server.SimpleHTTPRequestHandler): + """HTTP handler for remote view and WebSockets.""" + + protocol_version = 'HTTP/1.1' + + def __init__(self, request, client_address, server_inst): + # Set the directory before calling super().__init__ because it calls do_GET + self.directory = server_inst.static_dir + super().__init__( + request, client_address, server_inst, directory=server_inst.static_dir + ) + + def do_GET(self): # pylint: disable=invalid-name + """Handles GET requests for HTML and WebSocket handshake.""" + server_inst = self.server + if not isinstance(server_inst, PlaygroundServer): + self.send_error(500, 'Unexpected server type') + return + state = server_inst.playground_state + parsed_url = urllib.parse.urlparse(self.path) + if parsed_url.path == '/ws': + websocket_lib.handle_websocket(self, state) + else: + # Fallback to SimpleHTTPRequestHandler to serve static files + super().do_GET() diff --git a/android_env/playground/server_test.py b/android_env/playground/server_test.py new file mode 100644 index 00000000..55d1077d --- /dev/null +++ b/android_env/playground/server_test.py @@ -0,0 +1,232 @@ +# coding=utf-8 +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for playground server.""" + +import http.server +import os +import socket +import socketserver +from unittest import mock +from absl.testing import absltest +from android_env.playground import server +from android_env.playground import state as state_lib +from android_env.playground import websocket as websocket_lib + + +class ServerTest(absltest.TestCase): + + def setUp(self): + super().setUp() + self.mock_request = mock.create_autospec(socket.socket, instance=True) + self.mock_client_address = ('127.0.0.1', 12345) + self.mock_state = mock.create_autospec( + state_lib.PlaygroundState, instance=True + ) + self.mock_server = mock.create_autospec( + server.PlaygroundServer, instance=True + ) + self.mock_server.playground_state = self.mock_state + self.mock_server.static_dir = '/my/static/dir' + + @mock.patch.object( + socketserver.ThreadingTCPServer, '__init__', return_value=None + ) + def test_server_init(self, mock_super_init): + server_inst = server.PlaygroundServer( + ('127.0.0.1', 8080), + server.StreamHandler, + self.mock_state, + '/my/static/dir', + ) + self.assertEqual(server_inst.playground_state, self.mock_state) + self.assertEqual(server_inst.static_dir, '/my/static/dir') + mock_super_init.assert_called_once_with( + ('127.0.0.1', 8080), server.StreamHandler + ) + + @mock.patch.object( + http.server.SimpleHTTPRequestHandler, '__init__', return_value=None + ) + def test_handler_init(self, mock_super_init): + # Verify that handler init sets directory and calls super init + handler = server.StreamHandler( + self.mock_request, self.mock_client_address, self.mock_server + ) + self.assertEqual(handler.directory, '/my/static/dir') + mock_super_init.assert_called_once_with( + self.mock_request, + self.mock_client_address, + self.mock_server, + directory='/my/static/dir', + ) + + @mock.patch.object( + http.server.SimpleHTTPRequestHandler, '__init__', return_value=None + ) + @mock.patch.object( + http.server.SimpleHTTPRequestHandler, 'do_GET', autospec=True + ) + def test_do_get_fallback_to_static( + self, mock_super_do_get, unused_mock_super_init + ): + handler = server.StreamHandler( + self.mock_request, self.mock_client_address, self.mock_server + ) + handler.server = self.mock_server + handler.path = '/index.html' + + handler.do_GET() + + mock_super_do_get.assert_called_once_with(handler) + + @mock.patch.object( + http.server.SimpleHTTPRequestHandler, '__init__', return_value=None + ) + @mock.patch.object(websocket_lib, 'handle_websocket', autospec=True) + @mock.patch.object( + http.server.SimpleHTTPRequestHandler, 'do_GET', autospec=True + ) + def test_do_get_ws( + self, mock_super_do_get, mock_handle_ws, unused_mock_super_init + ): + handler = server.StreamHandler( + self.mock_request, self.mock_client_address, self.mock_server + ) + handler.server = self.mock_server + handler.path = '/ws' + + handler.do_GET() + + mock_handle_ws.assert_called_once_with(handler, self.mock_state) + mock_super_do_get.assert_not_called() + + @mock.patch.object( + http.server.SimpleHTTPRequestHandler, '__init__', return_value=None + ) + def test_do_get_invalid_server_type(self, unused_mock_super_init): + handler = server.StreamHandler( + self.mock_request, self.mock_client_address, self.mock_server + ) + # Set server to a mock that is NOT PlaygroundServer + handler.server = mock.create_autospec(http.server.HTTPServer, instance=True) + handler.send_error = mock.create_autospec(handler.send_error) + + handler.do_GET() + + handler.send_error.assert_called_once_with(500, 'Unexpected server type') + + def test_get_default_static_dir_returns_existing_dir(self): + static_dir = server.get_default_static_dir() + self.assertIsNotNone(static_dir) + + @mock.patch.object(os.path, 'exists', autospec=True) + def test_get_default_static_dir_resource_path(self, mock_exists): + mock_exists.return_value = True + with mock.patch.object( + server, + '_get_resource_filename', + return_value='/fake/resource/dir/index.html', + ): + static_dir = server.get_default_static_dir() + self.assertEqual(static_dir, '/fake/resource/dir') + + @mock.patch.object(os.path, 'exists', autospec=True) + def test_get_default_static_dir_resource_path_does_not_exist( + self, mock_exists + ): + expected_dir = os.path.join(os.path.dirname(server.__file__), 'frontend') + + def side_effect(path): + return path == os.path.join(expected_dir, 'index.html') + + mock_exists.side_effect = side_effect + with ( + mock.patch.object( + server, + '_get_resource_filename', + return_value='/nonexistent/resource/index.html', + ), + mock.patch.dict(os.environ, {}, clear=True), + ): + static_dir = server.get_default_static_dir() + self.assertEqual(static_dir, expected_dir) + + @mock.patch.object(os.path, 'exists', autospec=True) + def test_get_default_static_dir_fallback_to_runfiles_dir(self, mock_exists): + def side_effect(path): + return path == '/fake/runfiles/android_env/playground/frontend/index.html' + + mock_exists.side_effect = side_effect + with ( + mock.patch.object(server, '_get_resource_filename', return_value=''), + mock.patch.dict( + os.environ, {'RUNFILES_DIR': '/fake/runfiles'}, clear=True + ), + ): + static_dir = server.get_default_static_dir() + self.assertEqual( + static_dir, '/fake/runfiles/android_env/playground/frontend' + ) + + @mock.patch.object(os.path, 'exists', autospec=True) + def test_get_default_static_dir_fallback_to_python_runfiles( + self, mock_exists + ): + def side_effect(path): + return ( + path == '/fake/py_runfiles/android_env/playground/frontend/index.html' + ) + + mock_exists.side_effect = side_effect + with ( + mock.patch.object(server, '_get_resource_filename', return_value=''), + mock.patch.dict( + os.environ, {'PYTHON_RUNFILES': '/fake/py_runfiles'}, clear=True + ), + ): + static_dir = server.get_default_static_dir() + self.assertEqual( + static_dir, '/fake/py_runfiles/android_env/playground/frontend' + ) + + @mock.patch.object(os.path, 'exists', autospec=True) + def test_get_default_static_dir_fallback_to_package_dir(self, mock_exists): + expected_dir = os.path.join(os.path.dirname(server.__file__), 'frontend') + + def side_effect(path): + return path == os.path.join(expected_dir, 'index.html') + + mock_exists.side_effect = side_effect + with ( + mock.patch.object(server, '_get_resource_filename', return_value=''), + mock.patch.dict(os.environ, {}, clear=True), + ): + static_dir = server.get_default_static_dir() + self.assertEqual(static_dir, expected_dir) + + @mock.patch.object(os.path, 'exists', return_value=False, autospec=True) + def test_get_default_static_dir_raises_file_not_found(self, mock_exists): + del mock_exists + with ( + mock.patch.object(server, '_get_resource_filename', return_value=''), + mock.patch.dict(os.environ, {}, clear=True), + ): + with self.assertRaises(FileNotFoundError): + server.get_default_static_dir() + + +if __name__ == '__main__': + absltest.main() diff --git a/android_env/playground/state.py b/android_env/playground/state.py new file mode 100644 index 00000000..04bc3c70 --- /dev/null +++ b/android_env/playground/state.py @@ -0,0 +1,226 @@ +# coding=utf-8 +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared state for AndroidEnv Playground.""" + +from collections.abc import Mapping +import queue +import threading +from typing import Any + + +class PlaygroundState: + """Holds the shared state for the AndroidEnv playground. + + This class acts as a thread-safe container for state variables shared between + the main server thread and the background AndroidEnv execution thread. + """ + + def __init__(self): + """Initializes the shared state.""" + # The most recent processed screen frame (JPEG bytes) + self.current_frame = None + # Lock protecting updates to current_frame + self.frame_lock = threading.Lock() + # Queue for actions received from client to be consumed by AndroidEnv + self.action_queue = queue.Queue() + # Flag to control the background AndroidEnv run loop + self.running = True + # Active WebSocket connection socket for streaming frames + self.active_ws = None + # Lock protecting access to active_ws + self.active_ws_lock = threading.RLock() + # Flag indicating if we are waiting for an ACK from client for sent frame + self.waiting_for_ack = False + # Timestamp of when the last frame was sent to the client + self.last_frame_sent_time = 0.0 + + # Log history buffer + self.logs_lock = threading.RLock() + self.logs: list[dict[str, Any]] = [] + self.max_logs = 10000 + + # Status message + self._status: str | None = None + self._status_lock = threading.Lock() + + # Task reload flags + self._active_task_id: str | None = None + self._reload_lock = threading.Lock() + + # Episode reset flag and ADB command queue + self._reset_requested = False + self._reset_lock = threading.Lock() + self.adb_queue: queue.Queue[str] = queue.Queue() + + # Stats accumulation + self.stats_lock = threading.Lock() + self.pending_stats = { + 'ack_latency': [], + 'env_step_ms': [], + 'action_transit_ms': [], + 'action_queue_ms': [], + 'action_total_pre_step_ms': [], + } + self.battery_power_stats: dict[str, Any] | None = None + self.episode_stats: dict[str, Any] | None = None + + def set_active_ws(self, sock): + """Thread-safely sets the active WebSocket connection. + + Args: + sock: The socket object representing the active WebSocket client. + """ + with self.active_ws_lock: + self.active_ws = sock + + def get_active_ws(self): + """Thread-safely retrieves the active WebSocket connection. + + Returns: + The active socket object, or None if no client is connected. + """ + with self.active_ws_lock: + return self.active_ws + + def set_status(self, status: str) -> None: + """Thread-safely sets the latest status message.""" + with self._status_lock: + self._status = status + + def get_status(self) -> str | None: + """Thread-safely retrieves the latest status message.""" + with self._status_lock: + return self._status + + def add_stat(self, category: str, value: float): + """Thread-safely adds a stat value to the pending stats.""" + with self.stats_lock: + if category in self.pending_stats: + self.pending_stats[category].append(value) + + def set_battery_power_stats(self, metrics: Any): + """Thread-safely sets the latest battery and power telemetry metrics.""" + + def _get_val(obj: Any, name: str) -> Any: + if isinstance(obj, (dict, Mapping)): + val = obj.get(name) + else: + val = getattr(obj, name, None) + if hasattr(val, 'item'): + val = val.item() + return val + + extracted = {} + for key in ( + 'level_percent', + 'voltage_mv', + 'current_now_ua', + 'power_mw', + 'temperature_celsius', + 'status', + 'plugged_source', + 'charge_counter_uah', + 'energy_counter_nwh', + 'thermal_status', + ): + v = _get_val(metrics, key) + if v is not None: + extracted[key] = v + + with self.stats_lock: + self.battery_power_stats = extracted + + def get_battery_power_stats(self) -> dict[str, Any] | None: + """Thread-safely retrieves the latest battery and power telemetry metrics.""" + with self.stats_lock: + return ( + dict(self.battery_power_stats) + if self.battery_power_stats is not None + else None + ) + + def set_episode_stats(self, episode_stats: dict[str, Any]) -> None: + """Thread-safely sets the latest RL episode telemetry metrics.""" + with self.stats_lock: + self.episode_stats = dict(episode_stats) + + def get_episode_stats(self) -> dict[str, Any] | None: + """Thread-safely retrieves the latest RL episode telemetry metrics.""" + with self.stats_lock: + return ( + dict(self.episode_stats) if self.episode_stats is not None else None + ) + + def flush_stats(self) -> dict[str, Any]: + """Thread-safely flushes and returns the averaged pending stats.""" + with self.stats_lock: + stats = {} + for k, v in self.pending_stats.items(): + if v: + stats[k] = sum(v) / len(v) + v.clear() + if self.battery_power_stats is not None: + stats['battery_power'] = dict(self.battery_power_stats) + if self.episode_stats is not None: + stats['episode'] = dict(self.episode_stats) + return stats + + def set_reload(self, task_id: str): + """Thread-safely sets the active task id for reload.""" + with self._reload_lock: + self._active_task_id = task_id + + def get_reload(self) -> str | None: + """Thread-safely retrieves the active task id.""" + with self._reload_lock: + return self._active_task_id + + def clear_reload(self): + """Thread-safely clears the active task id reload flag.""" + with self._reload_lock: + self._active_task_id = None + + def set_reset(self) -> None: + """Thread-safely requests an episode reset.""" + with self._reset_lock: + self._reset_requested = True + + def get_reset(self) -> bool: + """Thread-safely returns whether an episode reset is requested.""" + with self._reset_lock: + return self._reset_requested + + def clear_reset(self) -> None: + """Thread-safely clears the episode reset flag.""" + with self._reset_lock: + self._reset_requested = False + + def add_log(self, log_entry: dict[str, Any]): + """Thread-safely appends a log entry to history.""" + with self.logs_lock: + self.logs.append(log_entry) + if len(self.logs) > self.max_logs: + self.logs.pop(0) + + def get_logs(self) -> list[dict[str, Any]]: + """Thread-safely returns a copy of log history.""" + with self.logs_lock: + return list(self.logs) + + def clear_logs(self): + """Thread-safely clears log history.""" + with self.logs_lock: + self.logs.clear() diff --git a/android_env/playground/state_test.py b/android_env/playground/state_test.py new file mode 100644 index 00000000..c4cd94e1 --- /dev/null +++ b/android_env/playground/state_test.py @@ -0,0 +1,160 @@ +# coding=utf-8 +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for playground state.""" + +import socket +from absl.testing import absltest +from android_env.playground import state as state_lib + + +class StateTest(absltest.TestCase): + + def test_playground_state_init(self): + state = state_lib.PlaygroundState() + self.assertIsNone(state.current_frame) + self.assertTrue(state.running) + self.assertIsNone(state.get_active_ws()) + self.assertFalse(state.waiting_for_ack) + + def test_set_active_ws(self): + state = state_lib.PlaygroundState() + # Use a dummy socket object or mock + # A real socket is not connected, but we just want to verify it stores it + mock_socket = socket.socket() + state.set_active_ws(mock_socket) + self.assertEqual(state.get_active_ws(), mock_socket) + mock_socket.close() + + def test_reload_flags(self): + state = state_lib.PlaygroundState() + self.assertIsNone(state.get_reload()) + + state.set_reload('my_task') + self.assertEqual(state.get_reload(), 'my_task') + + state.clear_reload() + self.assertIsNone(state.get_reload()) + + def test_logs(self): + state = state_lib.PlaygroundState() + self.assertEqual(state.get_logs(), []) + + log1 = {'message': 'msg1', 'level': 'INFO', 'timestamp': 100.0} + state.add_log(log1) + self.assertEqual(state.get_logs(), [log1]) + + log2 = {'message': 'msg2', 'level': 'ERROR', 'timestamp': 101.0} + state.add_log(log2) + self.assertEqual(state.get_logs(), [log1, log2]) + + state.clear_logs() + self.assertEqual(state.get_logs(), []) + + def test_logs_max_capacity(self): + state = state_lib.PlaygroundState() + state.max_logs = 3 + for i in range(5): + state.add_log({'message': f'msg{i}'}) + logs = state.get_logs() + self.assertLen(logs, 3) + self.assertEqual(logs[0]['message'], 'msg2') + self.assertEqual(logs[2]['message'], 'msg4') + + def test_battery_power_stats(self): + state = state_lib.PlaygroundState() + self.assertIsNone(state.get_battery_power_stats()) + + battery_stats = { + 'level_percent': 85, + 'power_mw': 1500.0, + } + state.set_battery_power_stats(battery_stats) + self.assertEqual(state.get_battery_power_stats(), battery_stats) + + stats = state.flush_stats() + self.assertEqual(stats['battery_power'], battery_stats) + + class FakeScalar: + + def __init__(self, v): + self._v = v + + def item(self): + return self._v + + class FakeBatteryMetrics: + level_percent = FakeScalar(90) + status = 'CHARGING' + power_mw = 2000.0 + + state.set_battery_power_stats(FakeBatteryMetrics()) + self.assertEqual( + state.get_battery_power_stats(), + {'level_percent': 90, 'status': 'CHARGING', 'power_mw': 2000.0}, + ) + + dict_metrics = { + 'level_percent': 85, + 'status': 'DISCHARGING', + 'power_mw': 1500.0, + } + state.set_battery_power_stats(dict_metrics) + self.assertEqual( + state.get_battery_power_stats(), + {'level_percent': 85, 'status': 'DISCHARGING', 'power_mw': 1500.0}, + ) + + def test_status(self): + state = state_lib.PlaygroundState() + self.assertIsNone(state.get_status()) + state.set_status('Ready! Orientation: 0') + self.assertEqual(state.get_status(), 'Ready! Orientation: 0') + + def test_add_and_flush_stats(self): + state = state_lib.PlaygroundState() + self.assertEqual(state.flush_stats(), {}) + + state.add_stat('ack_latency', 10.0) + state.add_stat('ack_latency', 20.0) + state.add_stat('env_step_ms', 5.0) + state.add_stat('unknown_stat', 99.0) + + flushed = state.flush_stats() + self.assertEqual(flushed, {'ack_latency': 15.0, 'env_step_ms': 5.0}) + self.assertEqual(state.flush_stats(), {}) + + def test_reset_flags(self): + state = state_lib.PlaygroundState() + self.assertFalse(state.get_reset()) + state.set_reset() + self.assertTrue(state.get_reset()) + state.clear_reset() + self.assertFalse(state.get_reset()) + + def test_episode_stats(self): + state = state_lib.PlaygroundState() + self.assertIsNone(state.get_episode_stats()) + + ep_stats = {'step': 4, 'reward': 1.5, 'cumulative_reward': 6.0} + state.set_episode_stats(ep_stats) + self.assertEqual(state.get_episode_stats(), ep_stats) + + flushed = state.flush_stats() + self.assertEqual(flushed['episode'], ep_stats) + + +if __name__ == '__main__': + absltest.main() diff --git a/android_env/playground/websocket.py b/android_env/playground/websocket.py new file mode 100644 index 00000000..5871955c --- /dev/null +++ b/android_env/playground/websocket.py @@ -0,0 +1,460 @@ +# coding=utf-8 +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""WebSocket handler for AndroidEnv Playground.""" + +import base64 +import hashlib +import http.server +import json +import socket +import time +from typing import Any + +from absl import logging +from android_env.playground import actions as actions_lib +from android_env.playground import state as state_lib + + +def _build_ws_frame(first_byte: int, payload: bytes) -> bytes: + """Builds a WebSocket frame with the given first byte (FIN + opcode).""" + length = len(payload) + header = bytearray([first_byte]) + if length < 126: + header.append(length) + elif length < 65536: + header.append(126) + header.extend(length.to_bytes(2, byteorder='big')) + else: + header.append(127) + header.extend(length.to_bytes(8, byteorder='big')) + return bytes(header) + payload + + +def send_ws_binary_frame(sock: socket.socket, data: bytes): + """Sends a binary WebSocket frame to the client.""" + # 0x82: FIN bit is set (0x80) to indicate this is the final fragment, and + # opcode is set to 2 (0x02) for a binary frame. + frame = _build_ws_frame(0x82, data) + sock.sendall(frame) + + +def send_ws_text_frame(sock: socket.socket, data: str): + """Sends a text WebSocket frame to the client.""" + # 0x81: FIN bit is set (0x80) and opcode is 1 (0x01) for text frame. + frame = _build_ws_frame(0x81, data.encode('utf-8')) + sock.sendall(frame) + + +def _try_send_json(state: state_lib.PlaygroundState, payload: dict[str, Any]): + """Serializes payload as JSON and sends it to the active WebSocket client.""" + msg = json.dumps(payload) + with state.active_ws_lock: + ws_sock = state.active_ws + if ws_sock is None: + return + try: + send_ws_text_frame(ws_sock, msg) + except OSError: + state.active_ws = None + + +def try_send_status(state: state_lib.PlaygroundState, status_text: str): + """Records and sends status text to the client.""" + state.set_status(status_text) + _try_send_json(state, {'type': 'status', 'message': status_text}) + + +def try_send_stats(state: state_lib.PlaygroundState, stats: dict[str, Any]): + """Sends stats to the client.""" + _try_send_json(state, {'type': 'stats', 'stats': stats}) + + +def try_send_logs(state: state_lib.PlaygroundState, logs: list[dict[str, Any]]): + """Sends log history or real-time log entries to the client.""" + _try_send_json(state, {'type': 'logs', 'logs': logs}) + + +def try_send_adb_response( + state: state_lib.PlaygroundState, + command: str, + status: str, + output: str, + error: str = '', + duration_ms: float = 0.0, +): + """Sends an ADB command response to the client.""" + _try_send_json( + state, + { + 'type': 'adb_response', + 'command': command, + 'status': status, + 'output': output, + 'error': error, + 'duration_ms': duration_ms, + }, + ) + + +def try_send_frame(state: state_lib.PlaygroundState, jpeg_bytes: bytes): + """Sends a frame if client is ready, enforcing backpressure.""" + ack_timed_out = False + with state.active_ws_lock: + if state.waiting_for_ack: + if time.time() - state.last_frame_sent_time > 1.0: + ack_timed_out = True + state.waiting_for_ack = False + else: + return + ws_sock = state.active_ws + if ws_sock is None: + return + state.waiting_for_ack = True + state.last_frame_sent_time = time.time() + try: + send_ws_binary_frame(ws_sock, jpeg_bytes) + except OSError: + state.active_ws = None + state.waiting_for_ack = False + + if ack_timed_out: + logging.warning('ACK timeout, resetting waiting_for_ack') + + +def _recvall(rfile, n: int) -> bytes | None: + """Receives exactly n bytes from the rfile reader.""" + try: + data = rfile.read(n) + if len(data) < n: + return None # EOF or incomplete read + return data + except OSError: + return None + + +def _perform_handshake(handler: http.server.BaseHTTPRequestHandler) -> bool: + """Performs the WebSocket handshake. + + Args: + handler: The HTTP request handler. + + Returns: + True if handshake succeeded, False otherwise. + """ + key = handler.headers.get('Sec-WebSocket-Key') + if not key: + handler.send_error(400, 'Missing Sec-WebSocket-Key') + return False + + # Globally unique identifier (GUID) defined in RFC 6455 Sec 1.3 for the + # WebSocket handshake. It is appended to the Sec-WebSocket-Key, SHA-1 hashed, + # and base64 encoded to produce the Sec-WebSocket-Accept response header. + guid = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11' + accept = base64.b64encode( + hashlib.sha1((key + guid).encode('utf-8')).digest() + ).decode('utf-8') + + # Send handshake response directly to bypass http.server header formatting + response = ( + b'HTTP/1.1 101 Switching Protocols\r\n' + b'Upgrade: websocket\r\n' + b'Connection: Upgrade\r\n' + b'Sec-WebSocket-Accept: ' + + accept.encode('utf-8') + + b'\r\n\r\n' + ) + try: + handler.wfile.write(response) + handler.wfile.flush() + return True + except OSError as e: + logging.warning('Handshake write failed: %s', e) + return False + + +def _read_frame(rfile) -> tuple[int, bytes] | None: + """Reads a single WebSocket frame from the rfile reader. + + Args: + rfile: The file-like object to read from. + + Returns: + A tuple of (opcode, payload_bytes) if a frame was successfully read and + parsed, or None on error or EOF. + """ + # Byte 0 contains: FIN bit (1 bit), RSV1-3 (3 bits), Opcode (4 bits) + header = _recvall(rfile, 2) + if not header: + return None + + fin_opcode = header[0] + mask_payload = header[1] + + # Extract the 4-bit opcode (lower 4 bits of the first byte) + opcode = fin_opcode & 0x0F + + has_mask = (mask_payload & 0x80) != 0 + if not has_mask: + logging.warning('WebSocket received unmasked frame.') + return None + + payload_len = mask_payload & 0x7F + + if payload_len == 126: + len_bytes = _recvall(rfile, 2) + if not len_bytes: + return None + payload_len = int.from_bytes(len_bytes, byteorder='big') + elif payload_len == 127: + len_bytes = _recvall(rfile, 8) + if not len_bytes: + return None + payload_len = int.from_bytes(len_bytes, byteorder='big') + + mask_key = _recvall(rfile, 4) + if not mask_key: + return None + + payload = _recvall(rfile, payload_len) + if payload is None: + return None + + # Client-to-server frames MUST be masked. Unmask the payload by XORing + # each byte with the corresponding byte of the 4-byte mask key (cycling). + unmasked = bytearray(payload_len) + for i in range(payload_len): + unmasked[i] = payload[i] ^ mask_key[i % 4] + + return opcode, bytes(unmasked) + + +def _handle_message(msg: str, state: state_lib.PlaygroundState): + """Handles a decoded WebSocket message.""" + if msg == 'ack': + with state.active_ws_lock: + state.waiting_for_ack = False + ack_latency = (time.time() - state.last_frame_sent_time) * 1000 + logging.vlog(1, 'ACK received in %.2f ms', ack_latency) + state.add_stat('ack_latency', ack_latency) + return + + if msg == 'reset': + logging.info('Received episode reset request.') + state.set_reset() + return + + if msg.startswith('adb,'): + cmd = msg[4:].strip() + if cmd: + logging.info('Received ADB command request: %s', cmd) + state.adb_queue.put(cmd) + else: + logging.warning('Received empty ADB command request.') + return + + parts = msg.split(',') + if len(parts) == 2 and parts[0] == 'load_task': + task_id = parts[1].strip() + if task_id: + logging.info('Received reload task request for: %s', task_id) + state.set_reload(task_id) + else: + logging.warning('Received load_task request with empty task_id.') + return + + if parts[0] in ('touch_down', 'touch_move', 'touch_up') and len(parts) == 4: + act_type, x_str, y_str, ts_str = parts + recv_ts = time.time() * 1000 + try: + x = float(x_str) + y = float(y_str) + client_ts = float(ts_str) + except ValueError: + logging.warning('Failed to parse %s message: %s', act_type, msg) + return + state.action_queue.put( + actions_lib.ClientAction( + action_type=act_type, + x=x, + y=y, + client_timestamp_ms=client_ts, + recv_timestamp_ms=recv_ts, + ) + ) + return + + if parts[0] == 'text': + recv_ts = time.time() * 1000 + prefix = 'text,' + payload = msg[len(prefix) :] + last_comma = payload.rfind(',') + if last_comma != -1: + text_str = payload[:last_comma] + ts_str = payload[last_comma + 1 :] + try: + client_ts = float(ts_str) + except ValueError: + client_ts = recv_ts + state.action_queue.put( + actions_lib.ClientAction( + action_type='text', + x=0.0, + y=0.0, + client_timestamp_ms=client_ts, + recv_timestamp_ms=recv_ts, + text=text_str, + ) + ) + return + + if parts[0] == 'keydown' and len(parts) == 3: + recv_ts = time.time() * 1000 + try: + keycode = int(parts[1]) + client_ts = float(parts[2]) + state.action_queue.put( + actions_lib.ClientAction( + action_type='keydown', + x=0.0, + y=0.0, + client_timestamp_ms=client_ts, + recv_timestamp_ms=recv_ts, + keycode=keycode, + ) + ) + except ValueError: + logging.warning('Failed to parse keydown message: %s', msg) + return + + if parts[0] == 'set_selection' and len(parts) == 4: + recv_ts = time.time() * 1000 + try: + start = int(parts[1]) + end = int(parts[2]) + client_ts = float(parts[3]) + state.action_queue.put( + actions_lib.ClientAction( + action_type='set_selection', + x=0.0, + y=0.0, + client_timestamp_ms=client_ts, + recv_timestamp_ms=recv_ts, + selection=(start, end), + ) + ) + except ValueError: + logging.warning('Failed to parse set_selection message: %s', msg) + return + + if parts[0] in ('cut', 'copy', 'paste') and len(parts) == 2: + recv_ts = time.time() * 1000 + try: + act_type = parts[0] + client_ts = float(parts[1]) + state.action_queue.put( + actions_lib.ClientAction( + action_type=act_type, + x=0.0, + y=0.0, + client_timestamp_ms=client_ts, + recv_timestamp_ms=recv_ts, + ) + ) + except ValueError: + logging.warning('Failed to parse %s message: %s', parts[0], msg) + return + + +def handle_websocket( + handler: http.server.BaseHTTPRequestHandler, + state: state_lib.PlaygroundState, +): + """Handles a WebSocket handshake and frame reading loop.""" + if not _perform_handshake(handler): + return + + sock = handler.request + sock.setblocking(True) + + # Disable Nagle's algorithm for low latency + try: + sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + except OSError as e: + logging.warning('Failed to set TCP_NODELAY: %s', e) + + logging.info('WebSocket connection established') + + # Close any existing active connection to prevent multiple active clients. + old_sock = state.get_active_ws() + if old_sock and old_sock != sock: + logging.info('Closing old connection') + try: + old_sock.shutdown(socket.SHUT_RDWR) + old_sock.close() + except OSError as e: + logging.warning('Failed to close old socket: %s', e) + + state.set_active_ws(sock) + with state.active_ws_lock: + state.waiting_for_ack = False + + # Send current status if available + current_status = state.get_status() + if current_status is not None: + try_send_status(state, current_status) + + # Send log history if available + logs_history = state.get_logs() + if logs_history: + try_send_logs(state, logs_history) + + # Send initial frame if available to show current status/screen immediately + with state.frame_lock: + init_frame = state.current_frame + if init_frame is not None: + try_send_frame(state, init_frame) + + rfile = handler.rfile + try: + while state.running: + # Ensure this thread is still handling the active connection. + if state.get_active_ws() != sock: + logging.info('Connection superseded, exiting loop.') + break + + frame = _read_frame(rfile) + if frame is None: + break + + opcode, payload = frame + if opcode == 8: # Close frame + logging.info('WebSocket close frame received') + break + + if opcode == 1: # Text frame + msg = payload.decode('utf-8') + _handle_message(msg, state) + else: + logging.warning('Unexpected WebSocket opcode: %d', opcode) + + except (OSError, UnicodeDecodeError, ValueError) as e: + logging.warning('WebSocket error: %s', e) + finally: + logging.info('WebSocket connection closed') + with state.active_ws_lock: + if state.active_ws == sock: + state.active_ws = None + state.waiting_for_ack = False diff --git a/android_env/playground/websocket_test.py b/android_env/playground/websocket_test.py new file mode 100644 index 00000000..b2660379 --- /dev/null +++ b/android_env/playground/websocket_test.py @@ -0,0 +1,667 @@ +# coding=utf-8 +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for playground websocket.""" + +import http.server +import io +import socket +from unittest import mock +from absl.testing import absltest +from android_env.playground import state as state_lib +from android_env.playground import websocket + + +class WebsocketTest(absltest.TestCase): + + def test_recvall_complete(self): + mock_rfile = mock.create_autospec(io.BufferedReader, instance=True) + mock_rfile.read.return_value = b'abcd' + res = websocket._recvall(mock_rfile, 4) + self.assertEqual(res, b'abcd') + mock_rfile.read.assert_called_once_with(4) + + def test_recvall_eof(self): + mock_rfile = mock.create_autospec(io.BufferedReader, instance=True) + mock_rfile.read.return_value = b'ab' + res = websocket._recvall(mock_rfile, 4) + self.assertIsNone(res) + mock_rfile.read.assert_called_once_with(4) + + def test_recvall_oserror(self): + mock_rfile = mock.create_autospec(io.BufferedReader, instance=True) + mock_rfile.read.side_effect = OSError('read failed') + self.assertIsNone(websocket._recvall(mock_rfile, 4)) + + def test_send_ws_binary_frame_small(self): + mock_sock = mock.create_autospec(socket.socket, instance=True) + websocket.send_ws_binary_frame(mock_sock, b'hello') + mock_sock.sendall.assert_called_once_with(b'\x82\x05hello') + + def test_send_ws_binary_frame_medium(self): + mock_sock = mock.create_autospec(socket.socket, instance=True) + data = b'x' * 200 + websocket.send_ws_binary_frame(mock_sock, data) + expected_header = b'\x82\x7e\x00\xc8' + mock_sock.sendall.assert_called_once_with(expected_header + data) + + def test_send_ws_binary_frame_large(self): + mock_sock = mock.create_autospec(socket.socket, instance=True) + data = b'x' * 70000 + websocket.send_ws_binary_frame(mock_sock, data) + expected_header = b'\x82\x7f\x00\x00\x00\x00\x00\x01\x11\x70' + mock_sock.sendall.assert_called_once_with(expected_header + data) + + def test_send_ws_binary_frame_error(self): + mock_sock = mock.create_autospec(socket.socket, instance=True) + mock_sock.sendall.side_effect = OSError('send error') + with self.assertRaises(OSError): + websocket.send_ws_binary_frame(mock_sock, b'hello') + + def test_read_frame_text(self): + mock_rfile = mock.create_autospec(io.BufferedReader, instance=True) + mock_rfile.read.side_effect = [ + b'\x81\x85', # Header: FIN=1, Opcode=1 (Text), Mask=1, Length=5 + b'1234', # Mask Key + b'YW_X^', # Masked Payload "hello" XOR "1234" + ] + res = websocket._read_frame(mock_rfile) + self.assertEqual(res, (1, b'hello')) + mock_rfile.read.assert_has_calls([ + mock.call(2), + mock.call(4), + mock.call(5), + ]) + + def test_read_frame_unmasked(self): + mock_rfile = mock.create_autospec(io.BufferedReader, instance=True) + mock_rfile.read.return_value = b'\x81\x05' + with mock.patch.object(websocket.logging, 'warning') as mock_warn: + res = websocket._read_frame(mock_rfile) + self.assertIsNone(res) + mock_warn.assert_called_once() + mock_rfile.read.assert_called_once_with(2) + + def test_read_frame_eof(self): + mock_rfile = mock.create_autospec(io.BufferedReader, instance=True) + mock_rfile.read.return_value = b'' + res = websocket._read_frame(mock_rfile) + self.assertIsNone(res) + mock_rfile.read.assert_called_once_with(2) + + def test_read_frame_medium_payload(self): + mock_rfile = mock.create_autospec(io.BufferedReader, instance=True) + masked_payload = b'IJKL' * 50 + mock_rfile.read.side_effect = [ + b'\x81\xFE', # Header + b'\x00\xC8', # Extra length bytes + b'1234', # Mask Key + masked_payload, + ] + res = websocket._read_frame(mock_rfile) + self.assertEqual(res, (1, b'x' * 200)) + mock_rfile.read.assert_has_calls([ + mock.call(2), + mock.call(2), + mock.call(4), + mock.call(200), + ]) + + def test_read_frame_large_payload(self): + mock_rfile = mock.create_autospec(io.BufferedReader, instance=True) + masked_payload = b'IJKL' * 17500 + mock_rfile.read.side_effect = [ + b'\x81\xFF', # Header + b'\x00\x00\x00\x00\x00\x01\x11\x70', # Extra length bytes + b'1234', # Mask Key + masked_payload, + ] + res = websocket._read_frame(mock_rfile) + self.assertEqual(res, (1, b'x' * 70000)) + mock_rfile.read.assert_has_calls([ + mock.call(2), + mock.call(8), + mock.call(4), + mock.call(70000), + ]) + + def test_read_frame_127_eof_len_bytes(self): + mock_rfile = mock.create_autospec(io.BufferedReader, instance=True) + mock_rfile.read.side_effect = [ + b'\x81\xFF', + b'', + ] + self.assertIsNone(websocket._read_frame(mock_rfile)) + + def test_read_frame_incomplete_header(self): + mock_rfile = mock.create_autospec(io.BufferedReader, instance=True) + mock_rfile.read.return_value = b'\x81' + res = websocket._read_frame(mock_rfile) + self.assertIsNone(res) + + def test_read_frame_incomplete_extended_length(self): + mock_rfile = mock.create_autospec(io.BufferedReader, instance=True) + mock_rfile.read.side_effect = [ + b'\x81\xFE', # Header + b'\x00', # incomplete length + ] + res = websocket._read_frame(mock_rfile) + self.assertIsNone(res) + + def test_read_frame_incomplete_mask(self): + mock_rfile = mock.create_autospec(io.BufferedReader, instance=True) + mock_rfile.read.side_effect = [ + b'\x81\x85', # Header (masked, len 5) + b'123', # incomplete mask + ] + res = websocket._read_frame(mock_rfile) + self.assertIsNone(res) + + def test_read_frame_incomplete_payload(self): + mock_rfile = mock.create_autospec(io.BufferedReader, instance=True) + mock_rfile.read.side_effect = [ + b'\x81\x85', # Header (masked, len 5) + b'1234', # Mask Key + b'abcd', # incomplete payload + ] + res = websocket._read_frame(mock_rfile) + self.assertIsNone(res) + + def test_try_send_frame_success(self): + state = state_lib.PlaygroundState() + mock_sock = mock.create_autospec(socket.socket, instance=True) + state.active_ws = mock_sock + with mock.patch.object( + websocket.logging, 'warning', autospec=True + ) as mock_warning: + websocket.try_send_frame(state, b'frame_data') + mock_warning.assert_not_called() + expected_header = b'\x82\x0a' + mock_sock.sendall.assert_called_once_with(expected_header + b'frame_data') + self.assertTrue(state.waiting_for_ack) + + def test_try_send_frame_client_not_ready_backpressure(self): + state = state_lib.PlaygroundState() + state.active_ws = mock.create_autospec(socket.socket, instance=True) + state.waiting_for_ack = True + state.last_frame_sent_time = 100.0 + with mock.patch.object(websocket.time, 'time', return_value=100.5): + websocket.try_send_frame(state, b'frame_data') + state.active_ws.sendall.assert_not_called() + self.assertTrue(state.waiting_for_ack) + + def test_try_send_frame_client_not_ready_timeout(self): + state = state_lib.PlaygroundState() + mock_sock = mock.create_autospec(socket.socket, instance=True) + state.active_ws = mock_sock + state.waiting_for_ack = True + state.last_frame_sent_time = 100.0 + with mock.patch.object( + websocket.logging, 'warning', autospec=True + ) as mock_warning: + with mock.patch.object(websocket.time, 'time', return_value=102.0): + websocket.try_send_frame(state, b'frame_data') + mock_warning.assert_called_once_with( + 'ACK timeout, resetting waiting_for_ack' + ) + expected_header = b'\x82\x0a' + mock_sock.sendall.assert_called_once_with(expected_header + b'frame_data') + self.assertTrue(state.waiting_for_ack) + self.assertEqual(state.last_frame_sent_time, 102.0) + + def test_try_send_frame_reentrant_active_ws_lock(self): + state = state_lib.PlaygroundState() + mock_sock = mock.create_autospec(socket.socket, instance=True) + state.active_ws = mock_sock + state.waiting_for_ack = True + state.last_frame_sent_time = 100.0 + + # Acquire lock in current thread to ensure reentrancy does not deadlock. + with state.active_ws_lock: + with mock.patch.object(websocket.time, 'time', return_value=102.0): + websocket.try_send_frame(state, b'frame_data') + + expected_header = b'\x82\x0a' + mock_sock.sendall.assert_called_once_with(expected_header + b'frame_data') + self.assertTrue(state.waiting_for_ack) + + def test_try_send_frame_no_client(self): + state = state_lib.PlaygroundState() + state.active_ws = None + websocket.try_send_frame(state, b'frame_data') + + def test_try_send_frame_send_error(self): + state = state_lib.PlaygroundState() + mock_sock = mock.create_autospec(socket.socket, instance=True) + mock_sock.sendall.side_effect = OSError('conn reset') + state.active_ws = mock_sock + state.waiting_for_ack = False + with mock.patch.object(websocket.time, 'time', return_value=100.0): + websocket.try_send_frame(state, b'frame_data') + self.assertIsNone(state.active_ws) + self.assertFalse(state.waiting_for_ack) + + def test_perform_handshake_success(self): + mock_handler = mock.create_autospec( + http.server.BaseHTTPRequestHandler, instance=True + ) + mock_handler.headers = {'Sec-WebSocket-Key': 'dGhlIHNhbXBsZSBub25jZQ=='} + mock_handler.wfile = io.BytesIO() + res = websocket._perform_handshake(mock_handler) + self.assertTrue(res) + expected_response = ( + b'HTTP/1.1 101 Switching Protocols\r\n' + b'Upgrade: websocket\r\n' + b'Connection: Upgrade\r\n' + b'Sec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=\r\n\r\n' + ) + self.assertEqual(mock_handler.wfile.getvalue(), expected_response) + + def test_perform_handshake_missing_key(self): + mock_handler = mock.create_autospec( + http.server.BaseHTTPRequestHandler, instance=True + ) + mock_handler.headers = {} + res = websocket._perform_handshake(mock_handler) + self.assertFalse(res) + mock_handler.send_error.assert_called_once_with( + 400, 'Missing Sec-WebSocket-Key' + ) + + def test_perform_handshake_write_failed(self): + mock_handler = mock.create_autospec( + http.server.BaseHTTPRequestHandler, instance=True + ) + mock_handler.headers = {'Sec-WebSocket-Key': 'dGhlIHNhbXBsZSBub25jZQ=='} + mock_wfile = mock.create_autospec(io.BufferedWriter, instance=True) + mock_wfile.write.side_effect = OSError('write error') + mock_handler.wfile = mock_wfile + with mock.patch.object(websocket.logging, 'warning') as mock_warn: + res = websocket._perform_handshake(mock_handler) + self.assertFalse(res) + mock_warn.assert_called_once() + + def test_handle_message_ack(self): + state = state_lib.PlaygroundState() + state.waiting_for_ack = True + state.last_frame_sent_time = 100.0 + with mock.patch.object(websocket.time, 'time', return_value=100.5): + websocket._handle_message('ack', state) + self.assertFalse(state.waiting_for_ack) + self.assertEqual(state.pending_stats['ack_latency'], [500.0]) + + def test_handle_message_action(self): + state = state_lib.PlaygroundState() + msg = 'touch_down,0.5,0.6,123456789.0' + with mock.patch.object(websocket.time, 'time', return_value=1700000000.0): + websocket._handle_message(msg, state) + self.assertEqual(state.action_queue.qsize(), 1) + act = state.action_queue.get() + self.assertEqual(act.action_type, 'touch_down') + self.assertEqual(act.x, 0.5) + self.assertEqual(act.y, 0.6) + self.assertEqual(act.client_timestamp_ms, 123456789.0) + self.assertEqual(act.recv_timestamp_ms, 1700000000000.0) + + def test_handle_message_text(self): + state = state_lib.PlaygroundState() + msg = 'text,hello world,123456789.0' + with mock.patch.object(websocket.time, 'time', return_value=1700000000.0): + websocket._handle_message(msg, state) + self.assertEqual(state.action_queue.qsize(), 1) + act = state.action_queue.get() + self.assertEqual(act.action_type, 'text') + self.assertEqual(act.text, 'hello world') + self.assertEqual(act.client_timestamp_ms, 123456789.0) + + def test_handle_message_text_invalid_timestamp(self): + state = state_lib.PlaygroundState() + msg = 'text,hello world,not_a_number' + with mock.patch.object(websocket.time, 'time', return_value=1700000000.0): + websocket._handle_message(msg, state) + self.assertEqual(state.action_queue.qsize(), 1) + act = state.action_queue.get() + self.assertEqual(act.action_type, 'text') + self.assertEqual(act.client_timestamp_ms, 1700000000000.0) + + def test_handle_message_keydown(self): + state = state_lib.PlaygroundState() + msg = 'keydown,66,123456789.0' + with mock.patch.object(websocket.time, 'time', return_value=1700000000.0): + websocket._handle_message(msg, state) + self.assertEqual(state.action_queue.qsize(), 1) + act = state.action_queue.get() + self.assertEqual(act.action_type, 'keydown') + self.assertEqual(act.keycode, 66) + self.assertEqual(act.client_timestamp_ms, 123456789.0) + self.assertEqual(act.recv_timestamp_ms, 1700000000000.0) + + def test_handle_message_set_selection(self): + state = state_lib.PlaygroundState() + msg = 'set_selection,2,5,123456789.0' + with mock.patch.object(websocket.time, 'time', return_value=1700000000.0): + websocket._handle_message(msg, state) + self.assertEqual(state.action_queue.qsize(), 1) + act = state.action_queue.get() + self.assertEqual(act.action_type, 'set_selection') + self.assertEqual(act.selection, (2, 5)) + self.assertEqual(act.client_timestamp_ms, 123456789.0) + self.assertEqual(act.recv_timestamp_ms, 1700000000000.0) + + def test_handle_message_clipboard(self): + for clip_act in ('cut', 'copy', 'paste'): + state = state_lib.PlaygroundState() + msg = f'{clip_act},123456789.0' + with mock.patch.object(websocket.time, 'time', return_value=1700000000.0): + websocket._handle_message(msg, state) + self.assertEqual(state.action_queue.qsize(), 1) + act = state.action_queue.get() + self.assertEqual(act.action_type, clip_act) + self.assertEqual(act.client_timestamp_ms, 123456789.0) + self.assertEqual(act.recv_timestamp_ms, 1700000000000.0) + + def test_handle_message_load_task(self): + state = state_lib.PlaygroundState() + msg = 'load_task,classic_2048' + websocket._handle_message(msg, state) + self.assertEqual(state.get_reload(), 'classic_2048') + + def test_handle_message_load_task_empty_id(self): + for msg in ('load_task,', 'load_task, '): + state = state_lib.PlaygroundState() + websocket._handle_message(msg, state) + self.assertIsNone(state.get_reload()) + + def test_handle_message_load_task_whitespace(self): + state = state_lib.PlaygroundState() + msg = 'load_task, classic_2048 ' + websocket._handle_message(msg, state) + self.assertEqual(state.get_reload(), 'classic_2048') + + def test_handle_message_load_task_invalid(self): + state = state_lib.PlaygroundState() + websocket._handle_message('load_task', state) + self.assertIsNone(state.get_reload()) + + def test_handle_message_invalid(self): + state = state_lib.PlaygroundState() + for msg in ( + 'invalid,msg', + 'touch_down,not_a_float,0.5,123.0', + 'keydown,not_an_int,123.0', + 'keydown,66', # Incomplete keydown (2 parts instead of 3) + 'unknown,1,2', # 3 parts, but not keydown + 'set_selection,a,b,123.0', + 'set_selection,1,2', # 3 parts instead of 4 + 'set_selection,1,2,3,4', # 5 parts instead of 4 + 'unknown,1,2,3,4', # 5 parts, not set_selection + 'cut,not_a_float', + ): + websocket._handle_message(msg, state) + self.assertEqual(state.action_queue.qsize(), 0) + + @mock.patch.object(websocket, '_perform_handshake', return_value=True) + @mock.patch.object(websocket, '_read_frame') + @mock.patch.object(websocket, '_handle_message') + def test_handle_websocket_loop( + self, mock_handle_msg, mock_read_frame, mock_handshake + ): + mock_sock = mock.create_autospec(socket.socket, instance=True) + mock_handler = mock.create_autospec( + http.server.BaseHTTPRequestHandler, instance=True + ) + mock_handler.request = mock_sock + mock_handler.rfile = io.BytesIO() + state = state_lib.PlaygroundState() + state.set_status('Connected') + state.current_frame = b'init_frame' + mock_read_frame.side_effect = [ + (1, b'msg1'), + (8, b''), + ] + websocket.handle_websocket(mock_handler, state) + mock_handshake.assert_called_once_with(mock_handler) + mock_sock.setsockopt.assert_called_once_with( + socket.IPPROTO_TCP, socket.TCP_NODELAY, 1 + ) + self.assertIsNone(state.active_ws) + self.assertEqual(mock_sock.sendall.call_count, 2) + self.assertIn(b'Connected', mock_sock.sendall.call_args_list[0][0][0]) + expected_header = b'\x82\x0a' + self.assertEqual( + mock_sock.sendall.call_args_list[1][0][0], + expected_header + b'init_frame', + ) + self.assertEqual(mock_read_frame.call_count, 2) + mock_handle_msg.assert_called_once_with('msg1', state) + + @mock.patch.object(websocket, '_perform_handshake', return_value=True) + @mock.patch.object(websocket, '_read_frame', return_value=None) + def test_handle_websocket_supersedes_old_connection( + self, mock_read_frame, mock_handshake + ): + mock_old_sock = mock.create_autospec(socket.socket, instance=True) + mock_new_sock = mock.create_autospec(socket.socket, instance=True) + mock_handler = mock.create_autospec( + http.server.BaseHTTPRequestHandler, instance=True + ) + mock_handler.request = mock_new_sock + mock_handler.rfile = io.BytesIO() + state = state_lib.PlaygroundState() + state.set_active_ws(mock_old_sock) + websocket.handle_websocket(mock_handler, state) + mock_old_sock.shutdown.assert_called_once_with(socket.SHUT_RDWR) + mock_old_sock.close.assert_called_once() + mock_handshake.assert_called_once_with(mock_handler) + mock_read_frame.assert_called_once() + + @mock.patch.object(websocket, '_perform_handshake', return_value=True) + @mock.patch.object(websocket, '_read_frame') + def test_handle_websocket_loop_exits_when_superseded( + self, mock_read_frame, mock_handshake + ): + mock_sock = mock.create_autospec(socket.socket, instance=True) + mock_handler = mock.create_autospec( + http.server.BaseHTTPRequestHandler, instance=True + ) + mock_handler.request = mock_sock + mock_handler.rfile = io.BytesIO() + state = state_lib.PlaygroundState() + + def change_state(*_): + state.set_active_ws(mock.create_autospec(socket.socket, instance=True)) + return (1, b'dummy') + + mock_read_frame.side_effect = change_state + websocket.handle_websocket(mock_handler, state) + mock_read_frame.assert_called_once() + self.assertIsNotNone(state.get_active_ws()) + self.assertNotEqual(state.get_active_ws(), mock_sock) + mock_handshake.assert_called_once_with(mock_handler) + + @mock.patch.object(websocket, '_perform_handshake', return_value=False) + def test_handle_websocket_handshake_fails(self, mock_handshake): + mock_handler = mock.create_autospec( + http.server.BaseHTTPRequestHandler, instance=True + ) + state = state_lib.PlaygroundState() + websocket.handle_websocket(mock_handler, state) + mock_handshake.assert_called_once_with(mock_handler) + self.assertIsNone(state.get_active_ws()) + + @mock.patch.object(websocket, '_perform_handshake', return_value=True) + @mock.patch.object(websocket, '_read_frame', return_value=None) + def test_handle_websocket_setsockopt_and_shutdown_oserror( + self, mock_read_frame, mock_handshake + ): + mock_old_sock = mock.create_autospec(socket.socket, instance=True) + mock_old_sock.shutdown.side_effect = OSError('shutdown failed') + mock_new_sock = mock.create_autospec(socket.socket, instance=True) + mock_new_sock.setsockopt.side_effect = OSError('nodelay failed') + mock_handler = mock.create_autospec( + http.server.BaseHTTPRequestHandler, instance=True + ) + mock_handler.request = mock_new_sock + mock_handler.rfile = io.BytesIO() + state = state_lib.PlaygroundState() + state.set_active_ws(mock_old_sock) + websocket.handle_websocket(mock_handler, state) + mock_handshake.assert_called_once_with(mock_handler) + mock_read_frame.assert_called_once() + self.assertIsNone(state.get_active_ws()) + + @mock.patch.object(websocket, '_perform_handshake', return_value=True) + @mock.patch.object(websocket, '_read_frame') + def test_handle_websocket_sends_logs_history_and_handles_unexpected_opcode( + self, mock_read_frame, mock_handshake + ): + mock_sock = mock.create_autospec(socket.socket, instance=True) + mock_handler = mock.create_autospec( + http.server.BaseHTTPRequestHandler, instance=True + ) + mock_handler.request = mock_sock + mock_handler.rfile = io.BytesIO() + state = state_lib.PlaygroundState() + state.add_log({'message': 'hello log', 'level': 'INFO', 'timestamp': 100.0}) + mock_read_frame.side_effect = [ + (2, b'binary data'), # unexpected opcode + (8, b''), # close frame + ] + websocket.handle_websocket(mock_handler, state) + mock_handshake.assert_called_once_with(mock_handler) + self.assertEqual(mock_read_frame.call_count, 2) + + @mock.patch.object(websocket, '_perform_handshake', return_value=True) + @mock.patch.object( + websocket, '_read_frame', side_effect=OSError('read error') + ) + def test_handle_websocket_catches_read_error( + self, mock_read_frame, mock_handshake + ): + mock_sock = mock.create_autospec(socket.socket, instance=True) + mock_handler = mock.create_autospec( + http.server.BaseHTTPRequestHandler, instance=True + ) + mock_handler.request = mock_sock + mock_handler.rfile = io.BytesIO() + state = state_lib.PlaygroundState() + websocket.handle_websocket(mock_handler, state) + mock_handshake.assert_called_once_with(mock_handler) + mock_read_frame.assert_called_once() + self.assertIsNone(state.get_active_ws()) + + def test_try_send_status_persists_without_client(self): + state = state_lib.PlaygroundState() + websocket.try_send_status(state, 'Starting environment...') + self.assertEqual(state.get_status(), 'Starting environment...') + + def test_try_send_logs(self): + state = state_lib.PlaygroundState() + mock_sock = mock.create_autospec(socket.socket, instance=True) + state.active_ws = mock_sock + + logs = [{'message': 'log1'}, {'message': 'log2'}] + websocket.try_send_logs(state, logs) + + mock_sock.sendall.assert_called_once() + sent_data = mock_sock.sendall.call_args[0][0] + self.assertIn(b'log1', sent_data) + self.assertIn(b'log2', sent_data) + + def test_try_send_stats_battery_power(self): + state = state_lib.PlaygroundState() + mock_sock = mock.create_autospec(socket.socket, instance=True) + state.active_ws = mock_sock + + stats = { + 'ack_latency': 10.0, + 'battery_power': { + 'level_percent': 85, + 'power_mw': 1500.0, + }, + } + websocket.try_send_stats(state, stats) + + mock_sock.sendall.assert_called_once() + sent_data = mock_sock.sendall.call_args[0][0] + self.assertIn(b'85', sent_data) + self.assertIn(b'1500.0', sent_data) + + def test_build_ws_frame_boundaries(self): + self.assertEqual( + websocket._build_ws_frame(0x81, b'a' * 125)[:2], b'\x81\x7d' + ) + self.assertEqual( + websocket._build_ws_frame(0x81, b'a' * 126)[:4], b'\x81\x7e\x00\x7e' + ) + self.assertEqual( + websocket._build_ws_frame(0x81, b'a' * 65535)[:4], b'\x81\x7e\xff\xff' + ) + self.assertEqual( + websocket._build_ws_frame(0x81, b'a' * 65536)[:10], + b'\x81\x7f\x00\x00\x00\x00\x00\x01\x00\x00', + ) + + def test_try_send_json_send_error_clears_active_ws(self): + state = state_lib.PlaygroundState() + mock_sock = mock.create_autospec(socket.socket, instance=True) + mock_sock.sendall.side_effect = OSError('broken pipe') + state.active_ws = mock_sock + + websocket.try_send_status(state, 'Connected') + + self.assertIsNone(state.active_ws) + self.assertEqual(state.get_status(), 'Connected') + + def test_try_send_adb_response(self): + state = state_lib.PlaygroundState() + mock_sock = mock.create_autospec(socket.socket, instance=True) + state.active_ws = mock_sock + + websocket.try_send_adb_response( + state, + command='shell wm size', + status='OK', + output='Physical size: 1080x1920', + error='', + duration_ms=4.5, + ) + + mock_sock.sendall.assert_called_once() + sent_data = mock_sock.sendall.call_args[0][0] + self.assertIn(b'adb_response', sent_data) + self.assertIn(b'shell wm size', sent_data) + self.assertIn(b'Physical size: 1080x1920', sent_data) + + def test_handle_message_reset(self): + state = state_lib.PlaygroundState() + self.assertFalse(state.get_reset()) + websocket._handle_message('reset', state) + self.assertTrue(state.get_reset()) + + def test_handle_message_adb(self): + state = state_lib.PlaygroundState() + websocket._handle_message('adb, shell dumpsys battery ', state) + self.assertEqual(state.adb_queue.qsize(), 1) + self.assertEqual(state.adb_queue.get(), 'shell dumpsys battery') + + def test_handle_message_adb_empty(self): + state = state_lib.PlaygroundState() + websocket._handle_message('adb, ', state) + self.assertTrue(state.adb_queue.empty()) + + +if __name__ == '__main__': + absltest.main() diff --git a/android_env/proto/BUILD.bazel b/android_env/proto/BUILD.bazel index 7144fbfa..a1637881 100644 --- a/android_env/proto/BUILD.bazel +++ b/android_env/proto/BUILD.bazel @@ -85,3 +85,31 @@ py_proto_library( name = "state_py_pb2", srcs = ["state.proto"], ) + +proto_library( + name = "device_env_service_proto", + srcs = ["device_env_service.proto"], + deps = [ + "//android_env/proto/a11y:a11y_proto", + "//android_env/proto/a11y:android_accessibility_forest_proto", + ], +) + +py_proto_library( + name = "device_env_service_py_pb2", + srcs = ["device_env_service.proto"], + proto_srcs = [ + "//android_env/proto/a11y:a11y.proto", + "//android_env/proto/a11y:android_accessibility_action.proto", + "//android_env/proto/a11y:android_accessibility_forest.proto", + "//android_env/proto/a11y:android_accessibility_node_info.proto", + "//android_env/proto/a11y:android_accessibility_node_info_clickable_span.proto", + "//android_env/proto/a11y:android_accessibility_tree.proto", + "//android_env/proto/a11y:android_accessibility_window_info.proto", + "//android_env/proto/a11y:rect.proto", + ], + deps = [ + "//android_env/proto/a11y:a11y_py_pb2", + "//android_env/proto/a11y:android_accessibility_forest_py_pb2", + ], +) diff --git a/android_env/proto/device_env_service.proto b/android_env/proto/device_env_service.proto new file mode 100644 index 00000000..f869459c --- /dev/null +++ b/android_env/proto/device_env_service.proto @@ -0,0 +1,305 @@ +// Copyright 2026 DeepMind Technologies Limited. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +syntax = "proto3"; + +package android_env.control; + +import "android_env/proto/a11y/a11y.proto"; +import "android_env/proto/a11y/android_accessibility_forest.proto"; + +option java_package = "com.androidenv.proto"; +option java_multiple_files = true; +option java_outer_classname = "DeviceEnvServiceProto"; + +// Enumerates all supported sensory streams and telemetry signals that can +// be observed from the Android device. +enum DeviceSignal { + DEVICE_SIGNAL_UNSPECIFIED = 0; + // A raw image of the device display (typically PNG or JPEG). + DEVICE_SIGNAL_SCREENSHOT = 1; + // The current UI accessibility hierarchy tree. + DEVICE_SIGNAL_ACCESSIBILITY_TREE = 2; + // Stream of UI accessibility events. + DEVICE_SIGNAL_ACCESSIBILITY_EVENTS = 3; + // The package name of the foreground application (e.g. "com.android.chrome"). + DEVICE_SIGNAL_ACTIVE_PACKAGE = 4; + // The class name of the current foreground activity. + DEVICE_SIGNAL_ACTIVE_ACTIVITY = 5; + // Recent logcat system logs. + DEVICE_SIGNAL_SYSTEM_LOGS = 6; + // List of currently running process names. + DEVICE_SIGNAL_RUNNING_PROCESSES = 7; + // Key-value map of system settings (e.g. screen_brightness, volume). + DEVICE_SIGNAL_DEVICE_SETTINGS = 8; + // Linear acceleration force along the X, Y, and Z axes (in m/s^2). + DEVICE_SIGNAL_SENSOR_ACCELEROMETER = 9; + // Rate of rotation around the X, Y, and Z axes (in rad/s). + DEVICE_SIGNAL_SENSOR_GYROSCOPE = 10; + // Ambient magnetic field along the X, Y, and Z axes (in micro-Tesla). + DEVICE_SIGNAL_SENSOR_MAGNETOMETER = 11; + // Atmospheric pressure (in hectopascals / hPa). + DEVICE_SIGNAL_SENSOR_BAROMETER = 12; + // Ambient room temperature (in degrees Celsius). + DEVICE_SIGNAL_SENSOR_AMBIENT_TEMP = 13; + // Ambient light level (in lux). + DEVICE_SIGNAL_SENSOR_LIGHT = 14; + // Proximity sensor distance (in centimeters). + DEVICE_SIGNAL_SENSOR_PROXIMITY = 15; + // Heart rate from wearable sensors (in beats per minute). + DEVICE_SIGNAL_SENSOR_HEART_RATE = 16; + // Current GPS coordinates (Latitude, Longitude, Altitude). + DEVICE_SIGNAL_LOCATION_GPS = 17; + // Continuous audio frames captured from the device speaker output. + DEVICE_SIGNAL_AUDIO_OUTPUT = 18; + // Battery and power consumption metrics (current, voltage, power, temp, + // status). + DEVICE_SIGNAL_BATTERY_POWER = 19; +} + +// Enumerates the types of interactive input actions that can be injected. +enum ActionType { + ACTION_TYPE_UNSPECIFIED = 0; + // Pressing down on the screen at a specific coordinate. + ACTION_TYPE_TOUCH_DOWN = 1; + // Lifting up from the screen. + ACTION_TYPE_TOUCH_UP = 2; + // Dragging/moving a pointer across the screen. + ACTION_TYPE_TOUCH_MOVE = 3; + // Injecting a physical hardware keycode event (e.g. BACK, HOME). + ACTION_TYPE_KEY_EVENT = 4; + // Native text insertion via Accessibility ACTION_SET_TEXT. + ACTION_TYPE_SET_TEXT = 5; + // Setting text selection or cursor position via Accessibility + // ACTION_SET_SELECTION. + ACTION_TYPE_SET_SELECTION = 6; + // Native clipboard cut via Accessibility ACTION_CUT. + ACTION_TYPE_CUT = 7; + // Native clipboard copy via Accessibility ACTION_COPY. + ACTION_TYPE_COPY = 8; + // Native clipboard paste via Accessibility ACTION_PASTE. + ACTION_TYPE_PASTE = 9; +} + +// Represents a 2D coordinate point on the screen. +message Point { + // X-coordinate normalized between 0.0 (left edge) and 1.0 (right edge). + float x = 1; + // Y-coordinate normalized between 0.0 (top edge) and 1.0 (bottom edge). + float y = 2; +} + +// Represents an input action to be executed on the target Android device. +message Action { + // The type of action to perform. Always required. + ActionType action_type = 1; + // The screen position for touch actions. + // REQUIRED if action_type is TOUCH_DOWN or TOUCH_MOVE. Omitted otherwise. + optional Point touch_position = 2; + // The Android keycode integer (e.g. 4 for BACK, 3 for HOME). + // REQUIRED if action_type is KEY_EVENT. Omitted otherwise. + optional int32 keycode = 3; + // Identifies the pointer (finger) for multi-touch support. + // Optional. Defaults to 0 (primary pointer) if unset. Value range: [0, 9]. + optional int32 pointer_id = 4; + // Text to insert into target text field. + // REQUIRED if action_type is SET_TEXT. Omitted otherwise. + optional string text = 5; + // Selection start index for text field cursor/highlight. + // REQUIRED if action_type is SET_SELECTION. Omitted otherwise. + optional int32 selection_start = 6; + // Selection end index for text field cursor/highlight. + // REQUIRED if action_type is SET_SELECTION. Omitted otherwise. + optional int32 selection_end = 7; +} + +// Metadata accompanying a streamed signal value, enabling synchronization and +// latency tracking. +message StreamedSignalMetadata { + // The device wall-clock time when the signal was captured (microseconds since + // epoch). + int64 timestamp_us = 1; + // Monotonically increasing sequence identifier to detect frame drops or + // out-of-order delivery. + int64 sequence_id = 2; +} + +// Represents raw decoded screenshot image pixel data. +message DecodedScreenshot { + // Uncompressed raw RGB pixel bytes. + bytes raw_pixels = 1; + // Frame width in pixels. + int32 width = 2; + // Frame height in pixels. + int32 height = 3; +} + +// Wraps screenshot image bytes. Exactly one representation format is set. +message ScreenshotValue { + oneof format { + // Raw compressed image bytes (typically JPEG or PNG format). + bytes encoded_bytes = 1; + // Uncompressed raw RGB pixel payload. + DecodedScreenshot decoded = 2; + } +} + +// GPS geolocation coordinates. +message GpsCoordinates { + // Latitude in degrees. Value range: [-90.0, 90.0]. + double latitude = 1; + // Longitude in degrees. Value range: [-180.0, 180.0]. + double longitude = 2; + // Altitude in meters above the WGS 84 reference ellipsoid. + double altitude = 3; +} + +// A single block of raw audio data captured from the device. +message AudioFrame { + // Raw PCM audio bytes. + bytes raw_bytes = 1; + // Audio sampling rate in Hertz. Value range: [8000, 48000] (typically 44100). + int32 sample_rate_hz = 2; + // Number of audio channels. Value range: [1, 2] (1 = mono, 2 = stereo). + int32 channel_count = 3; + // PCM encoding format string (e.g., "AUDIO_FORMAT_PCM_16BIT"). + optional string sample_format = 4; +} + +// Telemetry metrics for device battery state, current draw, voltage, and +// calculated power. +message BatteryPowerMetrics { + // Battery charge level percentage [0-100]. + int32 level_percent = 1; + // Instantaneous current draw in microamperes (uA). Positive for discharge, + // negative for charging. + int64 current_now_ua = 2; + // Voltage in millivolts (mV). + int32 voltage_mv = 3; + // Calculated instantaneous power in milliwatts (mW). + double power_mw = 4; + // Battery temperature in degrees Celsius (°C). + float temperature_celsius = 5; + // Battery health/status string (e.g., "CHARGING", "DISCHARGING", "FULL"). + string status = 6; + // Plugged power source ("AC", "USB", "WIRELESS", "UNPLUGGED"). + string plugged_source = 7; + // Total remaining charge counter in microampere-hours (uAh). + int64 charge_counter_uah = 8; + // Current thermal status (0 = NONE to 6 = SHUTDOWN). + int32 thermal_status = 9; + reserved 10; + reserved "is_synthetic_emulator_fallback"; + // Total remaining battery energy in nanowatt-hours (nWh). + int64 energy_counter_nwh = 11; +} + +// Wrapper for an array of floats (e.g. sensor readings). +message FloatArray { + repeated float values = 1; +} + +// Wrapper for a list of strings (e.g. log lines). +message StringList { + repeated string values = 1; +} + +// Wrapper for key-value maps (e.g. system settings). +message StringMap { + map entries = 1; +} + +// Represents a single streamed telemetry signal packaging metadata and the +// value. +message StreamedSignal { + // The type of signal this message represents. + DeviceSignal type = 1; + // Metadata for this signal capture. + StreamedSignalMetadata metadata = 2; + + // The actual signal value. + oneof value { + ScreenshotValue screenshot = 3; + AndroidAccessibilityForest accessibility_forest = 4; + EventRequestList accessibility_events = 5; + string active_package = 6; + string active_activity = 7; + StringList system_logs = 8; + StringList running_processes = 9; + StringMap device_settings = 10; + FloatArray accelerometer = 11; + FloatArray gyroscope = 12; + FloatArray magnetometer = 13; + float barometer = 14; + float ambient_temp = 15; + float light_level = 16; + float proximity = 17; + GpsCoordinates location_gps = 18; + float heart_rate = 19; + AudioFrame audio_output = 20; + BatteryPowerMetrics battery_power = 21; + } +} + +// Represents the aggregated state of all subscribed signals on the device. +message DeviceState { + // List of active signal observations currently subscribed. + repeated StreamedSignal signals = 1; + // Wall-clock time on the device when this state snapshot was assembled + // (microseconds since epoch). + int64 snapshot_timestamp_us = 2; + // Current screen rotation (0: ROTATION_0, 1: ROTATION_90, 2: ROTATION_180, 3: + // ROTATION_270). + int32 orientation = 3; +} + +// Top-level wrapper message sent from the Python Client to the Android Server. +message ClientMessage { + oneof payload { + // Configure which telemetry signals the server should stream/capture. + UpdateSubscriptionsRequest update_subscriptions = 1; + // Inject a physical user action. + InjectActionRequest inject_action = 2; + // Query the current state of the device synchronously. + GetStateRequest get_state = 3; + } +} + +// Request to modify the list of active telemetry subscriptions. +message UpdateSubscriptionsRequest { + repeated DeviceSignal active_signals = 1; +} + +// Request to inject an input action. +message InjectActionRequest { + Action action = 1; +} + +// Empty request to trigger a manual state fetch. +message GetStateRequest {} + +// Top-level wrapper message sent from the Android Server back to the Client. +message ServerMessage { + oneof payload { + // Contains the requested device state telemetry. + DeviceState state = 1; + // Error details if an operation failed on the server. + ErrorResponse error = 2; + } +} + +// Details of an error encountered on the server. +message ErrorResponse { + string message = 1; +} diff --git a/android_env/rl/BUILD.bazel b/android_env/rl/BUILD.bazel new file mode 100644 index 00000000..1b3f6fcc --- /dev/null +++ b/android_env/rl/BUILD.bazel @@ -0,0 +1,38 @@ +load("@rules_python//python:defs.bzl", "py_library", "py_test") + +package(default_visibility = ["//visibility:public"]) + +py_library( + name = "environment", + srcs = ["environment.py"], + deps = [ + "//android_env:env_interface", + "//android_env/components:action_type", + "//android_env/components:android_device", + "//android_env/components:specs", + "//android_env/components:telemetry", + "//android_env/proto:adb_py_pb2", + "//android_env/proto:device_env_service_py_pb2", + "@pip//absl_py", + "@pip//dm_env", + "@pip//numpy", + ], +) + +py_test( + name = "environment_test", + srcs = ["environment_test.py"], + legacy_create_init = False, + deps = [ + ":environment", + "//android_env/components:action_type", + "//android_env/components:android_device", + "//android_env/components:config_classes", + "//android_env/components:device_connection", + "//android_env/components:telemetry", + "//android_env/proto:adb_py_pb2", + "//android_env/proto:device_env_service_py_pb2", + "@pip//absl_py", + "@pip//numpy", + ], +) diff --git a/android_env/rl/environment.py b/android_env/rl/environment.py new file mode 100644 index 00000000..92b98ce0 --- /dev/null +++ b/android_env/rl/environment.py @@ -0,0 +1,292 @@ +# coding=utf-8 +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Task-Free AndroidEnv implementation connecting to on-device server. + +This environment connects to an on-device server and exposes device signals +including structured telemetry via `env.task_extras()`. +""" + +import time +from typing import Any, Self + +from absl import logging +from android_env import env_interface +from android_env.components import action_type +from android_env.components import android_device +from android_env.components import telemetry +from android_env.components.specs import base_action_spec +from android_env.components.specs import base_observation_spec +from android_env.proto import adb_pb2 +from android_env.proto import device_env_service_pb2 +import dm_env +from dm_env import specs +import numpy as np + + +class AndroidEnv(env_interface.AndroidEnvInterface): + """Task-Free AndroidEnv connected to a real device via WebSockets.""" + + def __init__(self, device: android_device.AndroidDevice): + """Initializes the AndroidEnv. + + Args: + device: AndroidDevice handle instance. + """ + self._device = device + self._latest_observation_time = 0.0 + state = self._device.get_state() + h, w = 1080, 1920 + for signal in state.signals: + if ( + signal.type == device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT + and signal.HasField('screenshot') + and signal.screenshot.HasField('decoded') + ): + h = signal.screenshot.decoded.height or 1080 + w = signal.screenshot.decoded.width or 1920 + break + self._pixel_shape = (h, w, 3) + logging.info('Detected video resolution: %s', self._pixel_shape) + + self._active_signals = { + device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT, + device_env_service_pb2.DEVICE_SIGNAL_ACTIVE_PACKAGE, + device_env_service_pb2.DEVICE_SIGNAL_SYSTEM_LOGS, + device_env_service_pb2.DEVICE_SIGNAL_ACCESSIBILITY_TREE, + device_env_service_pb2.DEVICE_SIGNAL_BATTERY_POWER, + } + self._device.update_subscriptions(self._active_signals) + self._latest_state = None + self._latest_pixels = None + self._is_touching = False + self._last_touch_position = None + self._last_non_touch_action = None + + def action_spec(self) -> dict[str, specs.Array]: + return base_action_spec(num_fingers=1, enable_key_events=True) + + def observation_spec(self) -> dict[str, specs.Array]: + return base_observation_spec( + height=self._pixel_shape[0], width=self._pixel_shape[1] + ) + + def execute_adb_call(self, call: adb_pb2.AdbRequest) -> adb_pb2.AdbResponse: + """Executes an ADB call using the underlying AndroidDevice.""" + return self._device.execute_adb_call(call) + + def task_extras(self, latest_only: bool = True) -> dict[str, Any]: + """Returns extras from the latest state signals. + + Args: + latest_only: Whether to return only the latest extra value per signal. + + Returns: + A dictionary mapping signal names to telemetry payloads. + """ + extras: dict[str, Any] = {} + if self._latest_state is not None: + for sig in self._latest_state.signals: + if sig.type == device_env_service_pb2.DEVICE_SIGNAL_SYSTEM_LOGS: + extras['logs'] = list(sig.system_logs.values) + elif ( + sig.type + == device_env_service_pb2.DEVICE_SIGNAL_ACCESSIBILITY_EVENTS + ): + extras['accessibility_events'] = list(sig.accessibility_events.events) + elif ( + sig.type == device_env_service_pb2.DEVICE_SIGNAL_ACCESSIBILITY_TREE + ): + extras['accessibility_tree'] = sig.accessibility_forest + elif ( + sig.type == device_env_service_pb2.DEVICE_SIGNAL_BATTERY_POWER + and sig.HasField('battery_power') + ): + extras['battery_power'] = telemetry.BatteryPowerMetrics.from_proto( + sig.battery_power + ) + return extras + + def _build_observation( + self, state: device_env_service_pb2.DeviceState, timedelta_us: int + ) -> dict[str, Any]: + pixels, active_package, audio, _ = self._parse_signals(state) + if pixels is not None: + self._latest_pixels = pixels + else: + pixels = ( + self._latest_pixels + if self._latest_pixels is not None + else np.zeros(self._pixel_shape, dtype=np.uint8) + ) + + orient_one_hot = np.zeros(4, dtype=np.uint8) + orient_val = state.orientation + if 0 <= orient_val < 4: + orient_one_hot[orient_val] = 1 + else: + orient_one_hot[0] = 1 + + obs: dict[str, Any] = { + 'pixels': pixels, + 'timedelta': np.int64(timedelta_us), + 'active_package': active_package or '', + 'orientation': orient_one_hot, + } + if audio is not None: + obs['audio'] = audio + return obs + + def reset(self) -> dm_env.TimeStep: + """Resets the environment.""" + logging.info('Resetting AndroidEnv...') + self._latest_pixels = None + self._latest_state = self._device.get_state() + obs = self._build_observation(self._latest_state, timedelta_us=0) + self._latest_observation_time = time.time() + return dm_env.restart(obs) + + def step(self, action: dict[str, Any]) -> dm_env.TimeStep: + """Takes a step in the environment by executing an action.""" + env_action_type = action.get('action_type') + + proto_action = device_env_service_pb2.Action() + if env_action_type is not None: + if env_action_type == action_type.ActionType.TOUCH: + proto_action.action_type = device_env_service_pb2.ACTION_TYPE_TOUCH_DOWN + pos = action.get('touch_position', [0.0, 0.0]) + proto_action.touch_position.x = float(pos[0]) + proto_action.touch_position.y = float(pos[1]) + self._is_touching = True + self._last_touch_position = (float(pos[0]), float(pos[1])) + elif env_action_type == action_type.ActionType.REPEAT: + if self._is_touching: + assert self._last_touch_position is not None + proto_action.action_type = ( + device_env_service_pb2.ACTION_TYPE_TOUCH_MOVE + ) + proto_action.touch_position.x = self._last_touch_position[0] + proto_action.touch_position.y = self._last_touch_position[1] + elif env_action_type == action_type.ActionType.LIFT: + proto_action.action_type = device_env_service_pb2.ACTION_TYPE_TOUCH_UP + if 'touch_position' in action: + pos = action['touch_position'] + proto_action.touch_position.x = float(pos[0]) + proto_action.touch_position.y = float(pos[1]) + elif self._last_touch_position is not None: + proto_action.touch_position.x = self._last_touch_position[0] + proto_action.touch_position.y = self._last_touch_position[1] + self._is_touching = False + self._last_touch_position = None + elif env_action_type == action_type.ActionType.KEYDOWN: + proto_action.action_type = device_env_service_pb2.ACTION_TYPE_KEY_EVENT + if 'keycode' in action: + proto_action.keycode = int(action['keycode']) + self._last_non_touch_action = device_env_service_pb2.Action() + self._last_non_touch_action.CopyFrom(proto_action) + elif env_action_type == action_type.ActionType.KEYUP: + proto_action.action_type = ( + device_env_service_pb2.ACTION_TYPE_UNSPECIFIED + ) + self._last_non_touch_action = None + elif env_action_type == action_type.ActionType.TEXT: + proto_action.action_type = device_env_service_pb2.ACTION_TYPE_SET_TEXT + if 'text' in action: + proto_action.text = str(action['text']) + self._last_non_touch_action = device_env_service_pb2.Action() + self._last_non_touch_action.CopyFrom(proto_action) + elif env_action_type == action_type.ActionType.SET_SELECTION: + proto_action.action_type = ( + device_env_service_pb2.ACTION_TYPE_SET_SELECTION + ) + if 'selection' in action: + sel = action['selection'] + proto_action.selection_start = int(sel[0]) + proto_action.selection_end = int(sel[1]) + self._last_non_touch_action = device_env_service_pb2.Action() + self._last_non_touch_action.CopyFrom(proto_action) + elif env_action_type == action_type.ActionType.CUT: + proto_action.action_type = device_env_service_pb2.ACTION_TYPE_CUT + self._last_non_touch_action = device_env_service_pb2.Action() + self._last_non_touch_action.CopyFrom(proto_action) + elif env_action_type == action_type.ActionType.COPY: + proto_action.action_type = device_env_service_pb2.ACTION_TYPE_COPY + self._last_non_touch_action = device_env_service_pb2.Action() + self._last_non_touch_action.CopyFrom(proto_action) + elif env_action_type == action_type.ActionType.PASTE: + proto_action.action_type = device_env_service_pb2.ACTION_TYPE_PASTE + self._last_non_touch_action = device_env_service_pb2.Action() + self._last_non_touch_action.CopyFrom(proto_action) + else: + proto_action.action_type = ( + device_env_service_pb2.ACTION_TYPE_UNSPECIFIED + ) + + if ( + proto_action.action_type + != device_env_service_pb2.ACTION_TYPE_UNSPECIFIED + ): + self._device.send_action(proto_action) + self._latest_state = self._device.get_state() + + now = time.time() + timedelta_us = int((now - self._latest_observation_time) * 1e6) + self._latest_observation_time = now + + obs = self._build_observation(self._latest_state, timedelta_us=timedelta_us) + return dm_env.transition(reward=0.0, observation=obs) + + def _parse_signals( + self, state: device_env_service_pb2.DeviceState | None + ) -> tuple[ + np.ndarray | None, + str | None, + np.ndarray | None, + list[str] | None, + ]: + if state is None: + return None, None, None, None + pixels = None + active_package = None + audio = None + logs = None + + for sig in state.signals: + if sig.type == device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT: + if sig.screenshot.HasField('decoded'): + decoded = sig.screenshot.decoded + pixels = np.frombuffer(decoded.raw_pixels, dtype=np.uint8).reshape( + (decoded.height, decoded.width, 3) + ) + elif sig.type == device_env_service_pb2.DEVICE_SIGNAL_ACTIVE_PACKAGE: + active_package = sig.active_package + elif sig.type == device_env_service_pb2.DEVICE_SIGNAL_AUDIO_OUTPUT: + raw_bytes = sig.audio_output.raw_bytes + if raw_bytes: + audio = np.frombuffer(raw_bytes, dtype=np.int16).reshape(-1, 2) + elif sig.type == device_env_service_pb2.DEVICE_SIGNAL_SYSTEM_LOGS: + logs = list(sig.system_logs.values) + + return pixels, active_package, audio, logs + + def close(self) -> None: + logging.info('Closing AndroidEnv...') + self._device.close() + + def __enter__(self) -> Self: + return self + + def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: + self._device.__exit__(exc_type, exc_value, traceback) diff --git a/android_env/rl/environment_test.py b/android_env/rl/environment_test.py new file mode 100644 index 00000000..0203cecd --- /dev/null +++ b/android_env/rl/environment_test.py @@ -0,0 +1,454 @@ +# coding=utf-8 +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for task-free AndroidEnv in android_env.rl.environment.""" + +from unittest import mock +from absl.testing import absltest +from android_env.components import action_type +from android_env.components import android_device +from android_env.components import config_classes +from android_env.components import device_connection +from android_env.components import telemetry +from android_env.proto import adb_pb2 +from android_env.proto import device_env_service_pb2 +from android_env.rl import environment as android_device_env +import numpy as np + + +class FakeDeviceConnection(device_connection.DeviceConnection): + + def __init__( + self, config: config_classes.DeviceConnectionConfig | None = None + ): + config = config or config_classes.DeviceConnectionConfig() + self._config = config + self._video_codec: str = 'h264' + self._video_width: int = 640 + self._video_height: int = 480 + self.fake_frame: np.ndarray = np.ones((480, 640, 3), dtype=np.uint8) + self.device_state: device_env_service_pb2.DeviceState = ( + device_env_service_pb2.DeviceState() + ) + self.injected_actions: list[device_env_service_pb2.Action] = [] + self.subscriptions: set[int] = { + device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT + } + self.closed: bool = False + + def connect(self): + pass + + def get_video_metadata(self) -> tuple[str, int, int]: + return self._video_codec, self._video_width, self._video_height + + def get_device_state(self) -> device_env_service_pb2.DeviceState: + if device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT in self.subscriptions: + if not any( + sig.type == device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT + for sig in self.device_state.signals + ): + h, w, _ = self.fake_frame.shape + sig = self.device_state.signals.add() + sig.type = device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT + sig.screenshot.decoded.raw_pixels = self.fake_frame.tobytes() + sig.screenshot.decoded.width = w + sig.screenshot.decoded.height = h + return self.device_state + + def inject_action(self, action: device_env_service_pb2.Action): + self.injected_actions.append(action) + + def update_subscriptions(self, signals: set[int]): + self.subscriptions = set(signals) + + def send_message(self, msg: device_env_service_pb2.ClientMessage): + payload_type = msg.WhichOneof('payload') + if payload_type == 'inject_action': + self.injected_actions.append(msg.inject_action.action) + elif payload_type == 'update_subscriptions': + self.subscriptions = set(msg.update_subscriptions.active_signals) + elif msg.HasField('inject_action'): + self.injected_actions.append(msg.inject_action.action) + elif msg.HasField('update_subscriptions'): + self.subscriptions = set(msg.update_subscriptions.active_signals) + + def close(self): + self.closed = True + + +def _make_env( + fake_conn: FakeDeviceConnection, + on_close_callbacks=None, + on_error_callbacks=None, +) -> android_device_env.AndroidEnv: + dev = android_device.AndroidDevice( + connection=fake_conn, + on_close_callbacks=on_close_callbacks, + on_error_callbacks=on_error_callbacks, + ) + return android_device_env.AndroidEnv(device=dev) + + +class AndroidDeviceEnvTest(absltest.TestCase): + + def test_init(self): + fake_conn = FakeDeviceConnection() + _ = _make_env(fake_conn) + self.assertEqual( + fake_conn.subscriptions, + { + device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT, + device_env_service_pb2.DEVICE_SIGNAL_ACTIVE_PACKAGE, + device_env_service_pb2.DEVICE_SIGNAL_SYSTEM_LOGS, + device_env_service_pb2.DEVICE_SIGNAL_ACCESSIBILITY_TREE, + device_env_service_pb2.DEVICE_SIGNAL_BATTERY_POWER, + }, + ) + + def test_task_extras_battery_power(self): + fake_conn = FakeDeviceConnection() + env = _make_env(fake_conn) + state = device_env_service_pb2.DeviceState() + sig = state.signals.add() + sig.type = device_env_service_pb2.DEVICE_SIGNAL_BATTERY_POWER + sig.battery_power.level_percent = 85 + sig.battery_power.current_now_ua = 450000 + sig.battery_power.voltage_mv = 4200 + sig.battery_power.power_mw = 1890.0 + sig.battery_power.temperature_celsius = 28.5 + sig.battery_power.status = 'DISCHARGING' + sig.battery_power.plugged_source = 'UNPLUGGED' + sig.battery_power.charge_counter_uah = 3000000 + sig.battery_power.thermal_status = 0 + sig.battery_power.energy_counter_nwh = 15000000000 + fake_conn.device_state = state + env.reset() + extras = env.task_extras() + self.assertIn('battery_power', extras) + bp = extras['battery_power'] + self.assertIsInstance(bp, telemetry.BatteryPowerMetrics) + + # Verify clean structured attribute access. + self.assertEqual(bp.level_percent, 85) + self.assertEqual(bp.current_now_ua, 450000) + self.assertEqual(bp.voltage_mv, 4200) + self.assertAlmostEqual(bp.power_mw, 1890.0) + self.assertAlmostEqual(bp.temperature_celsius, 28.5) + self.assertEqual(bp.status, 'DISCHARGING') + self.assertEqual(bp.plugged_source, 'UNPLUGGED') + self.assertEqual(bp.charge_counter_uah, 3000000) + self.assertEqual(bp.thermal_status, 0) + self.assertEqual(bp.energy_counter_nwh, 15000000000) + + def test_specs(self): + fake_conn = FakeDeviceConnection() + env = _make_env(fake_conn) + action_spec = env.action_spec() + self.assertIn('action_type', action_spec) + self.assertIn('touch_position', action_spec) + self.assertIn('keycode', action_spec) + + obs_spec = env.observation_spec() + self.assertIn('pixels', obs_spec) + self.assertIn('orientation', obs_spec) + self.assertIn('timedelta', obs_spec) + + def test_reset(self): + fake_conn = FakeDeviceConnection() + env = _make_env(fake_conn) + timestep = env.reset() + self.assertTrue(timestep.first()) + self.assertIn('pixels', timestep.observation) + + def test_step_touch(self): + fake_conn = FakeDeviceConnection() + env = _make_env(fake_conn) + _ = env.reset() + action = { + 'action_type': np.array(action_type.ActionType.TOUCH, dtype=np.float32), + 'touch_position': np.array([0.5, 0.5], dtype=np.float32), + } + timestep = env.step(action) + self.assertTrue(timestep.mid()) + self.assertLen(fake_conn.injected_actions, 1) + + def test_close(self): + fake_conn = FakeDeviceConnection() + on_close = mock.MagicMock() + env = _make_env(fake_conn, on_close_callbacks=[on_close]) + env.close() + self.assertTrue(fake_conn.closed) + on_close.assert_called_once() + + def test_video_resolution_detection_from_screenshot_signal(self): + fake_conn = FakeDeviceConnection() + state = device_env_service_pb2.DeviceState() + sig = state.signals.add() + sig.type = device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT + sig.screenshot.decoded.height = 480 + sig.screenshot.decoded.width = 640 + fake_conn.device_state = state + + dev = android_device.AndroidDevice(connection=fake_conn) + env = android_device_env.AndroidEnv(device=dev) + obs_spec = env.observation_spec() + self.assertEqual(obs_spec['pixels'].shape, (480, 640, 3)) + + def test_video_resolution_fallback_default(self): + fake_conn = FakeDeviceConnection() + fake_conn.subscriptions = set() + fake_conn.device_state = device_env_service_pb2.DeviceState() + + dev = android_device.AndroidDevice(connection=fake_conn) + env = android_device_env.AndroidEnv(device=dev) + obs_spec = env.observation_spec() + self.assertEqual(obs_spec['pixels'].shape, (1080, 1920, 3)) + + def test_build_observation_pixels_and_cached(self): + fake_conn = FakeDeviceConnection() + raw_pixels = np.ones((2, 2, 3), dtype=np.uint8) + fake_conn.fake_frame = raw_pixels + state = device_env_service_pb2.DeviceState() + sig = state.signals.add() + sig.type = device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT + sig.screenshot.decoded.width = 2 + sig.screenshot.decoded.height = 2 + sig.screenshot.decoded.raw_pixels = raw_pixels.tobytes() + fake_conn.device_state = state + + dev = android_device.AndroidDevice(connection=fake_conn) + env = android_device_env.AndroidEnv(device=dev) + ts = env.reset() + np.testing.assert_array_equal(ts.observation['pixels'], raw_pixels) + self.assertEqual(ts.observation['active_package'], '') + + # Next step with no new screenshot reuses cached pixels. + fake_conn.device_state = device_env_service_pb2.DeviceState() + ts = env.step({}) + np.testing.assert_array_equal(ts.observation['pixels'], raw_pixels) + self.assertEqual(ts.observation['active_package'], '') + + def test_build_observation_orientation_bounds(self): + fake_conn = FakeDeviceConnection() + dev = android_device.AndroidDevice(connection=fake_conn) + env = android_device_env.AndroidEnv(device=dev) + + # Valid orientation 2 -> one-hot [0, 0, 1, 0] + fake_conn.device_state.orientation = 2 + ts = env.step({}) + np.testing.assert_array_equal( + ts.observation['orientation'], np.array([0, 0, 1, 0], dtype=np.uint8) + ) + + # Out-of-bounds orientation 99 -> fallback one-hot [1, 0, 0, 0] + fake_conn.device_state.orientation = 99 + ts = env.step({}) + np.testing.assert_array_equal( + ts.observation['orientation'], np.array([1, 0, 0, 0], dtype=np.uint8) + ) + + def test_task_extras(self): + fake_conn = FakeDeviceConnection() + state = device_env_service_pb2.DeviceState() + sig1 = state.signals.add() + sig1.type = device_env_service_pb2.DEVICE_SIGNAL_SYSTEM_LOGS + sig1.system_logs.values.append('log1') + sig2 = state.signals.add() + sig2.type = device_env_service_pb2.DEVICE_SIGNAL_ACCESSIBILITY_EVENTS + ev = sig2.accessibility_events.events.add() + sig3 = state.signals.add() + sig3.type = device_env_service_pb2.DEVICE_SIGNAL_ACCESSIBILITY_TREE + fake_conn.device_state = state + + dev = android_device.AndroidDevice(connection=fake_conn) + env = android_device_env.AndroidEnv(device=dev) + env.reset() + extras = env.task_extras() + self.assertEqual(extras['logs'], ['log1']) + self.assertEqual(extras['accessibility_events'], [ev]) + self.assertEqual(extras['accessibility_tree'], sig3.accessibility_forest) + + def test_step_action_types(self): + mock_dev = mock.create_autospec(android_device.AndroidDevice, instance=True) + mock_dev.get_state.return_value = device_env_service_pb2.DeviceState() + env = android_device_env.AndroidEnv(device=mock_dev) + + # 1. TOUCH + env.step({ + 'action_type': action_type.ActionType.TOUCH, + 'touch_position': [0.2, 0.8], + }) + self.assertTrue(mock_dev.send_action.called) + last_act = mock_dev.send_action.call_args[0][0] + self.assertEqual( + last_act.action_type, device_env_service_pb2.ACTION_TYPE_TOUCH_DOWN + ) + self.assertAlmostEqual(last_act.touch_position.x, 0.2) + self.assertAlmostEqual(last_act.touch_position.y, 0.8) + + # 2. REPEAT when touching + mock_dev.send_action.reset_mock() + env.step({'action_type': action_type.ActionType.REPEAT}) + last_act = mock_dev.send_action.call_args[0][0] + self.assertEqual( + last_act.action_type, device_env_service_pb2.ACTION_TYPE_TOUCH_MOVE + ) + + # 3. LIFT with position + mock_dev.send_action.reset_mock() + env.step({ + 'action_type': action_type.ActionType.LIFT, + 'touch_position': [0.5, 0.5], + }) + last_act = mock_dev.send_action.call_args[0][0] + self.assertEqual( + last_act.action_type, device_env_service_pb2.ACTION_TYPE_TOUCH_UP + ) + + # 4. REPEAT when not touching (no action sent) + mock_dev.send_action.reset_mock() + env.step({'action_type': action_type.ActionType.REPEAT}) + mock_dev.send_action.assert_not_called() + + # 5. KEYDOWN + mock_dev.send_action.reset_mock() + env.step({'action_type': action_type.ActionType.KEYDOWN, 'keycode': 66}) + last_act = mock_dev.send_action.call_args[0][0] + self.assertEqual( + last_act.action_type, device_env_service_pb2.ACTION_TYPE_KEY_EVENT + ) + self.assertEqual(last_act.keycode, 66) + + # 6. KEYUP + mock_dev.send_action.reset_mock() + env.step({'action_type': action_type.ActionType.KEYUP}) + + # 7. Unrecognized action type + mock_dev.send_action.reset_mock() + env.step({'action_type': 999}) + mock_dev.send_action.assert_not_called() + + # 8. TEXT action + mock_dev.send_action.reset_mock() + env.step( + {'action_type': action_type.ActionType.TEXT, 'text': 'hello world'} + ) + last_act = mock_dev.send_action.call_args[0][0] + self.assertEqual( + last_act.action_type, device_env_service_pb2.ACTION_TYPE_SET_TEXT + ) + self.assertEqual(last_act.text, 'hello world') + + def test_parse_signals_extra_signals(self): + fake_conn = FakeDeviceConnection() + state = device_env_service_pb2.DeviceState() + sig1 = state.signals.add() + sig1.type = device_env_service_pb2.DEVICE_SIGNAL_ACTIVE_PACKAGE + sig1.active_package = 'com.example.app' + sig2 = state.signals.add() + sig2.type = device_env_service_pb2.DEVICE_SIGNAL_AUDIO_OUTPUT + pcm = np.array([[10, 20]], dtype=np.int16) + sig2.audio_output.raw_bytes = pcm.tobytes() + fake_conn.device_state = state + + dev = android_device.AndroidDevice(connection=fake_conn) + env = android_device_env.AndroidEnv(device=dev) + ts = env.reset() + self.assertEqual(ts.observation['active_package'], 'com.example.app') + np.testing.assert_array_equal(ts.observation['audio'], pcm) + + def test_parse_signals_none_state(self): + fake_conn = FakeDeviceConnection() + dev = android_device.AndroidDevice(connection=fake_conn) + env = android_device_env.AndroidEnv(device=dev) + res = env._parse_signals(None) + self.assertEqual(res, (None, None, None, None)) + + def test_parse_signals_screenshot_undecoded(self): + fake_conn = FakeDeviceConnection() + fake_conn.subscriptions = set() + state = device_env_service_pb2.DeviceState() + sig = state.signals.add() + sig.type = device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT + sig.screenshot.encoded_bytes = b'fake_compressed_jpeg_bytes' + fake_conn.device_state = state + + dev = android_device.AndroidDevice(connection=fake_conn) + env = android_device_env.AndroidEnv(device=dev) + pixels, _, _, _ = env._parse_signals(state) + self.assertIsNone(pixels) + self.assertTrue(sig.screenshot.HasField('encoded_bytes')) + + def test_parse_signals_system_logs(self): + fake_conn = FakeDeviceConnection() + state = device_env_service_pb2.DeviceState() + sig = state.signals.add() + sig.type = device_env_service_pb2.DEVICE_SIGNAL_SYSTEM_LOGS + sig.system_logs.values.append('log line 1') + sig.system_logs.values.append('log line 2') + fake_conn.device_state = state + + dev = android_device.AndroidDevice(connection=fake_conn) + env = android_device_env.AndroidEnv(device=dev) + _, _, _, logs = env._parse_signals(state) + self.assertEqual(logs, ['log line 1', 'log line 2']) + + def test_parse_signals_screenshot_empty_signal(self): + fake_conn = FakeDeviceConnection() + fake_conn.subscriptions = set() + state = device_env_service_pb2.DeviceState() + sig = state.signals.add() + sig.type = device_env_service_pb2.DEVICE_SIGNAL_SCREENSHOT + fake_conn.device_state = state + + dev = android_device.AndroidDevice(connection=fake_conn) + env = android_device_env.AndroidEnv(device=dev) + pixels, _, _, _ = env._parse_signals(state) + self.assertIsNone(pixels) + + def test_parse_signals_logs_none_for_non_logs_signals(self): + fake_conn = FakeDeviceConnection() + state = device_env_service_pb2.DeviceState() + sig1 = state.signals.add() + sig1.type = device_env_service_pb2.DEVICE_SIGNAL_ACTIVE_PACKAGE + sig1.active_package = 'com.example.app' + sig2 = state.signals.add() + sig2.type = device_env_service_pb2.DEVICE_SIGNAL_AUDIO_OUTPUT + fake_conn.device_state = state + + dev = android_device.AndroidDevice(connection=fake_conn) + env = android_device_env.AndroidEnv(device=dev) + _, _, _, logs = env._parse_signals(state) + self.assertIsNone(logs) + + def test_context_manager(self): + mock_dev = mock.create_autospec(android_device.AndroidDevice, instance=True) + with android_device_env.AndroidEnv(device=mock_dev) as env: + self.assertIsInstance(env, android_device_env.AndroidEnv) + mock_dev.__exit__.assert_called_once() + + def test_execute_adb_call_delegates_to_device(self): + mock_dev = mock.create_autospec(android_device.AndroidDevice, instance=True) + env = android_device_env.AndroidEnv(device=mock_dev) + req = adb_pb2.AdbRequest() + mock_dev.execute_adb_call.return_value = adb_pb2.AdbResponse() + res = env.execute_adb_call(req) + mock_dev.execute_adb_call.assert_called_once_with(req) + self.assertIsInstance(res, adb_pb2.AdbResponse) + + +if __name__ == '__main__': + absltest.main() diff --git a/android_env/wrappers/BUILD.bazel b/android_env/wrappers/BUILD.bazel index 1759c1a1..7bef5002 100644 --- a/android_env/wrappers/BUILD.bazel +++ b/android_env/wrappers/BUILD.bazel @@ -290,3 +290,36 @@ py_test( "@pip//numpy", ], ) + +py_library( + name = "task_wrapper", + srcs = ["task_wrapper.py"], + deps = [ + ":base_wrapper", + "//android_env:env_interface", + "//android_env/components:adb_call_parser", + "//android_env/components:config_classes", + "//android_env/components:log_stream", + "//android_env/components:task_manager", + "//android_env/proto:adb_py_pb2", + "//android_env/proto:task_py_pb2", + "@pip//dm_env", + "@pip//numpy", + ], +) + +py_test( + name = "task_wrapper_test", + srcs = ["task_wrapper_test.py"], + legacy_create_init = False, + deps = [ + ":task_wrapper", + "//android_env:env_interface", + "//android_env/components:fake_log_stream", + "//android_env/proto:adb_py_pb2", + "//android_env/proto:task_py_pb2", + "@pip//absl_py", + "@pip//dm_env", + "@pip//numpy", + ], +) diff --git a/android_env/wrappers/task_wrapper.py b/android_env/wrappers/task_wrapper.py new file mode 100644 index 00000000..6f3fff79 --- /dev/null +++ b/android_env/wrappers/task_wrapper.py @@ -0,0 +1,174 @@ +# coding=utf-8 +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Wraps AndroidEnv to add task support (rewards, setup/reset steps, extras).""" + +from collections.abc import Generator, Iterable +import queue +import re +from typing import Any + +from android_env import env_interface +from android_env.components import adb_call_parser as adb_call_parser_lib +from android_env.components import config_classes +from android_env.components import log_stream as log_stream_lib +from android_env.components import task_manager as task_manager_lib +from android_env.proto import adb_pb2 +from android_env.proto import task_pb2 +from android_env.wrappers import base_wrapper +import dm_env +import numpy as np + +_LOGCAT_PREFIX_RE = re.compile( + r'^\s*(?:\d{2}-\d{2}\s+\d{2}:\d{2}:\d{2}\.\d+|\d+\.\d+)' + r'\s+\d+\s+\d+\s+\S\s+[^:]*:\s*' +) + + +def _normalize_log_line(line: str) -> str: + """Formats a device log line into the epoch format expected by LogcatThread.""" + msg = _LOGCAT_PREFIX_RE.sub('', line) + return f' 0.000 1 1 I AndroidEnv: {msg}' + + +class _SignalLogStream(log_stream_lib.LogStream): + """LogStream backed by a synchronous queue fed from device state signals.""" + + def __init__(self) -> None: + super().__init__(verbose=False) + self._queue: queue.Queue[str | None] = queue.Queue() + + def _get_stream_output(self) -> Generator[str, None, None]: + while (line := self._queue.get()) is not None: + yield line + self._queue.task_done() + + def push_logs(self, lines: Iterable[str]) -> None: + """Enqueues log lines and blocks until LogcatThread processes all of them.""" + for line in lines: + self._queue.put(_normalize_log_line(line)) + self._queue.join() + + def stop_stream(self) -> None: + self._queue.put(None) + + +class _AdbCallParserAdapter(adb_call_parser_lib.AdbCallParser): + """Adapts AndroidEnvInterface.execute_adb_call to AdbCallParser interface.""" + + def __init__(self, env: env_interface.AndroidEnvInterface): + self._env = env + + def parse(self, request: adb_pb2.AdbRequest) -> adb_pb2.AdbResponse: + return self._env.execute_adb_call(request) + + +class TaskWrapper(base_wrapper.BaseWrapper): + """Wraps an AndroidEnv instance to support task-based RL execution.""" + + def __init__( + self, + env: env_interface.AndroidEnvInterface, + task: task_pb2.Task, + config: config_classes.TaskManagerConfig | None = None, + log_stream: log_stream_lib.LogStream | None = None, + ) -> None: + """Initializes TaskWrapper. + + Args: + env: Base environment implementing AndroidEnvInterface. + task: Task definition proto. + config: TaskManager configuration. + log_stream: Optional LogStream instance for log parsing. + """ + super().__init__(env) + self._task = task + self._task_manager = task_manager_lib.TaskManager(task=task, config=config) + self._signal_log_stream: _SignalLogStream | None = None + self._log_stream = log_stream + self._task_manager_started = False + self._reset_next_step = True + self._latest_extras: dict[str, np.ndarray] = {} + + def _ensure_task_manager_started(self) -> None: + if self._task_manager_started: + return + if self._log_stream is None: + self._signal_log_stream = _SignalLogStream() + self._log_stream = self._signal_log_stream + self._task_manager.start( + adb_call_parser_factory=lambda: _AdbCallParserAdapter(self._env), + log_stream=self._log_stream, + ) + self._task_manager.setup_task() + self._task_manager_started = True + + def _forward_device_logs(self) -> None: + if self._signal_log_stream is None: + return + logs = self._env.task_extras(latest_only=False).get('logs', ()) + self._signal_log_stream.push_logs(logs) + + def _extract_timestep_extras( + self, timestep: dm_env.TimeStep + ) -> dm_env.TimeStep: + if timestep.observation is not None: + self._latest_extras = timestep.observation.pop('extras', {}) + return timestep + + def reset(self) -> dm_env.TimeStep: + self._ensure_task_manager_started() + self._task_manager.reset_task() + timestep = self._env.reset() + self._forward_device_logs() + observation = dict(timestep.observation) + rl_timestep = self._task_manager.rl_reset(observation) + self._reset_next_step = False + return self._extract_timestep_extras(rl_timestep) + + def step(self, action: Any) -> dm_env.TimeStep: + if self._reset_next_step: + return self.reset() + processed_action = self._process_action(action) + timestep = self._env.step(processed_action) + self._forward_device_logs() + observation = dict(timestep.observation) + rl_timestep = self._task_manager.rl_step(observation) + if rl_timestep.last(): + self._reset_next_step = True + return self._extract_timestep_extras(rl_timestep) + + def task_extras(self, latest_only: bool = True) -> dict[str, Any]: + extras = dict(self._env.task_extras(latest_only=latest_only)) + for key, spec in self._latest_extras.items(): + extra_values = spec.astype(spec.dtype) + extras[key] = extra_values[-1] if latest_only else extra_values + return extras + + def stats(self) -> dict[str, Any]: + info = dict(self._env.stats()) + if self._task_manager_started: + info.update(self._task_manager.stats()) + return info + + def stop_task(self) -> None: + """Stops the current task manager without closing the wrapped environment.""" + if self._task_manager_started: + self._task_manager.stop() + self._task_manager_started = False + + def close(self) -> None: + self.stop_task() + super().close() diff --git a/android_env/wrappers/task_wrapper_test.py b/android_env/wrappers/task_wrapper_test.py new file mode 100644 index 00000000..aab8f2e2 --- /dev/null +++ b/android_env/wrappers/task_wrapper_test.py @@ -0,0 +1,380 @@ +# coding=utf-8 +# Copyright 2026 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for TaskWrapper environment wrapper.""" + +import threading +from unittest import mock + +from absl.testing import absltest +from android_env import env_interface +from android_env.components import fake_log_stream +from android_env.proto import adb_pb2 +from android_env.proto import task_pb2 +from android_env.wrappers import task_wrapper +import dm_env +import numpy as np + + +class TaskWrapperTest(absltest.TestCase): + + def setUp(self): + super().setUp() + self._thread_exceptions = [] + self._orig_excepthook = threading.excepthook + + def _excepthook(args): + self._thread_exceptions.append(args) + self._orig_excepthook(args) + + threading.excepthook = _excepthook + self.base_env = mock.create_autospec( + env_interface.AndroidEnvInterface, instance=True + ) + self.base_env.task_extras.return_value = {} + self.base_env.stats.return_value = {} + self.fake_log_stream = fake_log_stream.FakeLogStream() + self.task = task_pb2.Task( + name='test_task', + max_episode_steps=5, + ) + + def tearDown(self): + threading.excepthook = self._orig_excepthook + self.assertEmpty(self._thread_exceptions) + self.assertEqual(threading.active_count(), 1) + super().tearDown() + + def test_reset_and_step_with_default_log_stream(self): + self.base_env.reset.return_value = dm_env.TimeStep( + step_type=dm_env.StepType.FIRST, + reward=None, + discount=None, + observation={'pixels': np.zeros((10, 10, 3), dtype=np.uint8)}, + ) + self.base_env.step.return_value = dm_env.TimeStep( + step_type=dm_env.StepType.MID, + reward=0.0, + discount=1.0, + observation={'pixels': np.ones((10, 10, 3), dtype=np.uint8)}, + ) + + wrapper = task_wrapper.TaskWrapper( + env=self.base_env, + task=self.task, + ) + + ts_reset = wrapper.reset() + self.assertEqual(ts_reset.step_type, dm_env.StepType.FIRST) + self.assertNotIn('extras', ts_reset.observation) + self.base_env.reset.assert_called_once() + + ts_step = wrapper.step({'action_type': 0}) + self.assertEqual(ts_step.step_type, dm_env.StepType.MID) + self.assertNotIn('extras', ts_step.observation) + self.base_env.step.assert_called_once() + + wrapper.close() + + def test_device_signal_logs_compute_rewards_and_task_extras(self): + self.task.log_parsing_config.log_regexps.score = r'^Score: ([-+]?[0-9]+)$' + self.task.log_parsing_config.log_regexps.reward.append( + r'^Bonus: ([-+]?[0-9]+)$' + ) + self.task.log_parsing_config.log_regexps.extra.append( + r'^Extra (?P\w+): (?P.*)$' + ) + self.task.log_parsing_config.log_regexps.episode_end.append(r'^GAME_OVER$') + + self.base_env.reset.return_value = dm_env.TimeStep( + step_type=dm_env.StepType.FIRST, + reward=None, + discount=None, + observation={'pixels': np.zeros((10, 10, 3), dtype=np.uint8)}, + ) + self.base_env.step.return_value = dm_env.TimeStep( + step_type=dm_env.StepType.MID, + reward=0.0, + discount=1.0, + observation={'pixels': np.ones((10, 10, 3), dtype=np.uint8)}, + ) + reset_logs = [ + '07-26 12:00:00.123 100 200 I Game: Score: 10', + '07-26 12:00:00.124 100 200 I Game: Extra reset_extra: 1', + ] + step_logs = [ + '07-26 12:00:01.000 100 200 I Game: Score: 25', + ' 1553110400.424 100 200 I Game: Bonus: 5', + 'Extra level: 2', + 'Extra level: 3', + 'GAME_OVER', + ] + active_logs = [reset_logs] + + def _fake_task_extras(latest_only: bool = True): + logs = active_logs[0] + return { + 'logs': logs[-1] if latest_only else logs, + 'battery_power': {'level_percent': 90}, + } + + self.base_env.task_extras.side_effect = _fake_task_extras + + wrapper = task_wrapper.TaskWrapper( + env=self.base_env, + task=self.task, + ) + + wrapper.reset() + self.assertEqual(wrapper.task_extras(latest_only=True)['reset_extra'], 1) + + active_logs[0] = step_logs + ts_step = wrapper.step({'action_type': 0}) + + self.assertEqual(ts_step.step_type, dm_env.StepType.LAST) + self.assertEqual(ts_step.reward, 30.0) + + latest_extras = wrapper.task_extras(latest_only=True) + self.assertEqual(latest_extras['battery_power'], {'level_percent': 90}) + self.assertEqual(latest_extras['level'], 3) + + all_extras = wrapper.task_extras(latest_only=False) + np.testing.assert_array_equal(all_extras['level'], np.array([2, 3])) + + wrapper.close() + + def test_step_before_reset_or_after_episode_end_triggers_reset(self): + self.task.max_episode_steps = 1 + self.base_env.reset.return_value = dm_env.TimeStep( + step_type=dm_env.StepType.FIRST, + reward=None, + discount=None, + observation={'pixels': np.zeros((10, 10, 3), dtype=np.uint8)}, + ) + self.base_env.step.return_value = dm_env.TimeStep( + step_type=dm_env.StepType.MID, + reward=0.0, + discount=1.0, + observation={'pixels': np.ones((10, 10, 3), dtype=np.uint8)}, + ) + + wrapper = task_wrapper.TaskWrapper( + env=self.base_env, + task=self.task, + log_stream=self.fake_log_stream, + ) + + # First step() before explicit reset() automatically calls reset(). + ts_first = wrapper.step({'action_type': 0}) + self.assertEqual(ts_first.step_type, dm_env.StepType.FIRST) + self.assertEqual(self.base_env.reset.call_count, 1) + + # First real step is within max_episode_steps=1 (MID). + ts_mid = wrapper.step({'action_type': 0}) + self.assertEqual(ts_mid.step_type, dm_env.StepType.MID) + + # Second real step exceeds max_episode_steps=1 (LAST). + ts_last = wrapper.step({'action_type': 0}) + self.assertEqual(ts_last.step_type, dm_env.StepType.LAST) + + # Subsequent step() automatically resets the environment. + ts_auto_reset = wrapper.step({'action_type': 0}) + self.assertEqual(ts_auto_reset.step_type, dm_env.StepType.FIRST) + self.assertEqual(self.base_env.reset.call_count, 2) + + wrapper.close() + + def test_rewards_and_task_extras_from_custom_log_stream(self): + self.task.log_parsing_config.log_regexps.reward.append( + r'^Bonus: ([-+]?[0-9]+)$' + ) + self.task.log_parsing_config.log_regexps.extra.append( + r'^Extra (?P\w+): (?P.*)$' + ) + self.base_env.reset.return_value = dm_env.TimeStep( + step_type=dm_env.StepType.FIRST, + reward=None, + discount=None, + observation={'pixels': np.zeros((10, 10, 3), dtype=np.uint8)}, + ) + self.base_env.step.return_value = dm_env.TimeStep( + step_type=dm_env.StepType.MID, + reward=0.0, + discount=1.0, + observation={'pixels': np.ones((10, 10, 3), dtype=np.uint8)}, + ) + self.base_env.task_extras.return_value = { + 'battery_power': {'level_percent': 90}, + 'logs': ['Bonus: 100'], + } + + ready_to_emit = threading.Event() + logs_processed = threading.Event() + + def _log_gen(): + ready_to_emit.wait(timeout=5.0) + yield ' 0.000 1 1 I Game: Bonus: 5' + yield ' 0.000 1 1 I Game: Extra level: 2' + yield ' 0.000 1 1 I Game: Extra level: 3' + logs_processed.set() + + wrapper = task_wrapper.TaskWrapper( + env=self.base_env, + task=self.task, + log_stream=fake_log_stream.FakeLogStream(log_generator=_log_gen()), + ) + + wrapper.reset() + ready_to_emit.set() + self.assertTrue(logs_processed.wait(timeout=5.0)) + + ts_step = wrapper.step({'action_type': 0}) + + self.assertEqual(ts_step.reward, 5.0) + latest_extras = wrapper.task_extras(latest_only=True) + self.assertEqual(latest_extras['battery_power'], {'level_percent': 90}) + self.assertEqual(latest_extras['level'], 3) + + all_extras = wrapper.task_extras(latest_only=False) + np.testing.assert_array_equal(all_extras['level'], np.array([2, 3])) + + wrapper.close() + + def test_adb_call_delegation_runs_setup_steps_only_once(self): + self.task.setup_steps.add( + adb_request=adb_pb2.AdbRequest( + press_button=adb_pb2.AdbRequest.PressButton( + button=adb_pb2.AdbRequest.PressButton.Button.HOME + ) + ) + ) + self.base_env.execute_adb_call.return_value = adb_pb2.AdbResponse( + status=adb_pb2.AdbResponse.Status.OK + ) + self.base_env.reset.return_value = dm_env.TimeStep( + step_type=dm_env.StepType.FIRST, + reward=None, + discount=None, + observation={'pixels': np.zeros((10, 10, 3), dtype=np.uint8)}, + ) + + wrapper = task_wrapper.TaskWrapper( + env=self.base_env, + task=self.task, + log_stream=self.fake_log_stream, + ) + + wrapper.reset() + wrapper.reset() + self.base_env.execute_adb_call.assert_called_once() + wrapper.close() + + def test_stats(self): + self.base_env.stats.return_value = {'env_stat': 42} + self.base_env.reset.return_value = dm_env.TimeStep( + step_type=dm_env.StepType.FIRST, + reward=None, + discount=None, + observation={'pixels': np.zeros((10, 10, 3), dtype=np.uint8)}, + ) + + wrapper = task_wrapper.TaskWrapper( + env=self.base_env, + task=self.task, + log_stream=self.fake_log_stream, + ) + + self.assertEqual(wrapper.stats(), {'env_stat': 42}) + + wrapper.reset() + stats = wrapper.stats() + self.assertIn('env_stat', stats) + self.assertIn('episode_steps', stats) + wrapper.close() + self.assertEqual(wrapper.stats(), {'env_stat': 42}) + + def test_active_package_exit_truncates_and_auto_resets(self): + self.task.expected_app_screen.activity = 'com.example.game/.MainActivity' + self.base_env.reset.return_value = dm_env.TimeStep( + step_type=dm_env.StepType.FIRST, + reward=None, + discount=None, + observation={ + 'pixels': np.zeros((10, 10, 3), dtype=np.uint8), + 'active_package': 'com.example.game', + }, + ) + self.base_env.step.side_effect = [ + dm_env.TimeStep( + step_type=dm_env.StepType.MID, + reward=0.0, + discount=1.0, + observation={ + 'pixels': np.ones((10, 10, 3), dtype=np.uint8), + 'active_package': 'com.example.game', + }, + ), + dm_env.TimeStep( + step_type=dm_env.StepType.MID, + reward=0.0, + discount=1.0, + observation={ + 'pixels': np.ones((10, 10, 3), dtype=np.uint8), + 'active_package': 'com.android.launcher3', + }, + ), + ] + + wrapper = task_wrapper.TaskWrapper( + env=self.base_env, + task=self.task, + log_stream=self.fake_log_stream, + ) + + wrapper.reset() + ts_in_app = wrapper.step({'action_type': 0}) + self.assertEqual(ts_in_app.step_type, dm_env.StepType.MID) + self.assertEqual(wrapper.stats()['reset_count_user_exited'], 0) + + ts_exited = wrapper.step({'action_type': 0}) + self.assertEqual(ts_exited.step_type, dm_env.StepType.LAST) + self.assertEqual(wrapper.stats()['reset_count_user_exited'], 1) + self.base_env.execute_adb_call.assert_not_called() + + wrapper.close() + + def test_stop_task_stops_task_manager_without_closing_base_env(self): + self.base_env.reset.return_value = dm_env.TimeStep( + step_type=dm_env.StepType.FIRST, + reward=None, + discount=None, + observation={'pixels': np.zeros((10, 10, 3), dtype=np.uint8)}, + ) + self.base_env.stats.return_value = {'env_stat': 7} + wrapper = task_wrapper.TaskWrapper( + env=self.base_env, + task=self.task, + log_stream=self.fake_log_stream, + ) + wrapper.reset() + self.assertIn('episode_steps', wrapper.stats()) + + wrapper.stop_task() + self.assertEqual(wrapper.stats(), {'env_stat': 7}) + self.base_env.close.assert_not_called() + + +if __name__ == '__main__': + absltest.main() diff --git a/pyproject.toml b/pyproject.toml index ba7d11db..c8082d81 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,8 @@ dependencies = [ "protobuf>=6,<8", "pygame>=2.6.0", "pillow>=10.4.0", + "opencv-python-headless>=4.10.0", + "websocket-client>=1.8.0", ] [project.optional-dependencies] @@ -32,6 +34,11 @@ testing = [ ] acme = ["dm-acme"] gym = ["gymnasium"] +mcp = ["mcp>=1.26,<2"] +playground = ["ml-collections"] + +[tool.setuptools.package-data] +"android_env.components" = ["*.textproto"] [project.urls] repository = "https://github.com/deepmind/android_env" diff --git a/requirements.txt b/requirements.txt index b9a5d957..fef33114 100644 --- a/requirements.txt +++ b/requirements.txt @@ -11,6 +11,40 @@ grpcio-tools>=1.70.0 farama-notifications>=0.0.4 typing-extensions>=4.11.0 psutil>=5.9.8 -dm-tree +dm-tree>=0.1 attrs>=18.2.0 wrapt>=1.11.2 +opencv-python-headless>=4.10.0 +websocket-client>=1.8.0 +mcp>=1.26,<2 +ml-collections>=1.1 +# Transitive dependencies (pinned; Bazel pip.parse installs with --no-deps). +annotated-types==0.8.0 +anyio==4.15.1 +certifi==2026.7.22 +cffi==2.1.1 +click==8.5.0 +cloudpickle==3.1.2 +cryptography==50.0.2 +h11==0.16.0 +httpcore==1.0.9 +httpx==0.28.1 +httpx-sse==0.4.3 +idna==3.20 +jsonschema==4.26.0 +jsonschema-specifications==2025.9.1 +pycparser==3.11 +pydantic==2.13.5 +pydantic-core==2.46.5 +pydantic-settings==2.15.0 +pyjwt==2.15.1 +python-dotenv==1.2.4 +python-multipart==0.0.32 +pyyaml==6.0.3 +referencing==0.37.0 +rpds-py==2026.9.1 +setuptools==84.0.0 +sse-starlette==3.5.0 +starlette==1.7.0 +typing-inspection==0.4.4 +uvicorn==0.54.0