feat: imports duration breakdown

README.md

@@ -83,6 +83,7 @@ These options are resolved relative to the [workspace file](https://code.visuals
 - `vitest.debuggerPort`: Port that the debugger will be attached to. By default uses 9229 or tries to find a free port if it's not available.
 - `vitest.debuggerAddress`: TCP/IP address of process to be debugged. Default: localhost
 - `vitest.cliArguments`: Additional arguments to pass to the Vitest CLI. Note that some arguments will be ignored: `watch`, `reporter`, `api`, and `ui`. Example: `--mode=staging`
+- `vitest.showImportsDuration`: Show how long it took to import and transform the modules. When hovering, the extension provides more diagnostics.
 
 > 💡 The `vitest.nodeExecutable` and `vitest.nodeExecArgs` settings are used as `execPath` and `execArgv` when spawning a new `child_process`, and as `runtimeExecutable` and `runtimeArgs` when [debugging a test](https://github.com/microsoft/vscode-js-debug/blob/main/OPTIONS.md).
 > The `vitest.terminalShellPath` and `vitest.terminalShellArgs` settings are used as `shellPath` and `shellArgs` when creating a new [terminal](https://code.visualstudio.com/api/references/vscode-api#Terminal)
@@ -111,6 +112,16 @@ You can also type the same command in the quick picker while the file is open.
 
 ![Reveal test in explorer](./img/reveal-in-picker.png "Reveal test in explorer")
 
+### Import Breakdown
+
+If you use Vitest 4.0.15 or higher, the extension will show how long it took to load the module on the same line where the import is defined. This number includes transform time and evaluation time, including static imports.
+
+If you hover over it, you can get a more detailed diagnostic.
+
+![Import breakdown example](./img/import-breakdown.png "Import breakdown example")
+
+You can disable this feature by turning off `vitest.showImportsDuration`.
+
 ### Experimental
 
 If the extension hangs, consider enabling `vitest.experimentalStaticAstCollect` option to use static analysis instead of actually running the test file every time you make a change which can cause visible hangs if it takes a long time to setup the test.

package.json

@@ -262,6 +262,11 @@
           "description": "Show a squiggly line where the error was thrown. This also enables the error count in the File Tab.",
           "type": "boolean",
           "default": true
+        },
+        "vitest.showImportsDuration": {
+          "description": "Show how long it took to import the module during the last test run. If multiple isolated tests imported the module, the times will be aggregated.",
+          "type": "boolean",
+          "default": true
         }
       }
     }

packages/extension/src/api.ts

@@ -1,4 +1,4 @@
-import type { ExtensionTestSpecification } from 'vitest-vscode-shared'
+import type { ExtensionTestSpecification, ModuleDefinitionDurationsDiagnostic } from 'vitest-vscode-shared'
 import type { VitestPackage } from './api/pkg'
 import type { ExtensionWorkerEvents, VitestExtensionRPC } from './api/rpc'
 import type { ExtensionWorkerProcess } from './api/types'
@@ -44,6 +44,45 @@ export class VitestAPI {
     return this.api.forEach(callback)
   }
 
