chore(vscode): typing lineage and storybook

vscode/react/.gitignore

@@ -3,3 +3,6 @@ node_modules
 dist
 dist-ssr
 *.local
+
+*storybook.log
+storybook-static

vscode/react/.storybook/main.ts

@@ -0,0 +1,17 @@
+import type { StorybookConfig } from '@storybook/react-vite'
+
+const config: StorybookConfig = {
+  stories: ['../src/**/*.mdx', '../src/**/*.stories.@(js|jsx|mjs|ts|tsx)'],
+  addons: [
+    '@chromatic-com/storybook',
+    '@storybook/addon-docs',
+    '@storybook/addon-onboarding',
+    '@storybook/addon-a11y',
+    '@storybook/addon-vitest',
+  ],
+  framework: {
+    name: '@storybook/react-vite',
+    options: {},
+  },
+}
+export default config

vscode/react/.storybook/preview.ts

@@ -0,0 +1,22 @@
+import type { Preview } from '@storybook/react-vite'
+import './storybook.css'
+
+const preview: Preview = {
+  parameters: {
+    controls: {
+      matchers: {
+        color: /(background|color)$/i,
+        date: /Date$/i,
+      },
+    },
+
+    a11y: {
+      // 'todo' - show a11y violations in the test UI only
+      // 'error' - fail CI on a11y violations
+      // 'off' - skip a11y checks entirely
+      test: 'todo',
+    },
+  },
+}
+
+export default preview

vscode/react/.storybook/vitest.setup.ts

@@ -0,0 +1,7 @@
+import * as a11yAddonAnnotations from '@storybook/addon-a11y/preview'
+import { setProjectAnnotations } from '@storybook/react-vite'
+import * as projectAnnotations from './preview'
+
+// This is an important step to apply the right configuration when testing your stories.
+// More info at: https://storybook.js.org/docs/api/portable-stories/portable-stories-vitest#setprojectannotations
+setProjectAnnotations([a11yAddonAnnotations, projectAnnotations])

vscode/react/package.json

@@ -10,7 +10,9 @@
     "serve": "vite preview",
     "test": "vitest run",
     "generate:api": "orval --config ./orval.config.ts",
-    "lint": "tsc --noEmit"
+    "lint": "tsc --noEmit",
+    "storybook": "storybook dev -p 6006",
+    "build-storybook": "storybook build"
   },
   "dependencies": {
     "@headlessui/react": "^2.2.4",
@@ -35,15 +37,25 @@
     "vscode-uri": "^3.1.0"
   },
   "devDependencies": {
+    "@chromatic-com/storybook": "^4.0.1",
+    "@storybook/addon-a11y": "^9.0.15",
+    "@storybook/addon-docs": "^9.0.15",
+    "@storybook/addon-onboarding": "^9.0.15",
+    "@storybook/addon-vitest": "^9.0.15",
+    "@storybook/react-vite": "^9.0.15",
     "@testing-library/dom": "^10.4.0",
     "@testing-library/react": "^16.3.0",
     "@types/react": "^18.3.23",
     "@types/react-dom": "^18.3.7",
     "@vitejs/plugin-react": "^4.5.1",
     "jsdom": "^26.1.0",
+    "storybook": "^9.0.15",
     "typescript": "^5.8.3",
     "vite": "^6.3.5",
     "vitest": "^3.2.3",
-    "web-vitals": "^4.2.4"
+    "web-vitals": "^4.2.4",
+    "@vitest/browser": "3.2.3",
+    "playwright": "^1.53.2",
+    "@vitest/coverage-v8": "3.2.3"
   }
 }

vscode/react/src/components/graph/ModelColumns.tsx

@@ -8,10 +8,9 @@ import {
   isFalse,
   isNil,
   isNotNil,
-  toID,
   truncate,
 } from '@/utils/index'
-import { EnumSide, type Side } from './types'
+import { EnumSide, toID, type Side } from './types'
 import { NoSymbolIcon } from '@heroicons/react/24/solid'
 import { ClockIcon, ExclamationCircleIcon } from '@heroicons/react/24/outline'
 import clsx from 'clsx'
