fix: repaint Codex sessions after pane reattachment

Sources/Panels/AgentSessionWebHostView.swift

@@ -4,11 +4,14 @@ import WebKit
 @MainActor
 final class AgentSessionWebHostView: NSView {
     var onDidMoveToWindow: (() -> Void)?
+    /// Called when a retained agent web host is attached after its first pane attachment.
+    var onDidReattach: (() -> Void)?
     var onGeometryChanged: (() -> Void)?
     private(set) var geometryRevision: UInt64 = 0
     private var lastReportedAgentSessionWebHostGeometryState: AgentSessionWebHostGeometryState?
     private var hasPendingGeometryNotification = false
     private weak var hostedWebView: WKWebView?
+    private var hasBeenAttachedToSuperview = false
     private var sessionContentWidthPresentation = SessionContentWidthPresentation.disabled
     private var pendingScrollDelta = CGPoint.zero
     private var scrollFlushTask: Task<Void, Never>?
@@ -26,6 +29,12 @@ final class AgentSessionWebHostView: NSView {
 
     override func viewDidMoveToSuperview() {
         super.viewDidMoveToSuperview()
+        if superview != nil {
+            if hasBeenAttachedToSuperview {
+                onDidReattach?()
+            }
+            hasBeenAttachedToSuperview = true
+        }
         notifyGeometryChangedIfNeeded()
     }
 
@@ -163,10 +172,14 @@ final class AgentSessionWebHostView: NSView {
         if hostedWebView !== webView {
             resetPendingScroll()
         }
+        let wasAttachedToDifferentHost = webView.superview != nil && webView.superview !== self
         if webView.superview !== self {
             webView.removeFromSuperview()
             addSubview(webView, positioned: .above, relativeTo: nil)
         }
+        if wasAttachedToDifferentHost {
+            onDidReattach?()
+        }
         hostedWebView = webView
         webView.translatesAutoresizingMaskIntoConstraints = true
         webView.autoresizingMask = []

Sources/Panels/AgentSessionWebRenderer.swift

@@ -42,13 +42,17 @@ struct AgentSessionWebRenderer: NSViewRepresentable {
             isFocused: isFocused
         )
         let webView = context.coordinator.ensureWebView(onPointerDown: onRequestPanelFocus)
+        context.coordinator.attach(to: host)
         webView.onPointerDown = onRequestPanelFocus
         webView.navigationDelegate = context.coordinator
         webView.uiDelegate = context.coordinator
         applyBackground(to: host)
         applyBackground(to: webView)
         applyAppearance(to: webView)
         host.setSessionContentWidthPresentation(sessionContentWidthPresentation)
+        host.onDidReattach = { [weak coordinator = context.coordinator] in
+            coordinator?.invalidateVisiblePaintAfterReattachment()
+        }
         host.attachWebView(webView)
         host.onDidMoveToWindow = { [weak coordinator = context.coordinator] in
             coordinator?.loadShellIfNeeded()
@@ -68,6 +72,7 @@ struct AgentSessionWebRenderer: NSViewRepresentable {
         if let host = nsView as? AgentSessionWebHostView {
             host.detachHostedWebViewIfOwned(coordinator.webView)
             host.onDidMoveToWindow = nil
+            host.onDidReattach = nil
             host.onGeometryChanged = nil
         }
     }

Sources/Panels/AgentSessionWebRendererCoordinator.swift

@@ -8,6 +8,8 @@ final class AgentSessionWebRendererCoordinator: NSObject, WKNavigationDelegate,
     var webView: AgentSessionWebView?
     private var panelId = UUID()
     private var workspaceId = UUID()
+    private weak var attachedHost: AgentSessionWebHostView?
+    private var hasAttachedHost = false
     private var rendererKind: AgentSessionRendererKind = .react
     private var initialProviderID: AgentSessionProviderID = .codex
     private var workingDirectory: String?
@@ -18,6 +20,8 @@ final class AgentSessionWebRendererCoordinator: NSObject, WKNavigationDelegate,
     private var trustedShellURL: URL?
     private var hasFinishedNavigation = false
     private var hasCompletedVisiblePaintFlush = false
+    private var isVisiblePaintFlushInFlight = false
+    private(set) var visiblePaintGeneration: UInt64 = 0
     private var isPanelFocused = false
     private var isClosed = false
     private var isProviderStartPending = false
@@ -42,12 +46,15 @@ final class AgentSessionWebRendererCoordinator: NSObject, WKNavigationDelegate,
         isFocused: Bool
     ) {
         self.panelId = panelId
+        let workspaceChanged = self.workspaceId != workspaceId
         self.workspaceId = workspaceId
         if self.rendererKind != rendererKind {
             loadedRendererKind = nil
             trustedShellURL = nil
             hasFinishedNavigation = false
-            hasCompletedVisiblePaintFlush = false
+            resetVisiblePaintState()
+        } else if workspaceChanged {
+            resetVisiblePaintState()
         }
         self.rendererKind = rendererKind
         self.initialProviderID = initialProviderID
@@ -66,6 +73,16 @@ final class AgentSessionWebRendererCoordinator: NSObject, WKNavigationDelegate,
         }
     }
 
+    /// Records the host currently displaying this retained renderer session.
+    func attach(to host: AgentSessionWebHostView) {
+        guard attachedHost !== host else { return }
+        attachedHost = host
+        if hasAttachedHost {
+            invalidateVisiblePaintAfterReattachment()
+        }
+        hasAttachedHost = true
+    }
+
     func ensureWebView(onPointerDown: @escaping () -> Void) -> AgentSessionWebView {
         if let webView {
             webView.onPointerDown = onPointerDown
@@ -122,7 +139,7 @@ final class AgentSessionWebRendererCoordinator: NSObject, WKNavigationDelegate,
         webView.loadFileURL(indexURL, allowingReadAccessTo: Bundle.main.resourceURL ?? resourceDirectoryURL)
         loadedRendererKind = rendererKind
         hasFinishedNavigation = false
-        hasCompletedVisiblePaintFlush = false
+        resetVisiblePaintState()
     }
 
     func focus() {
@@ -157,7 +174,9 @@ final class AgentSessionWebRendererCoordinator: NSObject, WKNavigationDelegate,
         loadedRendererKind = nil
         trustedShellURL = nil
         hasFinishedNavigation = false
-        hasCompletedVisiblePaintFlush = false
+        attachedHost = nil
+        hasAttachedHost = false
+        resetVisiblePaintState()
     }
 
     func userContentController(
@@ -267,16 +286,33 @@ final class AgentSessionWebRendererCoordinator: NSObject, WKNavigationDelegate,
     func flushVisiblePaintIfReady() {
         guard hasFinishedNavigation,
               !hasCompletedVisiblePaintFlush,
+              !isVisiblePaintFlushInFlight,
               let webView,
               webView.window != nil,
               !webView.bounds.isEmpty else {
             return
         }
+        let generation = visiblePaintGeneration
+        isVisiblePaintFlushInFlight = true
         flushInitialPaint(for: webView) { [weak self] in
-            self?.hasCompletedVisiblePaintFlush = true
+            guard let self, self.visiblePaintGeneration == generation else { return }
+            self.isVisiblePaintFlushInFlight = false
+            self.hasCompletedVisiblePaintFlush = true
         }
     }
 
+    /// Reopens the visible paint gate after Bonsplit moves a retained web host.
+    func invalidateVisiblePaintAfterReattachment() {
+        resetVisiblePaintState()
+        flushVisiblePaintIfReady()
+    }
+
+    private func resetVisiblePaintState() {
+        visiblePaintGeneration &+= 1
+        isVisiblePaintFlushInFlight = false
+        hasCompletedVisiblePaintFlush = false
+    }
+
     private func flushInitialPaint(for webView: WKWebView, completion: (() -> Void)? = nil) {
         // Retained WKWebViews can finish loading before Bonsplit reattaches them
         // to a visible host. Reading layout after navigation forces WebKit to

cmuxTests/AgentSessionWebRendererTests.swift

@@ -1,5 +1,7 @@
 import Foundation
 import Testing
+import AppKit
+import WebKit
 
 #if canImport(cmux_DEV)
     @testable import cmux_DEV
@@ -9,6 +11,45 @@ import Testing
 
 @Suite(.serialized)
 struct AgentSessionWebRendererTests {
+    @Test
+    @MainActor
+    func testRetainedWebViewHostReportsReattachmentAfterPaneMove() {
+        let firstPane = NSView(frame: NSRect(x: 0, y: 0, width: 320, height: 240))
+        let secondPane = NSView(frame: NSRect(x: 0, y: 0, width: 320, height: 240))
+        let host = AgentSessionWebHostView(frame: firstPane.bounds)
+        let webView = WKWebView(frame: host.bounds)
+        var reattachmentCount = 0
+        host.onDidReattach = {
+            reattachmentCount += 1
+        }
+
+        firstPane.addSubview(host)
+        host.attachWebView(webView)
+        host.removeFromSuperview()
+        secondPane.addSubview(host)
+
+        #expect(reattachmentCount == 1)
+    }
+
+    @Test
+    @MainActor
+    func testRetainedCoordinatorReopensPaintGateForNewHost() {
+        let coordinator = AgentSessionWebRendererCoordinator()
+        var firstHost: AgentSessionWebHostView? = AgentSessionWebHostView()
+        let secondHost = AgentSessionWebHostView()
+        let initialGeneration = coordinator.visiblePaintGeneration
+
+        coordinator.attach(to: firstHost!)
+        #expect(coordinator.visiblePaintGeneration == initialGeneration)
+
+        firstHost = nil
+        coordinator.attach(to: secondHost)
+        #expect(coordinator.visiblePaintGeneration == initialGeneration + 1)
+
+        coordinator.attach(to: secondHost)
+        #expect(coordinator.visiblePaintGeneration == initialGeneration + 1)
+    }
+
     @Test
     @MainActor
     func testTerminalCommandQueuedBeforeCloseIsRejectedAfterClose() {