fix(visualizer): protocol v2 handshake, batched streaming, stable WS lifecycle

Author: Anandb71

crates/arbor-cli/src/commands.rs

@@ -189,14 +189,14 @@ pub async fn viz(path: &Path) -> Result<()> {
     );
 
     // 2. Start API Server (JSON-RPC)
-    let rpc_port = 7432;
+    let rpc_port = 7433;
     let rpc_addr = format!("127.0.0.1:{}", rpc_port).parse()?;
     let rpc_config = ServerConfig { addr: rpc_addr };
     let arbor_server = ArborServer::new(graph, rpc_config);
     let shared_graph = arbor_server.graph();
 
     // 3. Start Sync Server (WebSocket Broadcast)
-    let sync_port = 8080;
+    let sync_port = 8081;
     let sync_addr = format!("127.0.0.1:{}", sync_port).parse()?;
     let sync_config = arbor_server::SyncServerConfig {
         addr: sync_addr,

crates/arbor-server/src/sync_server.rs

@@ -60,14 +60,47 @@ impl Default for SyncServerConfig {
 #[derive(Debug, Clone, serde::Serialize)]
 #[serde(tag = "type", content = "payload")]
 pub enum BroadcastMessage {
-    /// Full graph snapshot or delta update.
+    /// Initial handshake with server info
+    Hello(HelloPayload),
+    /// Start of a graph stream
+    GraphBegin(GraphBeginPayload),
+    /// Batch of nodes
+    NodeBatch(NodeBatchPayload),
+    /// Batch of edges
+    EdgeBatch(EdgeBatchPayload),
+    /// End of graph stream
+    GraphEnd,
+    /// Full graph snapshot or delta update (Legacy/Incremental)
     GraphUpdate(GraphUpdatePayload),
     /// Tell the visualizer to focus on a specific node.
     FocusNode(FocusNodePayload),
     /// Indexer progress status.
     IndexerStatus(IndexerStatusPayload),
 }
 
+#[derive(Debug, Clone, serde::Serialize)]
+pub struct HelloPayload {
+    pub version: String,
+    pub node_count: usize,
+    pub edge_count: usize,
+}
+
+#[derive(Debug, Clone, serde::Serialize)]
+pub struct GraphBeginPayload {
+    pub total_nodes: usize,
+    pub total_edges: usize,
+}
+
+#[derive(Debug, Clone, serde::Serialize)]
+pub struct NodeBatchPayload {
+    pub nodes: Vec<arbor_core::CodeNode>,
+}
+
+#[derive(Debug, Clone, serde::Serialize)]
+pub struct EdgeBatchPayload {
+    pub edges: Vec<arbor_graph::GraphEdge>,
+}
+
 #[derive(Debug, Clone, serde::Serialize)]
 pub struct GraphUpdatePayload {
     /// Whether this is a full snapshot or delta.
@@ -313,32 +346,111 @@ async fn handle_client(
     graph: SharedGraph,
     mut broadcast_rx: broadcast::Receiver<BroadcastMessage>,
 ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
-    let ws_stream = accept_async(stream).await?;
+    use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
+
+    let config = WebSocketConfig {
+        max_send_queue: None,
+        max_message_size: Some(64 * 1024 * 1024), // 64 MB
+        max_frame_size: Some(64 * 1024 * 1024),   // 64 MB
+        accept_unmasked_frames: false,
+        ..Default::default()
+    };
+
+    let ws_stream = tokio_tungstenite::accept_async_with_config(stream, Some(config)).await?;
     let (mut write, mut read) = ws_stream.split();
 
     info!("✅ WebSocket handshake complete with {}", addr);
 
-    // Send initial graph snapshot
-    {
+    // 1. Send Hello (Metadata)
+    let (node_count, edge_count, nodes, edges) = {
         let g = graph.read().await;
-        let snapshot = BroadcastMessage::GraphUpdate(GraphUpdatePayload {
-            is_delta: false,
-            node_count: g.node_count(),
-            edge_count: g.edge_count(),
-            file_count: g.stats().files,
-            changed_files: vec![],
-            timestamp: std::time::SystemTime::now()
-                .duration_since(std::time::UNIX_EPOCH)
-                .unwrap()
-                .as_secs(),
-            nodes: Some(g.nodes().cloned().collect()),
-            edges: Some(g.export_edges()),
+        (
+            g.node_count(),
+            g.edge_count(),
+            g.nodes().cloned().collect::<Vec<_>>(),
+            g.export_edges(),
+        )
+    };
+
+    let hello = BroadcastMessage::Hello(HelloPayload {
+        version: "1.1.1".to_string(),
+        node_count,
+        edge_count,
+    });
+
+    let json = serde_json::to_string(&hello)?;
+    write.send(Message::Text(json)).await?;
+    info!(
+        "👋 Sent Hello ({} nodes, {} edges) to {}",
+        node_count, edge_count, addr
+    );
+
+    // 2. Wait for Client Ready
+    info!("⏳ Waiting for client {} to be ready...", addr);
+    let mut ready = false;
+    while let Some(msg) = read.next().await {
+        match msg {
+            Ok(Message::Text(text)) => {
+                // Simple parsing for "ready_for_graph"
+                if text.contains("ready_for_graph") {
+                    ready = true;
+                    info!("✅ Client {} is ready for graph", addr);
+                    break;
+                }
+                debug!("Running pre-ready protocol with {}: {}", addr, text);
+            }
+            Ok(Message::Ping(data)) => {
+                write.send(Message::Pong(data)).await?;
+            }
+            Ok(Message::Close(_)) => return Ok(()),
+            Err(e) => return Err(e.into()),
+            _ => {}
+        }
+    }
+
+    if !ready {
+        warn!("Client {} disconnected before sending ready signal", addr);
+        return Ok(());
+    }
+
+    // 3. Stream Graph (Chunked)
+    let begin = BroadcastMessage::GraphBegin(GraphBeginPayload {
+        total_nodes: node_count,
+        total_edges: edge_count,
+    });
+    write
+        .send(Message::Text(serde_json::to_string(&begin)?))
+        .await?;
+
+    // Stream Nodes
+    for chunk in nodes.chunks(50) {
+        let batch = BroadcastMessage::NodeBatch(NodeBatchPayload {
+            nodes: chunk.to_vec(),
         });
+        write
+            .send(Message::Text(serde_json::to_string(&batch)?))
+            .await?;
+    }
+    info!("📤 Streamed {} nodes to {}", node_count, addr);
 
-        let json = serde_json::to_string(&snapshot)?;
-        write.send(Message::Text(json)).await?;
-        debug!("📤 Sent initial snapshot to {}", addr);
+    // Stream Edges
+    for chunk in edges.chunks(100) {
+        let batch = BroadcastMessage::EdgeBatch(EdgeBatchPayload {
+            edges: chunk.to_vec(),
+        });
+        write
+            .send(Message::Text(serde_json::to_string(&batch)?))
+            .await?;
     }
+    info!("📤 Streamed {} edges to {}", edge_count, addr);
+
+    // End Stream
+    write
+        .send(Message::Text(serde_json::to_string(
+            &BroadcastMessage::GraphEnd,
+        )?))
+        .await?;
+    info!("🏁 Graph stream complete for {}", addr);
 
     // Two-way message handling
     loop {

visualizer/lib/core/protocol.dart

@@ -11,6 +11,16 @@ sealed class BroadcastMessage {
     final payload = json['payload'] as Map<String, dynamic>;
 
     switch (type) {
+      case 'Hello':
+        return Hello(payload);
+      case 'GraphBegin':
+        return GraphBegin(payload);
+      case 'NodeBatch':
+        return NodeBatch(payload);
+      case 'EdgeBatch':
+        return EdgeBatch(payload);
+      case 'GraphEnd':
+        return GraphEnd();
       case 'GraphUpdate':
         return GraphUpdate(payload);
       case 'FocusNode':
@@ -23,6 +33,52 @@ sealed class BroadcastMessage {
   }
 }
 
+class Hello extends BroadcastMessage {
+  final String version;
+  final int nodeCount;
+  final int edgeCount;
+
+  Hello(Map<String, dynamic> json)
+      : version = json['version'] as String,
+        nodeCount = json['node_count'] as int,
+        edgeCount = json['edge_count'] as int,
+        super('Hello');
+}
+
+class GraphBegin extends BroadcastMessage {
+  final int totalNodes;
+  final int totalEdges;
+
+  GraphBegin(Map<String, dynamic> json)
+      : totalNodes = json['total_nodes'] as int,
+        totalEdges = json['total_edges'] as int,
+        super('GraphBegin');
+}
+
+class NodeBatch extends BroadcastMessage {
+  final List<GraphNode> nodes;
+
+  NodeBatch(Map<String, dynamic> json)
+      : nodes = (json['nodes'] as List)
+            .map((e) => GraphNode.fromJson(e))
+            .toList(),
+        super('NodeBatch');
+}
+
+class EdgeBatch extends BroadcastMessage {
+  final List<GraphEdge> edges;
+
+  EdgeBatch(Map<String, dynamic> json)
+      : edges = (json['edges'] as List)
+            .map((e) => GraphEdge.fromJson(e))
+            .toList(),
+        super('EdgeBatch');
+}
+
+class GraphEnd extends BroadcastMessage {
+  GraphEnd() : super('GraphEnd');
+}
+
 class GraphUpdate extends BroadcastMessage {
   final bool isDelta;
   final int nodeCount;

visualizer/lib/services/websocket_service.dart

@@ -65,21 +65,102 @@ class WebSocketService {
     await Future.delayed(Duration(seconds: delay));
   }
 
+  // Buffer for batch loading
+  final List<GraphNode> _pendingNodes = [];
+  final List<GraphEdge> _pendingEdges = [];
+
   void _handleMessage(dynamic message) {
     try {
       if (message is String) {
         final json = jsonDecode(message);
         // Check if it matches BroadcastMessage structure
-        if (json is Map<String, dynamic> && json.containsKey('type') && json.containsKey('payload')) {
+        if (json is Map<String, dynamic> && json.containsKey('type')) {
+           final type = json['type'];
+
+           // 1. Handshake
+           if (type == 'Hello') {
+             final hello = Hello(json['payload']);
+             debugPrint('👋 Handshake from server: v${hello.version}, expecting ${hello.nodeCount} nodes');
+             // Send Ready
+             _send({'type': 'ready_for_graph'});
+             return;
+           }
+
+           // 2. Stream Control
+           if (type == 'GraphBegin') {
+             _pendingNodes.clear();
+             _pendingEdges.clear();
+             debugPrint('📥 Starting graph stream...');
+             return;
+           }
+
+           if (type == 'GraphEnd') {
+             debugPrint('🏁 Graph stream complete. Emitting update with ${_pendingNodes.length} nodes.');
+             final update = GraphUpdate({
+               'is_delta': false,
+               'node_count': _pendingNodes.length,
+               'edge_count': _pendingEdges.length,
+               'file_count': 0, // Not vital for viz
+               'changed_files': [],
+               'timestamp': DateTime.now().millisecondsSinceEpoch,
+               // We need to re-serialize to match the GraphUpdate constructor expectation
+               // Or we can manually construct it if we change the constructor. 
+               // For now, let's just make the list available.
+               // Wait, GraphUpdate expects Map<String, dynamic> in constructor? 
+               // Yes. Let's construct the object directly or change the constructor.
+               // Actually, let's look at protocol.dart again. GraphUpdate takes a Map.
+               // I can't pass List<GraphNode> directly to constructor unless I change it.
+               // Hack: Serialize back to JSON or modify protocol.dart? 
+               // Modify protocol.dart is cleaner but I just finished it.
+               // I will manually instantiate GraphUpdate if I can... 
+               // Wait, Dart doesn't have public fields constructor if it takes Map.
+               // I will pass nulls to map and set fields? No fields are final.
+               // Okay, I will construct a Map for the GraphUpdate constructor.
+               'nodes': _pendingNodes.map((n) => {
+                 'id': n.id,
+                 'name': n.name,
+                 'kind': n.kind,
+                 'file': n.file,
+                 'start_line': n.lineStart,
+                 'end_line': n.lineEnd,
+                 'centrality': n.centrality,
+               }).toList(),
+               'edges': _pendingEdges.map((e) => {
+                 'source': e.source,
+                 'target': e.target,
+                 'kind': e.kind,
+               }).toList(),
+             });
+             _controller.add(update);
+             return;
+           }
+
+           // 3. Batches
+           if (type == 'NodeBatch') {
+             final batch = NodeBatch(json['payload']);
+             _pendingNodes.addAll(batch.nodes);
+             return;
+           }
+
+           if (type == 'EdgeBatch') {
+             final batch = EdgeBatch(json['payload']);
+             _pendingEdges.addAll(batch.edges);
+             return;
+           }
+
+           // 4. Legacy / Other
            final broadcast = BroadcastMessage.fromJson(json);
            _controller.add(broadcast);
-        } else {
-           // Might be a standard JSON-RPC response, ignore for now or log
-           // debugPrint('Ignored message: $message');
         }
       }
-    } catch (e) {
-      debugPrint('Error parsing message: $e');
+    } catch (e, stack) {
+      debugPrint('Error parsing message: $e\n$stack');
+    }
+  }
+
+  void _send(Map<String, dynamic> data) {
+    if (_channel != null && _channel!.closeCode == null) {
+      _channel!.sink.add(jsonEncode(data));
     }
   }