import AppKit
|
import CoreImage
|
import CoreMedia
|
import CoreGraphics
|
import Foundation
|
@preconcurrency import ScreenCaptureKit
|
|
public enum ScreenshotProviderError: Error, Equatable, Sendable {
|
case windowNotFound(UInt32)
|
case captureTimedOut(windowID: UInt32, timeoutMilliseconds: Int)
|
case invalidCapture(windowID: UInt32, pixelWidth: Int, pixelHeight: Int, pointWidth: Double, pointHeight: Double)
|
}
|
|
public struct WindowScreenshotCapture: Sendable {
|
public let image: CGImage
|
public let pointSize: NSSize
|
|
public init(image: CGImage, pointSize: NSSize) {
|
self.image = image
|
self.pointSize = pointSize
|
}
|
}
|
|
public enum WindowServerWindowScreenshotCapturer {
|
public static func capture(windowID: UInt32) -> WindowScreenshotCapture? {
|
guard let image = captureImage(windowID: windowID) else {
|
return nil
|
}
|
|
return WindowScreenshotCapture(
|
image: image,
|
pointSize: NSSize(width: image.width, height: image.height)
|
)
|
}
|
|
public static func captureImage(windowID: UInt32) -> CGImage? {
|
var mutableWindowID = CGWindowID(windowID)
|
let images = CGSHWCaptureWindowList(
|
CGSMainConnectionID(),
|
&mutableWindowID,
|
1,
|
[.ignoreGlobalClipShape, .bestResolution, .fullSize]
|
).takeRetainedValue() as NSArray
|
|
guard let firstObject = images.firstObject else {
|
return nil
|
}
|
|
let cfObject = firstObject as CFTypeRef
|
guard CFGetTypeID(cfObject) == CGImage.typeID else {
|
return nil
|
}
|
|
return (firstObject as! CGImage)
|
}
|
}
|
|
@MainActor
|
public protocol WindowScreenshotCapturing: Sendable {
|
func captureWindow(windowID: UInt32) async throws -> WindowScreenshotCapture
|
}
|
|
public struct ScreenshotCapturePolicy: Equatable, Sendable {
|
public static let round01Default = ScreenshotCapturePolicy(
|
timeoutMilliseconds: 300,
|
maxRetriesPerSession: 1
|
)
|
|
public let timeoutMilliseconds: Int
|
public let maxRetriesPerSession: Int
|
|
public init(timeoutMilliseconds: Int, maxRetriesPerSession: Int) {
|
self.timeoutMilliseconds = max(1, timeoutMilliseconds)
|
self.maxRetriesPerSession = max(0, maxRetriesPerSession)
|
}
|
|
var timeoutNanoseconds: UInt64 {
|
UInt64(timeoutMilliseconds) * 1_000_000
|
}
|
}
|
|
public enum ScreenshotFallbackReason: Equatable, Sendable {
|
case screenRecordingDenied
|
case captureFailed(String)
|
case timedOut(windowID: UInt32, timeoutMilliseconds: Int)
|
case invalidCapture(windowID: UInt32, pixelWidth: Int, pixelHeight: Int, pointWidth: Double, pointHeight: Double)
|
case retryLimitReached(failedAttempts: Int)
|
case syntheticWindowID(windowID: UInt32)
|
}
|
|
public enum ScreenshotResolutionSource: Equatable, Sendable {
|
case realScreenshot
|
case skeletonFallback(ScreenshotFallbackReason)
|
}
|
|
@MainActor
|
public struct ScreenshotResolution {
|
public let image: NSImage
|
public let source: ScreenshotResolutionSource
|
|
public init(image: NSImage, source: ScreenshotResolutionSource) {
|
self.image = image
|
self.source = source
|
}
|
}
|
|
public enum ScreenshotDebugEvent: Equatable, Sendable {
|
case captureStarted(windowID: UInt32)
|
case captureSucceeded(windowID: UInt32)
|
case captureFailed(windowID: UInt32, reason: ScreenshotFallbackReason)
|
case retrySkipped(windowID: UInt32, failedAttempts: Int)
|
case fallbackUsed(windowID: UInt32, reason: ScreenshotFallbackReason)
|
}
|
|
public extension ScreenshotFallbackReason {
|
var diagnosticDescription: String {
|
switch self {
|
case .screenRecordingDenied:
|
return "screenRecordingDenied"
|
case .captureFailed(let summary):
|
return "captureFailed(\(summary))"
|
case .timedOut(let windowID, let timeoutMilliseconds):
|
return "timedOut(windowID: \(windowID), timeoutMilliseconds: \(timeoutMilliseconds))"
|
case .invalidCapture(let windowID, let pixelWidth, let pixelHeight, let pointWidth, let pointHeight):
|
return "invalidCapture(windowID: \(windowID), pixelSize: \(pixelWidth)x\(pixelHeight), pointSize: \(pointWidth)x\(pointHeight))"
|
case .retryLimitReached(let failedAttempts):
|
return "retryLimitReached(failedAttempts: \(failedAttempts))"
|
case .syntheticWindowID(let windowID):
|
return "syntheticWindowID(windowID: \(windowID))"
|
}
|
}
|
}
|
|
public extension ScreenshotDebugEvent {
|
var diagnosticDescription: String {
|
switch self {
|
case .captureStarted(let windowID):
|
return "captureStarted(\(windowID))"
|
case .captureSucceeded(let windowID):
|
return "captureSucceeded(\(windowID))"
|
case .captureFailed(let windowID, let reason):
|
return "captureFailed(\(windowID), \(reason.diagnosticDescription))"
|
case .retrySkipped(let windowID, let failedAttempts):
|
return "retrySkipped(\(windowID), failedAttempts: \(failedAttempts))"
|
case .fallbackUsed(let windowID, let reason):
|
return "fallbackUsed(\(windowID), \(reason.diagnosticDescription))"
|
}
|
}
|
}
|
|
@MainActor
|
public protocol ScreenshotDebugLogging: AnyObject {
|
func record(_ event: ScreenshotDebugEvent)
|
}
|
|
@MainActor
|
public final class NoopScreenshotDebugLogger: ScreenshotDebugLogging {
|
public init() {}
|
|
public func record(_ event: ScreenshotDebugEvent) {}
|
}
|
|
@MainActor
|
public final class ScreenshotCaptureSession {
|
private var failedAttemptsByWindowID: [UInt32: Int] = [:]
|
|
public init() {}
|
|
func failedAttempts(for windowID: UInt32) -> Int {
|
failedAttemptsByWindowID[windowID, default: 0]
|
}
|
|
func recordFailure(for windowID: UInt32) {
|
failedAttemptsByWindowID[windowID, default: 0] += 1
|
}
|
|
func recordSuccess(for windowID: UInt32) {
|
failedAttemptsByWindowID[windowID] = nil
|
}
|
}
|
|
@MainActor
|
private final class ScreenshotCaptureRace {
|
private var continuation: CheckedContinuation<WindowScreenshotCapture, any Error>?
|
private var result: Result<WindowScreenshotCapture, any Error>?
|
private var didFinish = false
|
|
func wait() async throws -> WindowScreenshotCapture {
|
if let result {
|
return try result.get()
|
}
|
|
return try await withCheckedThrowingContinuation { continuation in
|
self.continuation = continuation
|
}
|
}
|
|
func finish(_ result: Result<WindowScreenshotCapture, any Error>) {
|
guard !didFinish else {
|
return
|
}
|
|
didFinish = true
|
if let continuation {
|
self.continuation = nil
|
continuation.resume(with: result)
|
} else {
|
self.result = result
|
}
|
}
|
}
|
|
@MainActor
|
public final class ScreenCaptureKitWindowScreenshotCapturer: WindowScreenshotCapturing {
|
public init() {}
|
|
public func captureWindow(windowID: UInt32) async throws -> WindowScreenshotCapture {
|
if let privateCapture = WindowServerWindowScreenshotCapturer.capture(windowID: windowID) {
|
return privateCapture
|
}
|
|
let content = try await SCShareableContent.current
|
|
guard let captureWindow = content.windows.first(where: { $0.windowID == windowID }) else {
|
throw ScreenshotProviderError.windowNotFound(windowID)
|
}
|
|
let filter = SCContentFilter(desktopIndependentWindow: captureWindow)
|
let contentInfo = SCShareableContent.info(for: filter)
|
let configuration = configuration(for: contentInfo)
|
let image: CGImage
|
do {
|
image = try await SCScreenshotManager.captureImage(contentFilter: filter, configuration: configuration)
|
} catch {
|
if let sampleBufferImage = try? await captureSampleBufferImage(
|
filter: filter,
|
configuration: configuration,
|
windowID: windowID,
|
pointSize: NSSize(width: contentInfo.contentRect.width, height: contentInfo.contentRect.height)
|
) {
|
image = sampleBufferImage
|
} else {
|
throw error
|
}
|
}
|
|
return WindowScreenshotCapture(
|
image: image,
|
pointSize: NSSize(width: contentInfo.contentRect.width, height: contentInfo.contentRect.height)
|
)
|
}
|
|
private func configuration(for contentInfo: SCShareableContentInfo) -> SCStreamConfiguration {
|
let configuration = SCStreamConfiguration()
|
configuration.width = max(1, Int(contentInfo.contentRect.width * CGFloat(contentInfo.pointPixelScale)))
|
configuration.height = max(1, Int(contentInfo.contentRect.height * CGFloat(contentInfo.pointPixelScale)))
|
configuration.showsCursor = false
|
configuration.ignoreShadowsSingleWindow = true
|
configuration.includeChildWindows = true
|
return configuration
|
}
|
|
private func captureSampleBufferImage(
|
filter: SCContentFilter,
|
configuration: SCStreamConfiguration,
|
windowID: UInt32,
|
pointSize: NSSize
|
) async throws -> CGImage {
|
try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<CGImage, any Error>) in
|
SCScreenshotManager.captureSampleBuffer(
|
contentFilter: filter,
|
configuration: configuration
|
) { sampleBuffer, error in
|
if let error {
|
continuation.resume(throwing: error)
|
return
|
}
|
guard let sampleBuffer else {
|
continuation.resume(throwing: ScreenshotProviderError.invalidCapture(
|
windowID: windowID,
|
pixelWidth: 0,
|
pixelHeight: 0,
|
pointWidth: Double(pointSize.width),
|
pointHeight: Double(pointSize.height)
|
))
|
return
|
}
|
|
guard let pixelBuffer = CMSampleBufferGetImageBuffer(sampleBuffer) else {
|
continuation.resume(throwing: ScreenshotProviderError.invalidCapture(
|
windowID: windowID,
|
pixelWidth: 0,
|
pixelHeight: 0,
|
pointWidth: Double(pointSize.width),
|
pointHeight: Double(pointSize.height)
|
))
|
return
|
}
|
|
let ciImage = CIImage(cvPixelBuffer: pixelBuffer)
|
guard let cgImage = CIContext().createCGImage(ciImage, from: ciImage.extent) else {
|
continuation.resume(throwing: ScreenshotProviderError.invalidCapture(
|
windowID: windowID,
|
pixelWidth: CVPixelBufferGetWidth(pixelBuffer),
|
pixelHeight: CVPixelBufferGetHeight(pixelBuffer),
|
pointWidth: Double(pointSize.width),
|
pointHeight: Double(pointSize.height)
|
))
|
return
|
}
|
|
continuation.resume(returning: cgImage)
|
}
|
}
|
}
|
}
|
|
private typealias CGSConnectionID = UInt32
|
|
private struct CGSWindowCaptureOptions: OptionSet {
|
let rawValue: UInt32
|
|
static let ignoreGlobalClipShape = CGSWindowCaptureOptions(rawValue: 1 << 11)
|
static let bestResolution = CGSWindowCaptureOptions(rawValue: 1 << 8)
|
static let fullSize = CGSWindowCaptureOptions(rawValue: 1 << 19)
|
}
|
|
@_silgen_name("CGSMainConnectionID")
|
private func CGSMainConnectionID() -> CGSConnectionID
|
|
@_silgen_name("CGSHWCaptureWindowList")
|
private func CGSHWCaptureWindowList(
|
_ connectionID: CGSConnectionID,
|
_ windowList: UnsafeMutablePointer<CGWindowID>,
|
_ windowCount: UInt32,
|
_ options: CGSWindowCaptureOptions
|
) -> Unmanaged<CFArray>
|
|
@MainActor
|
public final class ScreenCaptureKitScreenshotProvider: ScreenshotProviderProtocol {
|
private let capturer: any WindowScreenshotCapturing
|
private let skeletonProvider: any SkeletonThumbnailProviderProtocol
|
private let policy: ScreenshotCapturePolicy
|
private let debugLogger: any ScreenshotDebugLogging
|
private let permissionChecker: any SystemPermissionChecking
|
private let defaultSession = ScreenshotCaptureSession()
|
|
public init(
|
capturer: any WindowScreenshotCapturing = ScreenCaptureKitWindowScreenshotCapturer(),
|
skeletonProvider: any SkeletonThumbnailProviderProtocol = NativeSkeletonThumbnailProvider(),
|
policy: ScreenshotCapturePolicy = .round01Default,
|
debugLogger: any ScreenshotDebugLogging = NoopScreenshotDebugLogger(),
|
permissionChecker: any SystemPermissionChecking = SystemPermissionChecker()
|
) {
|
self.capturer = capturer
|
self.skeletonProvider = skeletonProvider
|
self.policy = policy
|
self.debugLogger = debugLogger
|
self.permissionChecker = permissionChecker
|
}
|
|
public func screenshot(for window: AlignerWindow) async throws -> NSImage {
|
let capture = try await captureWindowWithinTimeout(windowID: window.id)
|
try validateCapture(capture, windowID: window.id)
|
return NSImage(cgImage: capture.image, size: capture.pointSize)
|
}
|
|
public func resolvedScreenshot(for window: AlignerWindow) async -> ScreenshotResolution {
|
await resolvedScreenshot(for: window, in: defaultSession)
|
}
|
|
public func resolvedScreenshot(for window: AlignerWindow, in session: ScreenshotCaptureSession) async -> ScreenshotResolution {
|
if MinimizedWindowFallbackPolicy.isSyntheticWindow(window) {
|
return fallbackResolution(for: window, reason: .syntheticWindowID(windowID: window.id))
|
}
|
if permissionChecker.status(for: .screenRecording) != .granted {
|
return fallbackResolution(for: window, reason: .screenRecordingDenied)
|
}
|
|
let failedAttempts = session.failedAttempts(for: window.id)
|
if failedAttempts > policy.maxRetriesPerSession {
|
let reason = ScreenshotFallbackReason.retryLimitReached(failedAttempts: failedAttempts)
|
debugLogger.record(.retrySkipped(windowID: window.id, failedAttempts: failedAttempts))
|
return fallbackResolution(for: window, reason: reason)
|
}
|
|
debugLogger.record(.captureStarted(windowID: window.id))
|
|
do {
|
let image = try await screenshot(for: window)
|
session.recordSuccess(for: window.id)
|
debugLogger.record(.captureSucceeded(windowID: window.id))
|
return ScreenshotResolution(image: image, source: .realScreenshot)
|
} catch {
|
let reason = fallbackReason(for: error)
|
session.recordFailure(for: window.id)
|
debugLogger.record(.captureFailed(windowID: window.id, reason: reason))
|
return fallbackResolution(for: window, reason: reason)
|
}
|
}
|
|
private func captureWindowWithinTimeout(windowID: UInt32) async throws -> WindowScreenshotCapture {
|
let capturer = self.capturer
|
let race = ScreenshotCaptureRace()
|
let captureTask = Task { @MainActor in
|
do {
|
let capture = try await capturer.captureWindow(windowID: windowID)
|
race.finish(.success(capture))
|
} catch {
|
race.finish(.failure(error))
|
}
|
}
|
let timeoutNanoseconds = policy.timeoutNanoseconds
|
let timeoutMilliseconds = policy.timeoutMilliseconds
|
let timeoutTask = Task {
|
do {
|
try await Task.sleep(nanoseconds: timeoutNanoseconds)
|
await MainActor.run {
|
race.finish(.failure(ScreenshotProviderError.captureTimedOut(
|
windowID: windowID,
|
timeoutMilliseconds: timeoutMilliseconds
|
)))
|
}
|
} catch {}
|
}
|
|
return try await withTaskCancellationHandler {
|
defer {
|
captureTask.cancel()
|
timeoutTask.cancel()
|
}
|
return try await race.wait()
|
} onCancel: {
|
captureTask.cancel()
|
timeoutTask.cancel()
|
Task { @MainActor in
|
race.finish(.failure(CancellationError()))
|
}
|
}
|
}
|
|
private func fallbackResolution(for window: AlignerWindow, reason: ScreenshotFallbackReason) -> ScreenshotResolution {
|
let title = window.title.isEmpty ? window.app.name : window.title
|
let image = skeletonProvider.withOverlayTitle(title, for: window.app)
|
debugLogger.record(.fallbackUsed(windowID: window.id, reason: reason))
|
return ScreenshotResolution(image: image, source: .skeletonFallback(reason))
|
}
|
|
private func fallbackReason(for error: Error) -> ScreenshotFallbackReason {
|
if case let ScreenshotProviderError.captureTimedOut(windowID, timeoutMilliseconds) = error {
|
return .timedOut(windowID: windowID, timeoutMilliseconds: timeoutMilliseconds)
|
}
|
if case let ScreenshotProviderError.invalidCapture(windowID, pixelWidth, pixelHeight, pointWidth, pointHeight) = error {
|
return .invalidCapture(
|
windowID: windowID,
|
pixelWidth: pixelWidth,
|
pixelHeight: pixelHeight,
|
pointWidth: pointWidth,
|
pointHeight: pointHeight
|
)
|
}
|
return .captureFailed(Self.sanitizedCaptureFailureSummary(for: error))
|
}
|
|
private static func sanitizedCaptureFailureSummary(for error: Error) -> String {
|
let nsError = error as NSError
|
return [
|
"domain=\(nsError.domain)",
|
"code=\(nsError.code)",
|
"descriptionHash=\(stableFingerprint(nsError.localizedDescription))",
|
"descriptionLength=\(nsError.localizedDescription.count)"
|
].joined(separator: " ")
|
}
|
|
private static func stableFingerprint(_ value: String) -> String {
|
var hash: UInt64 = 0xcbf29ce484222325
|
for byte in value.utf8 {
|
hash ^= UInt64(byte)
|
hash = hash &* 0x100000001b3
|
}
|
return String(format: "%016llx", hash)
|
}
|
|
private func validateCapture(_ capture: WindowScreenshotCapture, windowID: UInt32) throws {
|
let pointWidth = Double(capture.pointSize.width)
|
let pointHeight = Double(capture.pointSize.height)
|
guard capture.image.width > 0, capture.image.height > 0, pointWidth > 0, pointHeight > 0 else {
|
throw ScreenshotProviderError.invalidCapture(
|
windowID: windowID,
|
pixelWidth: capture.image.width,
|
pixelHeight: capture.image.height,
|
pointWidth: pointWidth,
|
pointHeight: pointHeight
|
)
|
}
|
}
|
}
|
