@wangyaoshen/remux 0.3.8-dev.29e114b

package/packages/RemuxKit/Sources/RemuxKit/Networking/RemuxConnection.swift

@@ -0,0 +1,395 @@
+import Foundation
+
+/// Connection status for the remux server
+public enum ConnectionStatus: Sendable, Equatable {
+    case disconnected
+    case connecting
+    case authenticating
+    case connected
+    case reconnecting(attempt: Int)
+}
+
+/// Credentials used to authenticate with the server
+public enum RemuxCredential: Sendable {
+    case token(String)
+    case password(String)
+    case resumeToken(String)
+}
+
+/// Delegate for connection events
+public protocol RemuxConnectionDelegate: AnyObject, Sendable {
+    func connectionDidChangeStatus(_ status: ConnectionStatus)
+    func connectionDidReceiveMessage(_ message: String)
+    func connectionDidReceiveData(_ data: Data)
+    func connectionDidAuthenticate(capabilities: ProtocolCapabilities)
+    func connectionDidFailAuth(reason: String)
+}
+
+/// WebSocket connection manager for a remux server.
+/// Uses URLSessionWebSocketTask (no third-party dependencies).
+/// Handles authentication, automatic reconnection with exponential backoff, and heartbeat.
+public final class RemuxConnection: NSObject, @unchecked Sendable {
+
+    public let serverURL: URL
+    private let credential: RemuxCredential
+    private let cols: Int
+    private let rows: Int
+
+    // Internals guarded by a lock for Sendable compliance
+    private let lock = NSLock()
+    private var _task: URLSessionWebSocketTask?
+    private var _status: ConnectionStatus = .disconnected
+    private var _reconnectAttempt: Int = 0
+    private var _reconnectTimer: DispatchWorkItem?
+    private var _heartbeatMissed: Int = 0
+    private var _heartbeatTimer: Timer?
+    private var _pendingMessages: [String] = []
+
+    public weak var delegate: RemuxConnectionDelegate?
+
+    private static let maxReconnectAttempts = 20
+    private static let heartbeatTimeout: TimeInterval = 45
+    private static let authTimeout: TimeInterval = 10
+
+    private lazy var urlSession: URLSession = {
+        URLSession(configuration: .default, delegate: nil, delegateQueue: nil)
+    }()
+
+    public init(serverURL: URL, credential: RemuxCredential, cols: Int = 80, rows: Int = 24) {
+        self.serverURL = serverURL
+        self.credential = credential
+        self.cols = cols
+        self.rows = rows
+        super.init()
+    }
+
+    // MARK: - Public API
+
+    public var status: ConnectionStatus {
+        lock.lock()
+        defer { lock.unlock() }
+        return _status
+    }
+
+    /// Connect to the server. Initiates WebSocket connection and authentication.
+    public func connect() {
+        lock.lock()
+        _reconnectAttempt = 0
+        lock.unlock()
+        startConnection()
+    }
+
+    /// Disconnect and stop reconnection attempts.
+    public func disconnect() {
+        lock.lock()
+        _reconnectTimer?.cancel()
+        _reconnectTimer = nil
+        _reconnectAttempt = Self.maxReconnectAttempts // prevent reconnect
+        let task = _task
+        _task = nil
+        lock.unlock()
+
+        stopHeartbeatTimer()
+        task?.cancel(with: .goingAway, reason: nil)
+        setStatus(.disconnected)
+    }
+
+    /// Send a JSON-encodable message to the server.
+    public func send<T: Encodable>(message: T) {
+        guard let data = try? JSONEncoder().encode(message),
+              let string = String(data: data, encoding: .utf8) else { return }
+        sendString(string)
+    }
+
+    /// Send raw text message.
+    public func sendString(_ string: String) {
+        lock.lock()
+        let task = _task
+        lock.unlock()
+
+        guard let task else {
+            // Buffer messages during reconnection
+            lock.lock()
+            _pendingMessages.append(string)
+            lock.unlock()
+            return
+        }
+        task.send(.string(string)) { _ in }
+    }
+
+    /// Send raw binary data (e.g. user terminal input forwarded as-is).
+    public func sendData(_ data: Data) {
+        lock.lock()
+        let task = _task
+        lock.unlock()
+        task?.send(.data(data)) { _ in }
+    }
+
+    // MARK: - Connection lifecycle
+
+    private func startConnection() {
+        setStatus(.connecting)
+
+        // Build WebSocket URL: ws(s)://host:port/ws
+        guard var components = URLComponents(url: serverURL, resolvingAgainstBaseURL: false) else {
+            setStatus(.disconnected)
+            return
+        }
+        // Convert http(s) to ws(s)
+        if components.scheme == "http" { components.scheme = "ws" }
+        else if components.scheme == "https" { components.scheme = "wss" }
+        components.path = "/ws"
+        guard let wsURL = components.url else { return }
+
+        NSLog("[RemuxConnection] connecting to %@", wsURL.absoluteString)
+        let task = urlSession.webSocketTask(with: wsURL)
+        lock.lock()
+        _task = task
+        lock.unlock()
+
+        task.resume()
+        receiveLoop(task: task)
+        authenticate(task: task)
+    }
+
+    private func authenticate(task: URLSessionWebSocketTask) {
+        setStatus(.authenticating)
+
+        var authDict: [String: Any] = [
+            "type": "auth",
+            "cols": cols,
+            "rows": rows,
+        ]
+
+        // Declare capabilities
+        authDict["capabilities"] = [
+            "envelope": true,
+            "inspectV2": true,
+            "deviceTrust": true,
+        ]
+
+        switch credential {
+        case .token(let t):
+            authDict["token"] = t
+        case .password(let p):
+            authDict["token"] = ""
+            authDict["password"] = p
+        case .resumeToken(let rt):
+            authDict["token"] = rt
+        }
+
+        if let data = try? JSONSerialization.data(withJSONObject: authDict),
+           let string = String(data: data, encoding: .utf8) {
+            task.send(.string(string)) { _ in }
+        }
+
+        // Auth timeout
+        DispatchQueue.global().asyncAfter(deadline: .now() + Self.authTimeout) { [weak self] in
+            guard let self else { return }
+            self.lock.lock()
+            let currentStatus = self._status
+            self.lock.unlock()
+
+            if case .authenticating = currentStatus {
+                self.handleDisconnect()
+            }
+        }
+    }
+
+    // MARK: - Receive loop
+
+    private func receiveLoop(task: URLSessionWebSocketTask) {
+        task.receive { [weak self] result in
+            guard let self else { return }
+
+            switch result {
+            case .success(let message):
+                switch message {
+                case .string(let text):
+                    self.handleTextMessage(text)
+                case .data(let data):
+                    let delegate = self.delegate
+                    DispatchQueue.main.async { delegate?.connectionDidReceiveData(data) }
+                @unknown default:
+                    break
+                }
+                // Continue receiving
+                self.receiveLoop(task: task)
+
+            case .failure:
+                self.handleDisconnect()
+            }
+        }
+    }
+
+    private func handleTextMessage(_ text: String) {
+        // PTY data from remux server is sent as raw text (not JSON).
+        // JSON control messages start with '{'. If text doesn't start with '{',
+        // it's PTY terminal data — forward as binary data for rendering.
+        guard text.hasPrefix("{") else {
+            if let data = text.data(using: .utf8) {
+                let delegate = self.delegate
+                DispatchQueue.main.async { delegate?.connectionDidReceiveData(data) }
+            }
+            return
+        }
+
+        // Try JSON parse for control messages
+        guard let data = text.data(using: .utf8),
+              let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else {
+            // Malformed JSON — treat as PTY data
+            if let data = text.data(using: .utf8) {
+                let delegate = self.delegate
+                DispatchQueue.main.async { delegate?.connectionDidReceiveData(data) }
+            }
+            return
+        }
+
+        // Handle both envelope and legacy format
+        let msgType: String?
+        if let v = json["v"] as? Int, v >= 1 {
+            msgType = json["type"] as? String
+        } else {
+            msgType = json["type"] as? String
+        }
+
+        switch msgType {
+        case "auth_ok":
+            handleAuthOk(json)
+            return
+        case "auth_error":
+            let reason = json["reason"] as? String
+                ?? (json["payload"] as? [String: Any])?["reason"] as? String
+                ?? "Unknown error"
+            let d = self.delegate
+            DispatchQueue.main.async { d?.connectionDidFailAuth(reason: reason) }
+            return
+        case "ping":
+            resetHeartbeatTimer()
+            sendString("{\"type\":\"pong\"}")
+            return
+        default:
+            break
+        }
+
+        let delegate = self.delegate
+        DispatchQueue.main.async { delegate?.connectionDidReceiveMessage(text) }
+    }
+
+    private func handleAuthOk(_ json: [String: Any]) {
+        // Parse capabilities from envelope or legacy format
+        let capsDict: [String: Any]?
+        if let payload = json["payload"] as? [String: Any] {
+            capsDict = payload["capabilities"] as? [String: Any]
+        } else {
+            capsDict = json["capabilities"] as? [String: Any]
+        }
+
+        let capabilities = ProtocolCapabilities(
+            envelope: capsDict?["envelope"] as? Bool ?? false,
+            inspectV2: capsDict?["inspectV2"] as? Bool ?? false,
+            deviceTrust: capsDict?["deviceTrust"] as? Bool ?? false
+        )
+
+        setStatus(.connected)
+
+        lock.lock()
+        _reconnectAttempt = 0
+        let pending = _pendingMessages
+        _pendingMessages.removeAll()
+        lock.unlock()
+
+        // Flush buffered messages
+        for msg in pending {
+            sendString(msg)
+        }
+
+        startHeartbeatTimer()
+        let delegate = self.delegate
+        DispatchQueue.main.async { delegate?.connectionDidAuthenticate(capabilities: capabilities) }
+    }
+
+    // MARK: - Reconnection (exponential backoff, ref: remodex reconnection strategy)
+
+    private func handleDisconnect() {
+        lock.lock()
+        let task = _task
+        _task = nil
+        let attempt = _reconnectAttempt
+        lock.unlock()
+
+        stopHeartbeatTimer()
+        task?.cancel(with: .goingAway, reason: nil)
+
+        guard attempt < Self.maxReconnectAttempts else {
+            setStatus(.disconnected)
+            return
+        }
+
+        let nextAttempt = attempt + 1
+        lock.lock()
+        _reconnectAttempt = nextAttempt
+        lock.unlock()
+
+        setStatus(.reconnecting(attempt: nextAttempt))
+
+        // Exponential backoff: 1s, 2s, 4s, 8s, 16s, 30s max
+        let delay = min(pow(2.0, Double(attempt)), 30.0)
+        let workItem = DispatchWorkItem { [weak self] in
+            self?.startConnection()
+        }
+
+        lock.lock()
+        _reconnectTimer = workItem
+        lock.unlock()
+
+        DispatchQueue.global().asyncAfter(deadline: .now() + delay, execute: workItem)
+    }
+
+    // MARK: - Heartbeat (E07-A-06)
+
+    private func startHeartbeatTimer() {
+        stopHeartbeatTimer()
+        lock.lock()
+        _heartbeatMissed = 0
+        lock.unlock()
+
+        DispatchQueue.main.async { [weak self] in
+            guard let self else { return }
+            let timer = Timer.scheduledTimer(withTimeInterval: Self.heartbeatTimeout, repeats: false) { [weak self] _ in
+                // No ping received within timeout — force reconnect
+                self?.handleDisconnect()
+            }
+            self.lock.lock()
+            self._heartbeatTimer = timer
+            self.lock.unlock()
+        }
+    }
+
+    private func resetHeartbeatTimer() {
+        // Called when we receive a ping — restart the timeout
+        startHeartbeatTimer()
+    }
+
+    private func stopHeartbeatTimer() {
+        lock.lock()
+        let timer = _heartbeatTimer
+        _heartbeatTimer = nil
+        lock.unlock()
+        DispatchQueue.main.async {
+            timer?.invalidate()
+        }
+    }
+
+    // MARK: - Status
+
+    private func setStatus(_ status: ConnectionStatus) {
+        lock.lock()
+        _status = status
+        lock.unlock()
+        let delegate = self.delegate
+        DispatchQueue.main.async {
+            delegate?.connectionDidChangeStatus(status)
+        }
+    }
+}

