refactor callbacks (#334)

refactor callbacks

src/App.tsx

@@ -278,6 +278,8 @@ export function App() {
 
   useEffect(() => {
     if (!isStreamReady) return
+    if (!engineCommandManager) return
+    let unsubFn: any[] = []
     const asyncWrap = async () => {
       try {
         if (!code) {
@@ -288,11 +290,8 @@ export function App() {
         setAst(_ast)
         resetLogs()
         resetKCLErrors()
-        if (engineCommandManager) {
+        engineCommandManager.endSession()
-          engineCommandManager.endSession()
+        engineCommandManager.startNewSession()
-          engineCommandManager.startNewSession()
-        }
-        if (!engineCommandManager) return
         const programMemory = await _executor(
           _ast,
           {
@@ -326,22 +325,29 @@ export function App() {
           await engineCommandManager.waitForAllCommands()
 
         setArtifactMap({ artifactMap, sourceRangeMap })
-        engineCommandManager.onHover((id) => {
+        const unSubHover = engineCommandManager.subscribeToLossy({
-          if (!id) {
+          event: 'highlight_set_entity',
-            setHighlightRange([0, 0])
+          callback: ({ data }) => {
-          } else {
+            if (!data?.entity_id) {
-            const sourceRange = sourceRangeMap[id]
+              setHighlightRange([0, 0])
-            setHighlightRange(sourceRange)
+            } else {
-          }
+              const sourceRange = sourceRangeMap[data.entity_id]
+              setHighlightRange(sourceRange)
+            }
+          },
         })
-        engineCommandManager.onClick((selections) => {
+        const unSubClick = engineCommandManager.subscribeTo({
-          if (!selections) {
+          event: 'select_with_point',
-            setCursor2()
+          callback: ({ data }) => {
-            return
+            if (!data?.entity_id) {
-          }
+              setCursor2()
-          const { id, type } = selections
+              return
-          setCursor2({ range: sourceRangeMap[id], type })
+            }
+            const sourceRange = sourceRangeMap[data.entity_id]
+            setCursor2({ range: sourceRange, type: 'default' })
+          },
         })
+        unsubFn.push(unSubHover, unSubClick)
         if (programMemory !== undefined) {
           setProgramMemory(programMemory)
         }
@@ -358,7 +364,10 @@ export function App() {
       }
     }
     asyncWrap()
-  }, [code, isStreamReady])
+    return () => {
+      unsubFn.forEach((fn) => fn())
+    }
+  }, [code, isStreamReady, engineCommandManager])
 
   const debounceSocketSend = throttle<EngineCommand>((message) => {
     engineCommandManager?.sendSceneCommand(message)

src/lang/std/engineConnection.ts

@@ -509,6 +509,23 @@ export class EngineConnection {
 }
 
 export type EngineCommand = Models['WebSocketRequest_type']
+type ModelTypes = Models['OkModelingCmdResponse_type']['type']
+
+type LossyResponses = Extract<
+  Models['OkModelingCmdResponse_type'],
+  { type: 'highlight_set_entity' }
+>
+interface LossySubscription<T extends LossyResponses['type']> {
+  event: T
+  callback: (data: Extract<LossyResponses, { type: T }>) => void
+}
+
+interface Subscription<T extends ModelTypes> {
+  event: T
+  callback: (
+    data: Extract<Models['OkModelingCmdResponse_type'], { type: T }>
+  ) => void
+}
 
 export class EngineCommandManager {
   artifactMap: ArtifactMap = {}
@@ -518,10 +535,17 @@ export class EngineCommandManager {
   engineConnection?: EngineConnection
   waitForReady: Promise<void> = new Promise(() => {})
   private resolveReady = () => {}
-  onHoverCallback: (id?: string) => void = () => {}
+
-  onClickCallback: (selection?: SelectionsArgs) => void = () => {}
+  subscriptions: {
-  onCursorsSelectedCallback: (selections: CursorSelectionsArgs) => void =
+    [event: string]: {
-    () => {}
+      [localUnsubscribeId: string]: (a: any) => void
+    }
+  } = {} as any
+  lossySubscriptions: {
+    [event: string]: {
+      [localUnsubscribeId: string]: (a: any) => void
+    }
+  } = {} as any
   constructor({
     setMediaStream,
     setIsStreamReady,
@@ -554,17 +578,23 @@ export class EngineCommandManager {
           let lossyDataChannel = event.channel
 
           lossyDataChannel.addEventListener('message', (event) => {
-            const result: Models['OkModelingCmdResponse_type'] = JSON.parse(
+            const result: LossyResponses = JSON.parse(event.data)
-              event.data
+            Object.values(this.lossySubscriptions[result.type] || {}).forEach(
+              // TODO: There is only one response that uses the lossy channel atm,
+              // highlight_set_entity, if there are more it's likely they will all have the same
+              // sequence logic, but I'm not sure if we use a single global sequence or a sequence
+              // per lossy subscription.
+              (callback) => {
+                if (
+                  result?.data?.sequence &&
+                  result?.data.sequence > this.inSequence &&
+                  result.type === 'highlight_set_entity'
+                ) {
+                  this.inSequence = result.data.sequence
+                  callback(result)
+                }
+              }
             )
-            if (
-              result.type === 'highlight_set_entity' &&
-              result?.data?.sequence &&
-              result.data.sequence > this.inSequence
-            ) {
-              this.onHoverCallback(result.data.entity_id)
-              this.inSequence = result.data.sequence
-            }
           })
         })
 
@@ -611,18 +641,11 @@ export class EngineCommandManager {
       return
     }
     const modelingResponse = message.data.modeling_response
+    Object.values(this.subscriptions[modelingResponse.type] || {}).forEach(
+      (callback) => callback(modelingResponse)
+    )
 
     const command = this.artifactMap[id]
-    if (modelingResponse.type === 'select_with_point') {
-      if (modelingResponse?.data?.entity_id) {
-        this.onClickCallback({
-          id: modelingResponse?.data?.entity_id,
-          type: 'default',
-        })
-      } else {
-        this.onClickCallback()
-      }
-    }
     if (command && command.type === 'pending') {
       const resolve = command.resolve
       this.artifactMap[id] = {
@@ -631,6 +654,7 @@ export class EngineCommandManager {
       }
       resolve({
         id,
+        data: modelingResponse,
       })
     } else {
       this.artifactMap[id] = {
@@ -646,21 +670,46 @@ export class EngineCommandManager {
     this.artifactMap = {}
     this.sourceRangeMap = {}
   }
+  subscribeTo<T extends ModelTypes>({
+    event,
+    callback,
+  }: Subscription<T>): () => void {
+    const localUnsubscribeId = uuidv4()
+    const otherEventCallbacks = this.subscriptions[event]
+    if (otherEventCallbacks) {
+      otherEventCallbacks[localUnsubscribeId] = callback
+    } else {
+      this.subscriptions[event] = {
+        [localUnsubscribeId]: callback,
+      }
+    }
+    return () => this.unSubscribeTo(event, localUnsubscribeId)
+  }
+  private unSubscribeTo(event: ModelTypes, id: string) {
+    delete this.subscriptions[event][id]
+  }
+  subscribeToLossy<T extends LossyResponses['type']>({
+    event,
+    callback,
+  }: LossySubscription<T>): () => void {
+    const localUnsubscribeId = uuidv4()
+    const otherEventCallbacks = this.lossySubscriptions[event]
+    if (otherEventCallbacks) {
+      otherEventCallbacks[localUnsubscribeId] = callback
+    } else {
+      this.lossySubscriptions[event] = {
+        [localUnsubscribeId]: callback,
+      }
+    }
+    return () => this.unSubscribeToLossy(event, localUnsubscribeId)
+  }
+  private unSubscribeToLossy(event: LossyResponses['type'], id: string) {
+    delete this.lossySubscriptions[event][id]
+  }
   endSession() {
     // this.websocket?.close()
     // socket.off('command')
   }
-  onHover(callback: (id?: string) => void) {
-    // It's when the user hovers over a part in the 3d scene, and so the engine should tell the
-    // frontend about that (with it's id) so that the FE can highlight code associated with that id
-    this.onHoverCallback = callback
-  }
-  onClick(callback: (selection?: SelectionsArgs) => void) {
-    // It's when the user clicks on a part in the 3d scene, and so the engine should tell the
-    // frontend about that (with it's id) so that the FE can put the user's cursor on the right
-    // line of code
-    this.onClickCallback = callback
-  }
   cusorsSelected(selections: {
     otherSelections: Selections['otherSelections']
     idBasedSelections: { type: string; id: string }[]
@@ -685,12 +734,12 @@ export class EngineCommandManager {
       cmd_id: uuidv4(),
     })
   }
-  sendSceneCommand(command: EngineCommand) {
+  sendSceneCommand(command: EngineCommand): Promise<any> {
     if (!this.engineConnection?.isReady()) {
       console.log('socket not ready')
-      return
+      return Promise.resolve()
     }
-    if (command.type !== 'modeling_cmd_req') return
+    if (command.type !== 'modeling_cmd_req') return Promise.resolve()
     const cmd = command.cmd
     if (
       cmd.type === 'camera_drag_move' &&
@@ -699,7 +748,7 @@ export class EngineCommandManager {
       cmd.sequence = this.outSequence
       this.outSequence++
       this.engineConnection?.lossyDataChannel?.send(JSON.stringify(command))
-      return
+      return Promise.resolve()
     } else if (
       cmd.type === 'highlight_set_entity' &&
       this.engineConnection?.lossyDataChannel
@@ -707,10 +756,12 @@ export class EngineCommandManager {
       cmd.sequence = this.outSequence
       this.outSequence++
       this.engineConnection?.lossyDataChannel?.send(JSON.stringify(command))
-      return
+      return Promise.resolve()
     }
     console.log('sending command', command)
+    // since it's not mouse drag or highlighting send over TCP and keep track of the command
     this.engineConnection?.send(command)
+    return this.handlePendingCommand(command.cmd_id)
   }
   sendModelingCommand({
     id,
@@ -725,9 +776,12 @@ export class EngineCommandManager {
 
     if (!this.engineConnection?.isReady()) {
       console.log('socket not ready')
-      return new Promise(() => {})
+      return Promise.resolve()
     }
     this.engineConnection?.send(command)
+    return this.handlePendingCommand(id)
+  }
+  handlePendingCommand(id: string) {
     let resolve: (val: any) => void = () => {}
     const promise = new Promise((_resolve, reject) => {
       resolve = _resolve