diff --git a/README.md b/README.md
index b9471e1..0ca344c 100644
--- a/README.md
+++ b/README.md
@@ -37,7 +37,7 @@ require("lazy").setup({
         'brianhuster/live-preview.nvim',
         dependencies = {'brianhuster/autosave.nvim'}, -- Not required, but recomended for autosaving
         opts = {},
-    }
+   }
 })
 ```
 
diff --git a/doc/livepreview.txt b/doc/livepreview.txt
index d259a12..de6b5dc 100644
--- a/doc/livepreview.txt
+++ b/doc/livepreview.txt
@@ -48,8 +48,8 @@ Server class for live-preview.nvim To call this class, do >lua
       • {watch_dir}  (`fun(self: Server, func: function)`) Watch a directory
                      for changes and send a message "reload" to a WebSocket
                      client
-      • {start}      (`fun(self: Server, ip: string, port: number)`) Start the
-                     server
+      • {start}      (`fun(self: Server, ip: string, port: number, func: function, client: uv_tcp_t)`)
+                     Start the server
       • {stop}       (`fun(self: Server)`) Stop the server
 
 
@@ -68,12 +68,15 @@ Server:routes({path})                     *livepreview.server.Server:routes()*
     Return: ~
         (`string`) path to the file
 
-Server:start({ip}, {port})                 *livepreview.server.Server:start()*
+                                           *livepreview.server.Server:start()*
+Server:start({ip}, {port}, {func}, {client})
     Start the server
 
     Parameters: ~
-      • {ip}    (`string`) IP address to bind to
-      • {port}  (`number`) port to bind to
+      • {ip}      (`string`) IP address to bind to
+      • {port}    (`number`) port to bind to
+      • {func}    (`function`)
+      • {client}  (`uv_tcp_t`) uv_tcp client
 
 Server:stop()                               *livepreview.server.Server:stop()*
     Stop the server
diff --git a/lua/livepreview/init.lua b/lua/livepreview/init.lua
index 5a4f7c0..b3be806 100644
--- a/lua/livepreview/init.lua
+++ b/lua/livepreview/init.lua
@@ -49,7 +49,14 @@ function M.preview_file(filepath, port)
 	M.utils.kill_port(port)
 	server = M.server.Server:new(vim.fs.dirname(filepath))
 	vim.wait(50, function()
-		server:start("127.0.0.1", port)
+		server:start("127.0.0.1", port, function(client)
+			if M.utils.supported_filetype(filepath) == 'html' then
+				M.server.websocket.send_json(client, { type = "reload" })
+			else
+				local content = M.utils.uv_read_file(filepath)
+				M.server.websocket.send_json(client, { type = "update", content = content })
+			end
+		end)
 	end)
 end
 
diff --git a/lua/livepreview/server/Server.lua b/lua/livepreview/server/Server.lua
index b8bb233..1d06537 100644
--- a/lua/livepreview/server/Server.lua
+++ b/lua/livepreview/server/Server.lua
@@ -76,7 +76,6 @@ function Server:routes(path)
 		file_path = vim.fs.joinpath(self.webroot, path)
 	end
 
-
 	return file_path
 end
 
@@ -96,7 +95,9 @@ end
 --- Start the server
 --- @param ip string: IP address to bind to
 --- @param port number: port to bind to
-function Server:start(ip, port)
+--- @param func function(client)|nil: function to call when when there is a change in the watched directory
+--- 	@param client uv_tcp_t: uv_tcp client
+function Server:start(ip, port, func)
 	self.server:bind(ip, port)
 	self.server:listen(128, function(err)
 		if err then
@@ -124,11 +125,8 @@ function Server:start(ip, port)
 		end)
 		ws_client = client
 		self:watch_dir(function()
-			if supported_filetype(filepath) == 'html' then
-				websocket.send_json(client, { type = "reload" })
-			else
-				local content = utils.uv_read_file(filepath)
-				websocket.send_json(client, { type = "update", content = content })
+			if func then
+				func(client)
 			end
 		end)
 	end)
diff --git a/lua/livepreview/server/init.lua b/lua/livepreview/server/init.lua
index 27c8630..0b16244 100644
--- a/lua/livepreview/server/init.lua
+++ b/lua/livepreview/server/init.lua
@@ -4,4 +4,5 @@ local M = {}
 M.handler = require('livepreview.server.handler')
 M.utils = require('livepreview.server.utils')
 M.Server = require('livepreview.server.Server')
+M.websocket = require('livepreview.server.websocket')
 return M
diff --git a/static/ws-client.js b/static/ws-client.js
index d31ee7d..c81ffa2 100644
--- a/static/ws-client.js
+++ b/static/ws-client.js
@@ -22,7 +22,6 @@ async function connectWebSocket() {
 	};
 
 	socket.onmessage = (event) => {
-		console.log("Message received: ", event.data);
 		const message = JSON.parse(event.data);
 
 		if (message.type === "reload") {