diff --git a/abp/static/scripts/websocket.js b/abp/static/scripts/websocket.js
index 91c8825..919b9d2 100644
--- a/abp/static/scripts/websocket.js
+++ b/abp/static/scripts/websocket.js
@@ -1,7 +1,11 @@
 var websocket = {};
+websocket.update = undefined;
 
 websocket.connect = function(update) {
     websocket.ws = new WebSocket("ws://localhost:5000");
+    if (update){
+        websocket.update = update;
+    }
     websocket.ws.onopen = function(evt) {
         gui.serverMessage("Connected to server.");
     };
@@ -19,7 +23,7 @@ websocket.connect = function(update) {
                 json.node[i].vop = 0;
             }
         }
-        update(json);
+        websocket.update(json);
     };
 
     websocket.ws.onclose = function(evt) {
diff --git a/bin/abpserver b/bin/abpserver
index f50419d..585cd25 100755
--- a/bin/abpserver
+++ b/bin/abpserver
@@ -14,6 +14,7 @@ import argparse
 import abp
 import json
 from pkg_resources import resource_filename
+import time
 
 clients = []
 local_state = abp.GraphState()