refactor: cleanup vitest-transform schema support

Author: sheremet-va

packages/extension/src/extension.ts

@@ -57,15 +57,15 @@ class VitestExtension {
     this.testController.resolveHandler = item => this.resolveTestFile(item)
     this.loadingTestItem = this.testController.createTestItem('_resolving', 'Resolving Vitest...')
     this.loadingTestItem.sortText = '.0' // show it first
-    this.testTree = new TestTree(this.testController, this.loadingTestItem)
-    this.tagsManager = new TagsManager(this.testTree)
-    this.debugManager = new DebugManager()
     this.schemaProvider = new SchemaProvider(
       async (apiId, project, environment, file) => {
         const api = this.api?.folderAPIs.find(a => a.id === apiId)
         return api?.getTransformedModule(project, environment, file) ?? null
       },
     )
+    this.testTree = new TestTree(this.testController, this.loadingTestItem, this.schemaProvider)
+    this.tagsManager = new TagsManager(this.testTree)
+    this.debugManager = new DebugManager()
   }
 
   private _defineTestProfilePromise: Promise<void> | undefined
@@ -371,7 +371,7 @@ class VitestExtension {
               label += environment
               return {
                 label,
-                uriParts: [api.id, project.name, environment],
+                uriParts: [api.id, project.name, environment.name, environment.transformTimestamp],
               }
             })
           })
