diff --git a/macos/Sources/LiveCapture.swift b/macos/Sources/LiveCapture.swift new file mode 100644 index 0000000..ec1523c --- /dev/null +++ b/macos/Sources/LiveCapture.swift @@ -0,0 +1,330 @@ +import AppKit +import CoreMedia +import CoreVideo +import Foundation +import ScreenCaptureKit +import VideoToolbox + +// Live display capture for MT Code's Computer View. +// +// A one-shot screenshot sets up a capture, takes one image and tears it down, +// which costs ~100 ms a frame and caps a live viewer at a few frames a second. +// A ScreenCaptureKit stream sets up once, and the system hands over a frame +// only when something on screen changed, so the viewer waits for the next +// change instead of capturing and diffing on a timer. +// +// While the stream runs, macOS shows its purple screen-recording indicator in +// the menu bar. That is accurate: someone is watching this screen. The stream +// stops a few seconds after the last request, so a viewer that crashes or +// loses its connection cannot leave it (or the indicator) running. +// +// Requested as `screenshot { display, live: true, after?, wait_ms? }`. +// Deliberately not in the tool schema: it serves that viewer, not agents, the +// same as the Windows server's `cursor: true`. + +/// Answers one live screenshot call. The image result gains a `live:` text line +/// carrying the frame's sequence number; when nothing changed within `wait_ms` +/// the result is that line alone, with no image. +func liveScreenshotResult( + _ args: [String: Any], display: Int, maxWidth: Int, encoding: ImageEncoding +) -> [String: Any] { + let after = (args["after"] as? NSNumber).map { $0.uint64Value } + let waitMs = min(max((args["wait_ms"] as? Int) ?? 1000, 0), 5000) + switch LiveCapture.shared.frame( + display: display, maxWidth: maxWidth, encoding: encoding, after: after, waitMs: waitMs) + { + case .frame(let shot, let seq): + var result = imageResult(shot, label: "display \(display)") + if var content = result["content"] as? [[String: Any]] { + content.append(["type": "text", "text": liveLine(seq: seq, changed: true)]) + result["content"] = content + } + return result + case .unchanged(let seq): + return textResult(liveLine(seq: seq, changed: false), isError: false) + case .failed(let message): + return textResult(message, isError: true) + } +} + +private func liveLine(seq: UInt64, changed: Bool) -> String { + "live: {\"seq\":\(seq),\"changed\":\(changed)}" +} + +final class LiveCapture: NSObject, SCStreamOutput, SCStreamDelegate, @unchecked Sendable { + static let shared = LiveCapture() + + enum Outcome { + case frame(CaptureShot, seq: UInt64) + case unchanged(seq: UInt64) + case failed(String) + } + + /// No request for this long and the stream stops. + private static let idleStop: TimeInterval = 5 + /// Upper bound on delivered frames; the viewer paces itself below this. + private static let maxFramesPerSecond: Int32 = 30 + /// How long a fresh stream gets to deliver its first frame. + private static let firstFrameWait: TimeInterval = 3 + + private let sampleQueue = DispatchQueue(label: "munim-computer-use.live-capture") + private let condition = NSCondition() + + // Guarded by `condition`. + private var stream: SCStream? + private var display: SCDisplay? + private var displayIndex = -1 + private var streamWidth = 0 + private var latest: CVPixelBuffer? + /// Never reset, so a frame number from before a restart is always older. + private var seq: UInt64 = 0 + private var lastRequest = Date.distantPast + private var screensChanged = false + /// Set when the person at the Mac stopped sharing from the menu bar. Honoured + /// for the life of the process: a viewer must not quietly start it again. + private var userStopped = false + private var encoded: (seq: UInt64, encoding: ImageEncoding, shot: CaptureShot)? + private var reaper: DispatchSourceTimer? + + private override init() { + super.init() + NotificationCenter.default.addObserver( + forName: NSApplication.didChangeScreenParametersNotification, object: nil, queue: nil + ) { [weak self] _ in + guard let self else { return } + self.condition.lock() + self.screensChanged = true + self.condition.unlock() + } + } + + func frame( + display index: Int, maxWidth: Int, encoding: ImageEncoding, after: UInt64?, waitMs: Int + ) -> Outcome { + condition.lock() + lastRequest = Date() + if userStopped { + condition.unlock() + return .failed( + "error: screen sharing was stopped from this Mac's menu bar, so live capture stays off " + + "until the viewer is reopened") + } + let restart = stream == nil || displayIndex != index || screensChanged + let width = Self.streamWidth(display: display, maxWidth: maxWidth) + let resize = !restart && width != streamWidth + condition.unlock() + + if restart { + if let failure = start(index: index, maxWidth: maxWidth) { return .failed(failure) } + } else if resize { + resizeStream(maxWidth: maxWidth) + } + + condition.lock() + let hadFrame = latest != nil + let wait = hadFrame ? TimeInterval(waitMs) / 1000 : max(TimeInterval(waitMs) / 1000, Self.firstFrameWait) + let deadline = Date().addingTimeInterval(wait) + while stream != nil, latest == nil || after.map({ seq <= $0 }) == true { + if !condition.wait(until: deadline) { break } + } + guard let buffer = latest, let frame = display?.frame else { + let stopped = userStopped + condition.unlock() + return .failed(stopped + ? "error: screen sharing was stopped from this Mac's menu bar" + : "error: the live capture delivered no frame — check Screen Recording permission") + } + let current = seq + if let after, current <= after { + condition.unlock() + return .unchanged(seq: current) + } + if let cached = encoded, cached.seq == current, cached.encoding == encoding { + condition.unlock() + return .frame(cached.shot, seq: current) + } + condition.unlock() + + // Encode outside the lock so new frames keep landing meanwhile. + var image: CGImage? + VTCreateCGImageFromCVPixelBuffer(buffer, options: nil, imageOut: &image) + guard let image, let data = encoding.encode(image) else { + return .failed("error: could not encode the live frame") + } + let shot = CaptureShot( + data: data, mimeType: encoding.mimeType, frame: frame, + pixelWidth: image.width, pixelHeight: image.height, title: nil) + condition.lock() + encoded = (current, encoding, shot) + condition.unlock() + return .frame(shot, seq: current) + } + + // MARK: Stream lifecycle + + /// Starts a stream on display `index`, replacing any running one. Returns an + /// error message on failure. + private func start(index: Int, maxWidth: Int) -> String? { + stop() + let semaphore = DispatchSemaphore(value: 0) + let found = DisplayBox() + let lookup = Task.detached { + defer { semaphore.signal() } + found.displays = try? await DisplayCache.shared.current(fresh: true) + } + if semaphore.wait(timeout: .now() + 10) == .timedOut { + lookup.cancel() + return "error: could not enumerate displays for the live capture" + } + guard let displays = found.displays, index >= 0, index < displays.count else { + return "error: display \(index) does not exist — call list_displays" + } + let display = displays[index] + let config = Self.configuration(display: display, maxWidth: maxWidth) + let stream = SCStream( + filter: SCContentFilter(display: display, excludingWindows: []), + configuration: config, delegate: self) + do { + try stream.addStreamOutput(self, type: .screen, sampleHandlerQueue: sampleQueue) + } catch { + return "error: could not start the live capture: \(error.localizedDescription)" + } + + // Publish the stream before it starts so its first frame is not dropped + // as belonging to some other stream. + condition.lock() + self.stream = stream + self.display = display + displayIndex = index + streamWidth = config.width + screensChanged = false + latest = nil + encoded = nil + condition.unlock() + + let started = DispatchSemaphore(value: 0) + let failure = ErrorBox() + stream.startCapture { error in + failure.error = error + started.signal() + } + let timedOut = started.wait(timeout: .now() + 10) == .timedOut + if timedOut || failure.error != nil { + stop() + let reason = failure.error?.localizedDescription ?? "it did not start within 10 s" + return "error: could not start the live capture (\(reason)) — check Screen Recording permission" + } + startReaper() + fputs("munim-computer-use: live capture started on display \(index) at \(config.width)×\(config.height)\n", stderr) + return nil + } + + private func resizeStream(maxWidth: Int) { + condition.lock() + guard let stream, let display else { + condition.unlock() + return + } + let config = Self.configuration(display: display, maxWidth: maxWidth) + streamWidth = config.width + condition.unlock() + // Frames at the old size keep coming until this lands; both map back to + // screen points through the display frame, so either is correct. + stream.updateConfiguration(config) { _ in } + } + + /// Stops the stream, if any, without waiting for ScreenCaptureKit to finish. + private func stop(reason: String? = nil) { + condition.lock() + let running = stream + stream = nil + latest = nil + encoded = nil + condition.broadcast() + condition.unlock() + guard let running else { return } + running.stopCapture { _ in } + if let reason { fputs("munim-computer-use: live capture stopped (\(reason))\n", stderr) } + } + + private func startReaper() { + condition.lock() + defer { condition.unlock() } + guard reaper == nil else { return } + let timer = DispatchSource.makeTimerSource(queue: sampleQueue) + timer.schedule(deadline: .now() + 1, repeating: 1) + timer.setEventHandler { [weak self] in + guard let self else { return } + self.condition.lock() + let idle = self.stream != nil && -self.lastRequest.timeIntervalSinceNow > Self.idleStop + self.condition.unlock() + if idle { self.stop(reason: "no viewer for \(Int(Self.idleStop)) s") } + } + timer.resume() + reaper = timer + } + + private static func streamWidth(display: SCDisplay?, maxWidth: Int) -> Int { + guard let display else { return 0 } + let scale = maxWidth > 0 ? min(1.0, Double(maxWidth) / max(1.0, Double(display.width))) : 1.0 + return max(1, Int((Double(display.width) * scale).rounded())) + } + + /// Same size rules as a one-shot display capture, so a viewer gets the same + /// pixels-per-point either way. + private static func configuration(display: SCDisplay, maxWidth: Int) -> SCStreamConfiguration { + let config = SCStreamConfiguration() + let width = streamWidth(display: display, maxWidth: maxWidth) + config.width = width + config.height = max(1, Int((Double(display.height) * Double(width) / max(1.0, Double(display.width))).rounded())) + config.pixelFormat = kCVPixelFormatType_32BGRA + config.minimumFrameInterval = CMTime(value: 1, timescale: maxFramesPerSecond) + config.queueDepth = 5 + config.showsCursor = false + return config + } + + // MARK: SCStreamOutput / SCStreamDelegate + + func stream(_ stream: SCStream, didOutputSampleBuffer sampleBuffer: CMSampleBuffer, of type: SCStreamOutputType) { + // Idle frames repeat the previous image; only complete ones are news. + guard type == .screen, sampleBuffer.isValid, + let attachments = CMSampleBufferGetSampleAttachmentsArray( + sampleBuffer, createIfNecessary: false) as? [[SCStreamFrameInfo: Any]], + let rawStatus = attachments.first?[.status] as? Int, + SCFrameStatus(rawValue: rawStatus) == .complete, + let buffer = sampleBuffer.imageBuffer + else { return } + condition.lock() + if stream === self.stream { + latest = buffer + seq += 1 + condition.broadcast() + } + condition.unlock() + } + + func stream(_ stream: SCStream, didStopWithError error: Error) { + condition.lock() + if stream === self.stream { + self.stream = nil + latest = nil + encoded = nil + // The menu-bar indicator's Stop, or the system's own privacy stop. + if (error as NSError).code == SCStreamError.Code.userStopped.rawValue { + userStopped = true + } + fputs("munim-computer-use: live capture stopped (\(error.localizedDescription))\n", stderr) + condition.broadcast() + } + condition.unlock() + } +} + +private final class DisplayBox: @unchecked Sendable { + var displays: [SCDisplay]? +} + +private final class ErrorBox: @unchecked Sendable { + var error: Error? +} diff --git a/macos/Sources/main.swift b/macos/Sources/main.swift index 52c0530..0a47cc4 100644 --- a/macos/Sources/main.swift +++ b/macos/Sources/main.swift @@ -1316,11 +1316,57 @@ func toolActivateApp(_ args: [String: Any]) -> String { // MARK: - Screen capture -/// Synchronizes capture results so a timed-out waiter never races a late write. +/// How a capture is encoded. PNG by default: agents read UI text, and PNG keeps +/// it sharp. JPEG is for live remote viewing, where every frame is replaced a +/// moment later and bytes matter more than exact pixels. Mirrors the +/// Windows/Linux server's `format` and `quality`. +enum ImageEncoding: Equatable { + case png + /// JPEG at this quality (1-100). + case jpeg(Int) + + /// JPEG quality when the caller does not ask for one, as on Windows/Linux. + static let defaultJPEGQuality = 55 + + static func parse(_ args: [String: Any]) throws -> ImageEncoding { + let format = ((args["format"] as? String) ?? "png").lowercased() + switch format { + case "png": + return .png + case "jpeg", "jpg": + let quality = (args["quality"] as? Int).map { min(max($0, 1), 100) } ?? defaultJPEGQuality + return .jpeg(quality) + default: + throw "error: unsupported screenshot format '\(format)' — use png or jpeg" + } + } + + var mimeType: String { + switch self { + case .png: return "image/png" + case .jpeg: return "image/jpeg" + } + } + + func encode(_ image: CGImage) -> Data? { + let data = NSMutableData() + let type = (self == .png ? "public.png" : "public.jpeg") as CFString + guard let destination = CGImageDestinationCreateWithData(data, type, 1, nil) else { return nil } + var options: [CFString: Any] = [:] + if case .jpeg(let quality) = self { + options[kCGImageDestinationLossyCompressionQuality] = Double(quality) / 100 + } + CGImageDestinationAddImage(destination, image, options as CFDictionary) + guard CGImageDestinationFinalize(destination) else { return nil } + return data as Data + } +} + /// One capture plus the geometry a model needs to turn image pixels back into /// the screen coordinates that click/hover/zoom accept. struct CaptureShot { let data: Data + let mimeType: String /// Captured area in global screen points (the coordinate space of click x/y). let frame: CGRect let pixelWidth: Int @@ -1328,6 +1374,7 @@ struct CaptureShot { let title: String? } +/// Synchronizes capture results so a timed-out waiter never races a late write. private final class CaptureBox: @unchecked Sendable { private let lock = NSLock() private var value: CaptureShot? @@ -1343,6 +1390,47 @@ private final class CaptureBox: @unchecked Sendable { } } +/// The attached displays, kept between captures. Asking ScreenCaptureKit for +/// shareable content walks every window on screen, which costs more than the +/// display capture itself when a viewer asks for frames back to back. Displays +/// only change when one is attached, removed or rearranged, which AppKit +/// announces, so the list is dropped then (and after a short lifetime, in case +/// a notification is missed). Window captures still query fresh: windows open, +/// close and move all the time, and a stale frame would put clicks in the +/// wrong place. +final class DisplayCache: @unchecked Sendable { + static let shared = DisplayCache() + private static let lifetime: TimeInterval = 30 + private let lock = NSLock() + private var displays: [SCDisplay]? + private var fetchedAt = Date.distantPast + + private init() { + NotificationCenter.default.addObserver( + forName: NSApplication.didChangeScreenParametersNotification, object: nil, queue: nil + ) { [weak self] _ in self?.invalidate() } + } + + func invalidate() { + lock.withLock { displays = nil } + } + + func current(fresh: Bool = false) async throws -> [SCDisplay] { + if !fresh, let cached = lock.withLock({ + -fetchedAt.timeIntervalSinceNow < Self.lifetime ? displays : nil + }) { + return cached + } + let content = try await SCShareableContent.excludingDesktopWindows( + false, onScreenWindowsOnly: true) + lock.withLock { + displays = content.displays + fetchedAt = Date() + } + return content.displays + } +} + /// The text block that rides with every screenshot/zoom image. Models read the /// image in pixels but every pointer tool takes screen points; without this /// line a coordinate click from a downscaled or Retina capture lands off-target. @@ -1362,7 +1450,7 @@ func captureMappingText(_ shot: CaptureShot, label: String) -> String { func imageResult(_ shot: CaptureShot, label: String) -> [String: Any] { return [ "content": [ - ["type": "image", "data": shot.data.base64EncodedString(), "mimeType": "image/png"], + ["type": "image", "data": shot.data.base64EncodedString(), "mimeType": shot.mimeType], ["type": "text", "text": captureMappingText(shot, label: label)], ], "isError": false, @@ -1378,10 +1466,10 @@ func backingScale(for display: SCDisplay) -> CGFloat { }?.backingScaleFactor ?? 2 } -/// Capture a window as PNG. Runs the async ScreenCaptureKit call on a background +/// Capture a window. Runs the async ScreenCaptureKit call on a background /// executor and blocks the JSON-RPC loop until it lands, with a timeout so a /// wedged capture can never hang the server. -func captureWindowPNG(pid: pid_t, maxWidth: Int) -> CaptureShot? { +func captureWindow(pid: pid_t, maxWidth: Int, encoding: ImageEncoding) -> CaptureShot? { let semaphore = DispatchSemaphore(value: 0) let box = CaptureBox() @@ -1411,11 +1499,10 @@ func captureWindowPNG(pid: pid_t, maxWidth: Int) -> CaptureShot? { let image = try await SCScreenshotManager.captureImage( contentFilter: SCContentFilter(desktopIndependentWindow: window), configuration: config) - guard let png = NSBitmapImageRep(cgImage: image).representation(using: .png, properties: [:]) - else { return } + guard let data = encoding.encode(image) else { return } box.set(CaptureShot( - data: png, frame: window.frame, pixelWidth: image.width, pixelHeight: image.height, - title: window.title)) + data: data, mimeType: encoding.mimeType, frame: window.frame, + pixelWidth: image.width, pixelHeight: image.height, title: window.title)) } catch { box.set(nil) } @@ -1432,96 +1519,94 @@ func captureWindowPNG(pid: pid_t, maxWidth: Int) -> CaptureShot? { /// Capture a whole display. Window capture covers one app; this is for seeing /// the desktop as a whole, including every monitor the user has attached. -func captureDisplayPNG(index: Int, maxWidth: Int) -> CaptureShot? { +func captureDisplay(index: Int, maxWidth: Int, encoding: ImageEncoding) -> CaptureShot? { let semaphore = DispatchSemaphore(value: 0) - let lock = NSLock() - var result: CaptureShot? - Task.detached { + let box = CaptureBox() + let task = Task.detached { defer { semaphore.signal() } - do { - let content = try await SCShareableContent.excludingDesktopWindows( - false, onScreenWindowsOnly: true) - let displays = content.displays - guard index >= 0, index < displays.count else { return } - let display = displays[index] - let config = SCStreamConfiguration() - let scale: Double - if maxWidth > 0 { - scale = min(1.0, Double(maxWidth) / max(1.0, Double(display.width))) - } else { - scale = 1.0 - } - config.width = Int(Double(display.width) * scale) - config.height = Int(Double(display.height) * scale) - config.showsCursor = false - let image = try await SCScreenshotManager.captureImage( - contentFilter: SCContentFilter(display: display, excludingWindows: []), - configuration: config) - if let png = NSBitmapImageRep(cgImage: image).representation(using: .png, properties: [:]) { - lock.lock() - result = CaptureShot( - data: png, frame: display.frame, pixelWidth: image.width, pixelHeight: image.height, - title: nil) - lock.unlock() + // The cached display list first; if that display is gone or the + // capture fails, once more with a fresh list. + for fresh in [false, true] { + do { + let displays = try await DisplayCache.shared.current(fresh: fresh) + guard index >= 0, index < displays.count else { continue } + let display = displays[index] + let config = SCStreamConfiguration() + let scale: Double + if maxWidth > 0 { + scale = min(1.0, Double(maxWidth) / max(1.0, Double(display.width))) + } else { + scale = 1.0 + } + config.width = Int(Double(display.width) * scale) + config.height = Int(Double(display.height) * scale) + config.showsCursor = false + let image = try await SCScreenshotManager.captureImage( + contentFilter: SCContentFilter(display: display, excludingWindows: []), + configuration: config) + guard let data = encoding.encode(image) else { return } + box.set(CaptureShot( + data: data, mimeType: encoding.mimeType, frame: display.frame, + pixelWidth: image.width, pixelHeight: image.height, title: nil)) + return + } catch { + continue } - } catch { - lock.lock() - result = nil - lock.unlock() } } - // On timeout the task may still write `result` — do not read it. if semaphore.wait(timeout: .now() + 20) == .timedOut { + task.cancel() + // Do not read `box` after cancel — the task may still be writing. return nil } - lock.lock() - defer { lock.unlock() } - return result + return box.get() } /// Capture one region of the screen at full physical resolution. `rect` is in /// global screen points, the same space click/hover take, so a model can zoom /// straight from the coordinates it already knows. -func captureRegionPNG(rect: CGRect, maxWidth: Int) -> CaptureShot? { +func captureRegion(rect: CGRect, maxWidth: Int, encoding: ImageEncoding) -> CaptureShot? { let semaphore = DispatchSemaphore(value: 0) let box = CaptureBox() let task = Task.detached { defer { semaphore.signal() } - do { - let content = try await SCShareableContent.excludingDesktopWindows( - false, onScreenWindowsOnly: true) - let center = CGPoint(x: rect.midX, y: rect.midY) - guard let display = content.displays.first(where: { $0.frame.contains(center) }) - ?? content.displays.first - else { return } - let clipped = rect.intersection(display.frame) - guard clipped.width >= 4, clipped.height >= 4 else { return } - let scale = backingScale(for: display) - var width = clipped.width * scale - var height = clipped.height * scale - if maxWidth > 0, width > CGFloat(maxWidth) { - let ratio = CGFloat(maxWidth) / width - width *= ratio - height *= ratio + for fresh in [false, true] { + do { + let displays = try await DisplayCache.shared.current(fresh: fresh) + let center = CGPoint(x: rect.midX, y: rect.midY) + guard let display = displays.first(where: { $0.frame.contains(center) }) + ?? displays.first + else { continue } + let clipped = rect.intersection(display.frame) + guard clipped.width >= 4, clipped.height >= 4 else { continue } + let scale = backingScale(for: display) + var width = clipped.width * scale + var height = clipped.height * scale + if maxWidth > 0, width > CGFloat(maxWidth) { + let ratio = CGFloat(maxWidth) / width + width *= ratio + height *= ratio + } + let config = SCStreamConfiguration() + config.sourceRect = CGRect( + x: clipped.origin.x - display.frame.origin.x, + y: clipped.origin.y - display.frame.origin.y, + width: clipped.width, height: clipped.height) + config.width = max(1, Int(width.rounded())) + config.height = max(1, Int(height.rounded())) + config.showsCursor = false + if #available(macOS 14.0, *) { config.captureResolution = .best } + let image = try await SCScreenshotManager.captureImage( + contentFilter: SCContentFilter(display: display, excludingWindows: []), + configuration: config) + guard let data = encoding.encode(image) else { return } + box.set(CaptureShot( + data: data, mimeType: encoding.mimeType, frame: clipped, + pixelWidth: image.width, pixelHeight: image.height, title: nil)) + return + } catch { + continue } - let config = SCStreamConfiguration() - config.sourceRect = CGRect( - x: clipped.origin.x - display.frame.origin.x, - y: clipped.origin.y - display.frame.origin.y, - width: clipped.width, height: clipped.height) - config.width = max(1, Int(width.rounded())) - config.height = max(1, Int(height.rounded())) - config.showsCursor = false - if #available(macOS 14.0, *) { config.captureResolution = .best } - let image = try await SCScreenshotManager.captureImage( - contentFilter: SCContentFilter(display: display, excludingWindows: []), - configuration: config) - guard let png = NSBitmapImageRep(cgImage: image).representation(using: .png, properties: [:]) - else { return } - box.set(CaptureShot( - data: png, frame: clipped, pixelWidth: image.width, pixelHeight: image.height, title: nil)) - } catch { - box.set(nil) } } if semaphore.wait(timeout: .now() + 15) == .timedOut { @@ -1538,17 +1623,16 @@ func toolListDisplays(_ args: [String: Any]) -> String { let task = Task.detached { defer { semaphore.signal() } var collected: [String] = [] - if let content = try? await SCShareableContent.excludingDesktopWindows( - false, onScreenWindowsOnly: true) { - for (i, display) in content.displays.enumerated() { + // Fresh, and it refreshes the cache the captures read: a viewer lists + // displays once, then asks for frames by index. + if let displays = try? await DisplayCache.shared.current(fresh: true) { + for (i, display) in displays.enumerated() { let frame = display.frame collected.append("[\(i)] \(display.width)x\(display.height) " + "at (\(Int(frame.origin.x)), \(Int(frame.origin.y)))") } } - lock.lock() - lines = collected - lock.unlock() + lock.withLock { lines = collected } } if semaphore.wait(timeout: .now() + 20) == .timedOut { task.cancel() @@ -3032,6 +3116,17 @@ let toolDefs: [[String: Any]] = [ "type": "integer", "description": "Downscale the image to this width in pixels (default 1400). Lower it to save tokens.", ], + "format": [ + "type": "string", + "enum": ["png", "jpeg"], + "description": "Image encoding (default png). Use jpeg for live remote viewing.", + ], + "quality": [ + "type": "integer", + "minimum": 1, + "maximum": 100, + "description": "JPEG quality from 1 to 100 (default 55). Ignored for png. Raise it when small text must stay sharp.", + ], ], ], "annotations": [ @@ -4169,7 +4264,12 @@ while let line = readLine(strippingNewline: true) { break } let maxWidth = clampedMaxWidth(args) - guard let shot = captureRegionPNG(rect: rect, maxWidth: maxWidth) else { + let encoding: ImageEncoding + do { encoding = try ImageEncoding.parse(args) } catch { + respond(id: id, result: textResult("\(error)", isError: true)) + break + } + guard let shot = captureRegion(rect: rect, maxWidth: maxWidth, encoding: encoding) else { respond(id: id, result: textResult( "error: could not capture that region — check Screen Recording permission and " + "that the region lies on an attached display (list_displays).", isError: true)) @@ -4179,9 +4279,19 @@ while let line = readLine(strippingNewline: true) { break } if name == "screenshot" { + let encoding: ImageEncoding + do { encoding = try ImageEncoding.parse(args) } catch { + respond(id: id, result: textResult("\(error)", isError: true)) + break + } if let display = args["display"] as? Int { let maxWidth = clampedMaxWidth(args) - guard let shot = captureDisplayPNG(index: display, maxWidth: maxWidth) else { + if (args["live"] as? Bool) == true { + respond(id: id, result: liveScreenshotResult( + args, display: display, maxWidth: maxWidth, encoding: encoding)) + break + } + guard let shot = captureDisplay(index: display, maxWidth: maxWidth, encoding: encoding) else { respond(id: id, result: textResult( "error: could not capture display \(display) — check Screen Recording " + "permission, or call list_displays for valid indices.", isError: true)) @@ -4201,7 +4311,9 @@ while let line = readLine(strippingNewline: true) { break } let maxWidth = clampedMaxWidth(args) - guard let shot = captureWindowPNG(pid: resolved.app.processIdentifier, maxWidth: maxWidth) else { + guard let shot = captureWindow( + pid: resolved.app.processIdentifier, maxWidth: maxWidth, encoding: encoding) + else { respond(id: id, result: textResult( "error: screen capture failed. The host app may be missing Screen Recording " + "permission, or this app may have no on-screen window.", diff --git a/scripts/check-tool-parity.mjs b/scripts/check-tool-parity.mjs index cb5182a..47eecf0 100644 --- a/scripts/check-tool-parity.mjs +++ b/scripts/check-tool-parity.mjs @@ -11,8 +11,6 @@ // Deliberate differences, and the only ones allowed: // * get_app_state.window — macOS only (scope to one window or the agent's // Chrome window). -// * screenshot.format, screenshot.quality — Windows/Linux only (jpeg -// encoding and its quality). // // Usage: node scripts/check-tool-parity.mjs (exit 1 on any mismatch) import { readFileSync } from "node:fs"; @@ -25,10 +23,7 @@ const rustPath = path.join(root, "windows-linux/src/tools.rs"); const ALLOWED_EXTRAS = { swift: [["get_app_state", "window"]], - rust: [ - ["screenshot", "format"], - ["screenshot", "quality"], - ], + rust: [], }; // ---------------------------------------------------------------- tokenizer @@ -299,5 +294,5 @@ if (problems.length) { for (const problem of problems) console.error(` - ${problem}`); process.exit(1); } -console.log(`Tool parity OK: ${swiftNames.length} tools, identical in order, text and schema (allowed extras: get_app_state.window on macOS, screenshot.format and screenshot.quality on Windows/Linux).`); +console.log(`Tool parity OK: ${swiftNames.length} tools, identical in order, text and schema (allowed extras: get_app_state.window on macOS).`); } diff --git a/windows-linux/Cargo.toml b/windows-linux/Cargo.toml index 0c255fa..575cb5b 100644 --- a/windows-linux/Cargo.toml +++ b/windows-linux/Cargo.toml @@ -33,6 +33,12 @@ windows = { version = "0.62.2", features = [ "Win32_System_Threading", "Win32_System_LibraryLoader", "Win32_Graphics_Gdi", + # Desktop Duplication for the live view (live.rs). + "Win32_Graphics_Direct3D", + "Win32_Graphics_Direct3D11", + "Win32_Graphics_Dxgi", + "Win32_Graphics_Dxgi_Common", + "Win32_UI_HiDpi", ] } [target.'cfg(any(windows, target_os = "linux"))'.dependencies] diff --git a/windows-linux/src/capture.rs b/windows-linux/src/capture.rs index 8c7ada1..657683c 100644 --- a/windows-linux/src/capture.rs +++ b/windows-linux/src/capture.rs @@ -149,6 +149,7 @@ pub struct CaptureFrame { pub height: f64, } +#[derive(Clone)] pub struct Capture { pub bytes: Vec, pub frame: CaptureFrame, @@ -172,7 +173,7 @@ pub fn mapping_text(capture: &Capture, label: &str) -> String { ) } -fn finish(image: RgbaImage, frame: CaptureFrame, max_width: u32, format: CaptureFormat) -> Result { +pub(crate) fn finish(image: RgbaImage, frame: CaptureFrame, max_width: u32, format: CaptureFormat) -> Result { let image = if max_width > 0 && image.width() > max_width { let height = ((image.height() as f64) * (max_width as f64) / (image.width() as f64)) .round() diff --git a/windows-linux/src/live.rs b/windows-linux/src/live.rs new file mode 100644 index 0000000..6426be6 --- /dev/null +++ b/windows-linux/src/live.rs @@ -0,0 +1,441 @@ +#![cfg_attr(not(windows), allow(dead_code))] +//! Live display capture for MT Code's Computer View. +//! +//! A one-shot screenshot enumerates monitors, copies the whole screen and +//! encodes it, every frame, so a viewer polling it pays that cost even while +//! nothing moves and only learns that nothing moved by comparing bytes. On +//! Windows, Desktop Duplication keeps one duplication of the display open and +//! hands over a new image only when something on it changed, so the viewer +//! waits for the next change instead. +//! +//! Requested as `screenshot { display, live: true, after?, wait_ms? }`. +//! Deliberately not in the tool schema: it serves that viewer, not agents, the +//! same as `cursor: true`. Linux has no live path yet (the Wayland capture +//! portal asks the user every time), so there the caller falls back to an +//! ordinary screenshot. + +use crate::capture::{Capture, CaptureFormat}; +#[cfg(windows)] +use crate::platform::DesktopError; + +pub struct LiveRequest { + pub display: usize, + pub max_width: u32, + pub format: CaptureFormat, + /// The sequence number of the frame the viewer already has. When set, the + /// call waits up to `wait_ms` for a newer one. + pub after: Option, + pub wait_ms: u32, + /// The viewer draws the pointer itself (`cursor: true`), so a pointer that + /// moved over a still screen ends the wait early. + pub pointer: bool, +} + +pub enum LiveFrame { + Changed { capture: Capture, seq: u64 }, + /// Nothing new within the wait: the viewer keeps its frame. Also returned + /// early when only the pointer moved, so the viewer can redraw it. + Unchanged { seq: u64 }, +} + +pub enum LiveError { + /// No live path for this display; the caller takes an ordinary screenshot. + Unsupported, + #[cfg_attr(not(windows), allow(dead_code))] + Failed(String), +} + +/// `wait_ms` when the caller gives none, and the most it may ask for. +pub const DEFAULT_WAIT_MS: u32 = 1000; +pub const MAX_WAIT_MS: u32 = 5000; + +/// The text line a live result carries, which the viewer parses. +pub fn live_line(seq: u64, changed: bool) -> String { + format!("live: {{\"seq\":{seq},\"changed\":{changed}}}") +} + +pub fn frame(request: &LiveRequest) -> Result { + #[cfg(windows)] + { + duplication::frame(request) + } + #[cfg(not(windows))] + { + let _ = request; + Err(LiveError::Unsupported) + } +} + +#[cfg(windows)] +impl From for LiveError { + fn from(error: DesktopError) -> Self { + LiveError::Failed(error.to_string()) + } +} + +#[cfg(windows)] +mod duplication { + use std::cell::RefCell; + use std::time::{Duration, Instant}; + + use image::RgbaImage; + use windows::Win32::Foundation::{E_ACCESSDENIED, HMODULE}; + use windows::Win32::Graphics::Direct3D::D3D_DRIVER_TYPE_UNKNOWN; + use windows::Win32::Graphics::Direct3D11::{ + D3D11_CPU_ACCESS_READ, D3D11_CREATE_DEVICE_BGRA_SUPPORT, D3D11_MAP_READ, + D3D11_MAPPED_SUBRESOURCE, D3D11_SDK_VERSION, D3D11_TEXTURE2D_DESC, D3D11_USAGE_STAGING, + D3D11CreateDevice, ID3D11Device, ID3D11DeviceContext, ID3D11Texture2D, + }; + use windows::Win32::Graphics::Dxgi::Common::{ + DXGI_FORMAT_B8G8R8A8_UNORM, DXGI_MODE_ROTATION_IDENTITY, DXGI_MODE_ROTATION_UNSPECIFIED, + }; + use windows::Win32::Graphics::Dxgi::{ + CreateDXGIFactory1, DXGI_ERROR_ACCESS_LOST, DXGI_ERROR_WAIT_TIMEOUT, + DXGI_OUTDUPL_FRAME_INFO, IDXGIFactory1, IDXGIOutput1, IDXGIOutputDuplication, + IDXGIResource, + }; + use windows::Win32::UI::HiDpi::{ + DPI_AWARENESS_CONTEXT_PER_MONITOR_AWARE_V2, SetThreadDpiAwarenessContext, + }; + use windows::core::Interface; + use xcap::Monitor; + + use super::{LiveError, LiveFrame, LiveRequest}; + use crate::capture::{self, Capture, CaptureFormat, CaptureFrame}; + + /// A duplication nobody asked for in this long is closed and reopened on + /// the next request. + const IDLE_RELEASE: Duration = Duration::from_secs(5); + /// How long a fresh duplication gets to produce its first image before the + /// display is seeded with an ordinary capture. A still desktop may produce + /// nothing until something on it changes. + const FIRST_IMAGE_WAIT: Duration = Duration::from_millis(250); + /// After a display turns out to have no live path, how long to stop trying. + const UNSUPPORTED_RETRY: Duration = Duration::from_secs(10); + + struct Session { + display: usize, + frame: CaptureFrame, + device: ID3D11Device, + context: ID3D11DeviceContext, + duplication: IDXGIOutputDuplication, + staging: Option<(ID3D11Texture2D, u32, u32)>, + } + + enum Acquired { + Image, + Pointer, + Nothing, + } + + #[derive(Default)] + struct State { + session: Option, + last_request: Option, + latest: Option, + /// Never reset, so a frame number from before a reopen is always older. + seq: u64, + encoded: Option<(u64, u32, CaptureFormat, Capture)>, + unsupported: Option<(usize, Instant)>, + } + + thread_local! { + // Requests are served one at a time on the stdio thread. + static STATE: RefCell = RefCell::new(State::default()); + } + + pub fn frame(request: &LiveRequest) -> Result { + STATE.with(|state| state.borrow_mut().serve(request)) + } + + impl State { + fn serve(&mut self, request: &LiveRequest) -> Result { + if self + .unsupported + .is_some_and(|(display, at)| display == request.display && at.elapsed() < UNSUPPORTED_RETRY) + { + return Err(LiveError::Unsupported); + } + let idle = self.last_request.is_some_and(|at| at.elapsed() > IDLE_RELEASE); + self.last_request = Some(Instant::now()); + if idle || self.session.as_ref().is_some_and(|session| session.display != request.display) { + self.close(); + } + if self.session.is_none() { + match Session::open(request.display) { + Ok(session) => self.session = Some(session), + Err(LiveError::Unsupported) => { + self.unsupported = Some((request.display, Instant::now())); + return Err(LiveError::Unsupported); + } + Err(error) => return Err(error), + } + if !matches!(self.acquire(FIRST_IMAGE_WAIT)?, Acquired::Image) { + self.seed()?; + } + } + + let deadline = Instant::now() + Duration::from_millis(u64::from(request.wait_ms)); + let wants_newer = + |state: &Self| state.latest.is_none() || request.after.is_some_and(|after| state.seq <= after); + // First pass takes whatever is already pending without waiting, so + // even a call that does not wait returns the newest image. + let mut timeout = Duration::ZERO; + loop { + let acquired = self.acquire(timeout)?; + if matches!(acquired, Acquired::Pointer) + && request.pointer + && request.after.is_some() + && self.latest.is_some() + { + break; + } + if !wants_newer(self) { + break; + } + timeout = deadline.saturating_duration_since(Instant::now()); + if timeout.is_zero() { + break; + } + } + + if request.after.is_some_and(|after| self.seq <= after) { + return Ok(LiveFrame::Unchanged { seq: self.seq }); + } + let seq = self.seq; + if let Some((cached_seq, max_width, format, capture)) = &self.encoded + && *cached_seq == seq + && *max_width == request.max_width + && *format == request.format + { + return Ok(LiveFrame::Changed { capture: capture.clone(), seq }); + } + let (Some(image), Some(session)) = (self.latest.clone(), self.session.as_ref()) else { + return Err(LiveError::Failed("the live capture has no image yet".into())); + }; + let capture = capture::finish(image, session.frame, request.max_width, request.format)?; + self.encoded = Some((seq, request.max_width, request.format, capture.clone())); + Ok(LiveFrame::Changed { capture, seq }) + } + + fn close(&mut self) { + self.session = None; + self.latest = None; + self.encoded = None; + } + + /// A still desktop may give a fresh duplication nothing to hand over, so + /// the first image comes from an ordinary capture instead. + fn seed(&mut self) -> Result<(), LiveError> { + let Some(session) = self.session.as_ref() else { return Ok(()) }; + let monitors = Monitor::all().map_err(|error| LiveError::Failed(error.to_string()))?; + let monitor = monitors + .get(session.display) + .ok_or_else(|| LiveError::Failed("the display went away".into()))?; + let image = monitor + .capture_image() + .map_err(|error| LiveError::Failed(format!("failed to capture display: {error}")))?; + self.latest = Some(image); + self.seq += 1; + Ok(()) + } + + fn acquire(&mut self, timeout: Duration) -> Result { + let Some(session) = self.session.as_mut() else { return Ok(Acquired::Nothing) }; + let mut info = DXGI_OUTDUPL_FRAME_INFO::default(); + let mut resource: Option = None; + let millis = u32::try_from(timeout.as_millis()).unwrap_or(u32::MAX); + let acquired = unsafe { session.duplication.AcquireNextFrame(millis, &mut info, &mut resource) }; + if let Err(error) = acquired { + if error.code() == DXGI_ERROR_WAIT_TIMEOUT { + return Ok(Acquired::Nothing); + } + if error.code() == DXGI_ERROR_ACCESS_LOST { + // The desktop switched (lock screen, UAC prompt), the mode + // changed, or a full-screen app took over: reopen, which + // hands over a full image again. + let display = session.display; + self.session = None; + self.session = Some(Session::open(display)?); + return Ok(Acquired::Nothing); + } + return Err(LiveError::Failed(format!("the live capture failed: {error}"))); + } + let image = if info.LastPresentTime != 0 { + resource.map(|resource| session.copy(&resource)).transpose().map(Option::flatten) + } else { + Ok(None) + }; + unsafe { + let _ = session.duplication.ReleaseFrame(); + } + match image { + Ok(Some(image)) => { + self.latest = Some(image); + self.seq += 1; + Ok(Acquired::Image) + } + Ok(None) if info.LastMouseUpdateTime != 0 => Ok(Acquired::Pointer), + Ok(None) => Ok(Acquired::Nothing), + Err(error) => Err(LiveError::Failed(format!("could not read the live frame: {error}"))), + } + } + } + + impl Session { + fn open(display: usize) -> Result { + let monitors = Monitor::all().map_err(|error| LiveError::Failed(error.to_string()))?; + let monitor = monitors.get(display).ok_or_else(|| { + LiveError::Failed(format!( + "display {display} does not exist — call list_displays ({} attached)", + monitors.len() + )) + })?; + // The same geometry an ordinary screenshot reports, so both map + // image pixels to the same screen coordinates. + let frame = CaptureFrame { + x: f64::from(monitor.x().unwrap_or(0)), + y: f64::from(monitor.y().unwrap_or(0)), + width: f64::from(monitor.width().unwrap_or(0)), + height: f64::from(monitor.height().unwrap_or(0)), + }; + let handle = monitor.id().map_err(|error| LiveError::Failed(error.to_string()))?; + + unsafe { + let factory: IDXGIFactory1 = CreateDXGIFactory1().map_err(|_| LiveError::Unsupported)?; + // Duplication must run on the adapter that drives the display, + // which on a laptop with two GPUs is not always the default. + let mut adapter_index = 0; + while let Ok(adapter) = factory.EnumAdapters1(adapter_index) { + adapter_index += 1; + let mut output_index = 0; + while let Ok(output) = adapter.EnumOutputs(output_index) { + output_index += 1; + let Ok(desc) = output.GetDesc() else { continue }; + if desc.Monitor.0 as usize as u32 != handle { + continue; + } + let mut device = None; + let mut context = None; + D3D11CreateDevice( + &adapter, + D3D_DRIVER_TYPE_UNKNOWN, + HMODULE::default(), + D3D11_CREATE_DEVICE_BGRA_SUPPORT, + None, + D3D11_SDK_VERSION, + Some(&mut device), + None, + Some(&mut context), + ) + .map_err(|_| LiveError::Unsupported)?; + let (Some(device), Some(context)) = (device, context) else { + return Err(LiveError::Unsupported); + }; + let output1: IDXGIOutput1 = output.cast().map_err(|_| LiveError::Unsupported)?; + // Duplication refuses processes that are not DPI aware + // on some builds; scope the awareness to this call. + let previous = SetThreadDpiAwarenessContext(DPI_AWARENESS_CONTEXT_PER_MONITOR_AWARE_V2); + let duplicated = output1.DuplicateOutput(&device); + if !previous.is_invalid() { + SetThreadDpiAwarenessContext(previous); + } + let duplication = match duplicated { + Ok(duplication) => duplication, + Err(error) if error.code() == E_ACCESSDENIED => { + return Err(LiveError::Failed( + "the screen is locked or showing a secure prompt (UAC); the live view \ + resumes once it closes" + .into(), + )); + } + // Unsupported adapter, too many duplications of + // this display already, a remote session... + Err(_) => return Err(LiveError::Unsupported), + }; + let duplication_desc = duplication.GetDesc(); + let rotation = duplication_desc.Rotation; + if (rotation != DXGI_MODE_ROTATION_IDENTITY && rotation != DXGI_MODE_ROTATION_UNSPECIFIED) + || duplication_desc.ModeDesc.Format != DXGI_FORMAT_B8G8R8A8_UNORM + { + // A rotated or HDR desktop arrives in a layout the + // copy below does not convert. + return Err(LiveError::Unsupported); + } + return Ok(Session { display, frame, device, context, duplication, staging: None }); + } + } + } + Err(LiveError::Unsupported) + } + + /// Copies the duplicated desktop into CPU memory as RGBA. + fn copy(&mut self, resource: &IDXGIResource) -> windows::core::Result> { + let texture: ID3D11Texture2D = resource.cast()?; + let mut desc = D3D11_TEXTURE2D_DESC::default(); + unsafe { texture.GetDesc(&mut desc) }; + let (width, height) = (desc.Width, desc.Height); + if self.staging.as_ref().is_none_or(|(_, w, h)| *w != width || *h != height) { + let staging_desc = D3D11_TEXTURE2D_DESC { + MipLevels: 1, + ArraySize: 1, + BindFlags: 0, + MiscFlags: 0, + Usage: D3D11_USAGE_STAGING, + CPUAccessFlags: D3D11_CPU_ACCESS_READ.0 as u32, + ..desc + }; + let mut staging = None; + unsafe { self.device.CreateTexture2D(&staging_desc, None, Some(&mut staging))? }; + let Some(staging) = staging else { return Ok(None) }; + self.staging = Some((staging, width, height)); + } + let Some((staging, _, _)) = self.staging.as_ref() else { return Ok(None) }; + let mut mapped = D3D11_MAPPED_SUBRESOURCE::default(); + unsafe { + self.context.CopyResource(staging, &texture); + self.context.Map(staging, 0, D3D11_MAP_READ, 0, Some(&mut mapped))?; + } + let row_bytes = width as usize * 4; + let mut rgba = vec![0u8; row_bytes * height as usize]; + for (row, target) in rgba.chunks_exact_mut(row_bytes).enumerate() { + let source = unsafe { + std::slice::from_raw_parts( + (mapped.pData as *const u8).add(row * mapped.RowPitch as usize), + row_bytes, + ) + }; + // BGRA to RGBA. The duplicated alpha is not meaningful. + for (to, from) in target.chunks_exact_mut(4).zip(source.chunks_exact(4)) { + to.copy_from_slice(&[from[2], from[1], from[0], 255]); + } + } + unsafe { self.context.Unmap(staging, 0) }; + Ok(RgbaImage::from_raw(width, height, rgba)) + } + } + +} + +#[cfg(test)] +mod tests { + #[test] + fn live_line_is_the_json_the_viewer_parses() { + assert_eq!(super::live_line(12, true), r#"live: {"seq":12,"changed":true}"#); + assert_eq!(super::live_line(0, false), r#"live: {"seq":0,"changed":false}"#); + } + + #[cfg(not(windows))] + #[test] + fn other_platforms_fall_back_to_an_ordinary_screenshot() { + let request = super::LiveRequest { + display: 0, + max_width: 1400, + format: crate::capture::CaptureFormat::Png, + after: None, + wait_ms: 0, + pointer: false, + }; + assert!(matches!(super::frame(&request), Err(super::LiveError::Unsupported))); + } +} diff --git a/windows-linux/src/main.rs b/windows-linux/src/main.rs index 82a2c7b..919166e 100644 --- a/windows-linux/src/main.rs +++ b/windows-linux/src/main.rs @@ -23,6 +23,7 @@ mod install; // Only the Windows and Linux backends press keys; macOS builds compile this for tests. #[cfg_attr(not(any(windows, target_os = "linux")), allow(dead_code))] mod keys; +mod live; mod platform; mod policy; mod tools; @@ -371,6 +372,38 @@ fn call_tool( Ok(format) => format, Err(error) => return text_result(format!("error: {error}"), true), }; + let Ok(index) = usize::try_from(display) else { + return text_result("error: display index must be zero or greater", true); + }; + if args.get("live").and_then(Value::as_bool) == Some(true) { + let request = live::LiveRequest { + display: index, + max_width, + format, + after: args.get("after").and_then(Value::as_u64), + wait_ms: arg_i64(&args, "wait_ms") + .unwrap_or(i64::from(live::DEFAULT_WAIT_MS)) + .clamp(0, i64::from(live::MAX_WAIT_MS)) as u32, + pointer: args.get("cursor").and_then(Value::as_bool) == Some(true), + }; + match live::frame(&request) { + Ok(live::LiveFrame::Changed { capture, seq }) => { + let text = capture::mapping_text(&capture, &format!("display {index}")); + let mut result = image_result(capture.bytes, format.mime_type(), text); + if let Some(content) = result.get_mut("content").and_then(Value::as_array_mut) { + content.push(json!({ "type": "text", "text": live::live_line(seq, true) })); + } + return with_cursor(result, &args); + } + Ok(live::LiveFrame::Unchanged { seq }) => { + return with_cursor(text_result(live::live_line(seq, false), false), &args); + } + Err(live::LiveError::Failed(error)) => return text_result(format!("error: {error}"), true), + // No live path here: an ordinary capture, without the live line, + // tells the viewer to keep polling. + Err(live::LiveError::Unsupported) => {} + } + } return match usize::try_from(display) { Ok(index) => match capture::capture_display(index, max_width, format) { Ok(capture) => {