Feat: Add "Only Direct Neighbors" filter to VSCode lineage panel

Selecting a model in the VSCode lineage panel renders the full transitive
lineage, which can be dense for hub models. This adds a new toggle to the
existing settings menu that limits the view to the selected model plus
its direct parents and direct children.

The filter is implemented as a new withOnlyDirect boolean on
LineageFlowContext, paired with a directNeighbors set memoized from the
lineage map and mainNode. When enabled, getUpdatedNodes and
getUpdatedEdges hide anything outside that set, leaving existing
withConnected/withImpacted/withSecondary logic untouched.

Refs #5811

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Bart Schuijt <schuijt.bart@gmail.com>

Author: tjirab

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}
+            onlyDirect={withOnlyDirect}
+            onOnlyDirectChange={setWithOnlyDirect}
           />
         </Controls>
         <Background

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

@@ -6,12 +6,21 @@ import clsx from 'clsx'
 interface SettingsControlProps {
   showColumns: boolean
   onWithColumnsChange: (value: boolean) => void
+  onlyDirect: boolean
+  onOnlyDirectChange: (value: boolean) => void
 }
 
 export function SettingsControl({
   showColumns,
   onWithColumnsChange,
+  onlyDirect,
+  onOnlyDirectChange,
 }: SettingsControlProps): JSX.Element {
+  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)]',
+  )
   return (
     <Menu
       as="div"
@@ -29,11 +38,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 +49,19 @@ export function SettingsControl({
             />
           )}
         </MenuItem>
+        <MenuItem
+          as="button"
+          className={itemClass}
+          onClick={() => onOnlyDirectChange(!onlyDirect)}
+        >
+          <span className="flex-1 text-left">Only Direct Neighbors</span>
+          {onlyDirect && (
+            <CheckIcon
+              className="h-4 w-4 text-primary-500"
+              aria-hidden="true"
+            />
+          )}
+        </MenuItem>
       </MenuItems>
     </Menu>
   )

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,23 @@ export default function LineageFlowProvider({
     [nodesConnections],
   )
 
+  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, info] of Object.entries(lineage) as Array<
+      [ModelEncodedFQN, Lineage]
+    >) {
+      if (info?.models?.includes(mainNode)) {
+        set.add(child)
+      }
+    }
+    return set
+  }, [mainNode, lineage])
+
   const selectedEdges = useMemo(
     () =>
       Array.from(selectedNodes)
@@ -292,6 +317,8 @@ export default function LineageFlowProvider({
         withConnected,
         withImpacted,
         withSecondary,
+        withOnlyDirect,
+        directNeighbors,
         showControls,
         hasBackground,
         nodesMap,
@@ -304,6 +331,7 @@ export default function LineageFlowProvider({
         setWithConnected,
         setWithImpacted,
         setWithSecondary,
+        setWithOnlyDirect,
         setHasBackground,
         setSelectedNodes,
         setMainNode,

vscode/react/src/components/graph/help.ts

@@ -562,6 +562,8 @@ export function getUpdatedEdges(
   withConnected: boolean = false,
   withImpacted: boolean = false,
   withSecondary: boolean = false,
+  withOnlyDirect: boolean = false,
+  directNeighbors: Set<string> = new Set(),
 ): Edge[] {
   const tempEdges = edges.map(edge => {
     const isActiveEdge = hasActiveEdge(activeEdges, [
@@ -593,10 +595,15 @@ export function getUpdatedEdges(
         hasEdge(selectedEdges, edge.id) &&
         activeNodes.has(edge.source) &&
         activeNodes.has(edge.target)
+      const shouldHideNonDirect =
+        withOnlyDirect &&
+        (isFalse(directNeighbors.has(edge.source)) ||
+          isFalse(directNeighbors.has(edge.target)))
 
       if (
         isFalse(shouldHideImpacted) &&
         isFalse(shouldHideSecondary) &&
+        isFalse(shouldHideNonDirect) &&
         (isFalse(hasSelections) || isVisibleEdge)
       ) {
         edge.hidden = false
@@ -653,6 +660,8 @@ export function getUpdatedNodes(
   withConnected: boolean,
   withImpacted: boolean,
   withSecondary: boolean,
+  withOnlyDirect: boolean = false,
+  directNeighbors: Set<string> = new Set(),
 ): Node[] {
   return nodes.map(node => {
     node.hidden = true
@@ -667,8 +676,14 @@ export function getUpdatedNodes(
       isFalse(withSecondary) && isFalse(hasSelections)
     const shouldHideSecondary = isSecondaryNode && withoutSecondaryNodes
     const shouldHideImpacted = isImpactedNode && withoutImpactedNodes
+    const shouldHideNonDirect =
+      withOnlyDirect && isFalse(directNeighbors.has(node.id))
 
-    if (isFalse(shouldHideImpacted) && isFalse(shouldHideSecondary)) {
+    if (
+      isFalse(shouldHideImpacted) &&
+      isFalse(shouldHideSecondary) &&
+      isFalse(shouldHideNonDirect)
+    ) {
       node.hidden = isFalse(isActiveNode)
     }