diff --git a/src/main/kotlin/platform/mcp/at/format/AtBlock.kt b/src/main/kotlin/platform/mcp/at/format/AtBlock.kt
index e0d4e76e6..8fc9748be 100644
--- a/src/main/kotlin/platform/mcp/at/format/AtBlock.kt
+++ b/src/main/kotlin/platform/mcp/at/format/AtBlock.kt
@@ -21,6 +21,7 @@
 package com.demonwav.mcdev.platform.mcp.at.format
 
 import com.demonwav.mcdev.platform.mcp.at.gen.psi.AtTypes
+import com.demonwav.mcdev.util.children
 import com.intellij.formatting.Alignment
 import com.intellij.formatting.Block
 import com.intellij.formatting.Indent
@@ -28,7 +29,6 @@ import com.intellij.formatting.Spacing
 import com.intellij.formatting.SpacingBuilder
 import com.intellij.formatting.Wrap
 import com.intellij.lang.ASTNode
-import com.intellij.lang.tree.util.children
 import com.intellij.psi.TokenType
 import com.intellij.psi.codeStyle.CodeStyleSettings
 import com.intellij.psi.formatter.common.AbstractBlock
diff --git a/src/main/kotlin/util/ast-utils.kt b/src/main/kotlin/util/ast-utils.kt
new file mode 100644
index 000000000..099c70cbf
--- /dev/null
+++ b/src/main/kotlin/util/ast-utils.kt
@@ -0,0 +1,25 @@
+/*
+ * Minecraft Development for IntelliJ
+ *
+ * https://mcdev.io/
+ *
+ * Copyright (C) 2024 minecraft-dev
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Lesser General Public License as published
+ * by the Free Software Foundation, version 3.0 only.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+package com.demonwav.mcdev.util
+
+import com.intellij.lang.ASTNode
+
+fun ASTNode.children(): Sequence<ASTNode> = generateSequence(firstChildNode) { it.treeNext }
diff --git a/src/test/kotlin/platform/mcp/at/AtCompletionTest.kt b/src/test/kotlin/platform/mcp/at/AtCompletionTest.kt
index 329cff22e..6c6f08d8c 100644
--- a/src/test/kotlin/platform/mcp/at/AtCompletionTest.kt
+++ b/src/test/kotlin/platform/mcp/at/AtCompletionTest.kt
@@ -27,8 +27,8 @@ import com.demonwav.mcdev.platform.PlatformType
 import com.demonwav.mcdev.platform.mcp.McpModuleSettings
 import com.demonwav.mcdev.platform.mcp.McpModuleType
 import com.demonwav.mcdev.platform.mcp.at.AtElementFactory.Keyword
-import com.demonwav.mcdev.util.runWriteActionAndWait
 import com.intellij.codeInsight.lookup.Lookup
+import com.intellij.openapi.application.runWriteActionAndWait
 import org.intellij.lang.annotations.Language
 import org.junit.jupiter.api.Assertions.assertNotNull
 import org.junit.jupiter.api.BeforeEach