From 4ee18c4f91448fe511cb3a8930034cd15bb327dc Mon Sep 17 00:00:00 2001
From: Ruben Fiszel <ruben@rubenfiszel.com>
Date: Wed, 28 Feb 2024 13:22:12 +0100
Subject: [PATCH] improve bun auto-type fetching

---
 frontend/src/lib/ata/apis.ts              |  6 +++++-
 frontend/src/lib/ata/index.ts             | 10 ++++++++--
 frontend/src/lib/components/Editor.svelte |  2 +-
 3 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/frontend/src/lib/ata/apis.ts b/frontend/src/lib/ata/apis.ts
index 4d5d1163d1..22e255a79c 100644
--- a/frontend/src/lib/ata/apis.ts
+++ b/frontend/src/lib/ata/apis.ts
@@ -63,8 +63,12 @@ export interface ResLimit {
 	usage: number
 }
 
+export function isOverlimit(resLimit: ResLimit) {
+	return resLimit.usage > 500000
+}
+
 function api<T>(url: string, resLimit: ResLimit, init?: RequestInit): Promise<T | Error> {
-	if (resLimit.usage > 500000) {
+	if (isOverlimit(resLimit)) {
 		console.warn(
 			`Exceeded limit of types downloaded for the needs of the assistant fetching: ${url}`
 		)
diff --git a/frontend/src/lib/ata/index.ts b/frontend/src/lib/ata/index.ts
index ff550bb7f6..f1a5719fb8 100644
--- a/frontend/src/lib/ata/index.ts
+++ b/frontend/src/lib/ata/index.ts
@@ -3,6 +3,7 @@ import {
 	getFiletreeForModuleWithVersion,
 	getNPMVersionForModuleReference,
 	getNPMVersionsForModule,
+	isOverlimit,
 	type NPMTreeMeta,
 	type ResLimit
 } from './apis'
@@ -51,11 +52,13 @@ export const setupTypeAcquisition = (config: ATABootstrapConfig) => {
 	let estimatedToDownload = 0
 	let estimatedDownloaded = 0
 
+	let resLimit = { usage: 0 }
+
 	return (initialSourceFile: string) => {
 		estimatedToDownload = 0
 		estimatedDownloaded = 0
 
-		return resolveDeps(initialSourceFile, 0, { usage: 0 }).then((t) => {
+		return resolveDeps(initialSourceFile, 0, resLimit).then((t) => {
 			if (estimatedDownloaded > 0) {
 				config.delegate.finished?.(fsMap)
 			}
@@ -116,7 +119,6 @@ export const setupTypeAcquisition = (config: ATABootstrapConfig) => {
 		const trees = await Promise.all(
 			depsToGet.map((f) => getFileTreeForModuleWithTag(f.module, f.version, f.raw, resLimit))
 		)
-		console.log(trees, 'trees')
 		const treesOnly = trees.filter((t) => !('error' in t)) as NPMTreeMeta[]
 
 		// These are the modules which we can grab directly
@@ -175,6 +177,10 @@ export const setupTypeAcquisition = (config: ATABootstrapConfig) => {
 		// Grab all dts files
 		await Promise.all(
 			allDTSFiles.map(async (dts) => {
+				if (isOverlimit(resLimit)) {
+					console.warn('Exceeded limit of types downloaded for the needs of the assistant')
+					return
+				}
 				const dtsCode = await getDTSFileForModuleWithVersion(
 					dts.moduleName,
 					dts.moduleVersion,
diff --git a/frontend/src/lib/components/Editor.svelte b/frontend/src/lib/components/Editor.svelte
index edd352836a..be8d80fcae 100644
--- a/frontend/src/lib/components/Editor.svelte
+++ b/frontend/src/lib/components/Editor.svelte
@@ -789,7 +789,7 @@
 					}
 					await initWasm()
 					const root = await genRoot(hostname)
-
+					console.log('SETUP TYPE ACQUISITION', { root, path })
 					ata = setupTypeAcquisition({
 						projectName: 'Windmill',
 						depsParser: (c) => {