WebSocket url is now passed to browser when a kernel is started.

IPython/frontend/html/notebook/handlers.py

@@ -61,8 +61,10 @@ class MainKernelHandler(web.RequestHandler):
         km = self.application.kernel_manager
         notebook_id = self.get_argument('notebook', default=None)
         kernel_id = km.start_kernel(notebook_id)
+        ws_url = self.application.ipython_app.get_ws_url()
+        data = {'ws_url':ws_url,'kernel_id':kernel_id}
         self.set_header('Location', '/'+kernel_id)
-        self.finish(jsonapi.dumps(kernel_id))
+        self.finish(jsonapi.dumps(data))
 
 
 class KernelHandler(web.RequestHandler):
@@ -85,7 +87,10 @@ class KernelActionHandler(web.RequestHandler):
             self.set_status(204)
         if action == 'restart':
             new_kernel_id = km.restart_kernel(kernel_id)
-            self.write(jsonapi.dumps(new_kernel_id))
+            ws_url = self.application.ipython_app.get_ws_url()
+            data = {'ws_url':ws_url,'kernel_id':new_kernel_id}
+            self.set_header('Location', '/'+new_kernel_id)
+            self.write(jsonapi.dumps(data))
         self.finish()
 
 
@@ -126,21 +131,36 @@ class IOPubHandler(ZMQStreamHandler):
     def initialize(self, *args, **kwargs):
         self._kernel_alive = True
         self._beating = False
+        self.iopub_stream = None
+        self.hb_stream = None
 
     def open(self, kernel_id):
         km = self.application.kernel_manager
         self.kernel_id = kernel_id
         self.session = Session()
         self.time_to_dead = km.time_to_dead
-        self.iopub_stream = km.create_iopub_stream(kernel_id)
-        self.hb_stream = km.create_hb_stream(kernel_id)
-        self.iopub_stream.on_recv(self._on_zmq_reply)
-        self.start_hb(self.kernel_died)
+        try:
+            self.iopub_stream = km.create_iopub_stream(kernel_id)
+            self.hb_stream = km.create_hb_stream(kernel_id)
+        except web.HTTPError:
+            # WebSockets don't response to traditional error codes so we
+            # close the connection.
+            if not self.stream.closed():
+                self.stream.close()
+        else:
+            self.iopub_stream.on_recv(self._on_zmq_reply)
+            self.start_hb(self.kernel_died)
 
     def on_close(self):
+        # This method can be called twice, once by self.kernel_died and once 
+        # from the WebSocket close event. If the WebSocket connection is
+        # closed before the ZMQ streams are setup, they could be None.
         self.stop_hb()
-        self.iopub_stream.close()
-        self.hb_stream.close()
+        if self.iopub_stream is not None and not self.iopub_stream.closed():
+            self.iopub_stream.on_recv(None)
+            self.iopub_stream.close()
+        if self.hb_stream is not None and not self.hb_stream.closed():
+            self.hb_stream.close()
  
     def start_hb(self, callback):
         """Start the heartbeating and call the callback if the kernel dies."""
@@ -188,15 +208,22 @@ class IOPubHandler(ZMQStreamHandler):
 class ShellHandler(ZMQStreamHandler):
 
     def initialize(self, *args, **kwargs):
-        pass
+        self.shell_stream = None
 
     def open(self, kernel_id):
         km = self.application.kernel_manager
         self.max_msg_size = km.max_msg_size
         self.kernel_id = kernel_id
-        self.session = Session()
-        self.shell_stream = self.application.kernel_manager.create_shell_stream(kernel_id)
-        self.shell_stream.on_recv(self._on_zmq_reply)
+        try:
+            self.shell_stream = km.create_shell_stream(kernel_id)
+        except web.HTTPError:
+            # WebSockets don't response to traditional error codes so we
+            # close the connection.
+            if not self.stream.closed():
+                self.stream.close()
+        else:
+            self.session = Session()
+            self.shell_stream.on_recv(self._on_zmq_reply)
 
     def on_message(self, msg):
         if len(msg) < self.max_msg_size:
@@ -204,7 +231,9 @@ class ShellHandler(ZMQStreamHandler):
             self.session.send(self.shell_stream, msg)
 
     def on_close(self):
-        self.shell_stream.close()
+        # Make sure the stream exists and is not already closed.
+        if self.shell_stream is not None and not self.shell_stream.closed():
+            self.shell_stream.close()
 
 
 #-----------------------------------------------------------------------------

IPython/frontend/html/notebook/kernelmanager.py

@@ -300,7 +300,21 @@ class MappingKernelManager(KernelManager):
         # Now save the new kernel/notebook association. We have to save it
         # after the old kernel is killed as that will delete the mapping.
         self.set_kernel_for_notebook(notebook_id, new_kernel_id)
