[codex] Show initial frame while live preview starts

snapo-app-mac/Snap-O/LivePreview/LivePreviewSession.swift

@@ -1,19 +1,32 @@
 import AppKit
 @preconcurrency import AVFoundation
 import Foundation
+import ImageIO
 import SnapODeviceClient
 
 /// Owns one device-side live-preview process and its decoded frame stream.
 @MainActor
 final class LivePreviewSession {
+  private static let initialFrameTimeout: Duration = .seconds(2)
+
   let deviceID: String
+  private(set) var initialFrame: CGImage?
 
   var media: Media?
+  var initialFrameHandler: ((CGImage) -> Void)? {
+    didSet {
+      if let initialFrame {
+        initialFrameHandler?(initialFrame)
+      }
+    }
+  }
+
   var sampleBufferHandler: ((CMSampleBuffer) -> Void)?
 
   private let densityScale: CGFloat
   private let screenStream: ScreenStreamSession
   private var decoder: H264StreamDecoder?
+  private var initialFrameTask: Task<Void, Never>?
   private var streamTask: Task<Void, Never>?
   private var hasStopped = false
 
@@ -26,11 +39,42 @@ final class LivePreviewSession {
     self.deviceID = deviceID
 
     let exec = await adb.exec()
-    densityScale = try await CGFloat(exec.displayDensity(deviceID: deviceID))
-    screenStream = try await exec.startScreenStream(deviceID: deviceID)
+    async let densityValue = exec.displayDensity(deviceID: deviceID)
+    async let startedStream = exec.startScreenStream(deviceID: deviceID)
+    densityScale = try await CGFloat(densityValue)
+    screenStream = try await startedStream
 
     setupDecoder()
     startStreamTask()
+    startInitialFrameTask(using: exec)
+  }
+
+  private static func captureInitialFrame(deviceID: String, using exec: ADBClient) async throws -> CGImage? {
+    let timeout = initialFrameTimeout
+    let data = try await withThrowingTaskGroup(of: Data?.self, returning: Data?.self) { group in
+      group.addTask {
+        do {
+          return try await exec.screencapPNG(deviceID: deviceID)
+        } catch {
+          try Task.checkCancellation()
+          return nil
+        }
+      }
+      group.addTask {
+        try await Task.sleep(for: timeout)
+        return nil
+      }
+
+      defer { group.cancelAll() }
+      guard let result = try await group.next() else { return nil }
+      return result
+    }
+    return data.flatMap(decodeInitialFrame)
+  }
+
+  private static func decodeInitialFrame(from data: Data) -> CGImage? {
+    guard let source = CGImageSourceCreateWithData(data as CFData, nil) else { return nil }
+    return CGImageSourceCreateImageAtIndex(source, 0, nil)
   }
 
   func waitUntilReady() async throws -> Media {
@@ -51,6 +95,27 @@ final class LivePreviewSession {
     finish(with: nil)
   }
 
+  func discardInitialFrame() {
+    initialFrame = nil
+    initialFrameHandler = nil
+  }
+
+  private func startInitialFrameTask(using exec: ADBClient) {
+    initialFrameTask = Task { [weak self] in
+      guard let self else { return }
+      defer { initialFrameTask = nil }
+
+      do {
+        guard let frame = try await Self.captureInitialFrame(deviceID: deviceID, using: exec),
+              !hasStopped else { return }
+        initialFrame = frame
+        initialFrameHandler?(frame)
+      } catch {
+        // The initial frame is best-effort; live preview continues without it.
+      }
+    }
+  }
+
   private func setupDecoder() {
     let decoder = H264StreamDecoder { [weak self] sample in
       guard let self else { return }
@@ -103,8 +168,12 @@ final class LivePreviewSession {
 
     streamTask?.cancel()
     streamTask = nil
+    initialFrameTask?.cancel()
+    initialFrameTask = nil
     screenStream.close()
     decoder = nil
+    initialFrame = nil
+    initialFrameHandler = nil
     sampleBufferHandler = nil
 
     stopResult = error

snapo-app-mac/Snap-O/LivePreview/LivePreviewView.swift

@@ -38,7 +38,9 @@ struct LivePreviewRendererView: NSViewRepresentable {
 final class LivePreviewDisplayView: NSView {
   private var renderer: LivePreviewRenderer?
   private var trackingArea: NSTrackingArea?
+  private var readyForDisplayObserver: NSObjectProtocol?
   private let displayLayer = AVSampleBufferDisplayLayer()
+  private let initialFrameLayer = CALayer()
   private var endedLivePreviewTrace = false
 
   private var pointerState = PointerState()
@@ -73,32 +75,46 @@ final class LivePreviewDisplayView: NSView {
       detachSession()
     }
     self.renderer = renderer
-    attachSession()
+    if shouldDetach {
+      attachSession()
+    }
   }
 
   private func configureLayerIfNeeded() {
     guard displayLayer.superlayer == nil else { return }
     wantsLayer = true
     layer?.addSublayer(displayLayer)
+    layer?.addSublayer(initialFrameLayer)
     displayLayer.videoGravity = .resizeAspect
     displayLayer.backgroundColor = NSColor.black.cgColor
     displayLayer.frame = bounds
     displayLayer.autoresizingMask = [.layerWidthSizable, .layerHeightSizable]
+    initialFrameLayer.contentsGravity = .resizeAspect
+    initialFrameLayer.frame = bounds
+    initialFrameLayer.autoresizingMask = [.layerWidthSizable, .layerHeightSizable]
   }
 
   private func attachSession() {
-    guard let session = renderer?.session else { return }
+    guard let renderer else { return }
+    let session = renderer.session
+    let operationID = renderer.operation.id
+    session.initialFrameHandler = { [weak self] frame in
+      self?.displayInitialFrame(frame, operationID: operationID)
+    }
     session.sampleBufferHandler = { [weak self] sample in
       self?.enqueue(sample)
     }
     endedLivePreviewTrace = false
   }
 
   private func detachSession() {
+    stopObservingDisplayReadiness()
+    renderer?.session.initialFrameHandler = nil
     renderer?.session.sampleBufferHandler = nil
     displayLayer.sampleBufferRenderer.stopRequestingMediaData()
     displayLayer.sampleBufferRenderer.flush(removingDisplayedImage: true, completionHandler: nil)
     displayLayer.sampleBufferRenderer.requestMediaDataWhenReady(on: .main) {}
+    setInitialFrame(nil)
     endedLivePreviewTrace = false
   }
 
@@ -112,6 +128,52 @@ final class LivePreviewDisplayView: NSView {
     }
   }
 
+  private func displayInitialFrame(_ image: CGImage, operationID: UUID) {
+    guard renderer?.operation.id == operationID else { return }
+    observeDisplayReadiness(operationID: operationID)
+    guard !displayLayer.isReadyForDisplay else {
+      hideInitialFrameIfReady(operationID: operationID)
+      return
+    }
+    setInitialFrame(image)
+    hideInitialFrameIfReady(operationID: operationID)
+  }
+
+  private func observeDisplayReadiness(operationID: UUID) {
+    stopObservingDisplayReadiness()
+    readyForDisplayObserver = NotificationCenter.default.addObserver(
+      forName: .AVSampleBufferDisplayLayerReadyForDisplayDidChange,
+      object: displayLayer,
+      queue: .main
+    ) { [weak self] _ in
+      Task { @MainActor [weak self] in
+        self?.hideInitialFrameIfReady(operationID: operationID)
+      }
+    }
+  }
+
+  private func hideInitialFrameIfReady(operationID: UUID) {
+    guard renderer?.operation.id == operationID,
+          displayLayer.isReadyForDisplay else { return }
+    renderer?.session.discardInitialFrame()
+    setInitialFrame(nil)
+    stopObservingDisplayReadiness()
+  }
+
+  private func stopObservingDisplayReadiness() {
+    guard let readyForDisplayObserver else { return }
+    NotificationCenter.default.removeObserver(readyForDisplayObserver)
+    self.readyForDisplayObserver = nil
+  }
+
+  private func setInitialFrame(_ image: CGImage?) {
+    if image == nil, initialFrameLayer.contents == nil { return }
+    CATransaction.begin()
+    CATransaction.setDisableActions(true)
+    initialFrameLayer.contents = image
+    CATransaction.commit()
+  }
+
   override func updateTrackingAreas() {
     super.updateTrackingAreas()
     if let trackingArea { removeTrackingArea(trackingArea) }