Fix chat send button behavior

Rosetta/Core/Network/Protocol/ProtocolManager.swift

@@ -43,9 +43,13 @@ final class ProtocolManager: @unchecked Sendable {
 
     private let client = WebSocketClient()
     private var packetQueue: [any Packet] = []
+    private var queuedPacketKeys: Set<String> = []
     private var handshakeComplete = false
     private var heartbeatTask: Task<Void, Never>?
     private var handshakeTimeoutTask: Task<Void, Never>?
+    private let searchHandlersLock = NSLock()
+    private let packetQueueLock = NSLock()
+    private var searchResultHandlers: [UUID: (PacketSearch) -> Void] = [:]
 
     // Saved credentials for auto-reconnect
     private var savedPublicKey: String?
@@ -88,14 +92,30 @@ final class ProtocolManager: @unchecked Sendable {
     // MARK: - Sending
 
     func sendPacket(_ packet: any Packet) {
-        if !handshakeComplete && !(packet is PacketHandshake) {
-            Self.logger.info("Queueing packet \(type(of: packet).packetId)")
-            packetQueue.append(packet)
+        if (!handshakeComplete && !(packet is PacketHandshake)) || !client.isConnected {
+            enqueuePacket(packet)
             return
         }
         sendPacketDirect(packet)
     }
 
+    // MARK: - Search Handlers (Android-like wait/unwait)
+
+    @discardableResult
+    func addSearchResultHandler(_ handler: @escaping (PacketSearch) -> Void) -> UUID {
+        let id = UUID()
+        searchHandlersLock.lock()
+        searchResultHandlers[id] = handler
+        searchHandlersLock.unlock()
+        return id
+    }
+
+    func removeSearchResultHandler(_ id: UUID) {
+        searchHandlersLock.lock()
+        searchResultHandlers.removeValue(forKey: id)
+        searchHandlersLock.unlock()
+    }
+
     // MARK: - Private Setup
 
     private func setupClientCallbacks() {
@@ -212,6 +232,7 @@ final class ProtocolManager: @unchecked Sendable {
             if let p = packet as? PacketSearch {
                 Self.logger.debug("📥 Search result: \(p.users.count) users, callback=\(self.onSearchResult != nil)")
                 onSearchResult?(p)
+                notifySearchResultHandlers(p)
             }
         case 0x05:
             if let p = packet as? PacketOnlineState {
@@ -242,14 +263,23 @@ final class ProtocolManager: @unchecked Sendable {
         }
     }
 
+    private func notifySearchResultHandlers(_ packet: PacketSearch) {
+        searchHandlersLock.lock()
+        let handlers = Array(searchResultHandlers.values)
+        searchHandlersLock.unlock()
+
+        for handler in handlers {
+            handler(packet)
+        }
+    }
+
     private func handleHandshakeResponse(_ packet: PacketHandshake) {
-        // Set handshakeComplete BEFORE cancelling timeout to prevent race
-        handshakeComplete = true
         handshakeTimeoutTask?.cancel()
         handshakeTimeoutTask = nil
 
         switch packet.handshakeState {
         case .completed:
+            handshakeComplete = true
             Self.logger.info("Handshake completed. Protocol v\(packet.protocolVersion), heartbeat \(packet.heartbeatInterval)s")
 
             Task { @MainActor in
@@ -261,7 +291,9 @@ final class ProtocolManager: @unchecked Sendable {
             onHandshakeCompleted?(packet)
 
         case .needDeviceVerification:
+            handshakeComplete = false
             Self.logger.info("Server requires device verification")
+            clearPacketQueue()
             startHeartbeat(interval: packet.heartbeatInterval)
         }
     }
@@ -289,15 +321,71 @@ final class ProtocolManager: @unchecked Sendable {
     private func sendPacketDirect(_ packet: any Packet) {
         let data = PacketRegistry.encode(packet)
         Self.logger.info("Sending packet 0x\(String(type(of: packet).packetId, radix: 16)) (\(data.count) bytes)")
-        client.send(data)
+        if !client.send(data, onFailure: { [weak self] _ in
+            guard let self else { return }
+            Self.logger.warning("Send failed, re-queueing packet 0x\(String(type(of: packet).packetId, radix: 16))")
+            self.enqueuePacket(packet)
+        }) {
+            Self.logger.warning("WebSocket unavailable, re-queueing packet 0x\(String(type(of: packet).packetId, radix: 16))")
+            enqueuePacket(packet)
+        }
     }
 
     private func flushPacketQueue() {
-        Self.logger.info("Flushing \(self.packetQueue.count) queued packets")
-        let packets = packetQueue
-        packetQueue.removeAll()
+        let packets = drainPacketQueue()
+        Self.logger.info("Flushing \(packets.count) queued packets")
         for packet in packets {
             sendPacketDirect(packet)
         }
     }
+
+    private func enqueuePacket(_ packet: any Packet) {
+        packetQueueLock.lock()
+        if let key = packetQueueKey(packet), queuedPacketKeys.contains(key) {
+            packetQueueLock.unlock()
+            return
+        }
+        packetQueue.append(packet)
+        if let key = packetQueueKey(packet) {
+            queuedPacketKeys.insert(key)
+        }
+        let count = packetQueue.count
+        packetQueueLock.unlock()
+        Self.logger.info("Queueing packet 0x\(String(type(of: packet).packetId, radix: 16)) (queue=\(count))")
+    }
+
+    private func drainPacketQueue() -> [any Packet] {
+        packetQueueLock.lock()
+        let packets = packetQueue
+        packetQueue.removeAll()
+        queuedPacketKeys.removeAll()
+        packetQueueLock.unlock()
+        return packets
+    }
+
+    private func clearPacketQueue() {
+        packetQueueLock.lock()
+        packetQueue.removeAll()
+        queuedPacketKeys.removeAll()
+        packetQueueLock.unlock()
+    }
+
+    private func packetQueueKey(_ packet: any Packet) -> String? {
+        switch packet {
+        case let message as PacketMessage:
+            return message.messageId.isEmpty ? nil : "0x06:\(message.messageId)"
+        case let delivery as PacketDelivery:
+            return delivery.messageId.isEmpty ? nil : "0x08:\(delivery.toPublicKey):\(delivery.messageId)"
+        case let read as PacketRead:
+            guard !read.fromPublicKey.isEmpty, !read.toPublicKey.isEmpty else { return nil }
+            return "0x07:\(read.fromPublicKey):\(read.toPublicKey)"
+        case let typing as PacketTyping:
+            guard !typing.fromPublicKey.isEmpty, !typing.toPublicKey.isEmpty else { return nil }
+            return "0x0b:\(typing.fromPublicKey):\(typing.toPublicKey)"
+        case let sync as PacketSync:
+            return "0x19:\(sync.status.rawValue):\(sync.timestamp)"
+        default:
+            return nil
+        }
+    }
 }