Read the active display, and let a client name one

CompanionLib/SwiftServer/CompanionServiceProvider.swift

@@ -396,6 +396,12 @@ final class CompanionServiceProvider: Idb_CompanionService.SimpleServiceProtocol
     }
   }
 
+  func list_displays(request: Idb_ListDisplaysRequest, context: ServerContext) async throws -> Idb_ListDisplaysResponse {
+    return try await trackedUnaryCall(context, request: request) {
+      try await ListDisplaysMethodHandler(target: target).handle(request: request, context: context)
+    }
+  }
+
   func crash_delete(request: Idb_CrashLogQuery, context: ServerContext) async throws -> Idb_CrashLogResponse {
     return try await trackedUnaryCall(context, request: request) {
       try await CrashDeleteMethodHandler(commandExecutor: commandExecutor)

CompanionLib/SwiftServer/MethodHandlers/FramebufferStreamMethodHandler.swift

@@ -25,6 +25,9 @@ enum FramebufferStreamError: Error, LocalizedError {
   case sharedMemoryTooSmall(name: String, needed: Int, available: UInt64)
   case scaleFailed(status: Int)
   case scaleChanged(latched: Float?, requested: Float)
+  case configuredLate
+  case unknownDisplay(uniqueID: String)
+  case displaysUnavailable
 
   var errorDescription: String? {
     switch self {
@@ -42,6 +45,12 @@ enum FramebufferStreamError: Error, LocalizedError {
       return "Shared memory \(name) holds \(available) bytes, the frame needs \(needed)"
     case let .scaleFailed(status):
       return "Scaling the framebuffer failed (vImage status \(status))"
+    case .configuredLate:
+      return "The stream is already reading a display; Configure has to arrive before the first copy"
+    case let .unknownDisplay(uniqueID):
+      return "No display with unique id \(uniqueID); ask list_displays for the ones this simulator has"
+    case .displaysUnavailable:
+      return "This runtime cannot report its displays, so a display cannot be selected by id"
     case let .scaleChanged(latched, requested):
       let was = latched.map { "\($0)" } ?? "none"
       return "The stream was opened at scale \(was); a later request asked for \(requested). The scale is fixed for the life of the stream"
@@ -84,52 +93,109 @@ struct FramebufferStreamMethodHandler: @unchecked Sendable {
       throw RPCError(code: .failedPrecondition, message: FramebufferStreamError.simulatorRequired(targetDescription: String(describing: target)).localizedDescription)
     }
 
-    let framebuffer = try Framebuffer.mainScreenSurface(for: simulator, logger: targetLogger)
-    let attachment = try framebuffer.attach()
-    defer { attachment.cancel() }
-
-    let state = FramebufferStreamState(surface: attachment.initialSurface)
-
-    // The request loop returns on `stop`; the frame loop runs until the attachment finishes. Whichever
-    // ends first tears the other down, so a client disconnect and a display teardown both land here.
-    let requests = Task {
-      try await readRequests(requestStream, into: state, responseStream: responseStream)
-    }
-    let frames = Task {
-      try await serviceFrames(attachment, state: state)
-    }
+    // The display is not known until a Configure names one, so the attachment is made on the first
+    // request that needs it rather than up front. A stream that never configures attaches to the
+    // active integrated display, which is what every client got before Configure existed.
+    var attached: Attached?
+    var frames: Task<Void, any Error>?
     defer {
-      requests.cancel()
-      frames.cancel()
+      frames?.cancel()
+      attached?.attachment.cancel()
     }
-    _ = try await Task.select(requests, frames).value
-  }
-
-  // MARK: - Loops
 
-  /// Reads control frames. A zero-length copy request is a geometry probe and is answered at once,
-  /// so a client can learn the frame size without waiting for the display to render.
-  private func readRequests(
-    _ requestStream: RequestStreamReader<Idb_FramebufferStreamRequest>,
-    into state: FramebufferStreamState,
-    responseStream: RPCWriter<Idb_FramebufferStreamResponse>
-  ) async throws {
     for try await request in requestStream {
       switch request.control {
+      case let .configure(configure):
+        guard attached == nil else {
+          try await responseStream.send(
+            Self.failure(FramebufferStreamError.configuredLate, sharedMemoryName: ""))
+          continue
+        }
+        do {
+          let established = try await attach(
+            simulator: simulator,
+            displayUniqueID: configure.hasDisplayUniqueID ? configure.displayUniqueID : nil,
+            scaleFactor: configure.hasScaleFactor ? configure.scaleFactor : nil)
+          attached = established
+          frames = Task { try await serviceFrames(established.attachment, state: established.state) }
+          // Answered with the chosen display's geometry, so Configure doubles as a probe.
+          try await responseStream.send(
+            await probeResponse(for: .init(), state: established.state))
+        } catch {
+          try await responseStream.send(Self.failure(error, sharedMemoryName: ""))
+        }
+
       case let .copyFramebuffer(copy):
+        if attached == nil {
+          do {
+            let established = try await attach(simulator: simulator, displayUniqueID: nil, scaleFactor: nil)
+            attached = established
+            frames = Task { try await serviceFrames(established.attachment, state: established.state) }
+          } catch {
+            try await responseStream.send(Self.failure(error, sharedMemoryName: copy.sharedMemoryName))
+            continue
+          }
+        }
+        guard let state = attached?.state else { continue }
         guard copy.sharedMemoryLength > 0 else {
           try await responseStream.send(await probeResponse(for: copy, state: state))
           continue
         }
         // Answered from the current surface rather than on the next rendered frame: a client pulls
         // when it wants a frame, and a still display renders none to wait for.
         try await responseStream.send(await copyResponse(for: copy, state: state))
+
       case .stop, .none:
         return
       }
     }
   }
 
+  /// A stream's display, once chosen.
+  private struct Attached {
+    let attachment: FramebufferAttachment
+    let state: FramebufferStreamState
+  }
+
+  /// Resolves the display and attaches to it.
+  ///
+  /// Without an id this is the active integrated display -- the lit panel, which on a foldable is the
+  /// one the user is looking at. A runtime too old to report displays falls back to the main-screen
+  /// heuristic, which is all that was ever available there.
+  private func attach(
+    simulator: Simulator, displayUniqueID: String?, scaleFactor: Float?
+  ) async throws -> Attached {
+    let framebuffer: Framebuffer
+    if let displayUniqueID {
+      guard let displays = try? await simulator.displays.list() else {
+        throw FramebufferStreamError.displaysUnavailable
+      }
+      guard let display = displays.first(where: { $0.uniqueID == displayUniqueID }) else {
+        throw FramebufferStreamError.unknownDisplay(uniqueID: displayUniqueID)
+      }
+      framebuffer = try await Framebuffer.surface(for: display, simulator: simulator)
+    } else if let active = try await simulator.displays.activeIntegratedDisplayIfSupported() {
+      framebuffer = try await Framebuffer.surface(for: active, simulator: simulator)
+    } else {
+      framebuffer = try Framebuffer.mainScreenSurface(for: simulator, logger: targetLogger)
+    }
+    let attachment = try framebuffer.attach()
+    let state = FramebufferStreamState(surface: attachment.initialSurface)
+    if let scaleFactor {
+      _ = try state.scale(requesting: scaleFactor)
+    }
+    return Attached(attachment: attachment, state: state)
+  }
+
+  private static func failure(_ error: any Error, sharedMemoryName: String) -> Idb_FramebufferStreamResponse {
+    Idb_FramebufferStreamResponse.with {
+      $0.sharedMemoryName = sharedMemoryName
+      $0.error = error.localizedDescription
+    }
+  }
+
+  // MARK: - Loops
+
   /// Tracks the surface the display currently holds, and ends when the display tears down.
   private func serviceFrames(
     _ attachment: FramebufferAttachment,

CompanionLib/SwiftServer/MethodHandlers/ListDisplaysMethodHandler.swift

@@ -0,0 +1,43 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import CompanionUtilities
+import FBControlCore
+import FBSimulatorControl
+import Foundation
+import GRPCCore
+import IDBGRPCSwift
+
+/// Reports the simulator's displays, so a client can name one when it opens a framebuffer stream.
+///
+/// A runtime whose provider cannot report displays answers with none rather than guessing, which is
+/// the same signal `framebuffer_stream` uses to fall back to its main-screen heuristic.
+struct ListDisplaysMethodHandler {
+
+  let target: any Target
+
+  func handle(request: Idb_ListDisplaysRequest, context: ServerContext) async throws -> Idb_ListDisplaysResponse {
+    guard let simulator = target as? Simulator else {
+      throw RPCError(code: .failedPrecondition, message: "Displays can only be listed for a simulator")
+    }
+    let displays = (try? await simulator.displays.list()) ?? []
+    return Idb_ListDisplaysResponse.with {
+      $0.displays = displays.map { display in
+        Idb_Display.with {
+          $0.uniqueID = display.uniqueID
+          $0.name = display.name
+          $0.isActive = display.isActive
+          $0.isPrimary = display.isPrimary
+          $0.isIntegrated = display.isIntegrated
+          $0.width = UInt32(display.bounds.width)
+          $0.height = UInt32(display.bounds.height)
+          $0.scale = display.scale
+        }
+      }
+    }
+  }
+}

FBSimulatorControl/Commands/SimulatorDisplayCommands.swift

@@ -203,7 +203,7 @@ public final class SimulatorDisplayCommands: DisplayCommands, @unchecked Sendabl
 
   /// Returns nil only when the runtime lacks the feature, or the provider the fields, needed to
   /// select a display.
-  func activeIntegratedDisplayIfSupported() async throws -> SimulatorDisplay? {
+  public func activeIntegratedDisplayIfSupported() async throws -> SimulatorDisplay? {
     let snapshot = try await target().coreDevice.performIfSupported(
       action: SimulatorDisplayProtocol.action, service: SimulatorDisplayProtocol.service, input: CoreDeviceEmptyInput(),
       decode: SimulatorDisplayProtocol.snapshot)

FBSimulatorControl/Framebuffer/Framebuffer.swift

@@ -83,6 +83,11 @@ public final class Framebuffer: @unchecked Sendable {
   private let statsRecorder: FramebufferStatsRecorder
   private let logger: any ControlCoreLogger
 
+  /// The framebuffer of one named display, for a caller that has chosen from `displays.list()`.
+  public class func surface(for display: SimulatorDisplay, simulator: Simulator) async throws -> Framebuffer {
+    try await FramebufferSurfaceLocator.framebuffer(for: display, simulator: simulator)
+  }
+
   public class func mainScreenSurface(for simulator: Simulator, logger: any ControlCoreLogger) throws -> Framebuffer {
     let surface = try FramebufferSurfaceLocator.mainDisplaySurface(for: simulator, logger: logger)
     return Framebuffer(surface: surface, logger: logger)

proto/idb.proto

@@ -70,6 +70,7 @@ service CompanionService {
       returns (stream VideoStreamResponse) {}
   rpc framebuffer_stream(stream FramebufferStreamRequest)
       returns (stream FramebufferStreamResponse) {}
+  rpc list_displays(ListDisplaysRequest) returns (ListDisplaysResponse) {}
   // Crash Operations
   rpc crash_delete(CrashLogQuery) returns (CrashLogResponse) {}
   rpc crash_list(CrashLogQuery) returns (CrashLogResponse) {}
@@ -898,17 +899,33 @@ message VideoStreamResponse {
 }
 
 message FramebufferStreamRequest {
+  // Fixes what the stream reads for its lifetime. Optional, at most once, and only
+  // before the first copy: a client allocates against the geometry it is told, so
+  // neither the display nor the scale can change under it afterwards. Answered with
+  // the selected display's geometry, so a Configure doubles as a probe.
+  //
+  // Omitted entirely, the stream reads the active integrated display unscaled --
+  // what every client did before this message existed.
+  message Configure {
+    // From `list_displays`. Unset selects the active integrated display.
+    optional string display_unique_id = 1;
+    // In (0, 1]. Unset copies at the display's own size.
+    optional float scale_factor = 2;
+  }
   message CopyFramebuffer {
     string shared_memory_name = 1;
     // A zero length asks only for the framebuffer's geometry, answered from the
     // current surface without waiting for a frame.
     uint64 shared_memory_length = 2;
+    // Superseded by Configure.scale_factor. Still honoured when no Configure was
+    // sent, where the first request naming a scale fixes it for the stream.
     optional float scale_factor = 3;
   }
   message Stop {}
   oneof control {
     CopyFramebuffer copy_framebuffer = 1;
     Stop stop = 2;
+    Configure configure = 3;
   }
 }
 
@@ -920,6 +937,29 @@ message FramebufferInfo {
   string format = 5;
 }
 
+message ListDisplaysRequest {}
+
+// One of the simulator's displays. A foldable reports an integrated panel per
+// fold state, of which one is active; a runtime that cannot report displays
+// answers with none rather than guessing.
+message Display {
+  string unique_id = 1;
+  string name = 2;
+  // Lit and rendering. Exactly one integrated display is active at a time.
+  bool is_active = 3;
+  bool is_primary = 4;
+  // False for TVOut, CarPlay and the resizable scene display.
+  bool is_integrated = 5;
+  // Bounds in the display's own unrotated pixel space.
+  uint32 width = 6;
+  uint32 height = 7;
+  double scale = 8;
+}
+
+message ListDisplaysResponse {
+  repeated Display displays = 1;
+}
+
 message FramebufferStreamResponse {
   string shared_memory_name = 1;
   optional uint64 bytes_written = 2;