refactor: extract UI test socket sanity coordinator

Sources/AppDelegate.swift

@@ -601,12 +601,32 @@ final class AppDelegate: NSObject, NSApplicationDelegate, UNUserNotificationCent
     )
     private var didSetupPortalStatsUITestDiagnostics = false
     private var portalStatsUITestObservers: [NSObjectProtocol] = []
+    private lazy var uiTestSocketSanityCoordinator = UITestSocketSanityCoordinator(
+        dependencies: .init(
+            configuration: { [weak self] in self?.socketListenerConfigurationIfEnabled() },
+            activeSocketPath: { path in
+                TerminalController.shared.activeSocketPath(preferredPath: path)
+            },
+            health: { path in
+                TerminalController.shared.socketListenerHealth(expectedSocketPath: path)
+            },
+            probe: { [weak self] command, path, timeout in
+                self?.socketTransport.probeCommand(command, at: path, timeout: timeout)
+            },
+            restart: { [weak self] source in
+                self?.restartSocketListenerIfEnabled(source: source)
+            },
+            recordStage: { [weak self] stage in
+                self?.uiTestDiagnosticsWriter.write(stage: stage)
+            }
+        )
+    )
     private lazy var uiTestDiagnosticsWriter = UITestDiagnosticsWriter(
         isRunningUnderXCTest: { [weak self] environment in
             self?.isRunningUnderXCTest(environment) ?? false
         },
         socketDiagnostics: { [weak self] environment in
-            self?.uiTestSocketDiagnostics(environment: environment) ?? [:]
+            self?.uiTestSocketSanityCoordinator.diagnostics(environment: environment) ?? [:]
         },
         renderDiagnostics: { [weak self] in
             self?.currentUITestRenderDiagnosticsForWriter()
@@ -1297,50 +1317,6 @@ final class AppDelegate: NSObject, NSApplicationDelegate, UNUserNotificationCent
     }
 
 #if DEBUG
-    private func uiTestSocketDiagnostics(
-        environment env: [String: String]
-    ) -> [String: String] {
-        var payload: [String: String] = [:]
-        guard env["CMUX_UI_TEST_SOCKET_SANITY"] == "1" else { return payload }
-
-        guard let config = socketListenerConfigurationIfEnabled() else {
-            payload["socketExpectedPath"] = env["CMUX_SOCKET_PATH"] ?? ""
-            payload["socketMode"] = "off"
-            payload["socketReady"] = "0"
-            payload["socketPingResponse"] = ""
-            payload["socketIsRunning"] = "0"
-            payload["socketAcceptLoopAlive"] = "0"
-            payload["socketPathMatches"] = "0"
-            payload["socketPathExists"] = "0"
-            payload["socketPathOwnedByListener"] = "0"
-            payload["socketFailureSignals"] = "socket_disabled"
-            return payload
-        }
-
-        let socketPath = TerminalController.shared.activeSocketPath(preferredPath: config.preferredSocketPath)
-        let health = TerminalController.shared.socketListenerHealth(expectedSocketPath: socketPath)
-        let pingResponse = health.isHealthy
-            ? socketTransport.probeCommand("ping", at: socketPath, timeout: 1.0)
-            : nil
-        let isReady = health.isHealthy && pingResponse == "PONG"
-        var failureSignals = health.failureSignals
-        if health.isHealthy && pingResponse != "PONG" {
-            failureSignals.append("ping_timeout")
-        }
-
-        payload["socketExpectedPath"] = socketPath
-        payload["socketMode"] = config.accessMode.rawValue
-        payload["socketReady"] = isReady ? "1" : "0"
-        payload["socketPingResponse"] = pingResponse ?? ""
-        payload["socketIsRunning"] = health.isRunning ? "1" : "0"
-        payload["socketAcceptLoopAlive"] = health.acceptLoopAlive ? "1" : "0"
-        payload["socketPathMatches"] = health.socketPathMatches ? "1" : "0"
-        payload["socketPathExists"] = health.socketPathExists ? "1" : "0"
-        payload["socketPathOwnedByListener"] = health.socketPathOwnedByListener ? "1" : "0"
-        payload["socketFailureSignals"] = failureSignals.joined(separator: ",")
-        return payload
-    }
-
     private func currentUITestRenderDiagnosticsForWriter() -> UITestDiagnosticsWriter.RenderDiagnostics? {
         guard let tabManager,
               let tabId = tabManager.selectedTabId,
@@ -1856,7 +1832,7 @@ final class AppDelegate: NSObject, NSApplicationDelegate, UNUserNotificationCent
 
         let env = ProcessInfo.processInfo.environment
         if isRunningUnderXCTest(env) || env["CMUX_UI_TEST_MODE"] == "1" {
-            scheduleUITestSocketSanityCheckIfNeeded()
+            uiTestSocketSanityCoordinator.scheduleIfNeeded(environment: env)
         }
         // Best-effort one-time migration: a value previously stored in the
         // legacy ~/.config/cmux/dev-window-display file moves into the shared
@@ -1898,38 +1874,6 @@ final class AppDelegate: NSObject, NSApplicationDelegate, UNUserNotificationCent
 
 
 #if DEBUG
-    private func scheduleUITestSocketSanityCheckIfNeeded() {
-        let env = ProcessInfo.processInfo.environment
-        guard env["CMUX_UI_TEST_SOCKET_SANITY"] == "1" else { return payload }
-
-        DispatchQueue.main.asyncAfter(deadline: .now() + 0.75) { [weak self] in
-            guard let self else { return }
-            guard let config = self.socketListenerConfigurationIfEnabled() else {
-                self.uiTestDiagnosticsWriter.write(stage: "socketSanityDisabled")
-                return
-            }
-
-            let expectedPath = TerminalController.shared.activeSocketPath(
-                preferredPath: config.preferredSocketPath
-            )
-            let health = TerminalController.shared.socketListenerHealth(expectedSocketPath: expectedPath)
-            let pingResponse = health.isHealthy
-                ? socketTransport.probeCommand("ping", at: expectedPath, timeout: 1.0)
-                : nil
-            let isReady = health.isHealthy && pingResponse == "PONG"
-            if isReady {
-                self.uiTestDiagnosticsWriter.write(stage: "socketSanityReady")
-                return
-            }
-
-            self.uiTestDiagnosticsWriter.write(stage: "socketSanityRestart")
-            self.restartSocketListenerIfEnabled(source: "uiTest.socketSanity")
-            DispatchQueue.main.asyncAfter(deadline: .now() + 0.75) { [weak self] in
-                self?.uiTestDiagnosticsWriter.write(stage: "socketSanityPostRestart")
-            }
-        }
-    }
-
     private func setupDisplayResolutionUITestDiagnosticsIfNeeded() {
         let env = ProcessInfo.processInfo.environment
         guard env["CMUX_UI_TEST_DISPLAY_RENDER_STATS"] == "1" else { return }
@@ -11064,7 +11008,7 @@ final class AppDelegate: NSObject, NSApplicationDelegate, UNUserNotificationCent
         window2Id: UUID? = nil
     ) {
         let env = ProcessInfo.processInfo.environment
-        guard env["CMUX_UI_TEST_SOCKET_SANITY"] == "1" else { return payload }
+        guard env["CMUX_UI_TEST_SOCKET_SANITY"] == "1" else { return }
 
         guard let config = socketListenerConfigurationIfEnabled() else {
             writeMultiWindowNotificationTestData([

Sources/Debug/UITests/UITestSocketSanityCoordinator.swift

@@ -0,0 +1,111 @@
+#if DEBUG
+import AppKit
+import CmuxControlSocket
+
+/// Owns the opt-in UI-test socket health probe and restart sequence.
+@MainActor
+final class UITestSocketSanityCoordinator {
+    struct Dependencies {
+        let configuration: () -> SocketControlServerConfiguration?
+        let activeSocketPath: (String) -> String
+        let health: (String) -> SocketListenerHealth
+        let probe: (String, String, TimeInterval) -> String?
+        let restart: (String) -> Void
+        let recordStage: (String) -> Void
+    }
+
+    private let dependencies: Dependencies
+    private var scheduled = false
+    private var tasks: [Task<Void, Never>] = []
+
+    init(dependencies: Dependencies) {
+        self.dependencies = dependencies
+    }
+
+    deinit {
+        tasks.forEach { $0.cancel() }
+    }
+
+    func diagnostics(environment: [String: String]) -> [String: String] {
+        guard environment["CMUX_UI_TEST_SOCKET_SANITY"] == "1" else { return [:] }
+        guard let config = dependencies.configuration() else {
+            return [
+                "socketExpectedPath": environment["CMUX_SOCKET_PATH"] ?? "",
+                "socketMode": "off",
+                "socketReady": "0",
+                "socketPingResponse": "",
+                "socketIsRunning": "0",
+                "socketAcceptLoopAlive": "0",
+                "socketPathMatches": "0",
+                "socketPathExists": "0",
+                "socketPathOwnedByListener": "0",
+                "socketFailureSignals": "socket_disabled",
+            ]
+        }
+
+        let path = dependencies.activeSocketPath(config.preferredSocketPath)
+        let health = dependencies.health(path)
+        let pingResponse = health.isHealthy ? dependencies.probe("ping", path, 1.0) : nil
+        let isReady = health.isHealthy && pingResponse == "PONG"
+        var failureSignals = health.failureSignals
+        if health.isHealthy && pingResponse != "PONG" {
+            failureSignals.append("ping_timeout")
+        }
+
+        return [
+            "socketExpectedPath": path,
+            "socketMode": config.accessMode.rawValue,
+            "socketReady": isReady ? "1" : "0",
+            "socketPingResponse": pingResponse ?? "",
+            "socketIsRunning": health.isRunning ? "1" : "0",
+            "socketAcceptLoopAlive": health.acceptLoopAlive ? "1" : "0",
+            "socketPathMatches": health.socketPathMatches ? "1" : "0",
+            "socketPathExists": health.socketPathExists ? "1" : "0",
+            "socketPathOwnedByListener": health.socketPathOwnedByListener ? "1" : "0",
+            "socketFailureSignals": failureSignals.joined(separator: ","),
+        ]
+    }
+
+    func scheduleIfNeeded(environment: [String: String]) {
+        guard !scheduled,
+              environment["CMUX_UI_TEST_SOCKET_SANITY"] == "1" else { return }
+        scheduled = true
+        schedule(after: .milliseconds(750)) { [weak self] in
+            self?.runProbe()
+        }
+    }
+
+    private func runProbe() {
+        guard let config = dependencies.configuration() else {
+            dependencies.recordStage("socketSanityDisabled")
+            return
+        }
+        let path = dependencies.activeSocketPath(config.preferredSocketPath)
+        let health = dependencies.health(path)
+        let pingResponse = health.isHealthy ? dependencies.probe("ping", path, 1.0) : nil
+        guard health.isHealthy && pingResponse == "PONG" else {
+            dependencies.recordStage("socketSanityRestart")
+            dependencies.restart("uiTest.socketSanity")
+            schedule(after: .milliseconds(750)) { [weak self] in
+                self?.dependencies.recordStage("socketSanityPostRestart")
+            }
+            return
+        }
+        dependencies.recordStage("socketSanityReady")
+    }
+
+    private func schedule(after duration: Duration, operation: @escaping @MainActor () -> Void) {
+        let task = Task { @MainActor [weak self] in
+            do {
+                try await ContinuousClock().sleep(for: duration)
+            } catch {
+                return
+            }
+            guard !Task.isCancelled else { return }
+            operation()
+            self?.tasks.removeAll { $0.isCancelled }
+        }
+        tasks.append(task)
+    }
+}
+#endif

cmux.xcodeproj/project.pbxproj

@@ -2545,6 +2545,7 @@ C0DE71B10000000000000001 /* AppDelegate+AgentChatNotifications.swift in Sources
 		468100000000000000000003 /* TypingHotPathRegressionTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 468100000000000000000004 /* TypingHotPathRegressionTests.swift */; };
 		D9A1B2C3D4E5000000000001 /* UITestDiagnosticsWriter.swift in Sources */ = {isa = PBXBuildFile; fileRef = D9A1B2C3D4E5000000000002 /* UITestDiagnosticsWriter.swift */; };
 		A5001501 /* UITestRecorder.swift in Sources */ = {isa = PBXBuildFile; fileRef = A5001511 /* UITestRecorder.swift */; };
+		E7F100000000000000000001 /* UITestSocketSanityCoordinator.swift in Sources */ = {isa = PBXBuildFile; fileRef = E7F100000000000000000002 /* UITestSocketSanityCoordinator.swift */; };
 		A500120D /* UpdateLogStore.swift in Sources */ = {isa = PBXBuildFile; fileRef = A5001223 /* UpdateLogStore.swift */; };
 		F2000000A1B2C3D4E5F60718 /* UpdatePillReleaseVisibilityTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = F2000001A1B2C3D4E5F60718 /* UpdatePillReleaseVisibilityTests.swift */; };
 		C0B4D9B0A1B2C3D4E5F60718 /* UpdatePillUITests.swift in Sources */ = {isa = PBXBuildFile; fileRef = C0B4D9B1A1B2C3D4E5F60718 /* UpdatePillUITests.swift */; };
@@ -5316,6 +5317,7 @@ C0DE71B10000000000000002 /* AppDelegate+AgentChatNotifications.swift */ = {isa =
 		468100000000000000000004 /* TypingHotPathRegressionTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TypingHotPathRegressionTests.swift; sourceTree = "<group>"; };
 		D9A1B2C3D4E5000000000002 /* UITestDiagnosticsWriter.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = "Debug/UITests/UITestDiagnosticsWriter.swift"; sourceTree = "<group>"; };
 		A5001511 /* UITestRecorder.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = UITestRecorder.swift; sourceTree = "<group>"; };
+		E7F100000000000000000002 /* UITestSocketSanityCoordinator.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = "Debug/UITests/UITestSocketSanityCoordinator.swift"; sourceTree = "<group>"; };
 		A5001223 /* UpdateLogStore.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Update/UpdateLogStore.swift; sourceTree = "<group>"; };
 		F2000001A1B2C3D4E5F60718 /* UpdatePillReleaseVisibilityTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = UpdatePillReleaseVisibilityTests.swift; sourceTree = "<group>"; };
 		C0B4D9B1A1B2C3D4E5F60718 /* UpdatePillUITests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = UpdatePillUITests.swift; sourceTree = "<group>"; };
@@ -6892,6 +6894,7 @@ C0DE71B10000000000000002 /* AppDelegate+AgentChatNotifications.swift */ = {isa =
 				FEE012340000000000000004 /* FeedSidebarUITestPushClient.swift */,
 				FEE012340000000000000006 /* FeedSidebarUITestRecorder.swift */,
 				D9A1B2C3D4E5000000000002 /* UITestDiagnosticsWriter.swift */,
+			E7F100000000000000000002 /* UITestSocketSanityCoordinator.swift */,
 				C51A740000000000000000A2 /* AppDelegate+SimulatorShortcutRouting.swift */,
 				D35B00000000000000000011 /* TerminalController+BrowserDesignMode.swift */,
 				D35B00000000000000000014 /* TerminalController+AgentPromptDelivery.swift */,
@@ -10590,6 +10593,7 @@ C0DE71B10000000000000002 /* AppDelegate+AgentChatNotifications.swift */ = {isa =
 				E30760000000000000000002 /* TmuxWorkspacePaneOverlayView.swift in Sources */,
 				D9A1B2C3D4E5000000000001 /* UITestDiagnosticsWriter.swift in Sources */,
 				A5001501 /* UITestRecorder.swift in Sources */,
+			E7F100000000000000000001 /* UITestSocketSanityCoordinator.swift in Sources */,
 				A500120D /* UpdateLogStore.swift in Sources */,
 				A5001208 /* UpdateTitlebarAccessory.swift in Sources */,
 				C2035A070000000000000001 /* URLRequest+BrowserFailedNavigation.swift in Sources */,