diff --git a/server/src/api/controllers/commandController.js b/server/src/api/controllers/commandController.js
index def7b1b..fd583e1 100644
--- a/server/src/api/controllers/commandController.js
+++ b/server/src/api/controllers/commandController.js
@@ -7,7 +7,6 @@ import CommandHandler from '../../game/commandHandler.js'
 export async function newCommand(req, res) {
   console.log('new command:', req.body)
   const { token, command, pwd, output, error, ...additionalData } = req.body
-  console.log('additionalData:', additionalData)
 
   if (!command) {
     res.status(400).json({ message: 'command is required' })
@@ -50,7 +49,6 @@ export async function newCommand(req, res) {
   const commandHandler = new CommandHandler(session, c, additionalData)
 
   const response = await commandHandler.run()
-  console.log(response)
 
   res.status(201).json(response)
   await session.save()
diff --git a/server/src/api/socket.js b/server/src/api/socket.js
index 25d26d6..ead74bf 100644
--- a/server/src/api/socket.js
+++ b/server/src/api/socket.js
@@ -10,9 +10,18 @@ function listenToSession(sessionId, event, callback) {
   sessions.set((sessionId, event), callbacks.concat(callback))
 }
 
+function removeFromSession(sessionId, event, callback) {
+  const callbacks = sessions.get((sessionId, event)) || []
+  sessions.set(
+    (sessionId, event),
+    callbacks.filter((cb) => cb !== 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))
   }
@@ -65,22 +74,27 @@ export default (server) => {
     const container = await getAndStartContainer(session.containerId)
     const stream = await attachContainer(container, { token })
 
-    stream.on('data', (chunk) => {
+    stream.socket.on('data', (chunk) => {
       socket.emit('terminal', chunk.toString())
     })
 
     socket.on('message', console.log)
 
     socket.on('terminal', function incoming(message) {
-      stream.write(message)
+      stream.socket.write(message)
     })
 
-    socket.on('close', () => {
-      console.log('Disconnected from the client.')
-    })
-
-    listenToSession(session.id, 'graph', (data) => {
+    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`)
diff --git a/server/src/containers/docker.js b/server/src/containers/docker.js
index b687897..0b6d82f 100644
--- a/server/src/containers/docker.js
+++ b/server/src/containers/docker.js
@@ -63,7 +63,7 @@ export async function attachContainer(container, { token }) {
     Tty: true,
   })
 
-  return execOutput.socket
+  return execOutput
 }
 
 export async function deleteContainer(id) {
diff --git a/server/src/game/commandHandler.js b/server/src/game/commandHandler.js
index e1aa818..c5e6555 100644
--- a/server/src/game/commandHandler.js
+++ b/server/src/game/commandHandler.js
@@ -94,8 +94,6 @@ export default class CommandHandler {
       return {}
     }
 
-    console.log('stages', stages)
-
     this.handleEvent()
     this.handleCommand()