package/packages/RemuxKit/Sources/RemuxKit/State/RemuxState.swift

@@ -0,0 +1,188 @@
+import Foundation
+
+/// Central state container for remux client.
+/// Uses @Observable macro (ref: remodex's CodexService pattern).
+/// All state updates happen on MainActor.
+@Observable
+@MainActor
+public final class RemuxState {
+
+    // MARK: - Connection
+
+    public private(set) var connectionStatus: ConnectionStatus = .disconnected
+    public private(set) var serverURL: URL?
+    public private(set) var capabilities: ProtocolCapabilities?
+
+    // MARK: - Workspace
+
+    public private(set) var currentSession: String = ""
+    public private(set) var tabs: [WorkspaceTab] = []
+    public private(set) var activeTabIndex: Int = 0
+
+    // MARK: - Client role
+
+    public private(set) var clientRole: String = "active" // "active" or "observer"
+
+    // MARK: - Inspect
+
+    public private(set) var inspectSnapshot: InspectSnapshot?
+
+    // MARK: - Devices
+
+    public private(set) var devices: [DeviceInfo] = []
+
+    // MARK: - Connection reference
+
+    private var connection: RemuxConnection?
+    private let router = MessageRouter()
+
+    public init() {}
+
+    // MARK: - Connection management
+
+    public func connect(url: URL, credential: RemuxCredential) {
+        serverURL = url
+        let conn = RemuxConnection(serverURL: url, credential: credential)
+        conn.delegate = self
+        connection = conn
+        conn.connect()
+    }
+
+    public func disconnect() {
+        connection?.disconnect()
+        connection = nil
+    }
+
+    // MARK: - Actions
+
+    public func switchTab(id: String) {
+        sendJSON(["type": "attach_tab", "tabId": id])
+    }
+
+    public func createTab() {
+        sendJSON(["type": "new_tab"])
+    }
+
+    public func closeTab(id: String) {
+        sendJSON(["type": "close_tab", "tabId": id])
+    }
+
+    public func renameTab(id: String, name: String) {
+        sendJSON(["type": "rename_tab", "tabId": id, "name": name])
+    }
+
+    public func createSession(name: String) {
+        sendJSON(["type": "new_session", "name": name])
+    }
+
+    public func deleteSession(name: String) {
+        sendJSON(["type": "delete_session", "name": name])
+    }
+
+    public func requestInspect(tabIndex: Int? = nil, query: String? = nil) {
+        var dict: [String: Any] = ["type": "inspect"]
+        if let idx = tabIndex { dict["tabIndex"] = idx }
+        if let q = query { dict["query"] = q }
+        if let data = try? JSONSerialization.data(withJSONObject: dict),
+           let str = String(data: data, encoding: .utf8) {
+            connection?.sendString(str)
+        }
+    }
+
+    public func requestControl() {
+        sendJSON(["type": "request_control"])
+    }
+
+    public func releaseControl() {
+        sendJSON(["type": "release_control"])
+    }
+
+    public func sendTerminalInput(_ text: String) {
+        connection?.sendString(text)
+    }
+
+    /// Safe JSON message construction — prevents injection via string interpolation
+    public func sendJSON(_ dict: [String: Any]) {
+        guard let data = try? JSONSerialization.data(withJSONObject: dict),
+              let str = String(data: data, encoding: .utf8) else { return }
+        connection?.sendString(str)
+    }
+
+    public func sendTerminalData(_ data: Data) {
+        connection?.sendData(data)
+    }
+
+    // MARK: - Internal state updates
+
+    fileprivate func updateConnectionStatus(_ status: ConnectionStatus) {
+        connectionStatus = status
+    }
+
+    fileprivate func handleAuthenticated(capabilities: ProtocolCapabilities) {
+        self.capabilities = capabilities
+        connectionStatus = .connected
+    }
+
+    fileprivate func processMessage(_ text: String) {
+        guard let routed = router.route(text) else { return }
+
+        switch routed {
+        case .state(let ws):
+            currentSession = ws.session
+            tabs = ws.tabs
+            activeTabIndex = ws.activeTabIndex
+        case .inspectResult(let snapshot):
+            inspectSnapshot = snapshot
+        case .roleChanged(let role):
+            clientRole = role
+        case .deviceList(let list):
+            devices = list
+        case .pairResult:
+            // Handled by pairing UI flow
+            break
+        case .pushStatus:
+            break
+        case .error:
+            break
+        case .unknown:
+            break
+        }
+    }
+}
+
+// MARK: - RemuxConnectionDelegate
+
+extension RemuxState: @preconcurrency RemuxConnectionDelegate {
+
+    public func connectionDidChangeStatus(_ status: ConnectionStatus) {
+        updateConnectionStatus(status)
+    }
+
+    public func connectionDidReceiveMessage(_ message: String) {
+        processMessage(message)
+    }
+
+    public func connectionDidReceiveData(_ data: Data) {
+        NotificationCenter.default.post(
+            name: .remuxTerminalData,
+            object: nil,
+            userInfo: ["data": data]
+        )
+    }
+
+    public func connectionDidAuthenticate(capabilities: ProtocolCapabilities) {
+        handleAuthenticated(capabilities: capabilities)
+        // Auto-attach to first tab to start receiving PTY data
+        connection?.sendString("{\"type\":\"attach_first\"}")
+    }
+
+    public func connectionDidFailAuth(reason: String) {
+        connectionStatus = .disconnected
+    }
+}
+
+// MARK: - Notifications
+
+public extension Notification.Name {
+    static let remuxTerminalData = Notification.Name("remuxTerminalData")
+}