@@ -31,8 +30,7 @@ import { useLineageFlow } from './context'
 import { useApiColumnLineage } from '@/api/index'
 import SourceList from '@/components/sourceList/SourceList'
 import type { Lineage } from '@/domain/lineage'
-import type { ModelName } from '@/domain/models'
-import type { Column } from '@/domain/column'
+import type { Column, ColumnName } from '@/domain/column'
 
 export function ModelColumns({
   nodeId,
@@ -590,5 +588,5 @@ function getColumnFromLineage(
   nodeId: string,
   columnName: string,
 ): LineageColumn | undefined {
-  return lineage?.[nodeId]?.columns?.[encodeURI(columnName) as ModelName]
+  return lineage?.[nodeId]?.columns?.[columnName as ColumnName]
 }

vscode/react/src/components/graph/ModelLineage.tsx

@@ -35,9 +35,14 @@ import {
 import { Popover } from '@headlessui/react'
 import ModelLineageDetails from './ModelLineageDetails'
 import { Divider } from '@/components/divider/Divider'
-import { type ModelLineageApiLineageModelNameGet200 } from '@/api/client'
 import { SettingsControl } from '@/components/graph/SettingsControl'
+import {
+  toModelLineage,
+  type ModelLineage as ModelLineageType,
+} from '@/domain/lineage'
 import './Graph.css'
+import { toKeys } from './types'
+import { encode } from '@/domain/models'
 
 const WITH_COLUMNS_LIMIT = 30
 
@@ -81,7 +86,7 @@ export function ModelLineage({
 
   const [isMergingModels, setIsMergingModels] = useState(false)
   const [modelLineage, setModelLineage] = useState<
-    ModelLineageApiLineageModelNameGet200 | undefined
+    ModelLineageType | undefined
   >(undefined)
 
   useEffect(() => {
@@ -91,7 +96,7 @@ export function ModelLineage({
 
     getModelLineage()
       .then(({ data }) => {
-        setModelLineage(data)
+        setModelLineage(data ? toModelLineage(data) : undefined)
         if (isNil(data)) return
 
         setIsMergingModels(true)
@@ -129,13 +134,16 @@ export function ModelLineage({
   }, [model.name, model.hash])
 
   useEffect(() => {
-    Object.keys(modelLineage ?? {}).forEach(modelName => {
-      modelName = encodeURI(modelName)
-
-      if (isFalse(modelName in models) && isFalse(modelName in unknownModels)) {
-        unknownModels.add(modelName)
+    const modelNames = toKeys(modelLineage ?? {})
+    for (const modelName of modelNames) {
+      const encodedModelName = encode(modelName)
+      if (
+        isFalse(encodedModelName in models) &&
+        isFalse(encodedModelName in unknownModels)
+      ) {
+        unknownModels.add(encodedModelName)
       }
-    })
+    }
 
     setUnknownModels(new Set(unknownModels))
   }, [modelLineage, models])
@@ -329,6 +337,7 @@ function ModelColumnLineage(): JSX.Element {
 
     setEdges(newEdges)
     setNodes(newNodes)
+    console.log('newActiveNodes', newActiveNodes)
     setActiveNodes(newActiveNodes)
   }, [
     connections,

vscode/react/src/components/graph/ModelNode.tsx

@@ -8,6 +8,7 @@ import { Position, type NodeProps } from 'reactflow'
 import { ModelNodeHeaderHandles } from './ModelNodeHeaderHandles'
 import { ModelColumns } from './ModelColumns'
 import { fromAPIColumn, type Column } from '@/domain/column'
+import type { ModelEncodedFQN } from '@/domain/models'
 
 export const EnumLineageNodeModelType = {
   ...ModelType,
@@ -150,7 +151,7 @@ export default function ModelNode({
   const isActiveNode =
     selectedNodes.size > 0 || activeNodes.size > 0 || withConnected
       ? isSelected ||
-        activeNodes.has(id) ||
+        activeNodes.has(id as ModelEncodedFQN) ||
         (withConnected && connectedNodes.has(id))
       : connectedNodes.has(id)
   const isInteractive = true

vscode/react/src/components/graph/ModelNodeHeaderHandles.tsx

@@ -2,8 +2,8 @@ import { type MouseEvent } from 'react'
 import { Handle, Position } from 'reactflow'
 import 'reactflow/dist/base.css'
 import { getModelNodeTypeTitle } from './help'
-import { isNotNil, toID, truncate } from '@/utils/index'
-import { EnumSide } from './types'
+import { isNotNil, truncate } from '@/utils/index'
+import { EnumSide, toID } from './types'
 import { ArrowRightCircleIcon } from '@heroicons/react/24/solid'
 import clsx from 'clsx'
 import { type LineageNodeModelType } from './ModelNode'

vscode/react/src/components/graph/context.tsx

@@ -1,4 +1,3 @@
-import { type Model } from '@/api/client'
 import {
   createContext,
   useState,
@@ -12,14 +11,16 @@ import { type Node } from 'reactflow'
 import type { Lineage } from '@/domain/lineage'
 import type { ModelSQLMeshModel } from '@/domain/sqlmesh-model'
 import type { Column } from '@/domain/column'
+import type { ModelEncodedFQN } from '@/domain/models'
+import type { Model } from '@/api/client'
 
 export interface Connections {
   left: string[]
   right: string[]
 }
 export type ActiveColumns = Map<string, { ins: string[]; outs: string[] }>
 export type ActiveEdges = Map<string, Array<[string, string]>>
-export type ActiveNodes = Set<string>
+export type ActiveNodes = Set<ModelEncodedFQN>
 export type SelectedNodes = Set<string>
 export type HighlightedNodes = Record<string, string[]>
 
@@ -32,8 +33,8 @@ interface LineageFlow {
   activeNodes: ActiveNodes
   selectedNodes: SelectedNodes
   selectedEdges: any[]
-  models: Record<string, Model>
-  unknownModels: Set<string>
+  models: Record<string, ModelSQLMeshModel>
+  unknownModels: Set<ModelEncodedFQN>
   connections: Map<string, Connections>
   withConnected: boolean
   withColumns: boolean
@@ -58,7 +59,7 @@ interface LineageFlow {
   addActiveEdges: (edges: Array<[string, string]>) => void
   removeActiveEdges: (edges: Array<[string, string]>) => void
   setActiveEdges: React.Dispatch<React.SetStateAction<ActiveEdges>>
-  setUnknownModels: React.Dispatch<React.SetStateAction<Set<string>>>
+  setUnknownModels: React.Dispatch<React.SetStateAction<Set<ModelEncodedFQN>>>
   setLineage: React.Dispatch<React.SetStateAction<Record<string, Lineage>>>
   setLineageCache: React.Dispatch<
     React.SetStateAction<Record<string, Lineage> | undefined>
@@ -135,7 +136,7 @@ export default function LineageFlowProvider({
   models: Record<string, Model>
 }): JSX.Element {
   const [lineage, setLineage] = useState<Record<string, Lineage>>({})
-  const [unknownModels, setUnknownModels] = useState(new Set<string>())
+  const [unknownModels, setUnknownModels] = useState(new Set<ModelEncodedFQN>())
   const [lineageCache, setLineageCache] = useState<
     Record<string, Lineage> | undefined
   >(undefined)
@@ -162,6 +163,7 @@ export default function LineageFlowProvider({
     () =>
       getNodeMap({
         lineage,
+        // @ts-expect-error TODO: fix this, should move to internal representation
         models,
         unknownModels,
         withColumns,
@@ -278,6 +280,7 @@ export default function LineageFlowProvider({
         connections,
         lineage,
         lineageCache,
+        // @ts-expect-error TODO: fix this, should move to internal representation
         models,
         manuallySelectedColumn,
         withColumns,