improve bridge lifecycle

packages/backend/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@stack-spot/vscode-async-webview-backend",
-  "version": "0.8.2",
+  "version": "0.8.5",
   "repository": "github:stack-spot/vscode-async-webview",
   "main": "out/index.js",
   "module": "out/index.mjs",

packages/backend/rollup.config.mjs

@@ -1,6 +1,9 @@
+import { createRequire } from 'module'
 import dts from 'rollup-plugin-dts'
 import esbuild from 'rollup-plugin-esbuild'
-import packageJson from './package.json' assert { type: 'json' }
+
+const require = createRequire(import.meta.url)
+const packageJson = require('./package.json')
 
 const name = packageJson.main.replace(/\.js$/, '')
 

packages/backend/src/MessageHandler.ts

@@ -1,5 +1,5 @@
 import { uniqueId } from 'lodash'
-import { 
+import {
   messageType,
   WebviewRequestMessage,
   WebviewResponseMessage,
@@ -16,6 +16,10 @@ import {
 } from '@stack-spot/vscode-async-webview-shared'
 import { AnyFunction } from './types'
 
+interface MessageListenerDisposable {
+  dispose(): void,
+}
+
 interface Dependencies {
   /**
    * Function to send a message to the client app. It must return a promise that resolves to true if the message was successfully sent or
@@ -37,6 +41,7 @@ interface Dependencies {
  * Send/receive messages to/from the client.
  */
 export class MessageHandler {
+  private readonly MAX_QUEUE_SIZE = 100
   private readonly deps: Dependencies
   private readonly getStateCalls: Map<string, ManualPromise> = new Map()
   private readonly setStateCalls: Map<string, ManualPromise<void>> = new Map()
@@ -46,17 +51,28 @@ export class MessageHandler {
    */
   private queue: WebviewMessage[] = []
   private streaming = new Map<string, { index: number, pending?: Omit<WebviewStreamMessage, 'index' | 'type'>, result: string }>()
-
+  private messageListener?: MessageListenerDisposable
+  private disposed = false
+
   constructor(deps: Dependencies) {
     this.deps = deps
     this.listen()
   }
 
   private async sendMessageToClient(message: WebviewMessage, onNotSent?: () => void) {
+    // Don't send messages if already disposed
+    if (this.disposed) return
+
     const sent = await this.deps.sendMessageToClient(message)
     if (!sent) {
       if (onNotSent) onNotSent()
-      else this.queue.push(message)
+      else {
+        if (this.queue.length >= this.MAX_QUEUE_SIZE) {
+          this.queue.shift() // Remove oldest message if queue is full
+          logger.warn('Message queue full, dropping oldest message')
+        }
+        this.queue.push(message)
+      }
     }
   }
 
@@ -113,7 +129,11 @@ export class MessageHandler {
   }
 
   private listen() {
-    this.deps.listenToMessagesFromClient(async (message) => {
+    // Store the disposable if returned, to clean it up later
+    const result = this.deps.listenToMessagesFromClient(async (message) => {
+      // Don't process messages if already disposed
+      if (this.disposed) return
+
       switch (message?.type) {
         case messageType.bridge:
           await this.handleRequestToBridge(message)
@@ -128,17 +148,42 @@ export class MessageHandler {
           this.handleClientReadyness()
       }
     })
+
+    // Only store if a disposable was returned (not void/undefined)
+    // We need to handle this without testing void for truthiness
+    if (result !== undefined && result !== null) {
+      const disposable = result as MessageListenerDisposable
+      if ('dispose' in disposable && typeof disposable.dispose === 'function') {
+        this.messageListener = disposable
+      }
+    }
   }
 
   dispose() {
+    if (this.disposed) return // Prevent double disposal
+    this.disposed = true
+
+    // Reject all pending promises
     this.getStateCalls.forEach((value, key) => {
       value.reject(`The webview closed before the state "${String(key)}" could be retrieved.`)
     })
     this.setStateCalls.forEach((value) => {
       value.reject('The webview closed before the state could be set.')
     })
+
+    // Clear all maps
     this.getStateCalls.clear()
     this.setStateCalls.clear()
+    this.streaming.clear()
+
+    // Clear the queue
+    this.queue = []
+
+    // Dispose the message listener if it exists
+    if (this.messageListener) {
+      this.messageListener.dispose()
+      this.messageListener = undefined
+    }
   }
 
   getState(name: string): Promise<any> {
@@ -173,6 +218,18 @@ export class MessageHandler {
   }
 
   stream(message: Omit<WebviewStreamMessage, 'index' | 'type'>) {
+    if (this.disposed) return
+
+    if (this.streaming.size > 50) {
+      const toDelete: string[] = []
+      this.streaming.forEach((value, key) => {
+        if (!value.pending && value.result === '') {
+          toDelete.push(key)
+        }
+      })
+      toDelete.forEach(key => this.streaming.delete(key))
+    }
+
     const currentStreaming = this.streaming.get(message.id) ?? { index: -1, result: '' }
     this.streaming.set(message.id, currentStreaming)
     currentStreaming.result += message.content

packages/backend/src/VSCodeViewProvider.ts

@@ -1,4 +1,5 @@
 import { WebviewOptions, WebviewView, WebviewViewProvider, window } from 'vscode'
+import { logger } from '@stack-spot/vscode-async-webview-shared'
 import { ViewOptions } from './types'
 import { VSCodeWebview } from './VSCodeWebview'
 import { VSCodeWebviewBridge } from './VSCodeWebviewBridge'
@@ -10,6 +11,7 @@ export class VSCodeViewProvider<
   Bridge extends VSCodeWebviewBridge = VSCodeWebviewBridge
 > extends VSCodeWebview<Bridge> implements WebviewViewProvider {
   private view: WebviewView | undefined
+  private isResolving = false
 
   constructor(options: ViewOptions<Bridge>) {
     super(options)
@@ -19,17 +21,34 @@ export class VSCodeViewProvider<
   }
 
   async resolveWebviewView(view: WebviewView) {
-    this.view = view
-    const { webview } = view
-    const html = this.getHTML() ?? await this.buildHtml(webview.asWebviewUri(this.baseUri))
-    this.buildBridge(webview)
-    webview.options = this.options as WebviewOptions
-    webview.html = html
-    view.onDidDispose(() => {
-      this.view = undefined
-      this.bridge?.dispose()
-      this.bridge = undefined
-    })
+    if (this.isResolving) {
+      logger.warn('resolveWebviewView called while already resolving')
+      return
+    }
+
+    this.isResolving = true
+
+    try {
+      if (this.bridge) {
+        this.bridge.dispose()
+        this.bridge = undefined
+      }
+
+      this.view = view
+      const { webview } = view
+      const html = this.getHTML() ?? await this.buildHtml(webview.asWebviewUri(this.baseUri))
+      this.buildBridge(webview)
+      webview.options = this.options as WebviewOptions
+      webview.html = html
+
+      view.onDidDispose(() => {
+        this.view = undefined
+        this.bridge?.dispose()
+        this.bridge = undefined
+      })
+    } finally {
+      this.isResolving = false
+    } 
   }
 
   /**

packages/backend/src/VSCodeWebview.ts

@@ -48,6 +48,7 @@ export class VSCodeWebview<Bridge extends VSCodeWebviewBridge = VSCodeWebviewBri
   private html: string | undefined
   private readonly index: string
   private bridgePromise: ManualPromise<Bridge> | undefined
+  private static htmlCache = new Map<string, string>()
 
   constructor({
     path,
@@ -69,14 +70,23 @@ export class VSCodeWebview<Bridge extends VSCodeWebviewBridge = VSCodeWebviewBri
     this.index = index
     this.options = {
       enableScripts: true,
+      retainContextWhenHidden: true,
       localResourceRoots: [this.baseUri],
       ...options,
     }
   }
 
   protected async buildHtml(baseSrc: Uri) {
+    const cacheKey = `${this.baseUri.path}/${this.index}`
+
+    if (VSCodeWebview.htmlCache.has(cacheKey)) {
+      this.html = this.treatHTML(VSCodeWebview.htmlCache.get(cacheKey)!, baseSrc)
+      return this.html
+    }
+
     try {
       const htmlText = await this.htmlPromise
+      VSCodeWebview.htmlCache.set(cacheKey, htmlText)
       this.html = this.treatHTML(htmlText, baseSrc)
     } catch (error: any) {
       window.showErrorMessage('There was an error while loading the html for the webview. This is a bug, please report it to the team.')

packages/backend/src/VSCodeWebviewBridge.ts

@@ -50,13 +50,11 @@ export abstract class VSCodeWebviewBridge<StateType extends Record<string, any>
         return typeof member === 'function' ? (...args: any[]) => (member as AnyFunction).apply(this, args) : undefined
       },
       sendMessageToClient: async (message) => webview.postMessage.apply(webview, [message]),
-      listenToMessagesFromClient: (listener) => {
-        webview.onDidReceiveMessage((data) => {
-          logger.debug('received message from client:', data)
-          const message = asWebViewMessage(data)
-          if (message) listener(message)
-        })
-      },
+      listenToMessagesFromClient: (listener) => webview.onDidReceiveMessage((data) => {
+        logger.debug('received message from client:', data)
+        const message = asWebViewMessage(data)
+        if (message) listener(message)
+      }),
     })
   }
 

packages/backend/tsconfig.json

@@ -6,13 +6,9 @@
 		],
 		"sourceMap": true,
 		"rootDir": "src",
-		"strict": true,   /* enable all strict type-checking options */
+		"strict": true,
 		"declaration": true,
 		"noImplicitOverride": true,
 		"moduleResolution": "node",
-		/* Additional Checks */
-		// "noImplicitReturns": true, /* Report error when not all code paths in function return a value. */
-		// "noFallthroughCasesInSwitch": true, /* Report errors for fallthrough cases in switch statement. */
-		// "noUnusedParameters": true,  /* Report errors on unused parameters. */
 	}
 }

packages/client/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@stack-spot/vscode-async-webview-client",
-  "version": "0.7.4",
+  "version": "0.7.6",
   "repository": "github:stack-spot/vscode-async-webview",
   "main": "out/index.js",
   "module": "out/index.mjs",

packages/client/rollup.config.mjs

@@ -1,6 +1,9 @@
+import { createRequire } from 'module'
 import dts from 'rollup-plugin-dts'
 import esbuild from 'rollup-plugin-esbuild'
-import packageJson from './package.json' assert { type: 'json' }
+
+const require = createRequire(import.meta.url)
+const packageJson = require('./package.json')
 
 const name = packageJson.main.replace(/\.js$/, '')