fix: resume WebSocket trace from cursor instead of replaying from start

crates/cli/src/commands/serve.rs

@@ -149,10 +149,11 @@ async fn handle_ws_connection(socket: WebSocket, network: Arc<NetworkConfig>) {
             if let Ok(request) = serde_json::from_str::<TraceRequest>(&text) {
                 let (tx, mut rx) = tokio::sync::broadcast::channel::<TraceStreamMessage>(100);
                 let tx_hash = request.tx_hash.clone();
+                let resume_from = request.resume_from;
                 let network = Arc::clone(&network);
 
                 tokio::spawn(async move {
-                    let _ = stream_trace_replay(&tx_hash, &network, tx).await;
+                    let _ = stream_trace_replay(&tx_hash, &network, tx, resume_from).await;
                 });
 
                 while let Ok(update) = rx.recv().await {
@@ -174,12 +175,19 @@ fn get_static_assets_path() -> std::path::PathBuf {
 #[derive(Debug, serde::Deserialize)]
 struct TraceRequest {
     tx_hash: String,
+    /// Synchronization cursor: the index of the last trace node the client has
+    /// already received. When present, the server fast-forwards the replay and
+    /// only streams nodes that occurred after this index (used for resuming a
+    /// trace across a dropped/reconnected WebSocket instead of replaying from 0).
+    #[serde(default)]
+    resume_from: Option<usize>,
 }
 
 async fn stream_trace_replay(
     tx_hash: &str,
     network: &NetworkConfig,
     sender: tokio::sync::broadcast::Sender<TraceStreamMessage>,
+    resume_from: Option<usize>,
 ) -> anyhow::Result<()> {
     use std::time::Instant;
 
@@ -218,6 +226,17 @@ async fn stream_trace_replay(
 
     let mut node_count = 0;
     for (idx, event) in result.events.iter().enumerate() {
+        // Fast-forward: when the client resumes after a disconnect it tells us the
+        // last node index it already received. Skip re-broadcasting those nodes so
+        // the resumed stream contains only unseen trace data (no duplication, no
+        // wasted bandwidth). `node_count` still advances so totals stay accurate.
+        if let Some(cursor) = resume_from {
+            if idx <= cursor {
+                node_count += 1;
+                continue;
+            }
+        }
+
         let node_json = serde_json::to_value(event)?;
 
         let _ = sender.send(TraceStreamMessage::TraceNode {

examples/api-clients/websocket-client/index.js

@@ -12,78 +12,146 @@ if (!TX_HASH) {
 
 console.log(`Connecting to ${WS_URL}...`);
 
-const ws = new WebSocket(WS_URL);
+// Persistent synchronization cursor, declared OUTSIDE the connection scope so it
+// survives across reconnects. It tracks the index of the most recent trace node
+// the server has delivered to us. When the socket drops and we reconnect, we hand
+// this value back to the server via `resume_from` so the trace is fast-forwarded
+// to where it left off instead of replayed from the very beginning.
+let lastSeenNodeId = null;
 
+let ws = null;
 let nodeCount = 0;
-let startTime = Date.now();
+let startTime = null;
+let stopped = false; // set once the trace completes/errors so we stop reconnecting
+let reconnectAttempts = 0;
+const MAX_RECONNECT_DELAY = 5000; // cap exponential backoff at 5s
 
-ws.on('open', () => {
-  console.log('✓ Connected to Grat WebSocket server');
-  console.log(`Requesting trace for: ${TX_HASH}\n`);
-
-  ws.send(JSON.stringify({ tx_hash: TX_HASH }));
-});
+// Build the request payload, dynamically augmenting it with the resume cursor
+// once we have received at least one trace node.
+function buildRequest() {
+  const request = { tx_hash: TX_HASH };
+  if (lastSeenNodeId !== null) {
+    request.resume_from = lastSeenNodeId;
+  }
+  return request;
+}
+
+function connect() {
+  ws = new WebSocket(WS_URL);
 
-ws.on('message', (data) => {
-  try {
-    const message = JSON.parse(data.toString());
-
-    switch (message.type) {
-      case 'trace_started':
-        console.log('🚀 Trace started');
-        console.log(`   Transaction: ${message.tx_hash}`);
-        console.log(`   Ledger: ${message.ledger_sequence}\n`);
-        break;
-
-      case 'trace_node':
-        nodeCount++;
-        process.stdout.write(`\r📦 Received ${nodeCount} trace nodes...`);
-        break;
-
-      case 'resource_update':
-        const cpuPercent = (message.cpu_used / message.cpu_limit * 100).toFixed(1);
-        const memPercent = (message.memory_used / message.memory_limit * 100).toFixed(1);
-        console.log(`\n📊 Resources: CPU ${cpuPercent}%, Memory ${memPercent}%`);
-        break;
-
-      case 'state_diff_entry':
-        console.log(`\n📝 State change: ${message.key} (${message.change_type})`);
-        break;
-
-      case 'trace_completed':
-        const duration = Date.now() - startTime;
-        console.log('\n\n✅ Trace completed!');
-        console.log(`   Total nodes: ${message.total_nodes}`);
-        console.log(`   Server duration: ${message.duration_ms}ms`);
-        console.log(`   Client duration: ${duration}ms`);
-        ws.close();
-        break;
-
-      case 'trace_error':
-        console.error('\n\n❌ Trace error:', message.error);
-        ws.close();
-        process.exit(1);
-        break;
-
-      default:
-        console.log('\n⚠ Unknown message type:', message.type);
+  ws.on('open', () => {
+    reconnectAttempts = 0;
+    if (startTime === null) startTime = Date.now();
+
+    if (lastSeenNodeId === null) {
+      console.log('✓ Connected to Grat WebSocket server');
+    } else {
+      console.log(`✓ Reconnected to Grat WebSocket server (resuming from node ${lastSeenNodeId})`);
     }
-  } catch (err) {
-    console.error('Failed to parse message:', err);
-  }
-});
+    console.log(`Requesting trace for: ${TX_HASH}\n`);
 
-ws.on('error', (err) => {
-  console.error('WebSocket error:', err.message);
-  process.exit(1);
-});
+    ws.send(JSON.stringify(buildRequest()));
+  });
 
-ws.on('close', () => {
-  console.log('\nConnection closed');
-});
+  ws.on('message', (data) => {
+    try {
+      const message = JSON.parse(data.toString());
+
+      switch (message.type) {
+        case 'trace_started':
+          console.log('🚀 Trace started');
+          console.log(`   Transaction: ${message.tx_hash}`);
+          console.log(`   Ledger: ${message.ledger_sequence}\n`);
+          break;
+
+        case 'trace_node': {
+          // Defensive deduplication: if the server hasn't (yet) honored
+          // `resume_from`, or we somehow receive a node we already processed
+          // before the disconnect, skip it rather than re-processing it.
+          const nodeId =
+            Array.isArray(message.path) && message.path.length > 0
+              ? message.path[0]
+              : null;
+
+          if (
+            nodeId !== null &&
+            lastSeenNodeId !== null &&
+            nodeId <= lastSeenNodeId
+          ) {
+            break;
+          }
+
+          nodeCount++;
+          if (nodeId !== null) lastSeenNodeId = nodeId;
+          process.stdout.write(`\r📦 Received ${nodeCount} trace nodes...`);
+          break;
+        }
+
+        case 'resource_update':
+          const cpuPercent = (message.cpu_used / message.cpu_limit * 100).toFixed(1);
+          const memPercent = (message.memory_used / message.memory_limit * 100).toFixed(1);
+          console.log(`\n📊 Resources: CPU ${cpuPercent}%, Memory ${memPercent}%`);
+          break;
+
+        case 'state_diff_entry':
+          console.log(`\n📝 State change: ${message.key} (${message.change_type})`);
+          break;
+
+        case 'trace_completed':
+          const duration = Date.now() - startTime;
+          console.log('\n\n✅ Trace completed!');
+          console.log(`   Total nodes: ${message.total_nodes}`);
+          console.log(`   Server duration: ${message.duration_ms}ms`);
+          console.log(`   Client duration: ${duration}ms`);
+          stopped = true;
+          ws.close();
+          break;
+
+        case 'trace_error':
+          console.error('\n\n❌ Trace error:', message.error);
+          stopped = true;
+          ws.close();
+          process.exit(1);
+          break;
+
+        default:
+          console.log('\n⚠ Unknown message type:', message.type);
+      }
+    } catch (err) {
+      console.error('Failed to parse message:', err);
+    }
+  });
+
+  ws.on('error', (err) => {
+    console.error('WebSocket error:', err.message);
+    // Do not exit here: the 'close' handler below drives reconnection.
+  });
+
+  ws.on('close', () => {
+    console.log('\nConnection closed');
+    if (stopped) {
+      process.exit(0);
+    }
+    scheduleReconnect();
+  });
+}
+
+function scheduleReconnect() {
+  reconnectAttempts++;
+  // Exponential backoff (500ms, 1s, 2s, 4s, capped at MAX_RECONNECT_DELAY).
+  const delay = Math.min(
+    MAX_RECONNECT_DELAY,
+    500 * 2 ** (reconnectAttempts - 1)
+  );
+  console.log(`Reconnecting in ${delay}ms (attempt ${reconnectAttempts})...`);
+  setTimeout(connect, delay);
+}
 
 process.on('SIGINT', () => {
   console.log('\n\nClosing connection...');
-  ws.close();
+  stopped = true;
+  if (ws) ws.close();
   process.exit(0);
 });
+
+connect();