Automatic reconnect with Language Server. (#9691)

Fixes #8520

If the websocket is closed not by us, we automatically try to reconnect with it, and initialize the protocol again. **Restoring state (execution contexts, attached visualizations) is not part of this PR**.

It's a part of making IDE work after hibernation (or LS crash).

# Important Notes
It required somewhat heavy refactoring:
1. I decided to use an existing implementation of reconnecting websocket. Replaced (later discovered by me) our implementation.
2. The LanguageServer class now handles both reconnecting and re-initializing - that make usage of it simpler (no more `Promise<LanguageServer>` - each method will just wait for (re)connection and initialization.
3. The stuff in `net` src's module was partially moved to shared's counterpart (with tests). Merged `exponentialBackoff` implementations, which also brought me to
4. Rewriting LS client, so it returns Result instead of throwing, what is closer our desired state, and allows us using exponentialBackoff method without any wrappers.

app/gui2/e2e/componentBrowser.spec.ts

@@ -254,6 +254,7 @@ test('Visualization preview: user visualization selection', async ({ page }) =>
   await input.fill('4')
   await expect(input).toHaveValue('4')
   await expect(locate.jsonVisualization(page)).toExist()
+  await expect(locate.jsonVisualization(page)).toContainText('"visualizedExpr": "4"')
   await locate.showVisualizationSelectorButton(page).click()
   await page.getByRole('button', { name: 'Table' }).click()
   // The table visualization is not currently working with `executeExpression` (#9194), but we can test that the JSON

app/gui2/e2e/graphNodeVisualization.spec.ts

@@ -12,6 +12,7 @@ test('node can open and load visualization', async ({ page }) => {
   await expect(locate.circularMenu(page)).toExist()
   await locate.toggleVisualizationButton(page).click()
   await expect(locate.anyVisualization(page)).toExist()
+  await expect(locate.loadingVisualization(page)).toHaveCount(0)
   await locate.showVisualizationSelectorButton(page).click()
   await page.getByText('JSON').click()
   const vis = locate.jsonVisualization(page)
@@ -36,6 +37,7 @@ test('Warnings visualization', async ({ page }) => {
   await expect(locate.circularMenu(page)).toExist()
   await locate.toggleVisualizationButton(page).click()
   await expect(locate.anyVisualization(page)).toExist()
+  await expect(locate.loadingVisualization(page)).toHaveCount(0)
   await locate.showVisualizationSelectorButton(page).click()
   await page.locator('.VisualizationSelector').getByRole('button', { name: 'Warnings' }).click()
   await expect(locate.warningsVisualization(page)).toExist()

app/gui2/e2e/locate.ts

@@ -145,6 +145,7 @@ function componentLocator<T extends string>(className: SanitizeClassName<T>) {
 export const graphEditor = componentLocator('GraphEditor')
 // @ts-expect-error
 export const anyVisualization = componentLocator('GraphVisualization > *')
+export const loadingVisualization = componentLocator('LoadingVisualization')
 export const circularMenu = componentLocator('CircularMenu')
 export const addNewNodeButton = componentLocator('PlusButton')
 export const componentBrowser = componentLocator('ComponentBrowser')

app/gui2/mock/MockFSWrapper.vue

@@ -2,9 +2,9 @@
 import { useProjectStore } from '@/stores/project'
 import { mockFsDirectoryHandle } from '@/util/convert/fsAccess'
 import { MockWebSocket, type WebSocketHandler } from '@/util/net'
-import { mockDataWSHandler } from 'shared/dataServer/mock'
 import { type Path as LSPath } from 'shared/languageServerTypes'
 import { watchEffect } from 'vue'
+import { mockDataWSHandler } from './dataServer'
 
 const projectStore = useProjectStore()
 

app/gui2/shared/dataServer/mock.ts → app/gui2/mock/dataServer.ts

@@ -27,9 +27,9 @@ import {
   type AnyOutboundPayload,
   type Offset,
   type Table,
-} from '../binaryProtocol'
-import { LanguageServerErrorCode } from '../languageServerTypes'
-import { uuidToBits } from '../uuid'
+} from 'shared/binaryProtocol'
+import { LanguageServerErrorCode } from 'shared/languageServerTypes'
+import { uuidToBits } from 'shared/uuid'
 
 const sha3 = createSHA3(224)
 

app/gui2/mock/engine.ts

@@ -1,3 +1,4 @@
+import { Pattern } from '@/util/ast/match'
 import type { MockYdocProviderImpl } from '@/util/crdt'
 import * as random from 'lib0/random'
 import * as Ast from 'shared/ast'
@@ -9,7 +10,7 @@ import {
   VisualizationContext,
   VisualizationUpdate,
 } from 'shared/binaryProtocol'
-import { mockDataWSHandler as originalMockDataWSHandler } from 'shared/dataServer/mock'
+import { ErrorCode } from 'shared/languageServer'
 import type {
   ContextId,
   ExpressionId,
@@ -26,6 +27,7 @@ import type { QualifiedName } from 'src/util/qualifiedName'
 import * as Y from 'yjs'
 import { mockFsDirectoryHandle, type FileTree } from '../src/util/convert/fsAccess'
 import mockDb from '../stories/mockSuggestions.json' assert { type: 'json' }
+import { mockDataWSHandler as originalMockDataWSHandler } from './dataServer'
 
 const mockProjectId = random.uuidv4() as Uuid
 const standardBase = 'Standard.Base' as QualifiedName
@@ -93,19 +95,21 @@ const visualizationExprIds = new Map<Uuid, ExpressionId>()
 const encoder = new TextEncoder()
 const encodeJSON = (data: unknown) => encoder.encode(JSON.stringify(data))
 
-const scatterplotJson = encodeJSON({
-  axis: {
-    x: { label: 'x-axis label', scale: 'linear' },
-    y: { label: 'y-axis label', scale: 'logarithmic' },
-  },
-  points: { labels: 'visible' },
-  data: [
-    { x: 0.1, y: 0.7, label: 'foo', color: '#FF0000', shape: 'circle', size: 0.2 },
-    { x: 0.4, y: 0.2, label: 'baz', color: '#0000FF', shape: 'square', size: 0.3 },
-  ],
-})
+const scatterplotJson = (params: string[]) =>
+  encodeJSON({
+    visualizedExpr: params[0],
+    axis: {
+      x: { label: 'x-axis label', scale: 'linear' },
+      y: { label: 'y-axis label', scale: 'logarithmic' },
+    },
+    points: { labels: 'visible' },
+    data: [
+      { x: 0.1, y: 0.7, label: 'foo', color: '#FF0000', shape: 'circle', size: 0.2 },
+      { x: 0.4, y: 0.2, label: 'baz', color: '#0000FF', shape: 'square', size: 0.3 },
+    ],
+  })
 
-const mockVizData: Record<string, Uint8Array | ((params: string[]) => Uint8Array)> = {
+const mockVizPreprocessors: Record<string, Uint8Array | ((params: string[]) => Uint8Array)> = {
   // JSON
   'Standard.Visualization.Preprocessor.default_preprocessor': scatterplotJson,
   'Standard.Visualization.Scatter_Plot.process_to_json_text': scatterplotJson,
@@ -320,7 +324,6 @@ const mockVizData: Record<string, Uint8Array | ((params: string[]) => Uint8Array
         return encodeJSON([])
     }
   },
-  'Standard.Visualization.AI.build_ai_prompt': () => encodeJSON('Could you __$$GOAL$$__, please?'),
 
   // The following visualizations do not have unique transformation methods, and as such are only kept
   // for posterity.
@@ -353,9 +356,9 @@ function createId(id: Uuid) {
   return (builder: Builder) => EnsoUUID.createEnsoUUID(builder, low, high)
 }
 
-function sendVizData(id: Uuid, config: VisualizationConfiguration) {
+function sendVizData(id: Uuid, config: VisualizationConfiguration, expressionId?: Uuid) {
   const vizDataHandler =
-    mockVizData[
+    mockVizPreprocessors[
       typeof config.expression === 'string' ?
         `${config.visualizationModule}.${config.expression}`
       : `${config.expression.definedOnType}.${config.expression.name}`
@@ -365,12 +368,22 @@ function sendVizData(id: Uuid, config: VisualizationConfiguration) {
     vizDataHandler instanceof Uint8Array ? vizDataHandler : (
       vizDataHandler(config.positionalArgumentsExpressions ?? [])
     )
+  const exprId = expressionId ?? visualizationExprIds.get(id)
+  sendVizUpdate(id, config.executionContextId, exprId, vizData)
+}
+
+function sendVizUpdate(
+  id: Uuid,
+  executionCtxId: Uuid,
+  exprId: Uuid | undefined,
+  vizData: Uint8Array,
+) {
+  if (!sendData) return
   const builder = new Builder()
-  const exprId = visualizationExprIds.get(id)
   const visualizationContextOffset = VisualizationContext.createVisualizationContext(
     builder,
     createId(id),
-    createId(config.executionContextId),
+    createId(executionCtxId),
     exprId ? createId(exprId) : null,
   )
   const dataOffset = VisualizationUpdate.createDataVector(builder, vizData)
@@ -446,16 +459,27 @@ export const mockLSHandler: MockTransportData = async (method, data, transport)
         expressionId: ExpressionId
         expression: string
       }
-      const { func, args } = Ast.analyzeAppLike(Ast.parse(data_.expression))
-      if (!(func instanceof Ast.PropertyAccess && func.lhs)) return
-      const visualizationConfig: VisualizationConfiguration = {
-        executionContextId: data_.executionContextId,
-        visualizationModule: func.lhs.code(),
-        expression: func.rhs.code(),
-        positionalArgumentsExpressions: args.map((ast) => ast.code()),
+      const aiPromptPat = Pattern.parse('Standard.Visualization.AI.build_ai_prompt __ . to_json')
+      const exprAst = Ast.parse(data_.expression)
+      if (aiPromptPat.test(exprAst)) {
+        sendVizUpdate(
+          data_.visualizationId,
+          data_.executionContextId,
+          data_.expressionId,
+          encodeJSON('Could you __$$GOAL$$__, please?'),
+        )
+      } else {
+        // Check if there's existing preprocessor mock which matches our expression
+        const { func, args } = Ast.analyzeAppLike(exprAst)
+        if (!(func instanceof Ast.PropertyAccess && func.lhs)) return
+        const visualizationConfig: VisualizationConfiguration = {
+          executionContextId: data_.executionContextId,
+          visualizationModule: func.lhs.code(),
+          expression: func.rhs.code(),
+          positionalArgumentsExpressions: args.map((ast) => ast.code()),
+        }
+        sendVizData(data_.visualizationId, visualizationConfig, data_.expressionId)
       }
-      visualizationExprIds.set(data_.visualizationId, data_.expressionId)
-      sendVizData(data_.visualizationId, visualizationConfig)
       return
     }
     case 'search/getSuggestionsDatabase':
@@ -487,9 +511,16 @@ export const mockLSHandler: MockTransportData = async (method, data, transport)
           if (!child || typeof child === 'string' || child instanceof ArrayBuffer) break
         }
       }
-      if (!child) return Promise.reject(`Folder '/${data_.path.segments.join('/')}' not found.`)
+      if (!child)
+        return Promise.reject({
+          code: ErrorCode.FILE_NOT_FOUND,
+          message: `Folder '/${data_.path.segments.join('/')}' not found.`,
+        })
       if (typeof child === 'string' || child instanceof ArrayBuffer)
-        return Promise.reject(`File '/${data_.path.segments.join('/')}' is not a folder.`)
+        return Promise.reject({
+          code: ErrorCode.NOT_DIRECTORY,
+          message: `File '/${data_.path.segments.join('/')}' is not a folder.`,
+        })
       return {
         paths: Object.entries(child).map(([name, entry]) => ({
           type: typeof entry === 'string' || entry instanceof ArrayBuffer ? 'File' : 'Directory',
@@ -500,8 +531,7 @@ export const mockLSHandler: MockTransportData = async (method, data, transport)
     }
     case 'ai/completion': {
       const { prompt } = data
-      console.log(prompt)
-      const match = /^"Could you (.*), please\?"$/.exec(prompt)
+      const match = /^Could you (.*), please\?$/.exec(prompt)
       if (!match) {
         return { code: 'How rude!' }
       } else if (match[1] === 'convert to table') {

app/gui2/package.json

@@ -63,6 +63,7 @@
     "lib0": "^0.2.85",
     "magic-string": "^0.30.3",
     "murmurhash": "^2.0.1",
+    "partysocket": "^1.0.1",
     "pinia": "^2.1.7",
     "postcss-inline-svg": "^6.0.0",
     "postcss-nesting": "^12.0.1",