Feat: VSCode lineage "Only Direct Neighbors" filter (#5815)

tjirab · claude · web-flow · commit 0a5980cb6c1e · 2026-05-29T09:09:03.000-06:00
Signed-off-by: Bart Schuijt &lt;schuijt.bart@gmail.com&gt;
Co-authored-by: Claude Opus 4.7 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/vscode/extension/tests/lineage_settings.spec.ts b/vscode/extension/tests/lineage_settings.spec.ts
@@ -1,4 +1,5 @@
 import { test, expect } from './fixtures'
+import type { FrameLocator, Page } from '@playwright/test'
 import fs from 'fs-extra'
 import {
   openLineageView,
@@ -8,6 +9,31 @@ import {
 } from './utils'
 import { createPythonInterpreterSettingsSpecifier } from './utils_code_server'
 
+/**
+ * Find the iframe that hosts the lineage UI (the one containing the
+ * Settings cog button). Returns null if it can't be located.
+ */
+async function findLineageFrame(page: Page): Promise<FrameLocator | null> {
+  const iframes = page.locator('iframe')
+  const iframeCount = await iframes.count()
+
+  for (let i = 0; i < iframeCount; i++) {
+    const contentFrame = iframes.nth(i).contentFrame()
+    if (!contentFrame) continue
+    const activeFrame = contentFrame.locator('#active-frame').contentFrame()
+    if (!activeFrame) continue
+    try {
+      await activeFrame
+        .getByRole('button', { name: 'Settings' })
+        .waitFor({ timeout: 1000 })
+      return activeFrame
+    } catch {
+      continue
+    }
+  }
+  return null
+}
+
 test('Settings button is visible in the lineage view', async ({
   page,
   sharedCodeServer,
@@ -35,30 +61,64 @@ test('Settings button is visible in the lineage view', async ({
   // Open lineage
   await openLineageView(page)
 
-  const iframes = page.locator('iframe')
-  const iframeCount = await iframes.count()
-  let settingsCount = 0
+  const lineageFrame = await findLineageFrame(page)
+  expect(lineageFrame).not.toBeNull()
+})
 
-  for (let i = 0; i < iframeCount; i++) {
-    const iframe = iframes.nth(i)
-    const contentFrame = iframe.contentFrame()
-    if (contentFrame) {
-      const activeFrame = contentFrame.locator('#active-frame').contentFrame()
-      if (activeFrame) {
-        try {
-          await activeFrame
-            .getByRole('button', {
-              name: 'Settings',
-            })
-            .waitFor({ timeout: 1000 })
-          settingsCount++
-        } catch {
-          // Continue to next iframe if this one doesn't have the error
-          continue
-        }
-      }
-    }
-  }
+test('Only Direct Neighbors toggle filters the lineage graph', async ({
+  page,
+  sharedCodeServer,
+  tempDir,
+}) => {
+  await fs.copy(SUSHI_SOURCE_PATH, tempDir)
+  await createPythonInterpreterSettingsSpecifier(tempDir)
 
-  expect(settingsCount).toBeGreaterThan(0)
+  await openServerPage(page, tempDir, sharedCodeServer)
+  await page.waitForSelector('text=models')
+
+  await page
+    .getByRole('treeitem', { name: 'models', exact: true })
+    .locator('a')
+    .click()
+  await page
+    .getByRole('treeitem', { name: 'waiters.py', exact: true })
+    .locator('a')
+    .click()
+  await waitForLoadedSQLMesh(page)
+
+  await openLineageView(page)
+
+  const lineageFrame = await findLineageFrame(page)
+  expect(lineageFrame).not.toBeNull()
+  if (!lineageFrame) return
+
+  // Wait for the graph to render at least one node
+  await lineageFrame.locator('.react-flow__node').first().waitFor()
+  const nodesBefore = await lineageFrame.locator('.react-flow__node').count()
+  expect(nodesBefore).toBeGreaterThan(0)
+
+  // Open the settings menu and toggle "Only Direct Neighbors"
+  await lineageFrame.getByRole('button', { name: 'Settings' }).click()
+  const toggle = lineageFrame.getByRole('button', {
+    name: 'Only Direct Neighbors',
+  })
+  await toggle.waitFor()
+  await toggle.click()
+
+  // After enabling, the visible node set must be a subset of the original.
+  // We assert a strict drop only when the original graph had room to shrink
+  // (i.e. more than the worst-case direct-neighbor count of 1 + parents + children).
+  await page.waitForTimeout(250) // let React Flow re-layout
+  const nodesAfter = await lineageFrame.locator('.react-flow__node').count()
+  expect(nodesAfter).toBeLessThanOrEqual(nodesBefore)
+  expect(nodesAfter).toBeGreaterThan(0) // main node is always shown
+
+  // Toggle off → graph returns to the full size
+  await lineageFrame.getByRole('button', { name: 'Settings' }).click()
+  await lineageFrame
+    .getByRole('button', { name: 'Only Direct Neighbors' })
+    .click()
+  await page.waitForTimeout(250)
+  const nodesRestored = await lineageFrame.locator('.react-flow__node').count()
+  expect(nodesRestored).toBe(nodesBefore)
 })
diff --git a/vscode/react/src/components/graph/ModelLineage.tsx b/vscode/react/src/components/graph/ModelLineage.tsx
@@ -209,6 +209,8 @@ function ModelColumnLineage(): JSX.Element {
     withConnected,
     withImpacted,
     withSecondary,
+    withOnlyDirect,
+    directNeighbors,
     hasBackground,
     activeEdges,
     connectedNodes,
@@ -217,6 +219,7 @@ function ModelColumnLineage(): JSX.Element {
     handleError,
     setActiveNodes,
     setWithColumns,
+    setWithOnlyDirect,
   } = useLineageFlow()
 
   const { setCenter } = useReactFlow()
@@ -252,6 +255,8 @@ function ModelColumnLineage(): JSX.Element {
       withConnected,
       withImpacted,
       withSecondary,
+      withOnlyDirect,
+      directNeighbors,
     )
     const newEdges = getUpdatedEdges(
       allEdges,
@@ -264,6 +269,8 @@ function ModelColumnLineage(): JSX.Element {
       withConnected,
       withImpacted,
       withSecondary,
+      withOnlyDirect,
+      directNeighbors,
     )
     const createLayout = createGraphLayout({
       nodesMap,
@@ -324,6 +331,8 @@ function ModelColumnLineage(): JSX.Element {
       withConnected,
       withImpacted,
       withSecondary,
+      withOnlyDirect,
+      directNeighbors,
     )
 
     const newEdges = getUpdatedEdges(
@@ -337,6 +346,8 @@ function ModelColumnLineage(): JSX.Element {
       withConnected,
       withImpacted,
       withSecondary,
+      withOnlyDirect,
+      directNeighbors,
     )
 
     setEdges(newEdges)
@@ -353,6 +364,8 @@ function ModelColumnLineage(): JSX.Element {
     withConnected,
     withImpacted,
     withSecondary,
+    withOnlyDirect,
+    directNeighbors,
     withColumns,
     mainNode,
   ])
@@ -395,6 +408,8 @@ function ModelColumnLineage(): JSX.Element {
           <SettingsControl
             showColumns={withColumns}
             onWithColumnsChange={setWithColumns}
+            withOnlyDirect={withOnlyDirect}
+            onWithOnlyDirectChange={setWithOnlyDirect}
           />
         </Controls>
         <Background
diff --git a/vscode/react/src/components/graph/SettingsControl.tsx b/vscode/react/src/components/graph/SettingsControl.tsx
@@ -6,11 +6,21 @@ import clsx from 'clsx'
 interface SettingsControlProps {
   showColumns: boolean
   onWithColumnsChange: (value: boolean) => void
+  withOnlyDirect: boolean
+  onWithOnlyDirectChange: (value: boolean) => void
 }
 
+const itemClass = clsx(
+  'group flex w-full items-center px-2 py-1 text-sm',
+  'text-[var(--vscode-button-foreground)]',
+  'hover:bg-[var(--vscode-button-background)] bg-[var(--vscode-button-hoverBackground)]',
+)
+
 export function SettingsControl({
   showColumns,
   onWithColumnsChange,
+  withOnlyDirect,
+  onWithOnlyDirectChange,
 }: SettingsControlProps): JSX.Element {
   return (
     <Menu
@@ -29,11 +39,7 @@ export function SettingsControl({
       <MenuItems className="absolute bottom-0 left-full ml-2 w-56 origin-bottom-left divide-y bg-theme shadow-lg focus:outline-none z-50">
         <MenuItem
           as="button"
-          className={clsx(
-            'group flex w-full items-center px-2 py-1 text-sm',
-            'text-[var(--vscode-button-foreground)]',
-            'hover:bg-[var(--vscode-button-background)] bg-[var(--vscode-button-hoverBackground)]',
-          )}
+          className={itemClass}
           onClick={() => onWithColumnsChange(!showColumns)}
         >
           <span className="flex-1 text-left">Show Columns</span>
@@ -44,6 +50,19 @@ export function SettingsControl({
             />
           )}
         </MenuItem>
+        <MenuItem
+          as="button"
+          className={itemClass}
+          onClick={() => onWithOnlyDirectChange(!withOnlyDirect)}
+        >
+          <span className="flex-1 text-left">Only Direct Neighbors</span>
+          {withOnlyDirect && (
+            <CheckIcon
+              className="h-4 w-4 text-primary-500"
+              aria-hidden="true"
+            />
+          )}
+        </MenuItem>
       </MenuItems>
     </Menu>
   )
diff --git a/vscode/react/src/components/graph/context.tsx b/vscode/react/src/components/graph/context.tsx
@@ -7,6 +7,7 @@ import {
 } from 'react'
 import { getNodeMap, hasActiveEdge, hasActiveEdgeConnector } from './help'
 import { type Node } from 'reactflow'
+import { isNil } from '@/utils/index'
 import type { Lineage } from '@/domain/lineage'
 import type { ModelSQLMeshModel } from '@/domain/sqlmesh-model'
 import type { Column } from '@/domain/column'
@@ -43,6 +44,8 @@ interface LineageFlow {
   hasBackground: boolean
   withImpacted: boolean
   withSecondary: boolean
+  withOnlyDirect: boolean
+  directNeighbors: Set<ModelEncodedFQN>
   manuallySelectedColumn?: [ModelSQLMeshModel, Column]
   highlightedNodes: HighlightedNodes
   nodesMap: Record<ModelEncodedFQN, Node>
@@ -55,6 +58,7 @@ interface LineageFlow {
   setHasBackground: React.Dispatch<React.SetStateAction<boolean>>
   setWithImpacted: React.Dispatch<React.SetStateAction<boolean>>
   setWithSecondary: React.Dispatch<React.SetStateAction<boolean>>
+  setWithOnlyDirect: React.Dispatch<React.SetStateAction<boolean>>
   setConnections: React.Dispatch<React.SetStateAction<Map<string, Connections>>>
   hasActiveEdge: (edge: [string | undefined, string | undefined]) => boolean
   addActiveEdges: (edges: Array<[string, string]>) => void
@@ -85,6 +89,8 @@ export const LineageFlowContext = createContext<LineageFlow>({
   withConnected: false,
   withImpacted: true,
   withSecondary: false,
+  withOnlyDirect: false,
+  directNeighbors: new Set(),
   hasBackground: true,
   mainNode: undefined,
   activeEdges: new Map(),
@@ -103,6 +109,7 @@ export const LineageFlowContext = createContext<LineageFlow>({
   setWithImpacted: () => false,
   setWithSecondary: () => false,
   setWithConnected: () => false,
+  setWithOnlyDirect: () => false,
   hasActiveEdge: () => false,
   addActiveEdges: () => {},
   removeActiveEdges: () => {},
@@ -161,6 +168,7 @@ export default function LineageFlowProvider({
   const [hasBackground, setHasBackground] = useState(true)
   const [withImpacted, setWithImpacted] = useState(true)
   const [withSecondary, setWithSecondary] = useState(false)
+  const [withOnlyDirect, setWithOnlyDirect] = useState(false)
 
   const nodesMap = useMemo(
     () =>
@@ -264,6 +272,39 @@ export default function LineageFlowProvider({
     [nodesConnections],
   )
 
+  // Reverse adjacency index: parent -> children. Built once per `lineage`
+  // change so per-model `directNeighbors` lookups stay O(parents + children)
+  // instead of scanning the whole graph on every mainNode switch.
+  const childrenByParent = useMemo(() => {
+    const map = new Map<ModelEncodedFQN, ModelEncodedFQN[]>()
+    for (const [child, info] of Object.entries(lineage) as Array<
+      [ModelEncodedFQN, Lineage]
+    >) {
+      for (const parent of info?.models ?? []) {
+        const existing = map.get(parent)
+        if (existing) {
+          existing.push(child)
+        } else {
+          map.set(parent, [child])
+        }
+      }
+    }
+    return map
+  }, [lineage])
+
+  const directNeighbors = useMemo(() => {
+    const set = new Set<ModelEncodedFQN>()
+    if (isNil(mainNode)) return set
+    set.add(mainNode)
+    for (const parent of lineage[mainNode]?.models ?? []) {
+      set.add(parent)
+    }
+    for (const child of childrenByParent.get(mainNode) ?? []) {
+      set.add(child)
+    }
+    return set
+  }, [mainNode, lineage, childrenByParent])
+
   const selectedEdges = useMemo(
     () =>
       Array.from(selectedNodes)
@@ -292,6 +333,8 @@ export default function LineageFlowProvider({
         withConnected,
         withImpacted,
         withSecondary,
+        withOnlyDirect,
+        directNeighbors,
         showControls,
         hasBackground,
         nodesMap,
@@ -304,6 +347,7 @@ export default function LineageFlowProvider({
         setWithConnected,
         setWithImpacted,
         setWithSecondary,
+        setWithOnlyDirect,
         setHasBackground,
         setSelectedNodes,
         setMainNode,
diff --git a/vscode/react/src/components/graph/help.ts b/vscode/react/src/components/graph/help.ts
