diff --git a/src/analysis.zig b/src/analysis.zig
index 51b7521..f4adf86 100644
--- a/src/analysis.zig
+++ b/src/analysis.zig
@@ -1510,6 +1510,7 @@ pub const DocumentScope = struct {
         for (self.scopes) |scope| {
             scope.decls.deinit();
             allocator.free(scope.uses);
+            allocator.free(scope.tests);
         }
         allocator.free(self.scopes);
     }
diff --git a/src/document_store.zig b/src/document_store.zig
index ee056e8..88d7058 100644
--- a/src/document_store.zig
+++ b/src/document_store.zig
@@ -629,6 +629,7 @@ fn stdUriFromLibPath(allocator: *std.mem.Allocator, zig_lib_path: ?[]const u8) !
 pub fn deinit(self: *DocumentStore) void {
     var entry_iterator = self.handles.iterator();
     while (entry_iterator.next()) |entry| {
+        entry.value.document_scope.deinit(self.allocator);
         self.allocator.free(entry.value.document.mem);
 
         for (entry.value.import_uris.items) |uri| {
@@ -638,8 +639,6 @@ pub fn deinit(self: *DocumentStore) void {
         entry.value.import_uris.deinit();
         self.allocator.free(entry.key);
         self.allocator.destroy(entry.value);
-    
-        entry.value.document_scope.deinit(self.allocator);
     }
 
     self.handles.deinit();
diff --git a/src/main.zig b/src/main.zig
index f9bf5ca..587a3db 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -839,7 +839,9 @@ fn configFromUriOr(uri: []const u8, default: Config) Config {
     return default;
 }
 
-fn processJsonRpc(parser: *std.json.Parser, json: []const u8, config: Config) !void {
+// TODO Rewrite this, use a ComptimeStringMap that points to a fn pointer + Param type to decode into and pass to the function
+//      Split into multiple files?
+fn processJsonRpc(parser: *std.json.Parser, json: []const u8, config: Config, keep_running: *bool) !void {
     var tree = try parser.parse(json);
     defer tree.deinit();
 
@@ -920,6 +922,10 @@ fn processJsonRpc(parser: *std.json.Parser, json: []const u8, config: Config) !v
 
         std.debug.warn("{}\n", .{client_capabilities});
         try respondGeneric(id, initialize_response);
+    } else if (std.mem.eql(u8, method, "shutdown")) {
+        keep_running.* = false;
+        // Technically we shoudl deinitialize first and send possible errors to the client
+        try respondGeneric(id, null_result_response);
     } else if (std.mem.eql(u8, method, "initialized")) {
         // All gucci
     } else if (std.mem.eql(u8, method, "$/cancelRequest")) {
@@ -1240,6 +1246,11 @@ pub fn main() anyerror!void {
         allocator = &debug_alloc_state.allocator;
     }
 
+    defer if (debug_alloc) |dbg| {
+        std.debug.warn("Finished cleanup, last allocation info.\n", .{});
+        std.debug.warn("{}\n", .{dbg.info});
+    };
+
     // Init global vars
     const in_stream = std.io.getStdIn().inStream();
     stdout = std.io.bufferedOutStream(std.io.getStdOut().outStream());
@@ -1271,7 +1282,6 @@ pub fn main() anyerror!void {
 
     // Find the zig executable in PATH
     var zig_exe_path: ?[]const u8 = null;
-    defer if (zig_exe_path) |exe_path| allocator.free(exe_path);
 
     find_zig: {
         if (config.zig_exe_path) |exe_path| {
@@ -1343,7 +1353,8 @@ pub fn main() anyerror!void {
     var json_parser = std.json.Parser.init(allocator, false);
     defer json_parser.deinit();
 
-    while (true) {
+    var keep_running = true;
+    while (keep_running) {
         const headers = readRequestHeader(allocator, in_stream) catch |err| {
             std.debug.warn("{}; exiting!", .{@errorName(err)});
             return;
@@ -1352,7 +1363,7 @@ pub fn main() anyerror!void {
         const buf = try allocator.alloc(u8, headers.content_length);
         defer allocator.free(buf);
         try in_stream.readNoEof(buf);
-        try processJsonRpc(&json_parser, buf, config);
+        try processJsonRpc(&json_parser, buf, config, &keep_running);
         json_parser.reset();
 
         if (debug_alloc) |dbg| {