Restore state after reconnecting (#9755)

Fixes: #8522

Execution context is refactored slightly: now we have a single `sync` function to synchronize both visualization and execution stack.

Tested hibernation on Linux: I was able to continue my work  🎉

# Important Notes
The Refinement Notes state, that the execution mode should be set before updating the stack, but actually it makes an error on startup (changing context automatically re-executes programs, what fails if there's no frame on the stack).

app/gui2/mock/engine.ts

@@ -496,6 +496,7 @@ export const mockLSHandler: MockTransportData = async (method, data, transport)
     case 'executionContext/push':
     case 'executionContext/pop':
     case 'executionContext/recompute':
+    case 'executionContext/setExecutionEnvironment':
     case 'capability/acquire':
       return {}
     case 'file/list': {

app/gui2/src/components/GraphEditor.vue

@@ -267,7 +267,7 @@ function onRecordOnceButtonPress() {
 watch(
   () => projectStore.executionMode,
   (modeValue) => {
-    projectStore.executionContext.setExecutionEnvironment(modeValue === 'live' ? 'Live' : 'Design')
+    projectStore.executionContext.executionEnvironment = modeValue === 'live' ? 'Live' : 'Design'
   },
 )
 

app/gui2/src/composables/stackNavigator.ts

@@ -43,23 +43,7 @@ export function useStackNavigator() {
   }
 
   function handleBreadcrumbClick(index: number) {
-    const activeStack = projectStore.executionContext.desiredStack
+    projectStore.executionContext.desiredStack = breadcrumbs.value.slice(0, index + 1)
-    // Number of items in desired stack should be index + 1
-    if (index + 1 < activeStack.length) {
-      for (let i = activeStack.length; i > index + 1; i--) {
-        projectStore.executionContext.pop()
-      }
-    } else if (index + 1 > activeStack.length) {
-      for (let i = activeStack.length; i <= index; i++) {
-        const stackItem = breadcrumbs.value[i]
-        if (stackItem?.type === 'LocalCall') {
-          const exprId = stackItem.expressionId
-          projectStore.executionContext.push(exprId)
-        } else {
-          console.warn('Cannot enter non-local call.')
-        }
-      }
-    }
     graphStore.updateState()
   }
 

app/gui2/src/stores/project/executionContext.ts

@@ -1,3 +1,4 @@
+import { findIndexOpt } from '@/util/data/array'
 import { isSome, type Opt } from '@/util/data/opt'
 import { Err, Ok, type Result } from '@/util/data/result'
 import { AsyncQueue, type AbortScope } from '@/util/net'
@@ -6,16 +7,17 @@ import * as object from 'lib0/object'
 import { ObservableV2 } from 'lib0/observable'
 import * as random from 'lib0/random'
 import type { LanguageServer } from 'shared/languageServer'
-import type {
+import {
-  ContextId,
+  stackItemsEqual,
-  Diagnostic,
+  type ContextId,
-  ExecutionEnvironment,
+  type Diagnostic,
-  ExplicitCall,
+  type ExecutionEnvironment,
-  ExpressionId,
+  type ExplicitCall,
-  ExpressionUpdate,
+  type ExpressionId,
-  StackItem,
+  type ExpressionUpdate,
-  Uuid,
+  type StackItem,
-  VisualizationConfiguration,
+  type Uuid,
+  type VisualizationConfiguration,
 } from 'shared/languageServerTypes'
 import { exponentialBackoff } from 'shared/util/net'
 import type { ExternalId } from 'shared/yjsModel'
@@ -46,12 +48,14 @@ function visualizationConfigEqual(
   )
 }
 
-interface ExecutionContextState {
+type ExecutionContextState =
-  lsRpc: LanguageServer
+  | { status: 'not-created' }
-  created: boolean
+  | {
-  visualizations: Map<Uuid, NodeVisualizationConfiguration>
+      status: 'created'
-  stack: StackItem[]
+      visualizations: Map<Uuid, NodeVisualizationConfiguration>
-}
+      stack: StackItem[]
+      environment?: ExecutionEnvironment
+    } // | { status: 'broken'} TODO[ao] think about it
 
 type EntryPoint = Omit<ExplicitCall, 'type'>
 
@@ -80,146 +84,79 @@ type ExecutionContextNotification = {
  * run only when the previous call is done.
  */
 export class ExecutionContext extends ObservableV2<ExecutionContextNotification> {
-  id: ContextId = random.uuidv4() as ContextId
+  readonly id: ContextId = random.uuidv4() as ContextId
-  queue: AsyncQueue<ExecutionContextState>
+  private queue: AsyncQueue<ExecutionContextState>
-  taskRunning = false
+  private syncScheduled = false
-  visSyncScheduled = false
+  private clearScheduled = false
-  desiredStack: StackItem[] = reactive([])
+  private _desiredStack: StackItem[] = reactive([])
-  visualizationConfigs: Map<Uuid, NodeVisualizationConfiguration> = new Map()
+  private visualizationConfigs: Map<Uuid, NodeVisualizationConfiguration> = new Map()
+  private _executionEnvironment: ExecutionEnvironment = 'Design'
 
   constructor(
-    lsRpc: LanguageServer,
+    private lsRpc: LanguageServer,
     entryPoint: EntryPoint,
     private abort: AbortScope,
   ) {
     super()
     this.abort.handleDispose(this)
-
+    this.lsRpc.retain()
-    this.queue = new AsyncQueue<ExecutionContextState>(
+    this.queue = new AsyncQueue<ExecutionContextState>(Promise.resolve({ status: 'not-created' }))
-      Promise.resolve({
-        lsRpc,
-        created: false,
-        visualizations: new Map(),
-        stack: [],
-      }),
-    )
     this.registerHandlers()
-    this.create()
     this.pushItem({ type: 'ExplicitCall', ...entryPoint })
-    this.recompute()
   }
 
-  private async withBackoff<T>(f: () => Promise<Result<T>>, message: string): Promise<T> {
+  private registerHandlers() {
-    const result = await exponentialBackoff(f, {
+    this.abort.handleObserve(this.lsRpc, 'executionContext/expressionUpdates', (event) => {
-      onBeforeRetry: (error, _, delay) => {
+      if (event.contextId == this.id) this.emit('expressionUpdates', [event.updates])
-        if (this.abort.signal.aborted) return false
-        console.warn(`${error.message(message)}. Retrying after ${delay}ms...\n`)
-      },
     })
-    if (result.ok) return result.value
+    this.abort.handleObserve(this.lsRpc, 'executionContext/executionFailed', (event) => {
-    else throw result.error
+      if (event.contextId == this.id) this.emit('executionFailed', [event.message])
-  }
+    })
-
+    this.abort.handleObserve(this.lsRpc, 'executionContext/executionComplete', (event) => {
-  private syncVisualizations() {
+      if (event.contextId == this.id) this.emit('executionComplete', [])
-    if (this.visSyncScheduled || this.abort.signal.aborted) return
+    })
-    this.visSyncScheduled = true
+    this.abort.handleObserve(this.lsRpc, 'executionContext/executionStatus', (event) => {
-    this.queue.pushTask(async (state) => {
+      if (event.contextId == this.id) this.emit('executionStatus', [event.diagnostics])
-      this.visSyncScheduled = false
+    })
-      if (!state.created || this.abort.signal.aborted) return state
+    this.abort.handleObserve(
-      this.emit('newVisualizationConfiguration', [new Set(this.visualizationConfigs.keys())])
+      this.lsRpc,
-      const promises: Promise<void>[] = []
+      'executionContext/visualizationEvaluationFailed',
-
+      (event) => {
-      const attach = (id: Uuid, config: NodeVisualizationConfiguration) => {
+        if (event.contextId == this.id)
-        return this.withBackoff(
+          this.emit('visualizationEvaluationFailed', [
-          () =>
+            event.visualizationId,
-            state.lsRpc.attachVisualization(id, config.expressionId, {
+            event.expressionId,
-              executionContextId: this.id,
+            event.message,
-              expression: config.expression,
+            event.diagnostic,
-              visualizationModule: config.visualizationModule,
+          ])
-              ...(config.positionalArgumentsExpressions ?
+      },
-                { positionalArgumentsExpressions: config.positionalArgumentsExpressions }
+    )
-              : {}),
+    this.lsRpc.on('transport/closed', () => {
-            }),
+      // Connection closed: the created execution context is no longer available
-          'Failed to attach visualization',
+      // There is no point in any scheduled action until resynchronization
-        ).then(() => {
+      this.queue.clear()
-          state.visualizations.set(id, config)
+      this.syncScheduled = false
-        })
+      this.queue.pushTask(() => {
-      }
+        this.clearScheduled = false
-
+        this.sync()
-      const modify = (id: Uuid, config: NodeVisualizationConfiguration) => {
+        return Promise.resolve({ status: 'not-created' })
-        return this.withBackoff(
+      })
-          () =>
+      this.clearScheduled = true
-            state.lsRpc.modifyVisualization(id, {
-              executionContextId: this.id,
-              expression: config.expression,
-              visualizationModule: config.visualizationModule,
-              ...(config.positionalArgumentsExpressions ?
-                { positionalArgumentsExpressions: config.positionalArgumentsExpressions }
-              : {}),
-            }),
-          'Failed to modify visualization',
-        ).then(() => {
-          state.visualizations.set(id, config)
-        })
-      }
-
-      const detach = (id: Uuid, config: NodeVisualizationConfiguration) => {
-        return this.withBackoff(
-          () => state.lsRpc.detachVisualization(id, config.expressionId, this.id),
-          'Failed to detach visualization',
-        ).then(() => {
-          state.visualizations.delete(id)
-        })
-      }
-
-      // Attach new and update existing visualizations.
-      for (const [id, config] of this.visualizationConfigs) {
-        const previousConfig = state.visualizations.get(id)
-        if (previousConfig == null) {
-          promises.push(attach(id, config))
-        } else if (!visualizationConfigEqual(previousConfig, config)) {
-          if (previousConfig.expressionId === config.expressionId) {
-            promises.push(modify(id, config))
-          } else {
-            promises.push(detach(id, previousConfig).then(() => attach(id, config)))
-          }
-        }
-      }
-
-      // Detach removed visualizations.
-      for (const [id, config] of state.visualizations) {
-        if (!this.visualizationConfigs.get(id)) {
-          promises.push(detach(id, config))
-        }
-      }
-      const settled = await Promise.allSettled(promises)
-
-      // Emit errors for failed requests.
-      const errors = settled
-        .map((result) => (result.status === 'rejected' ? result.reason : null))
-        .filter(isSome)
-      if (errors.length > 0) {
-        console.error('Failed to synchronize visualizations:', errors)
-      }
-
-      this.emit('visualizationsConfigured', [new Set(this.visualizationConfigs.keys())])
-
-      // State object was updated in-place in each successful promise.
-      return state
     })
   }
 
   private pushItem(item: StackItem) {
-    this.desiredStack.push(item)
+    this._desiredStack.push(item)
-    this.queue.pushTask(async (state) => {
+    this.sync()
-      if (!state.created) return state
+  }
-      await this.withBackoff(
+
-        () => state.lsRpc.pushExecutionContextItem(this.id, item),
+  get desiredStack() {
-        'Failed to push item to execution context stack',
+    return this._desiredStack
-      )
+  }
-      state.stack.push(item)
+
-      return state
+  set desiredStack(stack: StackItem[]) {
-    })
+    this._desiredStack = stack
+    this.sync()
   }
 
   push(expressionId: ExpressionId) {
@@ -227,79 +164,21 @@ export class ExecutionContext extends ObservableV2<ExecutionContextNotification>
   }
 
   pop() {
-    if (this.desiredStack.length === 1) {
+    if (this._desiredStack.length === 1) {
       console.debug('Cannot pop last item from execution context stack')
       return
     }
-    this.desiredStack.pop()
+    this._desiredStack.pop()
-    this.queue.pushTask(async (state) => {
+    this.sync()
-      if (!state.created) return state
-      if (state.stack.length === 1) {
-        console.debug('Cannot pop last item from execution context stack')
-        return state
-      }
-      await this.withBackoff(
-        () => state.lsRpc.popExecutionContextItem(this.id),
-        'Failed to pop item from execution context stack',
-      )
-      state.stack.pop()
-      return state
-    })
   }
 
-  async setVisualization(id: Uuid, configuration: Opt<NodeVisualizationConfiguration>) {
+  setVisualization(id: Uuid, configuration: Opt<NodeVisualizationConfiguration>) {
     if (configuration == null) {
       this.visualizationConfigs.delete(id)
     } else {
       this.visualizationConfigs.set(id, configuration)
     }
-    this.syncVisualizations()
+    this.sync()
-  }
-
-  private create() {
-    this.queue.pushTask(async (state) => {
-      if (state.created) return state
-      return this.withBackoff(async () => {
-        const result = await state.lsRpc.createExecutionContext(this.id)
-        if (!result.ok) return result
-        if (result.value.contextId !== this.id) {
-          return Err('Unexpected Context ID returned by the language server.')
-        }
-        state.lsRpc.retain()
-        return Ok({ ...state, created: true })
-      }, 'Failed to create execution context')
-    })
-  }
-
-  private registerHandlers() {
-    this.queue.pushTask(async (state) => {
-      this.abort.handleObserve(state.lsRpc, 'executionContext/expressionUpdates', (event) => {
-        if (event.contextId == this.id) this.emit('expressionUpdates', [event.updates])
-      })
-      this.abort.handleObserve(state.lsRpc, 'executionContext/executionFailed', (event) => {
-        if (event.contextId == this.id) this.emit('executionFailed', [event.message])
-      })
-      this.abort.handleObserve(state.lsRpc, 'executionContext/executionComplete', (event) => {
-        if (event.contextId == this.id) this.emit('executionComplete', [])
-      })
-      this.abort.handleObserve(state.lsRpc, 'executionContext/executionStatus', (event) => {
-        if (event.contextId == this.id) this.emit('executionStatus', [event.diagnostics])
-      })
-      this.abort.handleObserve(
-        state.lsRpc,
-        'executionContext/visualizationEvaluationFailed',
-        (event) => {
-          if (event.contextId == this.id)
-            this.emit('visualizationEvaluationFailed', [
-              event.visualizationId,
-              event.expressionId,
-              event.message,
-              event.diagnostic,
-            ])
-        },
-      )
-      return state
-    })
   }
 
   recompute(
@@ -307,36 +186,220 @@ export class ExecutionContext extends ObservableV2<ExecutionContextNotification>
     executionEnvironment?: ExecutionEnvironment,
   ) {
     this.queue.pushTask(async (state) => {
-      if (!state.created) return state
+      if (state.status !== 'created') {
-      await state.lsRpc.recomputeExecutionContext(this.id, expressionIds, executionEnvironment)
+        this.sync()
+        return state
+      }
+      await this.lsRpc.recomputeExecutionContext(this.id, expressionIds, executionEnvironment)
       return state
     })
   }
 
   getStackBottom(): StackItem {
-    return this.desiredStack[0]!
+    return this._desiredStack[0]!
   }
 
   getStackTop(): StackItem {
-    return this.desiredStack[this.desiredStack.length - 1]!
+    return this._desiredStack[this._desiredStack.length - 1]!
   }
 
-  setExecutionEnvironment(mode: ExecutionEnvironment) {
+  get executionEnvironment() {
-    this.queue.pushTask(async (state) => {
+    return this._executionEnvironment
-      await state.lsRpc.setExecutionEnvironment(this.id, mode)
+  }
-      return state
+
-    })
+  set executionEnvironment(env: ExecutionEnvironment) {
+    this._executionEnvironment = env
+    this.sync()
   }
 
   dispose() {
     this.queue.pushTask(async (state) => {
-      if (!state.created) return state
+      if (state.status === 'created') {
-      const result = await state.lsRpc.destroyExecutionContext(this.id)
+        const result = await this.withBackoff(
-      if (!result.ok) {
+          () => this.lsRpc.destroyExecutionContext(this.id),
-        result.error.log('Failed to destroy execution context')
+          'Destroying execution context',
+        )
+        if (!result.ok) {
+          result.error.log('Failed to destroy execution context')
+        }
       }
-      state.lsRpc.release()
+      this.lsRpc.release()
-      return { ...state, created: false }
+      return { status: 'not-created' }
     })
   }
+
+  private sync() {
+    if (this.syncScheduled || this.abort.signal.aborted) return
+    this.syncScheduled = true
+    this.queue.pushTask(this.syncTask())
+  }
+
+  private withBackoff<T>(f: () => Promise<Result<T>>, message: string): Promise<Result<T>> {
+    return exponentialBackoff(f, {
+      onBeforeRetry: (error, _, delay) => {
+        if (this.abort.signal.aborted || this.clearScheduled) return false
+        console.warn(`${error.message(message)}. Retrying after ${delay}ms...\n`)
+      },
+      onFailure(error) {
+        error.log(message)
+      },
+    })
+  }
+
+  private syncTask() {
+    return async (state: ExecutionContextState) => {
+      this.syncScheduled = false
+      if (this.abort.signal.aborted) return state
+      let newState = { ...state }
+
+      const create = () => {
+        if (newState.status === 'created') return Ok()
+        // if (newState.status === 'broken') {
+        //   this.withBackoff(() => this.lsRpc.destroyExecutionContext(this.id), 'Failed to destroy broken execution context')
+        // }
+        return this.withBackoff(async () => {
+          const result = await this.lsRpc.createExecutionContext(this.id)
+          if (!result.ok) return result
+          if (result.value.contextId !== this.id) {
+            return Err('Unexpected Context ID returned by the language server.')
+          }
+          newState = { status: 'created', visualizations: new Map(), stack: [] }
+          return Ok()
+        }, 'Failed to create execution context')
+      }
+
+      const syncEnvironment = async () => {
+        const state = newState
+        if (state.status !== 'created')
+          return Err('Cannot sync execution environment when context is not created')
+        if (state.environment === this._executionEnvironment) return Ok()
+        const result = await this.lsRpc.setExecutionEnvironment(this.id, this._executionEnvironment)
+        if (!result.ok) return result
+        state.environment = this._executionEnvironment
+        return Ok()
+      }
+
+      const syncStack = async () => {
+        const state = newState
+        if (state.status !== 'created')
+          return Err('Cannot sync stack when execution context is not created')
+        const firstDifferent =
+          findIndexOpt(this._desiredStack, (item, index) => {
+            const stateStack = state.stack[index]
+            return stateStack == null || !stackItemsEqual(item, stateStack)
+          }) ?? this._desiredStack.length
+        for (let i = state.stack.length; i > firstDifferent; --i) {
+          const popResult = await this.withBackoff(
+            () => this.lsRpc.popExecutionContextItem(this.id),
+            'Failed to pop execution stack frame',
+          )
+          if (popResult.ok) state.stack.pop()
+          else return popResult
+        }
+        for (let i = state.stack.length; i < this._desiredStack.length; ++i) {
+          const newItem = this._desiredStack[i]!
+          const pushResult = await this.withBackoff(
+            () => this.lsRpc.pushExecutionContextItem(this.id, newItem),
+            'Failed to push execution stack frame',
+          )
+          if (pushResult.ok) state.stack.push(newItem)
+          else return pushResult
+        }
+        return Ok()
+      }
+
+      const syncVisualizations = async () => {
+        const state = newState
+        if (state.status !== 'created')
+          return Err('Cannot sync visualizations when execution context is not created')
+        const promises: Promise<void>[] = []
+
+        const attach = (id: Uuid, config: NodeVisualizationConfiguration) => {
+          return this.withBackoff(
+            () =>
+              this.lsRpc.attachVisualization(id, config.expressionId, {
+                executionContextId: this.id,
+                expression: config.expression,
+                visualizationModule: config.visualizationModule,
+                ...(config.positionalArgumentsExpressions ?
+                  { positionalArgumentsExpressions: config.positionalArgumentsExpressions }
+                : {}),
+              }),
+            'Failed to attach visualization',
+          ).then((result) => {
+            if (result.ok) state.visualizations.set(id, config)
+          })
+        }
+
+        const modify = (id: Uuid, config: NodeVisualizationConfiguration) => {
+          return this.withBackoff(
+            () =>
+              this.lsRpc.modifyVisualization(id, {
+                executionContextId: this.id,
+                expression: config.expression,
+                visualizationModule: config.visualizationModule,
+                ...(config.positionalArgumentsExpressions ?
+                  { positionalArgumentsExpressions: config.positionalArgumentsExpressions }
+                : {}),
+              }),
+            'Failed to modify visualization',
+          ).then((result) => {
+            if (result.ok) state.visualizations.set(id, config)
+          })
+        }
+
+        const detach = (id: Uuid, config: NodeVisualizationConfiguration) => {
+          return this.withBackoff(
+            () => this.lsRpc.detachVisualization(id, config.expressionId, this.id),
+            'Failed to detach visualization',
+          ).then((result) => {
+            if (result.ok) state.visualizations.delete(id)
+          })
+        }
+
+        // Attach new and update existing visualizations.
+        for (const [id, config] of this.visualizationConfigs) {
+          const previousConfig = state.visualizations.get(id)
+          if (previousConfig == null) {
+            promises.push(attach(id, config))
+          } else if (!visualizationConfigEqual(previousConfig, config)) {
+            if (previousConfig.expressionId === config.expressionId) {
+              promises.push(modify(id, config))
+            } else {
+              promises.push(detach(id, previousConfig).then(() => attach(id, config)))
+            }
+          }
+        }
+
+        // Detach removed visualizations.
+        for (const [id, config] of state.visualizations) {
+          if (!this.visualizationConfigs.get(id)) {
+            promises.push(detach(id, config))
+          }
+        }
+        const settled = await Promise.allSettled(promises)
+
+        // Emit errors for failed requests.
+        const errors = settled
+          .map((result) => (result.status === 'rejected' ? result.reason : null))
+          .filter(isSome)
+        if (errors.length > 0) {
+          console.error('Failed to synchronize visualizations:', errors)
+        }
+      }
+
+      const createResult = await create()
+      if (!createResult.ok) return newState
+      const syncStackResult = await syncStack()
+      if (!syncStackResult.ok) return newState
+      const syncEnvResult = await syncEnvironment()
+      if (!syncEnvResult.ok) return newState
+      this.emit('newVisualizationConfiguration', [new Set(this.visualizationConfigs.keys())])
+      await syncVisualizations()
+      this.emit('visualizationsConfigured', [
+        new Set(state.status === 'created' ? state.visualizations.keys() : []),
+      ])
+      return newState
+    }
+  }
 }

app/gui2/src/util/data/array.ts

@@ -6,7 +6,10 @@ import type { Opt } from '@/util/data/opt'
 export type NonEmptyArray<T> = [T, ...T[]]
 
 /** An equivalent of `Array.prototype.findIndex` method, but returns null instead of -1. */
-export function findIndexOpt<T>(arr: T[], pred: (elem: T) => boolean): number | null {
+export function findIndexOpt<T>(
+  arr: T[],
+  pred: (elem: T, index: number) => boolean,
+): number | null {
   const index = arr.findIndex(pred)
   return index >= 0 ? index : null
 }