diff --git a/Ensemble.xcodeproj/project.pbxproj b/Ensemble.xcodeproj/project.pbxproj index 1c4d9b0..c6bf666 100644 --- a/Ensemble.xcodeproj/project.pbxproj +++ b/Ensemble.xcodeproj/project.pbxproj @@ -28,6 +28,7 @@ 4978BAB12AD55E8B000C549C /* WindowView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 4978BAB02AD55E8B000C549C /* WindowView.swift */; }; 4989D3402B0B9393005E2E7A /* shut_up_logging.c in Sources */ = {isa = PBXBuildFile; fileRef = 4989D33F2B0B9393005E2E7A /* shut_up_logging.c */; }; 4989D3412B0B9393005E2E7A /* shut_up_logging.c in Sources */ = {isa = PBXBuildFile; fileRef = 4989D33F2B0B9393005E2E7A /* shut_up_logging.c */; }; + 4992A6012B68291900844A16 /* WindowManager.swift in Sources */ = {isa = PBXBuildFile; fileRef = 4992A6002B68291900844A16 /* WindowManager.swift */; }; 49B352C72AE53A9300BCE03D /* Frame.swift in Sources */ = {isa = PBXBuildFile; fileRef = 49B352C62AE53A9300BCE03D /* Frame.swift */; }; 49B352C82AE53A9300BCE03D /* Frame.swift in Sources */ = {isa = PBXBuildFile; fileRef = 49B352C62AE53A9300BCE03D /* Frame.swift */; }; 49B352CB2AE593C300BCE03D /* FrameView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 49B352C92AE593C300BCE03D /* FrameView.swift */; }; @@ -80,6 +81,7 @@ 4978BAAE2AD55D71000C549C /* WindowPreviewView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = WindowPreviewView.swift; sourceTree = ""; }; 4978BAB02AD55E8B000C549C /* WindowView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = WindowView.swift; sourceTree = ""; }; 4989D33F2B0B9393005E2E7A /* shut_up_logging.c */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.c; path = shut_up_logging.c; sourceTree = ""; }; + 4992A6002B68291900844A16 /* WindowManager.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = WindowManager.swift; sourceTree = ""; }; 49B352C62AE53A9300BCE03D /* Frame.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Frame.swift; sourceTree = ""; }; 49B352C92AE593C300BCE03D /* FrameView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = FrameView.swift; sourceTree = ""; }; 49E09B532AD2EE5000B56CD3 /* Ensemble.app */ = {isa = PBXFileReference; explicitFileType = wrapper.application; includeInIndex = 0; path = Ensemble.app; sourceTree = BUILT_PRODUCTS_DIR; }; @@ -172,6 +174,7 @@ 49E09BC42AD52CC900B56CD3 /* Remote.swift */, 49226A302AE447C10044CFC9 /* ScreenRecorder.swift */, 49226A2E2AE43EF50044CFC9 /* SPI.swift */, + 4992A6002B68291900844A16 /* WindowManager.swift */, 49E09B5A2AD2EE5100B56CD3 /* Assets.xcassets */, 49E09B5C2AD2EE5100B56CD3 /* Preview Content */, ); @@ -411,6 +414,7 @@ 49EDAA6E2B28E58A00546EAB /* Events.swift in Sources */, 4901A14C2B7246760040D2EE /* Preference.swift in Sources */, 49E09BC22AD52C7800B56CD3 /* macOSInterface.swift in Sources */, + 4992A6012B68291900844A16 /* WindowManager.swift in Sources */, 49E09BB32AD419CF00B56CD3 /* Messages.swift in Sources */, 4901A14A2B721EAB0040D2EE /* PermissionsView.swift in Sources */, 49226A312AE447C10044CFC9 /* ScreenRecorder.swift in Sources */, diff --git a/macOS/Local.swift b/macOS/Local.swift index 74c22db..b7df87f 100644 --- a/macOS/Local.swift +++ b/macOS/Local.swift @@ -15,6 +15,7 @@ class Local: LocalInterface, macOSInterface { let screenRecorder = ScreenRecorder() let eventDispatcher = EventDispatcher() + let windowManager = WindowManager() struct Mask { let mask: vImage.PixelBuffer @@ -108,7 +109,7 @@ class Local: LocalInterface, macOSInterface { func _windows(parameters: M.Windows.Request) async throws -> M.Windows.Reply { return try await .init( - windows: screenRecorder.windows.compactMap { + windows: windowManager.allWindows.compactMap { guard let application = $0.owningApplication?.applicationName, $0.isOnScreen else { @@ -119,7 +120,7 @@ class Local: LocalInterface, macOSInterface { } func _windowPreview(parameters: M.WindowPreview.Request) async throws -> M.WindowPreview.Reply { - guard let window = try await screenRecorder.lookup(windowID: parameters.windowID), + guard let window = try await windowManager.lookupWindow(byID: parameters.windowID), window.isOnScreen, let screenshot = try await screenRecorder.screenshot(window: window, size: M.WindowPreview.previewSize) else { @@ -130,7 +131,7 @@ class Local: LocalInterface, macOSInterface { } func _startCasting(parameters: M.StartCasting.Request) async throws -> M.StartCasting.Reply { - let window = try await screenRecorder.lookup(windowID: parameters.windowID)! + let window = try await windowManager.lookupWindow(byID: parameters.windowID)! let stream = try await screenRecorder.stream(window: window) Task { @@ -157,9 +158,11 @@ class Local: LocalInterface, macOSInterface { return .init() } + var childObservers = [CGWindowID: Task]() + func _startWatchingForChildWindows(parameters: M.StartWatchingForChildWindows.Request) async throws -> M.StartWatchingForChildWindows.Reply { - Task { - for await children in await screenRecorder.watchForChildren(windowID: parameters.windowID) { + childObservers[parameters.windowID] = Task { + for try await children in await windowManager.childrenOfWindow(idenitifiedBy: parameters.windowID) { try await remote.childWindows(parent: parameters.windowID, children: children) } } @@ -167,19 +170,19 @@ class Local: LocalInterface, macOSInterface { } func _stopWatchingForChildWindows(parameters: M.StopWatchingForChildWindows.Request) async throws -> M.StopWatchingForChildWindows.Reply { - await screenRecorder.stopWatchingForChildren(windowID: parameters.windowID) + childObservers.removeValue(forKey: parameters.windowID)!.cancel() return .init() } func _mouseMoved(parameters: M.MouseMoved.Request) async throws -> M.MouseMoved.Reply { - let window = try await screenRecorder.lookup(windowID: parameters.windowID)! + let window = try await windowManager.lookupWindow(byID: parameters.windowID)! await eventDispatcher.injectMouseMoved(to: .init(x: window.frame.minX + window.frame.width * parameters.x, y: window.frame.minY + window.frame.height * parameters.y)) return .init() } func _clicked(parameters: M.Clicked.Request) async throws -> M.Clicked.Reply { - let window = try await screenRecorder.lookup(windowID: parameters.windowID)! + let window = try await windowManager.lookupWindow(byID: parameters.windowID)! await eventDispatcher.injectClick(at: .init(x: window.frame.minX + window.frame.width * parameters.x, y: window.frame.minY + window.frame.height * parameters.y)) return .init() } @@ -203,21 +206,21 @@ class Local: LocalInterface, macOSInterface { } func _dragBegan(parameters: M.DragBegan.Request) async throws -> M.DragBegan.Reply { - let window = try await screenRecorder.lookup(windowID: parameters.windowID)! + let window = try await windowManager.lookupWindow(byID: parameters.windowID)! await eventDispatcher.injectDragBegan(at: .init(x: window.frame.minX + window.frame.width * parameters.x, y: window.frame.minY + window.frame.height * parameters.y)) return .init() } func _dragChanged(parameters: M.DragChanged.Request) async throws -> M.DragChanged.Reply { - let window = try await screenRecorder.lookup(windowID: parameters.windowID)! + let window = try await windowManager.lookupWindow(byID: parameters.windowID)! await eventDispatcher.injectDragChanged(to: .init(x: window.frame.minX + window.frame.width * parameters.x, y: window.frame.minY + window.frame.height * parameters.y)) return .init() } func _dragEnded(parameters: M.DragEnded.Request) async throws -> M.DragEnded.Reply { - let window = try await screenRecorder.lookup(windowID: parameters.windowID)! + let window = try await windowManager.lookupWindow(byID: parameters.windowID)! await eventDispatcher.injectDragEnded(at: .init(x: window.frame.minX + window.frame.width * parameters.x, y: window.frame.minY + window.frame.height * parameters.y)) return .init() diff --git a/macOS/ScreenRecorder.swift b/macOS/ScreenRecorder.swift index f96bbcd..c7cee4f 100644 --- a/macOS/ScreenRecorder.swift +++ b/macOS/ScreenRecorder.swift @@ -9,39 +9,6 @@ import AVFoundation import ScreenCaptureKit actor ScreenRecorder { - static let cacheDuration = Duration.seconds(1) - - var _windows = [CGWindowID: SCWindow]() - var _lastWindowFetch = ContinuousClock.Instant.now.advanced(by: ScreenRecorder.cacheDuration * -2) - - func _updateWindows(force: Bool = false) async throws { - guard ContinuousClock.Instant.now - _lastWindowFetch > Self.cacheDuration || force else { - return - } - - try await _windows = Dictionary( - uniqueKeysWithValues: SCShareableContent.excludingDesktopWindows(true, onScreenWindowsOnly: false).windows.map { - ($0.windowID, $0) - }) - _lastWindowFetch = ContinuousClock.Instant.now - } - - var windows: [SCWindow] { - get async throws { - try await _updateWindows() - return Array(_windows.values) - } - } - - func lookup(windowID: CGWindowID) async throws -> SCWindow? { - if let window = _windows[windowID] { - return window - } else { - try await _updateWindows(force: true) - return _windows[windowID] - } - } - static func streamConfiguration() -> SCStreamConfiguration { let configuration = SCStreamConfiguration() configuration.pixelFormat = kCVPixelFormatType_32BGRA @@ -110,39 +77,4 @@ actor ScreenRecorder { func stopStream(for windowID: CGWindowID) async { await streams.removeValue(forKey: windowID)!.stop() } - - var childObservers = Set() - - func watchForChildren(windowID: CGWindowID) -> AsyncStream<[CGWindowID]> { - let (stream, continuation) = AsyncStream.makeStream(of: [CGWindowID].self) - childObservers.insert(windowID) - Task { - while childObservers.contains(windowID) { - try await Task.sleep(for: .seconds(1)) - var childWindows = - if let SLSCopyAssociatedWindows, - let SLSMainConnectionID - { - Set(SLSCopyAssociatedWindows(SLSMainConnectionID(), windowID) as? [CGWindowID] ?? []) - } else { - Set() - } - childWindows.remove(windowID) - - let root = try await lookup(windowID: windowID)! - let overlays = try await windows.filter { - $0.owningApplication == root.owningApplication && $0.windowLayer > NSWindow.Level.normal.rawValue && $0.frame.intersects(root.frame) - }.map(\.windowID) - - continuation.yield(Array(childWindows) + overlays) - } - continuation.finish() - } - return stream - } - - func stopWatchingForChildren(windowID: CGWindowID) { - let result = childObservers.remove(windowID) - assert(result != nil) - } } diff --git a/macOS/WindowManager.swift b/macOS/WindowManager.swift new file mode 100644 index 0000000..7be7d36 --- /dev/null +++ b/macOS/WindowManager.swift @@ -0,0 +1,154 @@ +// +// WindowManager.swift +// macOS +// +// Created by Saagar Jha on 1/29/24. +// + +import ApplicationServices +import ScreenCaptureKit + +actor WindowManager { + var applications = [pid_t: Application]() + var windows = [CGWindowID: Window]() + + class Application { + let application: SCRunningApplication + var windows: [Window] + + var windowUpdates: AsyncStream { + if Permission.helper.supported && Permission.helper.enabled { + let stream = AXObserver.observe([kAXCreatedNotification, kAXMenuOpenedNotification, kAXUIElementDestroyedNotification], for: AXUIElementCreateApplication(application.processID)) + var iterator = stream.makeAsyncIterator() + + return AsyncStream { + _ = await iterator.next() + } + } else { + return AsyncStream { + try? await Task.sleep(for: .seconds(1)) + } + } + } + + init(application: SCRunningApplication) { + self.application = application + windows = [] + } + + func childWindows(of window: Window) -> [CGWindowID] { + var childWindows = + if let SLSCopyAssociatedWindows, + let SLSMainConnectionID + { + Set(SLSCopyAssociatedWindows(SLSMainConnectionID(), window.window.windowID) as? [CGWindowID] ?? []) + } else { + Set() + } + childWindows.remove(window.window.windowID) + + let overlays = windows.filter { + $0.window.windowLayer > NSWindow.Level.normal.rawValue && $0.window.frame.intersects(window.window.frame) + }.map(\.window.windowID) + + return Array(childWindows) + overlays + } + + static func sameApplication(lhs: SCRunningApplication, rhs: SCRunningApplication) -> Bool { + lhs.processID == rhs.processID && lhs.bundleIdentifier == rhs.bundleIdentifier && lhs.applicationName == rhs.applicationName + } + } + + struct Window { + weak var application: Application! + let window: SCWindow + } + + func updateWindows() async throws { + var newApplications = [pid_t: Application]() + for application in applications.values { + application.windows.removeAll() + } + windows.removeAll() + + func lookup(application: SCRunningApplication) -> Application { + if let _application = newApplications[application.processID] { + assert(Application.sameApplication(lhs: application, rhs: _application.application)) + return _application + } + + if let _application = applications[application.processID], + Application.sameApplication(lhs: application, rhs: _application.application) + { + newApplications[application.processID] = _application + return _application + } + + let _application = Application(application: application) + newApplications[application.processID] = _application + return _application + } + + for window in try await SCShareableContent.excludingDesktopWindows(true, onScreenWindowsOnly: false).windows where window.owningApplication != nil { + let _application = lookup(application: window.owningApplication!) + let _window = Window(application: _application, window: window) + _application.windows.append(_window) + windows[window.windowID] = _window + } + + applications = newApplications + } + + func childrenOfWindow(idenitifiedBy windowID: CGWindowID) -> AsyncThrowingStream<[CGWindowID], Error> { + let window = windows[windowID]! + let application = window.application! + var iterator = application.windowUpdates.makeAsyncIterator() + return AsyncThrowingStream { + await iterator.next() + try await self.updateWindows() + return application.childWindows(of: window) + } + } + + func lookupWindow(byID id: CGWindowID) async throws -> SCWindow? { + guard let window = windows[id]?.window else { + try await updateWindows() + return windows[id]?.window + } + return window + } + + var allWindows: [SCWindow] { + get async throws { + try await updateWindows() + return windows.values.map(\.window) + } + } +} + +extension AXObserver { + static func observe(_ notifications: [String], for element: AXUIElement) -> AsyncStream<(AXUIElement, String)> { + AsyncStream<(AXUIElement, String)> { continuation in + var pid: pid_t = 0 + AXUIElementGetPid(element, &pid) + var observer: AXObserver! + + AXObserverCreate( + pid, + { _, element, notification, refcon in + let continuation = Unmanaged.fromOpaque(refcon!).takeUnretainedValue() as! AsyncStream<(AXUIElement, String)>.Continuation + continuation.yield((element, notification as String)) + }, &observer) + for notification in notifications { + AXObserverAddNotification(observer, element, notification as CFString, Unmanaged.passRetained(continuation as AnyObject).toOpaque()) + } + CFRunLoopAddSource(CFRunLoopGetMain(), AXObserverGetRunLoopSource(observer), .defaultMode) + + // Retain the observer until the stream is finished + let _observer = observer + continuation.onTermination = { _ in + _ = _observer + } + } + } +}