diff --git a/server/src/api/socket.js b/server/src/api/socket.js
index ead74bf..72f60e9 100644
--- a/server/src/api/socket.js
+++ b/server/src/api/socket.js
@@ -19,14 +19,71 @@ function removeFromSession(sessionId, event, callback) {
 }
 
 export function pushToSession(sessionId, event, data) {
-  console.log('Push:', sessionId, event)
   const callbacks = sessions.get((sessionId, event))
-  console.debug('Callbacks:', sessionId, callbacks.length)
   if (callbacks) {
     callbacks.forEach((callback) => callback(data))
   }
 }
 
+async function connectClient(socket) {
+  console.log('Connected to the client.')
+  const sessionId = socket.handshake.query.session_id
+  if (!sessionId) {
+    socket.send('Session ID not found.')
+    socket.disconnect()
+    return
+  }
+
+  const session = await Session.findOne({
+    _id: sessionId,
+    user: socket.user,
+    status: 'active',
+  })
+  if (!session) {
+    socket.send('Session not found.')
+    socket.disconnect()
+    return
+  }
+
+  const token = await genSessionJWT(session)
+
+  let container
+  let stream
+  try {
+    container = await getAndStartContainer(session.containerId)
+    stream = await attachContainer(container, { token })
+  } catch (err) {
+    socket.send(err.message)
+    socket.disconnect()
+    return
+  }
+
+  stream.socket.on('data', (chunk) => {
+    socket.emit('terminal', chunk.toString())
+  })
+
+  socket.on('message', console.log)
+
+  socket.on('terminal', function incoming(message) {
+    stream.socket.write(message)
+  })
+
+  const graphCallback = (data) => {
+    socket.emit('graph', data)
+  }
+
+  listenToSession(session.id, 'graph', graphCallback)
+
+  socket.on('disconnect', () => {
+    console.log('Disconnected from the client.')
+    stream.socket.write('exit\n')
+    removeFromSession(session.id, 'graph', graphCallback)
+    stream.destroy()
+  })
+
+  socket.send(`${container.id}\n`)
+}
+
 export default (server) => {
   const io = new Server(server)
 
@@ -43,60 +100,5 @@ export default (server) => {
     }
   })
 
-  io.on('connection', async (socket) => {
-    console.log('Connected to the client.')
-    const sessionId = socket.handshake.query.session_id
-    if (!sessionId) {
-      socket.send('Session ID not found.')
-      socket.disconnect()
-      return
-    }
-    let session
-    try {
-      session = await Session.findOne({
-        _id: sessionId,
-        user: socket.user,
-        status: 'active',
-      })
-      if (!session) {
-        socket.send('Session not found.')
-        socket.disconnect()
-        return
-      }
-    } catch (err) {
-      socket.send(err.message)
-      socket.disconnect()
-      return
-    }
-
-    const token = await genSessionJWT(session)
-
-    const container = await getAndStartContainer(session.containerId)
-    const stream = await attachContainer(container, { token })
-
-    stream.socket.on('data', (chunk) => {
-      socket.emit('terminal', chunk.toString())
-    })
-
-    socket.on('message', console.log)
-
-    socket.on('terminal', function incoming(message) {
-      stream.socket.write(message)
-    })
-
-    const graphCallback = (data) => {
-      socket.emit('graph', data)
-    }
-
-    listenToSession(session.id, 'graph', graphCallback)
-
-    socket.on('disconnect', () => {
-      console.log('Disconnected from the client.')
-      stream.socket.write('exit\n')
-      removeFromSession(session.id, 'graph', graphCallback)
-      stream.destroy()
-    })
-
-    socket.send(`${container.id}\n`)
-  })
+  io.on('connection', connectClient)
 }
diff --git a/server/src/containers/docker.js b/server/src/containers/docker.js
index 0b6d82f..74e2d86 100644
--- a/server/src/containers/docker.js
+++ b/server/src/containers/docker.js
@@ -35,16 +35,11 @@ export async function createContainer(name) {
 
 export async function getAndStartContainer(id) {
   console.log(`Getting container: ${id}`)
-  let container
-  try {
-    container = engine.getContainer(id)
-    if (!(await container.inspect()).State.Running) {
-      await container.start()
-    }
-    return container
-  } catch (error) {
-    throw new Error(`Container not found: ${id}`)
+  const container = engine.getContainer(id)
+  if (!(await container.inspect()).State.Running) {
+    await container.start()
   }
+  return container
 }
 
 export async function attachContainer(container, { token }) {