-        self.log.debug("Kernel restarted: %s" % new_kernel_id)
+        self.log.info("Kernel restarted: %s" % new_kernel_id)
         return new_kernel_id
 
+    def create_iopub_stream(self, kernel_id):
+        if kernel_id not in self:
+            raise web.HTTPError(404)
+        return super(MappingKernelManager, self).create_iopub_stream(kernel_id)
+
+    def create_shell_stream(self, kernel_id):
+        if kernel_id not in self:
+            raise web.HTTPError(404)
+        return super(MappingKernelManager, self).create_shell_stream(kernel_id)
+
+    def create_hb_stream(self, kernel_id):
+        if kernel_id not in self:
+            raise web.HTTPError(404)
+        return super(MappingKernelManager, self).create_hb_stream(kernel_id)
 

IPython/frontend/html/notebook/notebookapp.py

@@ -72,7 +72,7 @@ ipython notebook --port=5555 --ip=*    # Listen on port 5555, all interfaces
 
 class NotebookWebApplication(web.Application):
 
-    def __init__(self, kernel_manager, notebook_manager, log):
+    def __init__(self, ipython_app, kernel_manager, notebook_manager, log):
         handlers = [
             (r"/", NBBrowserHandler),
             (r"/new", NewHandler),
@@ -95,6 +95,7 @@ class NotebookWebApplication(web.Application):
         self.kernel_manager = kernel_manager
         self.log = log
         self.notebook_manager = notebook_manager
+        self.ipython_app = ipython_app
 
 
 #-----------------------------------------------------------------------------
@@ -115,6 +116,7 @@ aliases.update({
     'port': 'IPythonNotebookApp.port',
     'keyfile': 'IPythonNotebookApp.keyfile',
     'certfile': 'IPythonNotebookApp.certfile',
+    'ws-hostname': 'IPythonNotebookApp.ws_hostname',
     'notebook-dir': 'NotebookManager.notebook_dir'
 })
 
@@ -160,6 +162,13 @@ class IPythonNotebookApp(BaseIPythonApplication):
         help="The port the notebook server will listen on."
     )
 
+    ws_hostname = Unicode(LOCALHOST, config=True,
+        help="""The FQDN or IP for WebSocket connections. The default will work
+                fine when the server is listening on localhost, but this needs to
+                be set if the ip option is used. It will be used as the hostname part
+                of the WebSocket url: ws://hostname/path."""
+    )
+
     certfile = Unicode(u'', config=True, 
         help="""The full path to an SSL/TLS certificate file."""
     )
@@ -168,6 +177,14 @@ class IPythonNotebookApp(BaseIPythonApplication):
         help="""The full path to a private key file for usage with SSL/TLS."""
     )
 
+    def get_ws_url(self):
+        """Return the WebSocket URL for this server."""
+        if self.certfile:
+            prefix = u'wss://'
+        else:
+            prefix = u'ws://'
+        return prefix + self.ws_hostname + u':' + unicode(self.port)
+
     def parse_command_line(self, argv=None):
         super(IPythonNotebookApp, self).parse_command_line(argv)
         if argv is None:
@@ -202,7 +219,7 @@ class IPythonNotebookApp(BaseIPythonApplication):
         super(IPythonNotebookApp, self).initialize(argv)
         self.init_configurables()
         self.web_app = NotebookWebApplication(
-            self.kernel_manager, self.notebook_manager, self.log
+            self, self.kernel_manager, self.notebook_manager, self.log
         )
         if self.certfile:
             ssl_options = dict(certfile=self.certfile)

IPython/frontend/html/notebook/static/js/kernel.js

@@ -61,9 +61,10 @@ var IPython = (function (IPython) {
     };
 
 
-    Kernel.prototype._handle_start_kernel = function (kernel_id, callback) {
+    Kernel.prototype._handle_start_kernel = function (json, callback) {
         this.running = true;
-        this.kernel_id = kernel_id;
+        this.kernel_id = json.kernel_id;
+        this.ws_url = json.ws_url;
         this.kernel_url = this.base_url + "/" + this.kernel_id;
         this.start_channels();
         callback();
@@ -73,7 +74,8 @@ var IPython = (function (IPython) {
 
     Kernel.prototype.start_channels = function () {
         this.stop_channels();
-        var ws_url = "ws://127.0.0.1:8888" + this.kernel_url;
+        var ws_url = this.ws_url + this.kernel_url;
+        console.log("Starting WS:", ws_url);
         this.shell_channel = new WebSocket(ws_url + "/shell");
         this.iopub_channel = new WebSocket(ws_url + "/iopub");
     };