diff --git a/frontend/src/lib/components/copilot/chat/flow/helperUtils.ts b/frontend/src/lib/components/copilot/chat/flow/helperUtils.ts
index 00c6a65ded..e07b4a6eb9 100644
--- a/frontend/src/lib/components/copilot/chat/flow/helperUtils.ts
+++ b/frontend/src/lib/components/copilot/chat/flow/helperUtils.ts
@@ -1,4 +1,5 @@
 import type { FlowModule, OpenFlow, RawScript } from '$lib/gen'
+import { forEachFlowModule } from '$lib/components/flows/dfs'
 import { dfs } from '$lib/components/flows/previousResults'
 import { SPECIAL_MODULE_IDS } from '../shared'
 import type { InlineScriptSession } from './inlineScriptsUtils'
@@ -140,47 +141,9 @@ function replaceNewInlineScriptRefsWithEmptyCode(
 		return ''
 	}
 
-	function visitModule(module: FlowModule) {
+	forEachFlowModule(modules, (module) => {
 		if (module.value.type === 'rawscript' && module.value.content) {
 			module.value.content = replaceInlineScriptRefWithEmptyCode(module.id, module.value.content)
-			return
 		}
-
-		if (module.value.type === 'forloopflow' || module.value.type === 'whileloopflow') {
-			module.value.modules?.forEach(visitModule)
-			return
-		}
-
-		if (module.value.type === 'branchone') {
-			module.value.branches?.forEach((branch) => branch.modules?.forEach(visitModule))
-			module.value.default?.forEach(visitModule)
-			return
-		}
-
-		if (module.value.type === 'branchall') {
-			module.value.branches?.forEach((branch) => branch.modules?.forEach(visitModule))
-			return
-		}
-
-		if (module.value.type === 'aiagent') {
-			for (const tool of module.value.tools ?? []) {
-				if (
-					tool.value &&
-					'tool_type' in tool.value &&
-					tool.value.tool_type === 'flowmodule' &&
-					'type' in tool.value &&
-					tool.value.type === 'rawscript' &&
-					'content' in tool.value &&
-					tool.value.content
-				) {
-					tool.value.content = replaceInlineScriptRefWithEmptyCode(
-						tool.id,
-						tool.value.content as string
-					)
-				}
-			}
-		}
-	}
-
-	modules.forEach(visitModule)
+	})
 }
diff --git a/frontend/src/lib/components/flows/dfs.ts b/frontend/src/lib/components/flows/dfs.ts
index d7627ae2ec..2c8bcffbd0 100644
--- a/frontend/src/lib/components/flows/dfs.ts
+++ b/frontend/src/lib/components/flows/dfs.ts
@@ -1,34 +1,60 @@
 import type { FlowModule } from '$lib/gen'
 
-export function dfs<T>(
+type FlowDfsOptions = { skipToolNodes?: boolean }
+
+type FlowModuleVisitor<T> = (
+	x: FlowModule,
 	modules: FlowModule[],
-	f: (x: FlowModule, modules: FlowModule[], branches: FlowModule[][]) => T,
-	opts: { skipToolNodes?: boolean } = {}
-): T[] {
-	let result: T[] = []
+	branches: FlowModule[][]
+) => T
+
+function traverseFlowModules(
+	modules: FlowModule[],
+	visit: FlowModuleVisitor<void>,
+	opts: FlowDfsOptions = {}
+): void {
 	for (const module of modules) {
 		if (module.value.type == 'forloopflow' || module.value.type == 'whileloopflow') {
-			result = result.concat(f(module, modules, [module.value.modules]))
-			result = result.concat(dfs(module.value.modules, f, opts))
+			visit(module, modules, [module.value.modules])
+			traverseFlowModules(module.value.modules, visit, opts)
 		} else if (module.value.type == 'branchone') {
 			const allBranches = [module.value.default, ...module.value.branches.map((b) => b.modules)]
-			result = result.concat(f(module, modules, allBranches))
+			visit(module, modules, allBranches)
 
 			for (const branch of allBranches) {
-				result = result.concat(dfs(branch, f, opts))
+				traverseFlowModules(branch, visit, opts)
 			}
 		} else if (module.value.type == 'branchall') {
 			const allBranches = module.value.branches.map((b) => b.modules)
-			result = result.concat(f(module, modules, allBranches))
+			visit(module, modules, allBranches)
 			for (const branch of allBranches) {
-				result = result.concat(dfs(branch, f, opts))
+				traverseFlowModules(branch, visit, opts)
 			}
 		} else if (module.value.type == 'aiagent' && !opts.skipToolNodes) {
-			result = result.concat(f(module, modules, [module.value.tools as FlowModule[]]))
-			result = result.concat(dfs(module.value.tools as FlowModule[], f, opts))
+			visit(module, modules, [module.value.tools as FlowModule[]])
+			traverseFlowModules(module.value.tools as FlowModule[], visit, opts)
 		} else {
-			result.push(f(module, modules, []))
+			visit(module, modules, [])
 		}
 	}
+}
+
+export function dfs<T>(
+	modules: FlowModule[],
+	f: FlowModuleVisitor<T>,
+	opts: FlowDfsOptions = {}
+): T[] {
+	let result: T[] = []
+	traverseFlowModules(modules, (module, parentModules, branches) => {
+		result.push(f(module, parentModules, branches))
+	}, opts)
 	return result
 }
+
+export function forEachFlowModule(
+	modules: FlowModule[],
+	f: FlowModuleVisitor<void>,
+	opts: FlowDfsOptions = {}
+): void {
+	traverseFlowModules(modules, f, opts)
+}