Refactor Zoodle into runtime scene extension

Author: franknoirot

src/components/MlEphantConversation.spec.tsx

@@ -8,6 +8,10 @@ import {
 } from '@testing-library/react'
 import { afterEach, beforeEach, describe, expect, test, vi } from 'vitest'
 
+const bootMockState = vi.hoisted(() => ({
+  registry: undefined as unknown,
+}))
+
 // Mock modules that access localStorage at import time
 vi.mock('@src/routes/utils', () => ({
   getAppVersion: () => 'test',
@@ -23,6 +27,13 @@ vi.mock('@src/lib/desktop', () => ({
 
 // Mock useSingletons which requires heavy initialization
 vi.mock('@src/lib/boot', () => ({
+  useApp: () => {
+    if (!bootMockState.registry) {
+      throw new Error('Test registry has not been configured')
+    }
+
+    return { registry: bootMockState.registry }
+  },
   useSingletons: () => ({
     kclManager: {
       astSignal: { value: null },
@@ -39,6 +50,8 @@ vi.mock('@src/lib/screenshot', async (importOriginal) => {
   }
 })
 
+import { Registry } from '@kittycad/registry'
+import { useSignals } from '@preact/signals-react/runtime'
 import { MAKEATHON_ANNOUNCEMENT_DISMISSED_STORAGE_KEY } from '@src/components/MakeathonAnnouncement'
 import { MlEphantConversation } from '@src/components/MlEphantConversation'
 import { takeViewportScreenshot } from '@src/lib/screenshot'
@@ -49,6 +62,40 @@ import type {
   MlCopilotModeId,
   MlCopilotModeOption,
 } from '@src/machines/mlEphantManagerMachine'
+import {
+  type EngineSceneExtensionContext,
+  engineSceneRuntimeExtensionsSlot,
+  engineSceneStreamLayersValueSpec,
+} from '@src/registry/contracts/engineScene'
+
+const configureTestRegistry = () => {
+  const registry = new Registry()
+  registry.configure([engineSceneRuntimeExtensionsSlot.of()])
+  bootMockState.registry = registry
+  return registry
+}
+
+const testEngineSceneContext = {
+  modelingState: { matches: () => false },
+  modelingSend: vi.fn(),
+  sketchSolveStreamDimming: 0.8,
+  setSketchSolveStreamDimming: vi.fn(),
+} as unknown as EngineSceneExtensionContext
+
+const TestEngineSceneStreamLayers = () => {
+  useSignals()
+  const registry = bootMockState.registry as Registry
+  const layers = registry.signal(engineSceneStreamLayersValueSpec).value
+
+  return (
+    <>
+      {layers.map((layer) => {
+        const Component = layer.Component
+        return <Component key={layer.id} {...testEngineSceneContext} />
+      })}
+    </>
+  )
+}
 
 const SERVER_MODE_OPTIONS: MlCopilotModeOption[] = [
   {
@@ -69,6 +116,7 @@ const SERVER_MODE_OPTIONS: MlCopilotModeOption[] = [
 
 describe('MlEphantConversation', () => {
   beforeEach(() => {
+    configureTestRegistry()
     window.localStorage.removeItem(MAKEATHON_ANNOUNCEMENT_DISMISSED_STORAGE_KEY)
   })
 
@@ -977,22 +1025,25 @@ describe('MlEphantConversation', () => {
 
     const renderConversation = (handleProcess = vi.fn(), disabled = false) => {
       return render(
-        <MlEphantConversation
-          isLoading={false}
-          conversation={{ exchanges: [] }}
-          onProcess={handleProcess}
-          onClickClearChat={() => {}}
-          onReconnect={() => {}}
-          onCancel={() => {}}
-          needsReconnect={false}
-          disabled={disabled}
-          hasPromptCompleted={true}
-          contexts={[]}
-          isProcessing={false}
-          queue={[]}
-          onRemoveFromQueue={() => {}}
-          onSteer={() => {}}
-        />
+        <>
+          <MlEphantConversation
+            isLoading={false}
+            conversation={{ exchanges: [] }}
+            onProcess={handleProcess}
+            onClickClearChat={() => {}}
+            onReconnect={() => {}}
+            onCancel={() => {}}
+            needsReconnect={false}
+            disabled={disabled}
+            hasPromptCompleted={true}
+            contexts={[]}
+            isProcessing={false}
+            queue={[]}
+            onRemoveFromQueue={() => {}}
+            onSteer={() => {}}
+          />
+          <TestEngineSceneStreamLayers />
+        </>
       )
     }
 

src/components/MlEphantConversation.tsx

@@ -5,8 +5,8 @@ import { isExternalFileDrag } from '@src/components/Explorer/utils'
 import Loading from '@src/components/Loading'
 import { MakeathonAnnouncement } from '@src/components/MakeathonAnnouncement'
 import Tooltip from '@src/components/Tooltip'
-import { ViewportAnnotationOverlay } from '@src/components/ViewportAnnotationOverlay'
 import { noAutofillInputProps } from '@src/lib/autofill'
+import { useApp } from '@src/lib/boot'
 import { dataUrlToFile, takeViewportScreenshot } from '@src/lib/screenshot'
 import { err } from '@src/lib/trap'
 import { isNonNullable } from '@src/lib/utils'
@@ -17,8 +17,12 @@ import type {
   MlCopilotModeOption,
 } from '@src/machines/mlEphantManagerMachine'
 import type { Selections } from '@src/machines/modelingSharedTypes'
+import {
+  activateZoodleRuntimeExtension,
+  deactivateZoodleRuntimeExtension,
+} from '@src/registry/extensions/engineScene/zoodleRuntimeExtension'
 import type { ChangeEvent, ReactNode } from 'react'
-import { useEffect, useRef, useState } from 'react'
+import { useCallback, useEffect, useRef, useState } from 'react'
 
 const noop = () => {}
 
@@ -263,6 +267,7 @@ interface MlEphantConversationInputProps {
 export const MlEphantConversationInput = (
   props: MlEphantConversationInputProps
 ) => {
+  const { registry } = useApp()
   const refDiv = useRef<HTMLTextAreaElement>(null)
   const fileInputRef = useRef<HTMLInputElement>(null)
   const [value, setValue] = useState<string>('')
@@ -273,9 +278,22 @@ export const MlEphantConversationInput = (
   const lastModeScopeKey = useRef(props.modeScopeKey)
   const [attachments, setAttachments] = useState<File[]>([])
   const [isDraggingOver, setIsDraggingOver] = useState(false)
-  const [annotationImageDataUrl, setAnnotationImageDataUrl] = useState<
-    string | null
-  >(null)
+  const zoodleRuntimeExtensionActive = useRef(false)
+
+  const stopZoodleRuntimeExtension = useCallback(() => {
+    if (!zoodleRuntimeExtensionActive.current) {
+      return
+    }
+
+    zoodleRuntimeExtensionActive.current = false
+    deactivateZoodleRuntimeExtension(registry)
+  }, [registry])
+
+  useEffect(() => {
+    return () => {
+      stopZoodleRuntimeExtension()
+    }
+  }, [stopZoodleRuntimeExtension])
 
   // Without this the cursor ends up at the start of the text
   useEffect(() => setValue(props.defaultPrompt || ''), [props.defaultPrompt])
@@ -395,7 +413,18 @@ export const MlEphantConversationInput = (
     try {
       const dataUrl = takeViewportScreenshot()
       if (!dataUrl) return
-      setAnnotationImageDataUrl(dataUrl)
+      zoodleRuntimeExtensionActive.current = true
+      activateZoodleRuntimeExtension(registry, {
+        imageDataUrl: dataUrl,
+        onCancel: stopZoodleRuntimeExtension,
+        onSend: (annotatedDataUrl) => {
+          appendDataUrlAttachment(
+            annotatedDataUrl,
+            'annotated-viewport-screenshot.png'
+          )
+          stopZoodleRuntimeExtension()
+        },
+      })
     } catch (e) {
       console.error('Failed to capture viewport screenshot for annotation', e)
     }
@@ -590,19 +619,6 @@ export const MlEphantConversationInput = (
         Zookeeper can make mistakes. We send selection context to help. Always
         verify information.
       </div>
-      {annotationImageDataUrl && (
-        <ViewportAnnotationOverlay
-          imageDataUrl={annotationImageDataUrl}
-          onCancel={() => setAnnotationImageDataUrl(null)}
-          onSend={(annotatedDataUrl) => {
-            appendDataUrlAttachment(
-              annotatedDataUrl,
-              'annotated-viewport-screenshot.png'
-            )
-            setAnnotationImageDataUrl(null)
-          }}
-        />
-      )}
     </div>
   )
 }

src/components/ViewportAnnotationOverlay.tsx

@@ -1,9 +1,7 @@
 import { TRIM_PREVIEW_LINE_COLOR } from '@src/lib/constants'
 import { getTrimPreviewLineWidth } from '@src/lib/freehandLineDrawing'
-import { getVisibleViewportRect } from '@src/lib/viewportElement'
-import type { PointerEvent as ReactPointerEvent } from 'react'
+import type { PointerEvent as ReactPointerEvent, SyntheticEvent } from 'react'
 import { useEffect, useRef, useState } from 'react'
-import { createPortal } from 'react-dom'
 
 const loadImage = (src: string): Promise<HTMLImageElement> =>
   new Promise((resolve, reject) => {
@@ -43,6 +41,10 @@ const getCanvasPoint = (
   }
 }
 
+const stopOverlayEventPropagation = (event: SyntheticEvent) => {
+  event.stopPropagation()
+}
+
 interface ViewportAnnotationOverlayProps {
   imageDataUrl: string
   onCancel: () => void
@@ -54,7 +56,6 @@ export const ViewportAnnotationOverlay = (
 ) => {
   const canvasRef = useRef<HTMLCanvasElement>(null)
   const lastPoint = useRef<{ x: number; y: number } | null>(null)
-  const [viewportRect, setViewportRect] = useState(getVisibleViewportRect)
   const [imageSize, setImageSize] = useState<{
     width: number
     height: number
@@ -90,27 +91,6 @@ export const ViewportAnnotationOverlay = (
     }
   }, [props.imageDataUrl])
 
-  useEffect(() => {
-    const updateViewportRect = () => setViewportRect(getVisibleViewportRect())
-    const viewport = document.querySelector('[data-engine]')
-    const resizeObserver =
-      typeof ResizeObserver === 'undefined'
-        ? undefined
-        : new ResizeObserver(updateViewportRect)
-
-    if (viewport instanceof HTMLElement) {
-      resizeObserver?.observe(viewport)
-    }
-    window.addEventListener('resize', updateViewportRect)
-    window.addEventListener('scroll', updateViewportRect, true)
-
-    return () => {
-      resizeObserver?.disconnect()
-      window.removeEventListener('resize', updateViewportRect)
-      window.removeEventListener('scroll', updateViewportRect, true)
-    }
-  }, [])
-
   const drawTo = (event: ReactPointerEvent<HTMLCanvasElement>) => {
     const canvas = canvasRef.current
     if (!canvas || !lastPoint.current) {
@@ -204,16 +184,20 @@ export const ViewportAnnotationOverlay = (
       })
   }
 
-  return createPortal(
+  return (
     <div
       data-testid="viewport-annotation-overlay"
-      className="fixed z-50 bg-chalkboard-100"
-      style={{
-        left: viewportRect.left,
-        top: viewportRect.top,
-        width: viewportRect.width,
-        height: viewportRect.height,
-      }}
+      className="absolute inset-0 z-50 bg-chalkboard-100"
+      onPointerDown={stopOverlayEventPropagation}
+      onPointerMove={stopOverlayEventPropagation}
+      onPointerUp={stopOverlayEventPropagation}
+      onPointerCancel={stopOverlayEventPropagation}
+      onMouseDown={stopOverlayEventPropagation}
+      onMouseMove={stopOverlayEventPropagation}
+      onMouseUp={stopOverlayEventPropagation}
+      onClick={stopOverlayEventPropagation}
+      onDoubleClick={stopOverlayEventPropagation}
+      onWheel={stopOverlayEventPropagation}
     >
       <img
         src={props.imageDataUrl}
@@ -257,7 +241,6 @@ export const ViewportAnnotationOverlay = (
           Send
         </button>
       </div>
-    </div>,
-    document.body
+    </div>
   )
 }

src/components/layout/areas/MlEphantConversationPane.spec.tsx

@@ -17,6 +17,11 @@ vi.mock('@src/lib/desktop', () => ({
 }))
 
 vi.mock('@src/lib/boot', () => ({
+  useApp: () => ({
+    registry: {
+      reconfigure: vi.fn(),
+    },
+  }),
   useSingletons: () => ({
     kclManager: {
       astSignal: { value: null },

src/lib/app.ts

@@ -76,6 +76,7 @@ import {
   commandSystemService,
   provideCommand,
 } from '@src/registry/contracts/commands'
+import { engineSceneRuntimeExtensionsSlot } from '@src/registry/contracts/engineScene'
 import { executingEditorService } from '@src/registry/contracts/executingEditor'
 import { keymapService } from '@src/registry/contracts/keymap'
 import { layoutContributionsValueSpec } from '@src/registry/contracts/layout'
@@ -163,6 +164,7 @@ function createAppRegistryItems({
       : []),
     appCommandsSlot.of(),
     appRegistryServicesSlot.of(),
+    engineSceneRuntimeExtensionsSlot.of(),
     ...coreRegistryItems,
   ]
 }

src/registry/contracts/engineScene.tsx

@@ -1,4 +1,4 @@
-import { defineContract, defineValueSpec } from '@kittycad/registry'
+import { Slot, defineContract, defineValueSpec } from '@kittycad/registry'
 import type { modelingMachine } from '@src/machines/modelingMachine'
 import type { ComponentType, Dispatch, SetStateAction } from 'react'
 import type { EventFrom, StateFrom } from 'xstate'
@@ -18,6 +18,8 @@ export const engineSceneViewExtensionZones = [
 export type EngineSceneViewExtensionZone =
   (typeof engineSceneViewExtensionZones)[number]
 
+export const engineSceneRuntimeExtensionsSlot = new Slot()
+
 export type EngineSceneExtensionContext = {
   modelingState: StateFrom<typeof modelingMachine>
   modelingSend: (event: EventFrom<typeof modelingMachine>) => void