fix(shower): capture virtual display screenshots without the overlay

app/src/main/java/com/ai/assistance/operit/core/tools/agent/AgentScreenshotReport.kt

@@ -0,0 +1,44 @@
+package com.ai.assistance.operit.core.tools.agent
+
+/**
+ * Screenshot handed to the vision model for one agent step.
+ *
+ * The two fields are mutually exclusive: either [link] points at the captured target display,
+ * or [diagnostic] explains why that display produced no image.
+ */
+data class AgentScreenshot(
+    val link: String?,
+    val diagnostic: String? = null
+)
+
+/**
+ * Describes a virtual-display capture failure to the vision model.
+ *
+ * The message names the display that was actually attempted and states that the main display was
+ * not used, so the model treats the missing image as a capture failure instead of reasoning about
+ * whatever happens to be on the physical screen.
+ */
+internal fun virtualDisplayScreenshotDiagnostic(
+    agentId: String,
+    displayId: Int?,
+    reason: String
+): String {
+    val target = displayId?.let { "displayId=$it" } ?: "displayId=unbound"
+    return "Virtual display capture failed (agentId=$agentId, $target, reason=$reason). " +
+        "The main display was NOT captured because it does not show this agent's screen, " +
+        "so no image is attached for this step and any earlier screenshot is stale."
+}
+
+/** Builds the screen section of the agent prompt for one step. */
+internal fun buildScreenInfo(screenshot: AgentScreenshot): String =
+    buildString {
+        if (screenshot.link != null) {
+            appendLine("[SCREENSHOT] Below is the latest screen image:")
+            appendLine(screenshot.link)
+        } else if (screenshot.diagnostic != null) {
+            appendLine("[SCREENSHOT UNAVAILABLE]")
+            appendLine(screenshot.diagnostic)
+        } else {
+            appendLine("No screenshot available for this step.")
+        }
+    }.trim()

app/src/main/java/com/ai/assistance/operit/core/tools/agent/PhoneAgent.kt

