diff --git a/assets/init.js b/assets/init.js
index eb4ca22..12bac7e 100644
--- a/assets/init.js
+++ b/assets/init.js
@@ -478,8 +478,8 @@ class HtmlDocument {
             key: this.key,
             query: query
         })
-        if(!k) return null;
-        return new HtmlElement(k);
+        if(k == null) return null;
+        return new HtmlElement(k, this.key);
     }
 
     /**
@@ -494,7 +494,19 @@ class HtmlDocument {
             key: this.key,
             query: query
         })
-        return ks.map(k => new HtmlElement(k));
+        return ks.map(k => new HtmlElement(k, this.key));
+    }
+
+    /**
+     * Dispose the HTML document.
+     * This should be called when the document is no longer needed.
+     */
+    dispose() {
+        sendMessage({
+            method: "html",
+            function: "dispose",
+            key: this.key
+        })
     }
 }
 
@@ -504,12 +516,16 @@ class HtmlDocument {
 class HtmlElement {
     key = 0;
 
+    doc = 0;
+
     /**
      * Constructor for HtmlDom.
      * @param {number} k - The key of the element.
+     * @param {number} doc - The key of the document.
      */
-    constructor(k) {
+    constructor(k, doc) {
         this.key = k;
+        this.doc = doc;
     }
 
     /**
@@ -520,7 +536,8 @@ class HtmlElement {
         return sendMessage({
             method: "html",
             function: "getText",
-            key: this.key
+            key: this.key,
+            doc: this.doc,
         })
     }
 
@@ -532,7 +549,8 @@ class HtmlElement {
         return sendMessage({
             method: "html",
             function: "getAttributes",
-            key: this.key
+            key: this.key,
+            doc: this.doc,
         })
     }
 
@@ -546,10 +564,11 @@ class HtmlElement {
             method: "html",
             function: "dom_querySelector",
             key: this.key,
-            query: query
+            query: query,
+            doc: this.doc,
         })
-        if(!k) return null;
-        return new HtmlElement(k);
+        if(k == null) return null;
+        return new HtmlElement(k, this.doc);
     }
 
     /**
@@ -562,9 +581,10 @@ class HtmlElement {
             method: "html",
             function: "dom_querySelectorAll",
             key: this.key,
-            query: query
+            query: query,
+            doc: this.doc,
         })
-        return ks.map(k => new HtmlElement(k));
+        return ks.map(k => new HtmlElement(k, this.doc));
     }
 
     /**
@@ -575,9 +595,10 @@ class HtmlElement {
         let ks = sendMessage({
             method: "html",
             function: "getChildren",
-            key: this.key
+            key: this.key,
+            doc: this.doc,
         })
-        return ks.map(k => new HtmlElement(k));
+        return ks.map(k => new HtmlElement(k, this.doc));
     }
 
     /**
@@ -588,9 +609,10 @@ class HtmlElement {
         let ks = sendMessage({
             method: "html",
             function: "getNodes",
-            key: this.key
+            key: this.key,
+            doc: this.doc,
         })
-        return ks.map(k => new HtmlNode(k));
+        return ks.map(k => new HtmlNode(k, this.doc));
     }
 
     /**
@@ -601,7 +623,8 @@ class HtmlElement {
         return sendMessage({
             method: "html",
             function: "getInnerHTML",
-            key: this.key
+            key: this.key,
+            doc: this.doc,
         })
     }
 
@@ -613,9 +636,10 @@ class HtmlElement {
         let k = sendMessage({
             method: "html",
             function: "getParent",
-            key: this.key
+            key: this.key,
+            doc: this.doc,
         })
-        if(!k) return null;
+        if(k == null) return null;
         return new HtmlElement(k);
     }
 }
@@ -623,8 +647,11 @@ class HtmlElement {
 class HtmlNode {
     key = 0;
 
-    constructor(k) {
+    doc = 0;
+
+    constructor(k, doc) {
         this.key = k;
+        this.doc = doc;
     }
 
     /**
@@ -635,7 +662,8 @@ class HtmlNode {
         return sendMessage({
             method: "html",
             function: "node_text",
-            key: this.key
+            key: this.key,
+            doc: this.doc,
         })
     }
 
@@ -647,7 +675,8 @@ class HtmlNode {
         return sendMessage({
             method: "html",
             function: "node_type",
-            key: this.key
+            key: this.key,
+            doc: this.doc,
         })
     }
 
@@ -659,10 +688,11 @@ class HtmlNode {
         let k = sendMessage({
             method: "html",
             function: "node_toElement",
-            key: this.key
+            key: this.key,
+            doc: this.doc,
         })
-        if(!k) return null;
-        return new HtmlElement(k);
+        if(k == null) return null;
+        return new HtmlElement(k, this.doc);
     }
 }
 
diff --git a/lib/foundation/js_engine.dart b/lib/foundation/js_engine.dart
index 371b1b7..f547224 100644
--- a/lib/foundation/js_engine.dart
+++ b/lib/foundation/js_engine.dart
@@ -83,7 +83,7 @@ class JsEngine with _JSEngineApi {
     }
   }
 
-  dynamic _messageReceiver(dynamic message) {
+  Object? _messageReceiver(dynamic message) {
     try {
       if (message is Map<dynamic, dynamic>) {
         String method = message["method"] as String;
@@ -166,6 +166,7 @@ class JsEngine with _JSEngineApi {
             }
         }
       }
+      return null;
     } catch (e, s) {
       Log.error("Failed to handle message: $message\n$e\n$s", "JsEngine");
       rethrow;
@@ -224,130 +225,57 @@ class JsEngine with _JSEngineApi {
 }
 
 mixin class _JSEngineApi {
-  final Map<int, dom.Document> _documents = {};
-  final Map<int, List<int>> _documentElements = {};
-  final Map<int, List<int>> _documentNodes = {};
-  final Map<int, dom.Element> _elements = {};
-  final Map<int, dom.Node> _nodes = {};
   CookieJarSql? _cookieJar;
 
-  int _elementKey = 0;
-  int _nodeKey = 0;
+  final _documents = <int, DocumentWrapper>{};
 
-  dynamic handleHtmlCallback(Map<String, dynamic> data) {
+  Object? handleHtmlCallback(Map<String, dynamic> data) {
+    print(data);
     switch (data["function"]) {
       case "parse":
-        _documents[data["key"]] = html.parse(data["data"]);
-        _documentElements[data["key"]] = [];
-        _documentNodes[data["key"]] = [];
-        Future.delayed(const Duration(seconds: 1), () {
-          handleHtmlCallback({"function": "dispose", "key": data["key"]});
-        });
+        _documents[data["key"]] = DocumentWrapper.parse(data["data"]);
         return null;
       case "querySelector":
-        var res = _documents[data["key"]]!.querySelector(data["query"]);
-        if (res == null) return null;
-        _elements[_elementKey] = res;
-        _elementKey++;
-        _documentElements[data["key"]]!.add(_elementKey - 1);
-        return _elementKey - 1;
+        var key = data["key"];
+        return _documents[key]!.querySelector(data["query"]);
       case "querySelectorAll":
-        var res = _documents[data["key"]]!.querySelectorAll(data["query"]);
-        var keys = <int>[];
-        for (var element in res) {
-          _elements[_elementKey] = element;
-          keys.add(_elementKey);
-          _documentElements[data["key"]]!.add(_elementKey);
-          _elementKey++;
-        }
-        return keys;
+        var key = data["key"];
+        return _documents[key]!.querySelectorAll(data["query"]);
       case "getText":
-        return _elements[data["key"]]!.text;
+        return _documents[data["doc"]]!.elementGetText(data["key"]);
       case "getAttributes":
-        return _elements[data["key"]]!.attributes;
+        var res = _documents[data["doc"]]!.elementGetAttributes(data["key"]);
+        return res;
       case "dom_querySelector":
-        var res = _elements[data["key"]]!.querySelector(data["query"]);
-        if (res == null) return null;
-        _elements[_elements.length] = res;
-        var docKey = _documentElements.keys
-            .firstWhere((key) => _documentElements[key]!.contains(data["key"]));
-        _documentElements[docKey]!.add(_elements.length - 1);
-        return _elements.length - 1;
+        var doc = _documents[data["doc"]]!;
+        return doc.elementQuerySelector(data["key"], data["query"]);
       case "dom_querySelectorAll":
-        var res = _elements[data["key"]]!.querySelectorAll(data["query"]);
-        var keys = <int>[];
-        var docKey = _documentElements.keys
-            .firstWhere((key) => _documentElements[key]!.contains(data["key"]));
-        for (var element in res) {
-          _elements[_elements.length] = element;
-          keys.add(_elements.length - 1);
-          _documentElements[docKey]!.add(_elements.length - 1);
-        }
-        return keys;
+        var doc = _documents[data["doc"]]!;
+        return doc.elementQuerySelectorAll(data["key"], data["query"]);
       case "getChildren":
-        var res = _elements[data["key"]]!.children;
-        var keys = <int>[];
-        var docKey = _documentElements.keys
-            .firstWhere((key) => _documentElements[key]!.contains(data["key"]));
-        for (var element in res) {
-          _elements[_elements.length] = element;
-          keys.add(_elements.length - 1);
-          _documentElements[docKey]!.add(_elements.length - 1);
-        }
-        return keys;
+        var doc = _documents[data["doc"]]!;
+        return doc.elementGetChildren(data["key"]);
       case "getNodes":
-        var res = _elements[data["key"]]!.nodes;
-        var keys = <int>[];
-        var docKey = _documentElements.keys
-            .firstWhere((key) => _documentElements[key]!.contains(data["key"]));
-        for (var node in res) {
-          _nodes[_nodeKey] = node;
-          keys.add(_nodeKey);
-          _documentNodes[docKey]!.add(_nodeKey);
-          _nodeKey++;
-        }
-        return keys;
+        var doc = _documents[data["doc"]]!;
+        return doc.elementGetNodes(data["key"]);
       case "getInnerHTML":
-        return _elements[data["key"]]!.innerHtml;
+        var doc = _documents[data["doc"]]!;
+        return doc.elementGetInnerHTML(data["key"]);
       case "getParent":
-        var res = _elements[data["key"]]!.parent;
-        if (res == null) return null;
-        _elements[_elementKey] = res;
-        _documentElements[data["key"]]!.add(_elementKey);
-        return _elementKey++;
+        var doc = _documents[data["doc"]]!;
+        return doc.elementGetParent(data["key"]);
       case "node_text":
-        return _nodes[data["key"]]!.text;
+        return _documents[data["doc"]]!.nodeGetText(data["key"]);
       case "node_type":
-        return switch (_nodes[data["key"]]!.nodeType) {
-          dom.Node.ELEMENT_NODE => "element",
-          dom.Node.TEXT_NODE => "text",
-          dom.Node.COMMENT_NODE => "comment",
-          dom.Node.DOCUMENT_NODE => "document",
-          _ => "unknown"
-        };
+        return _documents[data["doc"]]!.nodeType(data["key"]);
       case "node_to_element":
-        var node = _nodes[data["key"]]!;
-        if (node is dom.Element) {
-          _elements[_elementKey] = node;
-          var docKey = _documentNodes.keys
-              .firstWhere((key) => _documentElements[key]!.contains(data["key"]));
-          _documentElements[docKey]!.add(_elementKey);
-          return _elementKey++;
-        }
-        return null;
+        return _documents[data["doc"]]!.nodeToElement(data["key"]);
       case "dispose":
         var docKey = data["key"];
         _documents.remove(docKey);
-        for (var elementKey in _documentElements[docKey]!) {
-          _elements.remove(elementKey);
-        }
-        _documentElements.remove(docKey);
-        for (var nodeKey in _documentNodes[docKey]!) {
-          _nodes.remove(nodeKey);
-        }
-        _documentNodes.remove(docKey);
         return null;
     }
+    return null;
   }
 
   dynamic handleCookieCallback(Map<String, dynamic> data) {
@@ -384,11 +312,7 @@ mixin class _JSEngineApi {
     }
   }
 
-  void clearHtml() {
-    _documents.clear();
-    _elements.clear();
-    _nodes.clear();
-  }
+  void clearHtml() {}
 
   void clearCookies(List<String> domains) async {
     for (var domain in domains) {
@@ -398,7 +322,7 @@ mixin class _JSEngineApi {
     }
   }
 
-  dynamic _convert(Map<String, dynamic> data) {
+  Object? _convert(Map<String, dynamic> data) {
     String type = data["type"];
     var value = data["value"];
     bool isEncode = data["isEncode"];
@@ -535,3 +459,113 @@ mixin class _JSEngineApi {
     return (min + (max - min) * math.Random().nextDouble()).toInt();
   }
 }
+
+class DocumentWrapper {
+  final dom.Document doc;
+
+  DocumentWrapper.parse(String doc) : doc = html.parse(doc);
+
+  var elements = <dom.Element>[];
+
+  var nodes = <dom.Node>[];
+
+  int? querySelector(String query) {
+    var element = doc.querySelector(query);
+    if (element == null) return null;
+    elements.add(element);
+    return elements.length - 1;
+  }
+
+  List<int> querySelectorAll(String query) {
+    var res = doc.querySelectorAll(query);
+    var keys = <int>[];
+    for (var element in res) {
+      elements.add(element);
+      keys.add(elements.length - 1);
+    }
+    return keys;
+  }
+
+  String? elementGetText(int key) {
+    return elements[key].text;
+  }
+
+  Map<String, String> elementGetAttributes(int key) {
+    return elements[key].attributes.map(
+          (key, value) => MapEntry(
+            key.toString(),
+            value,
+          ),
+        );
+  }
+
+  String? elementGetInnerHTML(int key) {
+    return elements[key].innerHtml;
+  }
+
+  int? elementGetParent(int key) {
+    var res = elements[key].parent;
+    if (res == null) return null;
+    elements.add(res);
+    return elements.length - 1;
+  }
+
+  int? elementQuerySelector(int key, String query) {
+    var res = elements[key].querySelector(query);
+    if (res == null) return null;
+    elements.add(res);
+    return elements.length - 1;
+  }
+
+  List<int> elementQuerySelectorAll(int key, String query) {
+    var res = elements[key].querySelectorAll(query);
+    var keys = <int>[];
+    for (var element in res) {
+      elements.add(element);
+      keys.add(elements.length - 1);
+    }
+    return keys;
+  }
+
+  List<int> elementGetChildren(int key) {
+    var res = elements[key].children;
+    var keys = <int>[];
+    for (var element in res) {
+      elements.add(element);
+      keys.add(elements.length - 1);
+    }
+    return keys;
+  }
+
+  List<int> elementGetNodes(int key) {
+    var res = elements[key].nodes;
+    var keys = <int>[];
+    for (var node in res) {
+      nodes.add(node);
+      keys.add(nodes.length - 1);
+    }
+    return keys;
+  }
+
+  String? nodeGetText(int key) {
+    return nodes[key].text;
+  }
+
+  String nodeType(int key) {
+    return switch (nodes[key].nodeType) {
+      dom.Node.ELEMENT_NODE => "element",
+      dom.Node.TEXT_NODE => "text",
+      dom.Node.COMMENT_NODE => "comment",
+      dom.Node.DOCUMENT_NODE => "document",
+      _ => "unknown"
+    };
+  }
+
+  int? nodeToElement(int key) {
+    if (nodes[key] is dom.Element) {
+      elements.add(nodes[key] as dom.Element);
+      return elements.length - 1;
+    }
+    return null;
+  }
+}