Implement scrcpy screenshot backend in adb_vision

adb_vision/pyproject.toml

@@ -0,0 +1,7 @@
+[project]
+name = "adb-vision"
+version = "0.1.0"
+dependencies = [
+    "av>=16.1.0",
+    "pillow>=10.0.0",
+]

adb_vision/screenshot.py

@@ -0,0 +1,164 @@
+import asyncio
+import struct
+import base64
+import io
+import os
+import re
+from pathlib import Path
+from typing import Protocol, Optional
+
+# Type definition for the ADB run function
+class AdbRunFn(Protocol):
+    async def __call__(self, *args: str, timeout: float) -> bytes: ...
+
+async def _capture_scrcpy(*, adb_run: AdbRunFn, serial: str, adb_exe: str) -> str:
+    """Return base64-encoded PNG screenshot via scrcpy single-frame capture.
+
+    This implementation uses scrcpy-server v1.20 which is available in the repository.
+    It pushes the server to the device, starts it, captures one frame, and cleans up.
+    """
+    try:
+        import av
+        from av.codec import CodecContext
+    except ImportError:
+        raise ImportError("PyAV (av) is required for scrcpy capture. Please install it.")
+
+    from PIL import Image
+
+    # 1. Locate and push scrcpy-server.jar
+    # The file is expected at ALAS/alas_wrapped/bin/scrcpy/scrcpy-server-v1.20.jar
+    # This script is at ALAS/adb_vision/screenshot.py
+    base_dir = Path(__file__).resolve().parents[1]
+    local_jar = base_dir / "alas_wrapped" / "bin" / "scrcpy" / "scrcpy-server-v1.20.jar"
+
+    if not local_jar.exists():
+        # Fallback to current working directory relative path
+        local_jar = Path("alas_wrapped/bin/scrcpy/scrcpy-server-v1.20.jar")
+        if not local_jar.exists():
+            raise FileNotFoundError(f"scrcpy-server-v1.20.jar not found. Checked: {local_jar.absolute()}")
+
+    device_jar = "/data/local/tmp/scrcpy-server.jar"
+    await adb_run("-s", serial, "push", str(local_jar), device_jar, timeout=15.0)
+
+    # 2. Start scrcpy server
+    # Position arguments for scrcpy-server v1.20:
+    # 1: version (1.20)
+    # 2: log_level (info)
+    # 3: max_size (0 for no limit)
+    # 4: bit_rate (20000000)
+    # 5: max_fps (1 for single frame capture)
+    # 6: lock_video_orientation (-1: unlocked)
+    # 7: tunnel_forward (true: server listens on localabstract:scrcpy)
+    # 8: crop (-)
+    # 9: send_frame_meta (false)
+    # 10: control (true)
+    # 11: display_id (0)
+    # 12: show_touches (false)
+    # 13: stay_awake (false)
+    # 14: codec_options (i-frame-interval=0 to get I-frame immediately)
+    # 15: encoder_name (-)
+    # 16: power_off_on_close (false)
+
+    codec_options = "i-frame-interval=0"
+
+    cmd = [
+        "shell", f"CLASSPATH={device_jar}", "app_process", "/", "com.genymobile.scrcpy.Server",
+        "1.20", "info", "0", "20000000", "1", "-1", "true", "-", "false", "true", "0", "false", "false", codec_options, "-", "false"
+    ]
+
+    server_process = await asyncio.create_subprocess_exec(
+        adb_exe, "-s", serial, *cmd,
+        stdout=asyncio.subprocess.PIPE,
+        stderr=asyncio.subprocess.PIPE
+    )
+
+    local_port = None
+    try:
+        # 3. Setup port forwarding
+        forward_res = await adb_run("-s", serial, "forward", "tcp:0", "localabstract:scrcpy", timeout=5.0)
+        forward_out = forward_res.decode().strip()
+
+        # Parse assigned port (can be "21504" or "tcp:21504")
+        m = re.search(r"(\d+)", forward_out)
+        if m:
+            local_port = int(m.group(1))
+        else:
+            raise RuntimeError(f"Failed to parse port from adb forward: {forward_out}")
+
+        # 4. Connect to the server socket
+        reader, writer = None, None
+        for i in range(20):
+            try:
+                reader, writer = await asyncio.wait_for(
+                    asyncio.open_connection("127.0.0.1", local_port),
+                    timeout=2.0
+                )
+                break
+            except (ConnectionRefusedError, asyncio.TimeoutError, OSError):
+                if server_process.returncode is not None:
+                    stderr_data = await server_process.stderr.read()
+                    raise RuntimeError(f"scrcpy-server exited: {stderr_data.decode(errors='replace')}")
+                await asyncio.sleep(0.5)
+
+        if not reader or not writer:
+            raise RuntimeError("Failed to connect to scrcpy-server socket")
+
+        try:
+            # 5. Read header (69 bytes)
+            # Dummy byte(1), Device name(64), Resolution(4)
+            await reader.readexactly(69)
+
+            # 6. Read and decode video stream
+            codec = CodecContext.create("h264", "r")
+
+            img_b64 = None
+            # Read and parse chunks until a frame is decoded
+            for _ in range(100):
+                chunk = await reader.read(16384)
+                if not chunk:
+                    break
+
+                packets = codec.parse(chunk)
+                for packet in packets:
+                    frames = codec.decode(packet)
+                    for frame in frames:
+                        # Success! Convert to PNG
+                        img = frame.to_image()
+                        buf = io.BytesIO()
+                        img.save(buf, format="PNG")
+                        img_b64 = base64.b64encode(buf.getvalue()).decode("ascii")
+                        break
+                    if img_b64: break
+                if img_b64: break
+
+            if not img_b64:
+                raise RuntimeError("Failed to decode a frame from scrcpy stream")
+
+            return img_b64
+
+        finally:
+            if writer:
+                writer.close()
+                try:
+                    await writer.wait_closed()
+                except Exception:
+                    pass
+
+    finally:
+        # 7. Cleanup
+        if server_process:
+            try:
+                server_process.terminate()
+                await asyncio.wait_for(server_process.wait(), timeout=2.0)
+            except Exception:
+                try:
+                    server_process.kill()
+                    await server_process.wait()
+                except Exception:
+                    pass
+
+        if local_port:
+            try:
+                await adb_run("-s", serial, "forward", "--remove", f"tcp:{local_port}", timeout=5.0)
+            except Exception:
+                pass

