diff --git a/robotframework-ls/src/robotframework_ls/impl/ast_utils.py b/robotframework-ls/src/robotframework_ls/impl/ast_utils.py
index 2ff0097b87..69cb117f7e 100644
--- a/robotframework-ls/src/robotframework_ls/impl/ast_utils.py
+++ b/robotframework-ls/src/robotframework_ls/impl/ast_utils.py
@@ -259,6 +259,14 @@ def _iter_nodes(node, stack=None, recursive=True):
                 stack.pop()
 
 
+def iter_all_nodes_recursive(node):
+    """
+    This function will iterate over all the nodes. Use only if there's no
+    other way to implement it as iterating over all the nodes is slow...
+    """
+    yield from _iter_nodes(node)
+
+
 def _iter_nodes_filtered_not_recursive(
     ast, accept_class: Union[Tuple[str, ...], str]
 ) -> Iterator[Tuple[list, Any]]:
diff --git a/robotframework-ls/src/robotframework_ls/impl/semantic_tokens.py b/robotframework-ls/src/robotframework_ls/impl/semantic_tokens.py
index 69e8da6bab..9b985c4e77 100644
--- a/robotframework-ls/src/robotframework_ls/impl/semantic_tokens.py
+++ b/robotframework-ls/src/robotframework_ls/impl/semantic_tokens.py
@@ -356,18 +356,17 @@ def semantic_tokens_full(context: ICompletionContext):
         ast = context.doc.get_ast()
     except:
         return []
-    return semantic_tokens_full_from_ast(ast, context.monitor)
 
-
-def semantic_tokens_full_from_ast(ast, monitor: Optional[IMonitor]):
     from robotframework_ls.impl import ast_utils
 
+    monitor = context.monitor
+
     ret: List[int] = []
     append = ret.append
 
     last_line = 0
     last_column = 0
-    for _stack, node in ast_utils._iter_nodes(ast, recursive=True):
+    for _stack, node in ast_utils.iter_all_nodes_recursive(ast):
         if monitor:
             monitor.check_cancelled()
         tokens = getattr(node, "tokens", None)
diff --git a/robotframework-ls/src/robotframework_ls/rf_interactive_integration.py b/robotframework-ls/src/robotframework_ls/rf_interactive_integration.py
index d09133b5bb..7260cfd3f9 100644
--- a/robotframework-ls/src/robotframework_ls/rf_interactive_integration.py
+++ b/robotframework-ls/src/robotframework_ls/rf_interactive_integration.py
@@ -364,6 +364,7 @@ def run():
                     prefix=code["prefix"],
                     full_code=code["full_code"],
                     indent=code["indent"],
+                    uri=uri,
                     monitor=None,
                 )
         except:
diff --git a/robotframework-ls/src/robotframework_ls/server_api/client.py b/robotframework-ls/src/robotframework_ls/server_api/client.py
index 8d0d15f821..6dce41601c 100644
--- a/robotframework-ls/src/robotframework_ls/server_api/client.py
+++ b/robotframework-ls/src/robotframework_ls/server_api/client.py
@@ -133,7 +133,7 @@ def request_semantic_tokens_full(
         )
 
     def request_semantic_tokens_from_code_full(
-        self, prefix: str, full_code: str, indent: str
+        self, prefix: str, full_code: str, indent: str, uri: str
     ) -> Optional[IIdMessageMatcher]:
         """
         :Note: async complete.
@@ -144,6 +144,7 @@ def request_semantic_tokens_from_code_full(
                 prefix=prefix,
                 full_code=full_code,
                 indent=indent,
+                uri=uri,
             )
         )
 
diff --git a/robotframework-ls/src/robotframework_ls/server_api/server.py b/robotframework-ls/src/robotframework_ls/server_api/server.py
index 9555fe1c30..15774e7fd6 100644
--- a/robotframework-ls/src/robotframework_ls/server_api/server.py
+++ b/robotframework-ls/src/robotframework_ls/server_api/server.py
@@ -895,13 +895,14 @@ def threaded_monaco_completions_from_code_full(
         }
 
     def m_semantic_tokens_from_code_full(
-        self, prefix: str = "", full_code: str = "", indent: str = ""
+        self, prefix: str = "", full_code: str = "", indent: str = "", uri: str = ""
     ):
         func = partial(
             self.threaded_semantic_tokens_from_code_full,
             prefix=prefix,
             full_code=full_code,
             indent=indent,
+            uri=uri,
         )
         func = require_monitor(func)
         return func
@@ -911,17 +912,24 @@ def threaded_semantic_tokens_from_code_full(
         prefix: str,
         full_code: str,
         indent: str,
+        uri: str,
         monitor: Optional[IMonitor] = None,
     ):
-        from robotframework_ls.impl.semantic_tokens import semantic_tokens_full_from_ast
+        from robotframework_ls.impl.semantic_tokens import semantic_tokens_full
+        from robotframework_ls.impl.completion_context import CompletionContext
 
         try:
             from robotframework_ls.impl.robot_workspace import RobotDocument
 
-            doc = RobotDocument("")
-            doc.source = full_code
-            ast = doc.get_ast()
-            data = semantic_tokens_full_from_ast(ast, monitor)
+            document = RobotDocument(uri, full_code)
+            completion_context = CompletionContext(
+                document,
+                config=self.config,
+                monitor=monitor,
+                workspace=self.workspace,
+            )
+
+            data = semantic_tokens_full(completion_context)
             if not prefix:
                 return {"resultId": None, "data": data}
 
diff --git a/robotframework-ls/tests/robotframework_ls_tests/test_ast_utils.py b/robotframework-ls/tests/robotframework_ls_tests/test_ast_utils.py
index 6227f1ed98..a170b04fbc 100644
--- a/robotframework-ls/tests/robotframework_ls_tests/test_ast_utils.py
+++ b/robotframework-ls/tests/robotframework_ls_tests/test_ast_utils.py
@@ -6,7 +6,7 @@ def test_iter_nodes():
         "unused", source="*** settings ***\nResource    my_resource.resource"
     )
     lst = []
-    for stack, node in ast_utils._iter_nodes(doc.get_ast()):
+    for stack, node in ast_utils.iter_all_nodes_recursive(doc.get_ast()):
         lst.append(
             "%s - %s" % ([s.__class__.__name__ for s in stack], node.__class__.__name__)
         )