Files
clawdbot/apps/macos/Sources/Clawdis/Screenshotter.swift
2025-12-07 00:10:35 +01:00

81 lines
2.7 KiB
Swift

import AppKit
import CoreGraphics
import Foundation
@preconcurrency import ScreenCaptureKit
import VideoToolbox
enum Screenshotter {
@MainActor
static func capture(displayID: UInt32?, windowID: UInt32?) async -> Data? {
guard let content = try? await SCShareableContent.current else { return nil }
let targetDisplay: SCDisplay? = if let displayID {
content.displays.first(where: { $0.displayID == displayID })
} else {
content.displays.first
}
let filter: SCContentFilter
if let windowID, let win = content.windows.first(where: { $0.windowID == windowID }) {
filter = SCContentFilter(desktopIndependentWindow: win)
} else if let display = targetDisplay {
filter = SCContentFilter(display: display, excludingWindows: [])
} else {
return nil
}
let config = SCStreamConfiguration()
if let display = targetDisplay {
config.width = display.width
config.height = display.height
}
config.scalesToFit = true
config.colorSpaceName = CGColorSpace.displayP3
let stream = SCStream(filter: filter, configuration: config, delegate: nil)
let grabber = FrameGrabber()
try? stream.addStreamOutput(
grabber,
type: .screen,
sampleHandlerQueue: DispatchQueue(label: "com.steipete.clawdis.sshot"))
do {
try await stream.startCapture()
let data = await grabber.awaitPNG()
try? await stream.stopCapture()
return data
} catch {
return nil
}
}
}
final class FrameGrabber: NSObject, SCStreamOutput {
private var continuation: CheckedContinuation<Data?, Never>?
private var delivered = false
func awaitPNG() async -> Data? {
await withCheckedContinuation { cont in
self.continuation = cont
}
}
nonisolated func stream(
_ stream: SCStream,
didOutputSampleBuffer sampleBuffer: CMSampleBuffer,
of outputType: SCStreamOutputType)
{
guard outputType == .screen else { return }
if self.delivered { return }
guard let imageBuffer = sampleBuffer.imageBuffer else { return }
var cgImage: CGImage?
let result = VTCreateCGImageFromCVPixelBuffer(imageBuffer, options: nil, imageOut: &cgImage)
guard result == noErr, let cgImage else { return }
let rep = NSBitmapImageRep(cgImage: cgImage)
guard let data = rep.representation(using: .png, properties: [:]) else { return }
self.delivered = true
self.continuation?.resume(returning: data)
self.continuation = nil
}
}