@@ -603,15 +603,8 @@ class PhoneAgent(
         _stepCount++
         AppLogger.d("PhoneAgent", "[$agentId] _executeStep: begin, step=$_stepCount")
 
-        val screenshotLink = actionHandler.captureScreenshotForAgent()
-        val screenInfo = buildString {
-            if (screenshotLink != null) {
-                appendLine("[SCREENSHOT] Below is the latest screen image:")
-                appendLine(screenshotLink)
-            } else {
-                appendLine("No screenshot available for this step.")
-            }
-        }.trim()
+        val screenshot = actionHandler.captureScreenshotForAgent()
+        val screenInfo = buildScreenInfo(screenshot)
 
         val userMessage = if (isFirst) {
             "$userPrompt\n\n$screenInfo"
@@ -823,29 +816,45 @@ class ActionHandler(
         return ShowerUsageContext(isAdbOrHigher = permissionState.isAdbOrHigher, showerDisplayId = showerId)
     }
 
-    suspend fun captureScreenshotForAgent(): String? {
+    suspend fun captureScreenshotForAgent(): AgentScreenshot {
         val showerCtx = resolveShowerUsageContext()
         val floatingService = FloatingChatService.getInstance()
         val progressOverlay = UIAutomationProgressOverlay.getInstance(context)
 
         var screenshotLink: String? = null
         var dimensions: Pair<Int, Int>? = null
+        var diagnostic: String? = null
 
         try {
             // Keep screenshot captures clean: hide overlays first, then restore after capture.
             floatingService?.setStatusIndicatorVisible(false)
             progressOverlay.setOverlayVisible(false)
             delay(200)
 
-            if (showerCtx.canUseShowerForInput) {
-                val (link, dims) = captureScreenshotViaShower()
-                screenshotLink = link
-                dimensions = dims
+            val showerDisplayId = showerCtx.showerDisplayId
+            var showerFailure: String? = null
+            if (showerDisplayId != null) {
+                // The Shower display has its own size and input injection works in that display's
+                // pixels. Scaling the relative coordinates of later taps from the physical main
+                // display would aim them at the wrong place on the virtual screen.
+                ShowerController.getVideoSize(agentId)?.let { (width, height) ->
+                    if (width > 0 && height > 0) {
+                        screenWidth = width
+                        screenHeight = height
+                    }
+                }
+                val capture = captureScreenshotViaShower(showerDisplayId)
+                screenshotLink = capture.link
+                dimensions = capture.dimensions
+                showerFailure = capture.failure
             }
 
-            if (screenshotLink == null) {
-                val screenshotTool = buildScreenshotTool()
-                val (bitmap, fallbackDims) = toolImplementations.captureScreenshotBitmap(screenshotTool)
+            // A main-screen agent targets the physical display, so a MediaProjection capture is a
+            // valid second attempt there. A virtual-screen agent must never receive such an image:
+            // it shows another screen and its size would rescale the agent's tap coordinates.
+            if (screenshotLink == null && isMainScreenAgent()) {
+                val (bitmap, fallbackDims) =
+                    toolImplementations.captureScreenshotBitmap(buildScreenshotTool())
 
                 if (bitmap != null) {
                     val (compressedLink, rawDims) = saveCompressedScreenshotFromBitmap(bitmap)
@@ -854,6 +863,19 @@ class ActionHandler(
                     bitmap.recycle()
                 }
             }
+
+            if (screenshotLink == null && !isMainScreenAgent()) {
+                // A virtual-screen agent's content only exists on its Shower display. Returning a
+                // main-screen picture here would show the model another screen and scale its tap
+                // coordinates against that screen's size, so report the failure on the target
+                // display instead of substituting an image of a different screen.
+                diagnostic = virtualDisplayScreenshotDiagnostic(
+                    agentId = agentId,
+                    displayId = showerDisplayId,
+                    reason = showerFailure ?: "no Shower display is bound to this agent"
+                )
+                AppLogger.w("ActionHandler", "[$agentId] $diagnostic")
+            }
         } finally {
             val hasShowerDisplayNow = try {
                 ShowerController.getDisplayId(agentId) != null
@@ -871,7 +893,7 @@ class ActionHandler(
             screenWidth = dimensions.first
             screenHeight = dimensions.second
         }
-        return screenshotLink
+        return AgentScreenshot(link = screenshotLink, diagnostic = diagnostic)
     }
 
     private fun buildScreenshotTool(): AITool {
@@ -881,34 +903,63 @@ class ActionHandler(
         )
     }
 
-    private suspend fun captureScreenshotViaShower(): Pair<String?, Pair<Int, Int>?> {
+    /**
+     * Captures the display the agent drives through the Shower server.
+     *
+     * The server mirrors the requested display into an ImageReader, so displayId=0 and a
+     * virtual display take the same path and no overlay window has to be attached on the
+     * main screen for the capture to work.
+     *
+     * Returns the image link and its size, or a null link plus the reason the display produced
+     * no frame so the step can report a capture failure on the target display.
+     */
+    private suspend fun captureScreenshotViaShower(displayId: Int): ShowerCapture {
         return try {
-            val pngBytes = if (isMainScreenAgent()) {
-                ShowerController.requestScreenshot(agentId)
-            } else {
-                VirtualDisplayOverlay.getInstance(context, agentId).captureCurrentFramePng()
-            }
+            val pngBytes = ShowerController.requestScreenshot(agentId)
             if (pngBytes == null || pngBytes.isEmpty()) {
-                AppLogger.w("ActionHandler", "[$agentId] Shower WS screenshot returned no data")
-                Pair(null, null)
+                ShowerCapture(
+                    link = null,
+                    dimensions = null,
+                    failure = "Shower returned no frame for displayId=$displayId"
+                )
             } else {
                 val bitmap = BitmapFactory.decodeByteArray(pngBytes, 0, pngBytes.size)
                 if (bitmap == null) {
-                    AppLogger.e("ActionHandler", "[$agentId] Shower screenshot: failed to decode bytes")
-                    Pair(null, null)
+                    ShowerCapture(
+                        link = null,
+                        dimensions = null,
+                        failure = "captured frame for displayId=$displayId is not a decodable image"
+                    )
                 } else {
-                    val result = saveCompressedScreenshotFromBitmap(bitmap)
+                    val (link, dims) = saveCompressedScreenshotFromBitmap(bitmap)
                     bitmap.recycle()
-                    result
+                    ShowerCapture(
+                        link = link,
+                        dimensions = dims,
+                        failure = if (link == null) {
+                            "captured frame for displayId=$displayId could not be stored"
+                        } else {
+                            null
+                        }
+                    )
                 }
             }
         } catch (e: Exception) {
             if (e is CancellationException) throw e
-            AppLogger.e("ActionHandler", "[$agentId] Shower screenshot failed", e)
-            Pair(null, null)
+            ShowerCapture(
+                link = null,
+                dimensions = null,
+                failure = "Shower screenshot request failed for displayId=$displayId: ${e.message}"
+            )
         }
     }
 
+    private data class ShowerCapture(
+        val link: String?,
+        val dimensions: Pair<Int, Int>?,
+        val failure: String?
+    )
+
     private fun saveCompressedScreenshotFromBitmap(bitmap: Bitmap): Pair<String?, Pair<Int, Int>?> {
         return try {
             val originalWidth = bitmap.width

app/src/main/java/com/ai/assistance/operit/core/tools/agent/ShowerVideoRenderer.kt

@@ -21,7 +21,4 @@ object ShowerVideoRenderer {
     fun detach() = core.detach()
 
     fun onFrame(data: ByteArray) = core.onFrame(data)
-
-    suspend fun captureCurrentFramePng(): ByteArray? =
-        core.captureCurrentFramePng()
 }

app/src/main/java/com/ai/assistance/operit/ui/common/displays/VirtualDisplayOverlay.kt

@@ -125,8 +125,6 @@ class VirtualDisplayOverlay private constructor(private val context: Context, pr
     private var lifecycleOwner: ServiceLifecycleOwner? = null
     private var layoutParams: WindowManager.LayoutParams? = null
     private var displayId: Int? = null
-    @Volatile
-    private var surfaceView: ShowerSurfaceView? = null
     private var isFullscreen by mutableStateOf(false)
     private var isSnapped by mutableStateOf(false)
     private var snappedToRight by mutableStateOf(false)
@@ -252,8 +250,6 @@ class VirtualDisplayOverlay private constructor(private val context: Context, pr
         }
     }
 
-    suspend fun captureCurrentFramePng(): ByteArray? = surfaceView?.captureCurrentFramePng()
-
     fun showAutomationControls(
         totalSteps: Int,
         initialStatus: String,
@@ -315,7 +311,6 @@ class VirtualDisplayOverlay private constructor(private val context: Context, pr
                 windowManager = null
                 displayId = null
                 currentAppPackageName = null
-                surfaceView = null
             } catch (e: Exception) {
                 AppLogger.e("VirtualDisplayOverlay", "Error hiding overlay", e)
             }
@@ -744,11 +739,9 @@ class VirtualDisplayOverlay private constructor(private val context: Context, pr
                                         view.bindController(ShowerController.getInstance(agentId))
                                     } catch (_: Exception) {
                                     }
-                                    surfaceView = view
                                 }
                             },
                             update = { view ->
-                                surfaceView = view
                                 try {
                                     view.bindController(ShowerController.getInstance(agentId))
                                 } catch (_: Exception) {

app/src/test/java/com/ai/assistance/operit/core/tools/agent/AgentScreenshotReportTest.kt

@@ -0,0 +1,78 @@
+package com.ai.assistance.operit.core.tools.agent
+
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertFalse
+import org.junit.Assert.assertTrue
+import org.junit.Test
+
+/**
+ * Covers the virtual-display screenshot contract from issue #1371: the agent's screenshot comes
+ * from the display it drives, and a failed virtual-display capture is reported as a failure instead
+ * of being replaced by a main-screen picture.
+ */
+class AgentScreenshotReportTest {
+
+    @Test
+    fun virtualDisplayDiagnostic_namesAgentAndDisplay() {
+        val text = virtualDisplayScreenshotDiagnostic(
+            agentId = "agent-abc",
+            displayId = 7,
+            reason = "Shower returned no frame for displayId=7"
+        )
+
+        assertTrue(text.contains("agentId=agent-abc"))
+        assertTrue(text.contains("displayId=7"))
+        assertTrue(text.contains("Shower returned no frame for displayId=7"))
+    }
+
+    @Test
+    fun virtualDisplayDiagnostic_statesMainDisplayWasNotCaptured() {
+        val text = virtualDisplayScreenshotDiagnostic(
+            agentId = "agent-abc",
+            displayId = 7,
+            reason = "no frame returned"
+        )
+
+        // The model must not treat this as a usable image of its own screen.
+        assertTrue(text.contains("main display was NOT captured"))
+    }
+
+    @Test
+    fun virtualDisplayDiagnostic_marksUnboundDisplayExplicitly() {
+        val text = virtualDisplayScreenshotDiagnostic(
+            agentId = "agent-abc",
+            displayId = null,
+            reason = "no Shower display is bound to this agent"
+        )
+
+        assertTrue(text.contains("displayId=unbound"))
+        assertFalse(text.contains("displayId=null"))
+    }
+
+    @Test
+    fun buildScreenInfo_embedsImageLinkWhenCaptureSucceeded() {
+        val info = buildScreenInfo(AgentScreenshot(link = "<link type=\"image\" id=\"img-1\"></link>"))
+
+        assertTrue(info.contains("[SCREENSHOT]"))
+        assertTrue(info.contains("img-1"))
+        assertFalse(info.contains("UNAVAILABLE"))
+    }
+
+    @Test
+    fun buildScreenInfo_reportsDiagnosticInsteadOfImageOnFailure() {
+        val info = buildScreenInfo(
+            AgentScreenshot(link = null, diagnostic = "Virtual display capture failed (agentId=a1, displayId=7)")
+        )
+
+        assertTrue(info.contains("[SCREENSHOT UNAVAILABLE]"))
+        assertTrue(info.contains("displayId=7"))
+        assertFalse(info.contains("[SCREENSHOT] Below"))
+    }
+
+    @Test
+    fun buildScreenInfo_fallsBackToPlainMessageWithoutImageOrDiagnostic() {
+        val info = buildScreenInfo(AgentScreenshot(link = null))
+
+        assertEquals("No screenshot available for this step.", info)
+    }
+}