adb_vision/test_scrcpy.py

@@ -0,0 +1,130 @@
+import pytest
+import asyncio
+import base64
+import io
+from unittest.mock import MagicMock, patch, AsyncMock, ANY
+from adb_vision.screenshot import _capture_scrcpy
+
+@pytest.mark.asyncio
+async def test_capture_scrcpy_happy_path():
+    # Mock adb_run
+    adb_run = AsyncMock()
+    adb_run.side_effect = [
+        b"pushed", # push jar
+        b"21504",  # forward tcp:0
+        b"removed" # forward --remove
+    ]
+
+    # Mock asyncio.create_subprocess_exec
+    mock_process = AsyncMock()
+    mock_process.returncode = None
+    mock_process.terminate = MagicMock()
+    mock_process.wait = AsyncMock()
+
+    # Mock asyncio.open_connection
+    mock_reader = AsyncMock()
+    mock_writer = AsyncMock()
+    mock_writer.close = MagicMock()
+    mock_writer.wait_closed = AsyncMock()
+
+    # Header: 69 bytes
+    mock_reader.readexactly.return_value = b"\x00" + b"x" * 64 + b"\x05\x00\x02\xd0"
+
+    # Mock frame.to_image() to return a real PIL image so img.save works
+    from PIL import Image
+    real_img = Image.new("RGB", (1280, 720), color="red")
+
+    with patch("asyncio.create_subprocess_exec", return_value=mock_process), \
+         patch("asyncio.open_connection", return_value=(mock_reader, mock_writer)), \
+         patch("av.codec.CodecContext") as mock_codec_context_class:
+
+        mock_codec = MagicMock()
+        mock_codec_context_class.create.return_value = mock_codec
+
+        mock_packet = MagicMock()
+        mock_codec.parse.return_value = [mock_packet]
+
+        mock_frame = MagicMock()
+        mock_codec.decode.return_value = [mock_frame]
+        mock_frame.to_image.return_value = real_img
+
+        mock_reader.read.side_effect = [b"some h264 data", b""]
+
+        serial = "127.0.0.1:5555"
+        result = await _capture_scrcpy(
+            adb_run=adb_run,
+            serial=serial,
+            adb_exe="adb"
+        )
+
+        assert result is not None
+        assert len(result) > 0
+        # Verify it's valid base64
+        base64.b64decode(result)
+
+        # Verify ADB calls
+        assert adb_run.call_count >= 3
+        adb_run.assert_any_call("-s", serial, "push", ANY, "/data/local/tmp/scrcpy-server.jar", timeout=15.0)
+        adb_run.assert_any_call("-s", serial, "forward", "tcp:0", "localabstract:scrcpy", timeout=5.0)
+        adb_run.assert_any_call("-s", serial, "forward", "--remove", "tcp:21504", timeout=5.0)
+
+@pytest.mark.asyncio
+async def test_capture_scrcpy_server_fail():
+    adb_run = AsyncMock()
+    adb_run.side_effect = [
+        b"pushed", # push jar
+        b"21504",  # forward tcp:0
+        b"removed" # forward --remove
+    ]
+
+    mock_process = AsyncMock()
+    mock_process.returncode = 1
+    mock_process.stderr.read.return_value = b"some error"
+    mock_process.terminate = MagicMock()
+    mock_process.wait = AsyncMock()
+
+    with patch("asyncio.create_subprocess_exec", return_value=mock_process), \
+         patch("asyncio.open_connection", side_effect=ConnectionRefusedError()):
+
+        with pytest.raises(RuntimeError, match="scrcpy-server exited: some error"):
+            await _capture_scrcpy(
+                adb_run=adb_run,
+                serial="127.0.0.1:5555",
+                adb_exe="adb"
+            )
+
+@pytest.mark.asyncio
+async def test_capture_scrcpy_decode_fail():
+    adb_run = AsyncMock()
+    adb_run.side_effect = [
+        b"pushed", # push jar
+        b"21504",  # forward tcp:0
+        b"removed" # forward --remove
+    ]
+
+    mock_process = AsyncMock()
+    mock_process.returncode = None
+    mock_process.terminate = MagicMock()
+    mock_process.wait = AsyncMock()
+
+    mock_reader = AsyncMock()
+    mock_writer = AsyncMock()
+    mock_writer.close = MagicMock()
+    mock_writer.wait_closed = AsyncMock()
+    mock_reader.readexactly.return_value = b"\x00" + b"x" * 64 + b"\x05\x00\x02\xd0"
+    mock_reader.read.return_value = b"" # No data
+
+    with patch("asyncio.create_subprocess_exec", return_value=mock_process), \
+         patch("asyncio.open_connection", return_value=(mock_reader, mock_writer)), \
+         patch("av.codec.CodecContext") as mock_codec_context_class:
+
+        mock_codec = MagicMock()
+        mock_codec_context_class.create.return_value = mock_codec
+        mock_codec.parse.return_value = []
+
+        with pytest.raises(RuntimeError, match="Failed to decode a frame from scrcpy stream"):
+            await _capture_scrcpy(
+                adb_run=adb_run,
+                serial="127.0.0.1:5555",
+                adb_exe="adb"
+            )