@@ -385,9 +385,9 @@ class VitestExtension {
           return
         }
         try {
-          const [apiId, projectName, environment] = pick.uriParts
+          const [apiId, projectName, environment, t] = pick.uriParts
           const uri = vscode.Uri.parse(
-            `vitest-transform://${currentUri.fsPath}.js?apiId=${apiId}&project=${projectName}&environment=${environment}`,
+            `vitest-transform://${currentUri.fsPath}.js?apiId=${apiId}&project=${projectName}&environment=${environment}&t=${t}`,
           )
           const doc = await vscode.workspace.openTextDocument(uri)
           await vscode.window.showTextDocument(doc, { preview: false })

packages/extension/src/schemaProvider.ts

@@ -3,12 +3,28 @@ import * as vscode from 'vscode'
 export class SchemaProvider implements vscode.TextDocumentContentProvider, vscode.Disposable {
   private disposables: vscode.Disposable[] = []
 
+  private _onDidChangeEvents = new vscode.EventEmitter<vscode.Uri>()
+
   constructor(
     private getTransformedModule: (apiId: string, project: string, environment: string, file: string) => Promise<string | null>,
   ) {
     this.disposables.push(vscode.workspace.registerTextDocumentContentProvider('vitest-transform', this))
+    this.disposables.push(this._onDidChangeEvents)
+  }
+
+  public onDidChange = this._onDidChangeEvents.event
+
+  public emitChange(uri: vscode.Uri) {
+    const cachedFsPaths = this._cachedFsPaths.get(uri.fsPath)
+    if (cachedFsPaths) {
+      cachedFsPaths.forEach((uri) => {
+        this._onDidChangeEvents.fire(uri)
+      })
+    }
   }
 
+  private _cachedFsPaths = new Map<string, Set<vscode.Uri>>()
+
   async provideTextDocumentContent(uri: vscode.Uri): Promise<string> {
     const searchParams = new URLSearchParams(uri.query)
     const apiId = searchParams.get('apiId')
@@ -18,6 +34,24 @@ export class SchemaProvider implements vscode.TextDocumentContentProvider, vscod
       throw new Error(`Cannot parse the schema: ${uri.toString()}`)
     }
     const fsPath = uri.fsPath.slice(0, -3) // remove .js
+
+    let _cachedUris = this._cachedFsPaths.get(fsPath)
+    if (!_cachedUris) {
+      _cachedUris = new Set()
+    }
+    else {
+      // remove older files from the same environment
+      _cachedUris.forEach((uri) => {
+        const query = uri.query.replace(/&t=\d+/, '')
+        const currentQuery = uri.query.replace(/&t=\d+/, '')
+        if (query === currentQuery) {
+          _cachedUris?.delete(uri)
+        }
+      })
+    }
+    _cachedUris.add(uri)
+    this._cachedFsPaths.set(fsPath, _cachedUris)
+
     const content = await this.getTransformedModule(apiId, project, environment, fsPath)
     if (content == null) {
       throw new Error(`The file ${fsPath} was not processed by Vite yet.`)
@@ -26,6 +60,7 @@ export class SchemaProvider implements vscode.TextDocumentContentProvider, vscod
   }
 
   dispose() {
+    this._cachedFsPaths.clear()
     this.disposables.forEach(d => d.dispose())
   }
 }

packages/extension/src/testTree.ts

@@ -1,6 +1,7 @@
 import type { RunnerTask, RunnerTestFile } from 'vitest'
 import type { ExtensionTestFileSpecification } from 'vitest-vscode-shared'
 import type { VitestFolderAPI } from './api'
+import type { SchemaProvider } from './schemaProvider'
 import type { TestFileMetadata } from './testTreeData'
 import { realpathSync } from 'node:fs'
 import { resolve } from 'node:path'
@@ -30,6 +31,7 @@ export class TestTree extends vscode.Disposable {
   constructor(
     private readonly controller: vscode.TestController,
     private readonly loaderItem: vscode.TestItem,
+    schemaProvider: SchemaProvider,
   ) {
     super(() => {
       this.folderItems.clear()
@@ -38,7 +40,7 @@ export class TestTree extends vscode.Disposable {
       this.testItemsByFile.clear()
       this.watcher.dispose()
     })
-    this.watcher = new ExtensionWatcher(this)
+    this.watcher = new ExtensionWatcher(this, schemaProvider)
   }
 
   public getFileTestItems(fsPath: string) {

packages/extension/src/watcher.ts

@@ -1,4 +1,5 @@
 import type { VitestFolderAPI } from './api'
+import type { SchemaProvider } from './schemaProvider'
 import type { TestTree } from './testTree'
 import { relative } from 'node:path'
 import { normalize } from 'pathe'
@@ -10,7 +11,10 @@ export class ExtensionWatcher extends vscode.Disposable {
   private watcherByFolder = new Map<vscode.WorkspaceFolder, vscode.FileSystemWatcher>()
   private apisByFolder = new WeakMap<vscode.WorkspaceFolder, VitestFolderAPI[]>()
 
-  constructor(private readonly testTree: TestTree) {
+  constructor(
+    private readonly testTree: TestTree,
+    private readonly schemaProvider: SchemaProvider,
+  ) {
     super(() => {
       this.reset()
       log.verbose?.('[VSCODE] Watcher disposed')
@@ -42,9 +46,11 @@ export class ExtensionWatcher extends vscode.Disposable {
     watcher.onDidDelete((uri) => {
       log.verbose?.('[VSCODE] File deleted:', this.relative(api, uri))
       this.testTree.removeFile(normalize(uri.fsPath))
+      this.schemaProvider.emitChange(uri)
     })
 
     watcher.onDidChange(async (uri) => {
+      this.schemaProvider.emitChange(uri)
       const path = normalize(uri.fsPath)
       if (await this.shouldIgnoreFile(api, path, uri)) {
         return

packages/shared/src/index.ts

@@ -39,7 +39,10 @@ export interface ExtensionUserConsoleLog extends UserConsoleLog {
 
 export interface ExtensionEnvironment {
   name: string
-  environments: string[]
+  environments: {
+    name: string
+    transformTimestamp: number
+  }[]
 }
 
 export interface ExtensionWorkerTransport {

packages/worker/src/worker.ts

@@ -106,17 +106,20 @@ export class ExtensionWorker implements ExtensionWorkerTransport {
 
   getModuleEnvironments(moduleId: string): ExtensionEnvironment[] {
     return this.vitest.projects.map((project) => {
-      const environments = new Set<string>()
+      const environments = new Map<string, { timestamp: number }>()
       for (const name in project.vite.environments) {
         const environment = project.vite.environments[name]
         const nodes = [...environment.moduleGraph.getModulesByFile(moduleId) || []]
         if (nodes.some(n => n.transformResult)) {
-          environments.add(name)
+          environments.set(name, { timestamp: nodes[0].lastInvalidationTimestamp })
         }
       }
       return {
         name: project.name,
-        environments: Array.from(environments),
+        environments: Array.from(environments).map(([name, { timestamp }]) => ({
+          name,
+          transformTimestamp: timestamp,
+        })),
       }
     })
   }