diff --git a/app/gui2/src/stores/graph/imports.ts b/app/gui2/src/stores/graph/imports.ts
index ceeb1fac330..6d0c77795d1 100644
--- a/app/gui2/src/stores/graph/imports.ts
+++ b/app/gui2/src/stores/graph/imports.ts
@@ -144,6 +144,22 @@ export interface UnqualifiedImport {
   import: Identifier
 }
 
+/** Read imports from given module block */
+export function readImports(ast: Ast.Ast): Import[] {
+  const imports: Import[] = []
+  ast.visitRecursive((node) => {
+    if (node instanceof Ast.Import) {
+      const recognized = recognizeImport(node)
+      if (recognized) {
+        imports.push(recognized)
+      }
+      return false
+    }
+    return true
+  })
+  return imports
+}
+
 /** Insert the given imports into the given block at an appropriate location. */
 export function addImports(
   edit: MutableModule,
@@ -288,10 +304,10 @@ export function covers(existing: Import, required: RequiredImport): boolean {
 }
 
 export function filterOutRedundantImports(
-  existing: { import: Import }[],
+  existing: Import[],
   required: RequiredImport[],
 ): RequiredImport[] {
-  return required.filter((info) => !existing.some((existing) => covers(existing.import, info)))
+  return required.filter((info) => !existing.some((existing) => covers(existing, info)))
 }
 
 if (import.meta.vitest) {
diff --git a/app/gui2/src/stores/graph/index.ts b/app/gui2/src/stores/graph/index.ts
index fdcb176ded9..202d2200b0e 100644
--- a/app/gui2/src/stores/graph/index.ts
+++ b/app/gui2/src/stores/graph/index.ts
@@ -5,6 +5,7 @@ import { GraphDb } from '@/stores/graph/graphDatabase'
 import {
   addImports,
   filterOutRedundantImports,
+  readImports,
   recognizeImport,
   type Import,
   type RequiredImport,
@@ -95,7 +96,6 @@ export const useGraphStore = defineStore('graph', () => {
   )
   const portInstances = reactive(new Map<PortId, Set<PortViewInstance>>())
   const editedNodeInfo = ref<NodeEditInfo>()
-  const imports = ref<{ import: Import; span: SourceRange }[]>([])
   const methodAst = ref<Ast.Function>()
   const currentNodeIds = ref(new Set<ExprId>())
 
@@ -126,18 +126,6 @@ export const useGraphStore = defineStore('graph', () => {
       moduleRoot.value = newRoot.exprId
       module.doc.setIdMap(idMap_)
 
-      imports.value = []
-      newRoot.visitRecursive((node) => {
-        if (node instanceof Ast.Import) {
-          const recognized = recognizeImport(node)
-          if (recognized) {
-            imports.value.push({ import: recognized, span: node.span! })
-          }
-          return false
-        }
-        return true
-      })
-
       methodAst.value = methodAstInModule(astModule)
       if (methodAst.value) {
         currentNodeIds.value = db.readFunctionAst(methodAst.value, (id) => meta.get(id))
@@ -257,14 +245,15 @@ export const useGraphStore = defineStore('graph', () => {
 
   function addMissingImports(edit: MutableModule, newImports: RequiredImport[]) {
     if (!newImports.length) return
-    const topLevel_ = topLevel.value
-    if (!topLevel_) {
-      console.error(`BUG: Cannot add required imports: No module root.`)
+    const topLevel = moduleRoot.value ? edit.get(moduleRoot.value) : null
+    if (!topLevel || !(topLevel instanceof Ast.BodyBlock)) {
+      console.error(`BUG: Cannot add required imports: No BodyBlock module root.`)
       return
     }
-    const importsToAdd = filterOutRedundantImports(imports.value, newImports)
+    const existingImports = readImports(topLevel)
+    const importsToAdd = filterOutRedundantImports(existingImports, newImports)
     if (!importsToAdd.length) return
-    addImports(edit, topLevel_, importsToAdd)
+    addImports(edit, topLevel, importsToAdd)
   }
 
   function deleteNodes(ids: ExprId[]) {
@@ -525,7 +514,6 @@ export const useGraphStore = defineStore('graph', () => {
     transact,
     db: markRaw(db),
     mockExpressionUpdate,
-    imports,
     editedNodeInfo,
     unconnectedEdge,
     edges,