Playground stability

typescript/playground-common/src/baml_wasm_web/EventListener.tsx

@@ -10,19 +10,29 @@ import CustomErrorBoundary from '../utils/ErrorFallback'
 import { sessionStore, vscodeLocalStorageStore } from './JotaiProvider'
 import { availableProjectsAtom, projectFamilyAtom, projectFilesAtom, runtimeFamilyAtom } from './baseAtoms'
 import type { WasmDiagnosticError, WasmParam, WasmRuntime } from '@gloo-ai/baml-schema-wasm-web/baml_schema_build'
+const vscode = acquireVsCodeApi()
 
 // const wasm = await import("@gloo-ai/baml-schema-wasm-web/baml_schema_build");
 // const { WasmProject, WasmRuntime, WasmRuntimeContext, version: RuntimeVersion } = wasm;
+const postMessageToExtension = (message: any) => {
+  console.log(`Sending message to extension ${message.command}`)
+  vscode.postMessage(message)
+}
+
 const defaultEnvKeyValues: [string, string][] = (() => {
   if ((window as any).next?.version) {
     console.log('Running in nextjs')
+
     const domain = window?.location?.origin || ''
     if (domain.includes('localhost')) {
       // we can do somehting fancier here later if we want to test locally.
       return [['BOUNDARY_PROXY_URL', 'https://fiddle-proxy.fly.dev']]
     }
     return [['BOUNDARY_PROXY_URL', 'https://fiddle-proxy.fly.dev']]
   } else {
+    postMessageToExtension({ command: 'get_port' })
+    postMessageToExtension({ command: 'add_project' })
+
     console.log('Not running in a Next.js environment, set default value')
     // Not running in a Next.js environment, set default value
     return [['BOUNDARY_PROXY_URL', 'http://localhost:0000']]
@@ -121,8 +131,6 @@ export const selectedFunctionAtom = atom(
       const functions = get(availableFunctionsAtom)
       if (functions.find((f) => f.name === func)) {
         set(selectedFunctionStorageAtom, func)
-      } else {
-        // console.error(`Function ${func} not found in ${functions.map((f) => f.name).join(', ')}`)
       }
     }
   },
@@ -548,7 +556,6 @@ export const EventListener: React.FC<{ children: React.ReactNode }> = ({ childre
       >,
     ) => {
       const { command, content } = event.data
-      console.log('select Received message', command, content)
 
       switch (command) {
         case 'modify_file':
@@ -581,6 +588,13 @@ export const EventListener: React.FC<{ children: React.ReactNode }> = ({ childre
 
         case 'port_number':
           console.log('Setting port number', content.port)
+
+          if (content.port === 0) {
+            console.error('Port number is 0, cannot launch BAML extension')
+
+            return
+          }
+
           setEnvKeyValueStorage((prev) => {
             let keyExists = false
             const updated: [string, string][] = prev.map(([key, value]) => {

typescript/vscode-ext/packages/vscode/src/extension.ts

@@ -5,8 +5,6 @@ import { WebPanelView } from './panels/WebPanelView'
 import plugins from './plugins'
 import { requestDiagnostics } from './plugins/language-server'
 import { telemetry } from './plugins/language-server'
-import httpProxy from 'http-proxy'
-import express from 'express'
 import cors from 'cors'
 import { createProxyMiddleware } from 'http-proxy-middleware'
 import { type LanguageClient, type ServerOptions, TransportKind } from 'vscode-languageclient/node'
@@ -146,6 +144,8 @@ async function runDiagnostics(): Promise<void> {
   statusBarItem.hide()
 }
 
+import type { Express } from 'express'
+
 export function activate(context: vscode.ExtensionContext) {
   console.log('BAML extension activating')
 
@@ -164,12 +164,76 @@ export function activate(context: vscode.ExtensionContext) {
 
   context.subscriptions.push(codeActionProvider)
 
+  const app: Express = require('express')()
+  app.use(cors())
+  var port: number
+  const server = app.listen(0, () => {
+    console.log('Server started on port ' + getPort())
+    WebPanelView.currentPanel?.postMessage('port_number', {
+      port: port,
+    })
+  })
+
+  const getPort = () => {
+    let addr = server.address()
+    if (addr === null) {
+      vscode.window.showErrorMessage(
+        'Failed to start BAML extension server. Please try reloading the window, or restarting VSCode.',
+      )
+      return 0
+    }
+    if (typeof addr === 'string') {
+      return parseInt(addr)
+    }
+    return addr.port
+  }
+
+  app.use(
+    createProxyMiddleware({
+      changeOrigin: true,
+      pathRewrite: (path, req) => {
+        // Ensure the URL does not end with a slash
+        if (path.endsWith('/')) {
+          return path.slice(0, -1)
+        }
+        return path
+      },
+      router: (req) => {
+        // Extract the original target URL from the custom header
+        let originalUrl = req.headers['baml-original-url']
+        if (typeof originalUrl === 'string') {
+          delete req.headers['baml-original-url']
+          delete req.headers['baml-render-url']
+          req.headers['origin'] = `http://localhost:${port}`
+
+          // Ensure the URL does not end with a slash
+          if (originalUrl.endsWith('/')) {
+            originalUrl = originalUrl.slice(0, -1)
+          }
+          return originalUrl
+        } else {
+          throw new Error('baml-original-url header is missing or invalid')
+        }
+      },
+      logger: console,
+      on: {
+        proxyRes: (proxyRes, req, res) => {
+          proxyRes.headers['Access-Control-Allow-Origin'] = '*'
+        },
+        error: (error) => {
+          console.error('proxy error:', error)
+        },
+      },
+    }),
+  )
+
   const bamlPlaygroundCommand = vscode.commands.registerCommand(
     'baml.openBamlPanel',
     (args?: { projectId?: string; functionName?: string; implName?: string; showTests?: boolean }) => {
       const config = vscode.workspace.getConfiguration()
       config.update('baml.bamlPanelOpen', true, vscode.ConfigurationTarget.Global)
-      WebPanelView.render(context.extensionUri)
+
+      WebPanelView.render(context.extensionUri, getPort)
       if (telemetry) {
         telemetry.sendTelemetryEvent({
           event: 'baml.openBamlPanel',
@@ -183,29 +247,8 @@ export function activate(context: vscode.ExtensionContext) {
         root_path: 'default',
         function_name: args?.functionName ?? 'default',
       })
-      // to account for some delays (this is a hack, sorry)
-      setTimeout(() => {
-        WebPanelView.currentPanel?.postMessage('select_function', {
-          root_path: 'default',
-          function_name: args?.functionName ?? 'default',
-        })
-        WebPanelView.currentPanel?.postMessage('port_number', {
-          port: port,
-        })
-      }, 200)
-
-      // adding the select_function here causes glitches (cause it's way too delayed)
-      // for now lets resend only the port until we have a more reliable way to request this data.
-      setTimeout(() => {
-        WebPanelView.currentPanel?.postMessage('port_number', {
-          port: port,
-        })
-      }, 1000)
 
       console.info('Opening BAML panel')
-      WebPanelView.currentPanel?.postMessage('port_number', {
-        port: port,
-      })
     },
   )
 
@@ -243,6 +286,8 @@ export function activate(context: vscode.ExtensionContext) {
     }
   })
 
+  // Listen for messages from the webview
+
   plugins.map(async (plugin) => {
     const enabled = await plugin.enabled()
     if (enabled) {
@@ -260,61 +305,6 @@ export function activate(context: vscode.ExtensionContext) {
     vscode.commands.executeCommand('baml.openBamlPanel')
   }
 
-  try {
-    const app = require('express')()
-    app.use(cors())
-    var port: number
-    const server = app.listen(0, () => {
-      port = server.address().port
-      console.log('Server started on port ' + port)
-      WebPanelView.currentPanel?.postMessage('port_number', {
-        port: port,
-      })
-    })
-
-    app.use(
-      createProxyMiddleware({
-        changeOrigin: true,
-        pathRewrite: (path, req) => {
-          // Ensure the URL does not end with a slash
-          if (path.endsWith('/')) {
-            return path.slice(0, -1)
-          }
-          return path
-        },
-        router: (req) => {
-          // Extract the original target URL from the custom header
-          let originalUrl = req.headers['baml-original-url']
-          if (typeof originalUrl === 'string') {
-            delete req.headers['baml-original-url']
-            delete req.headers['baml-render-url']
-            req.headers['origin'] = `http://localhost:${port}`
-
-            // Ensure the URL does not end with a slash
-            if (originalUrl.endsWith('/')) {
-              originalUrl = originalUrl.slice(0, -1)
-            }
-            return originalUrl
-          } else {
-            throw new Error('baml-original-url header is missing or invalid')
-          }
-        },
-        logger: console,
-        on: {
-          proxyRes: (proxyRes, req, res) => {
-            proxyRes.headers['Access-Control-Allow-Origin'] = '*'
-          },
-          error: (error) => {
-            console.error('proxy error:', error)
-          },
-        },
-      }),
-    )
-  } catch (error) {
-    console.error('Failed to start proxy server:', error)
-    vscode.window.showErrorMessage('Failed to BAML localhost server. Contact support for help.')
-  }
-
   // TODO: Reactivate linter.
   // runDiagnostics();
 }

typescript/vscode-ext/packages/vscode/src/panels/WebPanelView.ts

@@ -6,6 +6,7 @@ import { getUri } from '../utils/getUri'
 
 import { type Config, adjectives, animals, colors, uniqueNamesGenerator } from 'unique-names-generator'
 import { URI } from 'vscode-uri'
+import { requestDiagnostics } from '../plugins/language-server'
 
 const customConfig: Config = {
   dictionaries: [adjectives, colors, animals],
@@ -27,15 +28,17 @@ export class WebPanelView {
   public static currentPanel: WebPanelView | undefined
   private readonly _panel: WebviewPanel
   private _disposables: Disposable[] = []
+  private _port: () => number
 
   /**
    * The WebPanelView class private constructor (called only from the render method).
    *
    * @param panel A reference to the webview panel
    * @param extensionUri The URI of the directory containing the extension
    */
-  private constructor(panel: WebviewPanel, extensionUri: Uri) {
+  private constructor(panel: WebviewPanel, extensionUri: Uri, portLoader: () => number) {
     this._panel = panel
+    this._port = portLoader
 
     // Set an event listener to listen for when the panel is disposed (i.e. when the user closes
     // the panel or when the panel is closed programmatically)
@@ -54,7 +57,7 @@ export class WebPanelView {
    *
    * @param extensionUri The URI of the directory containing the extension.
    */
-  public static render(extensionUri: Uri) {
+  public static render(extensionUri: Uri, portLoader: () => number) {
     if (WebPanelView.currentPanel) {
       // If the webview panel already exists reveal it
       WebPanelView.currentPanel._panel.reveal(ViewColumn.Beside)
@@ -81,7 +84,7 @@ export class WebPanelView {
         },
       )
 
-      WebPanelView.currentPanel = new WebPanelView(panel, extensionUri)
+      WebPanelView.currentPanel = new WebPanelView(panel, extensionUri, portLoader)
     }
   }
 
@@ -161,6 +164,17 @@ export class WebPanelView {
         const text = message.text
 
         switch (command) {
+          case 'get_port':
+            // Code that should run in response to the hello message command
+            console.log(`Sending port from WebPanelView: ${this._port()}`)
+            this.postMessage('port_number', {
+              port: this._port(),
+            })
+            return
+
+          case 'add_project':
+            requestDiagnostics()
+            return
           case 'receiveData':
             // Code that should run in response to the hello message command
             window.showInformationMessage(text)