+  async getSourceModuleDiagnostic(moduleId: string) {
+    const allDiagnostic = await Promise.all(
+      this.folderAPIs.map(api => api.getSourceModuleDiagnostic(moduleId)),
+    )
+    const modules = allDiagnostic[0]?.modules || []
+    const untrackedModules = allDiagnostic[0]?.untrackedModules || []
+
+    type TimeDiagnostic = Pick<ModuleDefinitionDurationsDiagnostic, 'selfTime' | 'totalTime' | 'transformTime' | 'resolvedId'>
+    const aggregateModules = (aggregatedModule: TimeDiagnostic, currentMod: TimeDiagnostic) => {
+      if (aggregatedModule.resolvedId === currentMod.resolvedId) {
+        aggregatedModule.selfTime += currentMod.selfTime
+        aggregatedModule.totalTime += currentMod.totalTime
+        if (aggregatedModule.transformTime != null && currentMod.transformTime != null) {
+          aggregatedModule.transformTime += currentMod.transformTime
+        }
+      }
+    }
+
+    // aggregate time from _other_ diagnostics that could've potentially imported this file
+    for (let i = 1; i < allDiagnostic.length; i++) {
+      const currentDiagnostic = allDiagnostic[i]
+      currentDiagnostic.modules.forEach((mod, index) => {
+        const aggregatedModule = modules[index]
+
+        aggregateModules(aggregatedModule, mod)
+      })
+      currentDiagnostic.untrackedModules.forEach((mod, index) => {
+        const aggregatedModule = untrackedModules[index]
+
+        aggregateModules(aggregatedModule, mod)
+      })
+    }
+
+    return {
+      modules,
+      untrackedModules,
+    }
+  }
+
   getModuleEnvironments(moduleId: string) {
     return Promise.all(
       this.api.map(async (api) => {
@@ -127,6 +166,10 @@ export class VitestFolderAPI {
     return this.meta.rpc.getTransformedModule(project, environment, moduleId)
   }
 
+  getSourceModuleDiagnostic(moduleId: string) {
+    return this.meta.rpc.getSourceModuleDiagnostic(moduleId)
+  }
+
   async getModuleEnvironments(moduleId: string) {
     return this.meta.rpc.getModuleEnvironments(moduleId)
   }

packages/extension/src/config.ts

@@ -97,6 +97,7 @@ export function getConfig(workspaceFolder?: WorkspaceFolder) {
     debuggerPort: get<number>('debuggerPort') || undefined,
     debuggerAddress: get<string>('debuggerAddress', undefined) || undefined,
     logLevel,
+    showImportsDuration: get<boolean>('showImportsDuration', true) ?? true,
   }
 }
 

packages/extension/src/debug.ts

@@ -2,6 +2,7 @@ import type { VitestPackage } from './api/pkg'
 import type { ExtensionWorkerProcess } from './api/types'
 import type { WsConnectionMetadata } from './api/ws'
 import type { ExtensionDiagnostic } from './diagnostic'
+import type { ImportsBreakdownProvider } from './importsBreakdownProvider'
 import type { TestTree } from './testTree'
 import crypto from 'node:crypto'
 import { createServer } from 'node:http'
@@ -26,6 +27,7 @@ export async function debugTests(
   tree: TestTree,
   pkg: VitestPackage,
   diagnostic: ExtensionDiagnostic | undefined,
+  importsBreakdown: ImportsBreakdownProvider,
 
   request: vscode.TestRunRequest,
   token: vscode.CancellationToken,
@@ -152,6 +154,7 @@ export async function debugTests(
             tree,
             api,
             diagnostic,
+            importsBreakdown,
           )
           disposables.push(api, runner)
 

packages/extension/src/extension.ts

@@ -10,6 +10,7 @@ import { configGlob, workspaceGlob } from './constants'
 import { coverageContext } from './coverage'
 import { DebugManager, debugTests } from './debug'
 import { ExtensionDiagnostic } from './diagnostic'
+import { ImportsBreakdownProvider } from './importsBreakdownProvider'
 import { log } from './log'
 import { TestRunner } from './runner'
 import { SchemaProvider } from './schemaProvider'
@@ -43,6 +44,7 @@ class VitestExtension {
   private diagnostic: ExtensionDiagnostic | undefined
   private debugManager: DebugManager
   private schemaProvider: SchemaProvider
+  private importsBreakdownProvider: ImportsBreakdownProvider
 
   /** @internal */
   _debugDisposable: vscode.Disposable | undefined
@@ -66,6 +68,12 @@ class VitestExtension {
     this.testTree = new TestTree(this.testController, this.loadingTestItem, this.schemaProvider)
     this.tagsManager = new TagsManager(this.testTree)
     this.debugManager = new DebugManager()
+    this.importsBreakdownProvider = new ImportsBreakdownProvider(
+      async (moduleId: string) => this.api?.getSourceModuleDiagnostic(moduleId) || {
+        modules: [],
+        untrackedModules: [],
+      },
+    )
   }
 
   private _defineTestProfilePromise: Promise<void> | undefined
@@ -158,6 +166,7 @@ class VitestExtension {
         this.testTree,
         api,
         this.diagnostic,
+        this.importsBreakdownProvider,
       )
       this.runners.push(runner)
 
@@ -200,6 +209,7 @@ class VitestExtension {
           this.testTree,
           api.package,
           this.diagnostic,
+          this.importsBreakdownProvider,
 
           request,
           token,
@@ -483,6 +493,7 @@ class VitestExtension {
     this.tagsManager.dispose()
     this.testController.dispose()
     this.schemaProvider.dispose()
+    this.importsBreakdownProvider.dispose()
     this.runProfiles.forEach(profile => profile.dispose())
     this.runProfiles.clear()
     this.disposables.forEach(d => d.dispose())

packages/extension/src/importsBreakdownProvider.ts

@@ -0,0 +1,163 @@
+import type { SourceModuleDiagnostic } from 'vitest-vscode-shared'
+import { relative } from 'node:path'
+import { pathToFileURL } from 'node:url'
+import * as vscode from 'vscode'
+import { getConfig } from './config'
+import { log } from './log'
+
+export class ImportsBreakdownProvider {
+  private disposables: vscode.Disposable[] = []
+  private decorationType: vscode.TextEditorDecorationType
+
+  private _decorations = new Map<string, vscode.DecorationOptions[]>()
+
+  private showDecorations = getConfig().showImportsDuration
+
+  constructor(
+    private getSourceModuleDiagnostic: (moduleId: string) => Promise<SourceModuleDiagnostic>,
+  ) {
+    // Create a decoration type with gray color
+    this.decorationType = vscode.window.createTextEditorDecorationType({
+      after: {
+        color: '#808080',
+        margin: '0 0 0 0.5em',
+      },
+    })
+
+    // Update decorations when the active editor changes
+    this.disposables.push(
+      vscode.window.onDidChangeActiveTextEditor((editor) => {
+        if (editor) {
+          this.updateDecorations(editor)
+        }
+      }),
+    )
+
+    // Update decorations when the document changes
+    this.disposables.push(
+      vscode.workspace.onDidChangeTextDocument((event) => {
+        const editor = vscode.window.activeTextEditor
+        if (editor && event.document === editor.document) {
+          this.updateDecorations(editor)
+        }
+      }),
+    )
+
+    this.disposables.push(
+      vscode.workspace.onDidChangeConfiguration((event) => {
+        if (event.affectsConfiguration('vitest.showImportsDuration')) {
+          this.showDecorations = getConfig().showImportsDuration
+
+          this.refreshCurrentDecorations()
+        }
+      }),
+    )
+
+    // Update decorations for the currently active editor
+    if (vscode.window.activeTextEditor) {
+      this.updateDecorations(vscode.window.activeTextEditor)
+    }
+  }
+
+  public refreshCurrentDecorations() {
+    log.info('[DECOR] Reset all decorations.')
+    this._decorations.clear()
+
+    // Update decorations for the currently active editor
+    if (vscode.window.activeTextEditor) {
+      this.updateDecorations(vscode.window.activeTextEditor)
+    }
+  }
+
+  private async updateDecorations(editor: vscode.TextEditor) {
+    const document = editor.document
+    if (!this.showDecorations || document.uri.scheme !== 'file' || !document.lineCount) {
+      editor.setDecorations(this.decorationType, [])
+      return
+    }
+    const fsPath = document.uri.fsPath
+    if (this._decorations.has(fsPath)) {
+      log.info('[DECOR] Decorations for', fsPath, 'are already cached. Displaying them.')
+      editor.setDecorations(this.decorationType, this._decorations.get(fsPath)!)
+      return
+    }
+
+    const diagnostic = await this.getSourceModuleDiagnostic(fsPath).catch(() => null)
+    if (!diagnostic || !diagnostic.modules) {
+      editor.setDecorations(this.decorationType, [])
+      return
+    }
+
+    const decorations: vscode.DecorationOptions[] = []
+
+    // TODO: untracked modules somehow?
+    diagnostic.modules.forEach((diagnostic) => {
+      const range = new vscode.Range(
+        diagnostic.start.line - 1,
+        diagnostic.start.column,
+        diagnostic.end.line - 1,
+        diagnostic.end.column,
+      )
+
+      const overallTime = diagnostic.totalTime + (diagnostic.transformTime || 0)
+      let color: string | undefined
+      if (overallTime >= 500) {
+        color = 'rgb(248 113 113 / 0.8)'
+      }
+      else if (overallTime >= 100) {
+        color = 'rgb(251 146 60 / 0.8)'
+      }
+
+      let diagnosticMessage = `
+### VITEST DIAGNOSTIC
+- It took **${formatPreciseTime(diagnostic.totalTime)}** to import this module, including static imports.
+- It took **${formatPreciseTime(diagnostic.selfTime)}** to import this modules, excluding static imports.
+- It took **${formatPreciseTime(diagnostic.transformTime || 0)}** to transform this module.`
+
+      if (diagnostic.external) {
+        diagnosticMessage += `\n- This module was **externalized** to [${diagnostic.resolvedUrl}](${pathToFileURL(diagnostic.resolvedId).toString()})`
+      }
+      if (diagnostic.importer && document.fileName !== diagnostic.importer) {
+        diagnosticMessage += `\n- This module was originally imported by [${relative(document.fileName, diagnostic.importer)}](${pathToFileURL(diagnostic.importer)})`
+      }
+
+      diagnosticMessage += `\n\nYou can disable diagnostic by setting [\`vitest.showImportsDuration\`](command:workbench.action.openSettings?%5B%22vitest.showImportsDuration%22%5D) option in your VSCode settings to \`false\`.`
+      const ms = new vscode.MarkdownString(diagnosticMessage)
+      ms.isTrusted = true
+
+      decorations.push({
+        range,
+        hoverMessage: ms,
+        renderOptions: {
+          after: {
+            color,
+            contentText: formatTime(overallTime),
+          },
+        },
+      })
+    })
+
+    this._decorations.set(fsPath, decorations)
+
+    editor.setDecorations(this.decorationType, decorations)
+  }
+
+  dispose() {
+    this.decorationType.dispose()
+    this.disposables.forEach(d => d.dispose())
+  }
+}
+
+function formatTime(time: number): string {
+  if (time > 1000) {
+    return `${(time / 1000).toFixed(2)}s`
+  }
+  return `${Math.round(time)}ms`
+}
+
+function formatPreciseTime(time: number): string {
+  if (time > 1000) {
+    return `${(time / 1000).toFixed(2)}s`
+  }
+  return `${time.toFixed(2)}ms`
+}

packages/extension/src/runner.ts

@@ -2,6 +2,7 @@ import type { ParsedStack, RunnerTaskResult, TestError } from 'vitest'
 import type { ExtensionTestSpecification } from 'vitest-vscode-shared'
 import type { VitestFolderAPI } from './api'
 import type { ExtensionDiagnostic } from './diagnostic'
+import type { ImportsBreakdownProvider } from './importsBreakdownProvider'
 import type { TestTree } from './testTree'
 import { rm } from 'node:fs/promises'
 import path from 'node:path'
@@ -34,6 +35,7 @@ export class TestRunner extends vscode.Disposable {
     private readonly tree: TestTree,
     private readonly api: VitestFolderAPI,
     private readonly diagnostic: ExtensionDiagnostic | undefined,
+    private readonly importsBreakdown: ImportsBreakdownProvider,
   ) {
     super(() => {
       log.verbose?.('Disposing test runner')
@@ -94,6 +96,8 @@ export class TestRunner extends vscode.Disposable {
       if (collecting)
         return
 
+      this.importsBreakdown.refreshCurrentDecorations()
+
       getTasks(file).forEach((task) => {
         const test = this.tree.getTestItemByTask(task)
         if (!test) {