fix(apple): bind widget snapshots to their document owner (#129666)

apps/shared/OpenClawKit/Sources/OpenClawChatUI/ChatInlineWidgetSnapshotCapture.swift

@@ -0,0 +1,79 @@
+#if canImport(WebKit) && (os(iOS) || os(macOS))
+import Foundation
+import WebKit
+
+#if os(iOS)
+import UIKit
+
+typealias ChatInlineWidgetSnapshotImage = UIImage
+#elseif os(macOS)
+import AppKit
+
+typealias ChatInlineWidgetSnapshotImage = NSImage
+#endif
+
+struct ChatInlineWidgetSnapshotRequest: Equatable {
+    enum Action: Equatable {
+        case copy
+        case save
+    }
+
+    let id = UUID()
+    let action: Action
+    let generation: UUID
+    let resource: OpenClawChatWidgetResource
+}
+
+enum ChatInlineWidgetSnapshotOutcome {
+    case success(ChatInlineWidgetSnapshotRequest, ChatInlineWidgetSnapshotImage)
+    case failure(ChatInlineWidgetSnapshotRequest)
+}
+
+@MainActor
+final class ChatInlineWidgetSnapshotCapture {
+    private var request: ChatInlineWidgetSnapshotRequest?
+    private weak var webView: WKWebView?
+    private var generation: UUID?
+    private var resource: OpenClawChatWidgetResource?
+
+    func capture(
+        _ request: ChatInlineWidgetSnapshotRequest?,
+        from webView: WKWebView,
+        generation: UUID,
+        resource: OpenClawChatWidgetResource,
+        onSnapshot: @escaping @MainActor @Sendable (ChatInlineWidgetSnapshotOutcome) -> Void)
+    {
+        if self.webView !== webView || self.generation != generation || self.resource != resource {
+            self.invalidate()
+        }
+        guard let request,
+              request.generation == generation,
+              request.resource == resource,
+              request.id != self.request?.id
+        else { return }
+        self.request = request
+        self.webView = webView
+        self.generation = generation
+        self.resource = resource
+
+        webView.takeSnapshot(with: WKSnapshotConfiguration()) { [weak self, weak webView] image, _ in
+            guard let self,
+                  let webView,
+                  self.request == request,
+                  self.webView === webView,
+                  self.generation == request.generation,
+                  self.resource == request.resource
+            else { return }
+            self.invalidate()
+            onSnapshot(image.map { .success(request, $0) } ?? .failure(request))
+        }
+    }
+
+    func invalidate() {
+        self.request = nil
+        self.webView = nil
+        self.generation = nil
+        self.resource = nil
+    }
+}
+#endif

apps/shared/OpenClawKit/Sources/OpenClawChatUI/ChatInlineWidgetView.swift

@@ -390,6 +390,7 @@ struct ChatInlineWidgetView: View {
     private func renderedWidget(resource: OpenClawChatWidgetResource) -> some View {
         ChatInlineWidgetWebView(
             resource: resource,
+            loadGeneration: self.loadGeneration,
             allowsScripts: self.preview.sandbox == "scripts",
             snapshotRequest: self.snapshotRequest,
             onFailure: { self.handleLoadFailure(resource: resource) },
@@ -429,16 +430,20 @@ struct ChatInlineWidgetView: View {
     }
 
     private func requestSnapshot(for action: ChatInlineWidgetSnapshotRequest.Action) {
-        self.snapshotRequest = ChatInlineWidgetSnapshotRequest(action: action)
+        guard let resource = self.resolvedResource else { return }
+        self.snapshotRequest = ChatInlineWidgetSnapshotRequest(
+            action: action,
+            generation: self.loadGeneration,
+            resource: resource)
     }
 
     private func handleSnapshot(_ outcome: ChatInlineWidgetSnapshotOutcome) {
         switch outcome {
         case let .failure(request):
-            self.clearSnapshotRequest(ifMatching: request)
+            guard self.consumeSnapshotRequest(request) else { return }
             self.exportErrorMessage = String(localized: "The widget image could not be captured.")
         case let .success(request, image):
-            self.clearSnapshotRequest(ifMatching: request)
+            guard self.consumeSnapshotRequest(request) else { return }
             switch request.action {
             case .copy:
                 self.copySnapshot(image)
@@ -448,9 +453,13 @@ struct ChatInlineWidgetView: View {
         }
     }
 
-    private func clearSnapshotRequest(ifMatching completedRequest: ChatInlineWidgetSnapshotRequest) {
-        guard self.snapshotRequest?.id == completedRequest.id else { return }
+    private func consumeSnapshotRequest(_ request: ChatInlineWidgetSnapshotRequest) -> Bool {
+        guard self.snapshotRequest == request,
+              request.generation == self.loadGeneration,
+              request.resource == self.resolvedResource
+        else { return false }
         self.snapshotRequest = nil
+        return true
     }
 
     #if os(iOS)
@@ -500,7 +509,7 @@ struct ChatInlineWidgetView: View {
     private func reset(path: String?) {
         self.loadGeneration = UUID()
         self.activePath = path
-        self.resolvedResource = nil
+        self.setResolvedResource(nil)
         self.recoveryAttempts = 0
         self.refreshInFlight = false
         self.unavailable = false
@@ -517,17 +526,26 @@ struct ChatInlineWidgetView: View {
               self.loadGeneration == generation
         else { return }
         let resource = candidate?.hasValidTLSBinding == true ? candidate : nil
-        self.resolvedResource = resource
+        self.setResolvedResource(resource)
         self.unavailable = resource == nil
     }
 
+    private func setResolvedResource(_ resource: OpenClawChatWidgetResource?) {
+        #if canImport(WebKit) && (os(iOS) || os(macOS))
+        if self.resolvedResource != resource || resource == nil {
+            self.snapshotRequest = nil
+        }
+        #endif
+        self.resolvedResource = resource
+    }
+
     private func handleLoadFailure(resource: OpenClawChatWidgetResource) {
         guard self.resolvedResource == resource,
               let path = self.activePath,
               !self.refreshInFlight
         else { return }
         guard self.recoveryAttempts < 3 else {
-            self.resolvedResource = nil
+            self.setResolvedResource(nil)
             self.unavailable = true
             return
         }
@@ -549,27 +567,6 @@ extension OpenClawChatWidgetResource {
 }
 
 #if canImport(WebKit) && (os(iOS) || os(macOS))
-#if os(iOS)
-private typealias ChatInlineWidgetSnapshotImage = UIImage
-#elseif os(macOS)
-private typealias ChatInlineWidgetSnapshotImage = NSImage
-#endif
-
-private struct ChatInlineWidgetSnapshotRequest: Equatable {
-    enum Action: Equatable {
-        case copy
-        case save
-    }
-
-    let id = UUID()
-    let action: Action
-}
-
-private enum ChatInlineWidgetSnapshotOutcome {
-    case success(ChatInlineWidgetSnapshotRequest, ChatInlineWidgetSnapshotImage)
-    case failure(ChatInlineWidgetSnapshotRequest)
-}
-
 enum ChatInlineWidgetTLSPin {
     static func normalize(_ raw: String) -> String? {
         let stripped = raw.replacingOccurrences(
@@ -618,6 +615,7 @@ private final class ChatInlineWidgetNavigationDelegate: NSObject, WKNavigationDe
         didSet {
             if self.resource != oldValue {
                 self.contentProcessRecovery.reset()
+                self.snapshotCapture.invalidate()
             }
         }
     }
@@ -625,7 +623,7 @@ private final class ChatInlineWidgetNavigationDelegate: NSObject, WKNavigationDe
     let onFailure: @MainActor @Sendable () -> Void
     var onSnapshot: @MainActor @Sendable (ChatInlineWidgetSnapshotOutcome) -> Void
     private var contentProcessRecovery = ChatInlineWidgetContentProcessRecovery()
-    private var lastSnapshotRequestID: UUID?
+    let snapshotCapture = ChatInlineWidgetSnapshotCapture()
 
     init(
         resource: OpenClawChatWidgetResource,
@@ -637,21 +635,6 @@ private final class ChatInlineWidgetNavigationDelegate: NSObject, WKNavigationDe
         self.onSnapshot = onSnapshot
     }
 
-    func captureSnapshot(_ request: ChatInlineWidgetSnapshotRequest?, from webView: WKWebView) {
-        guard let request, request.id != self.lastSnapshotRequestID else { return }
-        self.lastSnapshotRequestID = request.id
-
-        let configuration = WKSnapshotConfiguration()
-        webView.takeSnapshot(with: configuration) { [weak self] image, _ in
-            guard let self else { return }
-            if let image {
-                self.onSnapshot(.success(request, image))
-            } else {
-                self.onSnapshot(.failure(request))
-            }
-        }
-    }
-
     func webView(
         _: WKWebView,
         decidePolicyFor navigationAction: WKNavigationAction,
@@ -766,6 +749,7 @@ private func makeChatInlineWidgetWebView(
 #if os(iOS)
 private struct ChatInlineWidgetWebView: UIViewRepresentable {
     let resource: OpenClawChatWidgetResource
+    let loadGeneration: UUID
     let allowsScripts: Bool
     let snapshotRequest: ChatInlineWidgetSnapshotRequest?
     let onFailure: @MainActor @Sendable () -> Void
@@ -791,10 +775,16 @@ private struct ChatInlineWidgetWebView: UIViewRepresentable {
             context.coordinator.resource = self.resource
             webView.load(URLRequest(url: self.resource.url, cachePolicy: .reloadIgnoringLocalCacheData))
         }
-        context.coordinator.captureSnapshot(self.snapshotRequest, from: webView)
+        context.coordinator.snapshotCapture.capture(
+            self.snapshotRequest,
+            from: webView,
+            generation: self.loadGeneration,
+            resource: self.resource,
+            onSnapshot: context.coordinator.onSnapshot)
     }
 
     static func dismantleUIView(_ webView: WKWebView, coordinator: ChatInlineWidgetNavigationDelegate) {
+        coordinator.snapshotCapture.invalidate()
         webView.stopLoading()
         webView.navigationDelegate = nil
     }
@@ -802,6 +792,7 @@ private struct ChatInlineWidgetWebView: UIViewRepresentable {
 #elseif os(macOS)
 private struct ChatInlineWidgetWebView: NSViewRepresentable {
     let resource: OpenClawChatWidgetResource
+    let loadGeneration: UUID
     let allowsScripts: Bool
     let snapshotRequest: ChatInlineWidgetSnapshotRequest?
     let onFailure: @MainActor @Sendable () -> Void
@@ -827,10 +818,16 @@ private struct ChatInlineWidgetWebView: NSViewRepresentable {
             context.coordinator.resource = self.resource
             webView.load(URLRequest(url: self.resource.url, cachePolicy: .reloadIgnoringLocalCacheData))
         }
-        context.coordinator.captureSnapshot(self.snapshotRequest, from: webView)
+        context.coordinator.snapshotCapture.capture(
+            self.snapshotRequest,
+            from: webView,
+            generation: self.loadGeneration,
+            resource: self.resource,
+            onSnapshot: context.coordinator.onSnapshot)
     }
 
     static func dismantleNSView(_ webView: WKWebView, coordinator: ChatInlineWidgetNavigationDelegate) {
+        coordinator.snapshotCapture.invalidate()
         webView.stopLoading()
         webView.navigationDelegate = nil
     }

apps/shared/OpenClawKit/Tests/OpenClawKitTests/ChatInlineWidgetTests.swift

@@ -3,6 +3,11 @@ import OpenClawKit
 import Testing
 @testable import OpenClawChatUI
 
+#if os(macOS) && canImport(WebKit)
+import AppKit
+import WebKit
+#endif
+
 struct ChatInlineWidgetTests {
     @Test func `sanitizes widget export filenames`() {
         #expect(ChatInlineWidgetExport.filename(title: "  Sales / Q3\\Summary\u{0007}  ") ==
@@ -215,6 +220,111 @@ struct ChatInlineWidgetTests {
     }
 
     #if canImport(WebKit) && (os(iOS) || os(macOS))
+    #if os(macOS)
+    @Test @MainActor func `invalidated widget snapshots cannot publish a late completion`() throws {
+        let fixture = try ChatInlineWidgetSnapshotFixture()
+        fixture.capture(fixture.request())
+
+        fixture.owner.invalidate()
+        try fixture.complete()
+
+        #expect(fixture.outcomes.isEmpty)
+    }
+
+    @Test @MainActor func `widget snapshots reject a changed document or generation before capture`() throws {
+        let fixture = try ChatInlineWidgetSnapshotFixture()
+        let replacement = try OpenClawChatWidgetResource(
+            url: #require(URL(string: "https://gateway.example/replacement")))
+        let changedPin = OpenClawChatWidgetResource(
+            url: fixture.resource.url,
+            tlsFingerprintSHA256: String(repeating: "ab", count: 32))
+
+        fixture.capture(fixture.request(), resource: replacement)
+        fixture.capture(fixture.request(), resource: changedPin)
+        fixture.capture(fixture.request(generation: UUID()))
+
+        #expect(fixture.webView.completions.isEmpty)
+        #expect(fixture.outcomes.isEmpty)
+    }
+
+    @Test @MainActor func `widget snapshot completions cannot outlive document or TLS replacements`() throws {
+        let oldURL = try #require(URL(string: "https://gateway.example/old"))
+        let newURL = try #require(URL(string: "https://gateway.example/replacement"))
+        let replacements = [
+            OpenClawChatWidgetResource(url: newURL),
+            OpenClawChatWidgetResource(url: oldURL, tlsFingerprintSHA256: String(repeating: "ab", count: 32)),
+        ]
+
+        for replacement in replacements {
+            let fixture = try ChatInlineWidgetSnapshotFixture()
+            fixture.capture(fixture.request())
+            fixture.capture(nil, resource: replacement)
+
+            try fixture.complete()
+
+            #expect(fixture.outcomes.isEmpty)
+        }
+    }
+
+    @Test @MainActor func `widget snapshot completions cannot outlive a reset generation`() throws {
+        let fixture = try ChatInlineWidgetSnapshotFixture()
+        fixture.capture(fixture.request())
+        fixture.capture(nil, generation: UUID())
+
+        try fixture.complete()
+
+        #expect(fixture.outcomes.isEmpty)
+    }
+
+    @Test @MainActor func `widget snapshot completions cannot migrate to a replacement WebView`() throws {
+        let fixture = try ChatInlineWidgetSnapshotFixture()
+        let replacement = ChatInlineWidgetDeferredSnapshotWebView(
+            frame: .zero,
+            configuration: WKWebViewConfiguration())
+        fixture.capture(fixture.request())
+        fixture.capture(nil, from: replacement)
+
+        try fixture.complete()
+
+        #expect(fixture.outcomes.isEmpty)
+    }
+
+    @Test @MainActor func `newer widget snapshots supersede older asynchronous completions`() throws {
+        let fixture = try ChatInlineWidgetSnapshotFixture()
+        let first = fixture.request(action: .copy)
+        let second = fixture.request(action: .save)
+        fixture.capture(first)
+        fixture.capture(second)
+
+        try fixture.complete(at: 0)
+        #expect(fixture.outcomes.isEmpty)
+
+        try fixture.complete(at: 1)
+        guard case let .success(request, _) = try #require(fixture.outcomes.first) else {
+            Issue.record("The current widget snapshot did not produce an image")
+            return
+        }
+        #expect(request == second)
+    }
+
+    @Test @MainActor func `current widget snapshots publish one image for the exact document`() throws {
+        let fixture = try ChatInlineWidgetSnapshotFixture()
+        let request = fixture.request()
+        fixture.capture(request)
+        fixture.capture(request)
+
+        #expect(fixture.webView.completions.count == 1)
+        try fixture.complete()
+
+        guard case let .success(completed, _) = try #require(fixture.outcomes.first) else {
+            Issue.record("The current widget snapshot did not produce an image")
+            return
+        }
+        #expect(completed == request)
+        #expect(fixture.outcomes.count == 1)
+    }
+    #endif
+
     @Test func `bounds WebKit content process recovery per document`() {
         var recovery = ChatInlineWidgetContentProcessRecovery()
 
@@ -236,6 +346,70 @@ struct ChatInlineWidgetTests {
     #endif
 }
 
+#if os(macOS) && canImport(WebKit)
+@MainActor
+private final class ChatInlineWidgetSnapshotFixture {
+    let resource: OpenClawChatWidgetResource
+    let generation = UUID()
+    let webView = ChatInlineWidgetDeferredSnapshotWebView(
+        frame: .zero,
+        configuration: WKWebViewConfiguration())
+    let owner = ChatInlineWidgetSnapshotCapture()
+    var outcomes: [ChatInlineWidgetSnapshotOutcome] = []
+
+    init() throws {
+        self.resource = try OpenClawChatWidgetResource(
+            url: #require(URL(string: "https://gateway.example/old")))
+    }
+
+    func request(
+        action: ChatInlineWidgetSnapshotRequest.Action = .copy,
+        generation: UUID? = nil,
+        resource: OpenClawChatWidgetResource? = nil) -> ChatInlineWidgetSnapshotRequest
+    {
+        ChatInlineWidgetSnapshotRequest(
+            action: action,
+            generation: generation ?? self.generation,
+            resource: resource ?? self.resource)
+    }
+
+    func capture(
+        _ request: ChatInlineWidgetSnapshotRequest?,
+        from webView: WKWebView? = nil,
+        generation: UUID? = nil,
+        resource: OpenClawChatWidgetResource? = nil)
+    {
+        self.owner.capture(
+            request,
+            from: webView ?? self.webView,
+            generation: generation ?? self.generation,
+            resource: resource ?? self.resource)
+        { [weak self] outcome in
+            self?.outcomes.append(outcome)
+        }
+    }
+
+    func complete(at index: Int = 0) throws {
+        let completion = try #require(self.webView.completions.indices.contains(index)
+            ? self.webView.completions[index]
+            : nil)
+        completion(NSImage(size: NSSize(width: 1, height: 1)), nil)
+    }
+}
+
+@MainActor
+private final class ChatInlineWidgetDeferredSnapshotWebView: WKWebView {
+    var completions: [@MainActor (NSImage?, (any Error)?) -> Void] = []
+
+    override func takeSnapshot(
+        with _: WKSnapshotConfiguration?,
+        completionHandler: @escaping @MainActor (NSImage?, (any Error)?) -> Void)
+    {
+        self.completions.append(completionHandler)
+    }
+}
+#endif
+
 private actor ChatWidgetOperatorRouteRefreshProbe {
     private var route: GatewayCanvasHostRoute
     private let replacement: GatewayCanvasHostRoute