From f330bdabd228feeacc45f51b1d4e5321bc113e4e Mon Sep 17 00:00:00 2001 From: Bart Schuijt Date: Thu, 28 May 2026 09:32:34 +0200 Subject: [PATCH 1/2] 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) Signed-off-by: Bart Schuijt --- .../src/components/graph/ModelLineage.tsx | 15 ++++++++++ .../src/components/graph/SettingsControl.tsx | 28 +++++++++++++++---- vscode/react/src/components/graph/context.tsx | 28 +++++++++++++++++++ vscode/react/src/components/graph/help.ts | 17 ++++++++++- 4 files changed, 82 insertions(+), 6 deletions(-) diff --git a/vscode/react/src/components/graph/ModelLineage.tsx b/vscode/react/src/components/graph/ModelLineage.tsx index 3d157d3869..1ff9110f7e 100644 --- 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 { 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 ( onWithColumnsChange(!showColumns)} > Show Columns @@ -44,6 +49,19 @@ export function SettingsControl({ /> )} + onOnlyDirectChange(!onlyDirect)} + > + Only Direct Neighbors + {onlyDirect && ( + ) diff --git a/vscode/react/src/components/graph/context.tsx b/vscode/react/src/components/graph/context.tsx index 9ab4f0722e..5ae865a809 100644 --- 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 manuallySelectedColumn?: [ModelSQLMeshModel, Column] highlightedNodes: HighlightedNodes nodesMap: Record @@ -55,6 +58,7 @@ interface LineageFlow { setHasBackground: React.Dispatch> setWithImpacted: React.Dispatch> setWithSecondary: React.Dispatch> + setWithOnlyDirect: React.Dispatch> setConnections: React.Dispatch>> hasActiveEdge: (edge: [string | undefined, string | undefined]) => boolean addActiveEdges: (edges: Array<[string, string]>) => void @@ -85,6 +89,8 @@ export const LineageFlowContext = createContext({ 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({ 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() + 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, diff --git a/vscode/react/src/components/graph/help.ts b/vscode/react/src/components/graph/help.ts index 93e5c4db45..8b9c8881f9 100644 --- a/vscode/react/src/components/graph/help.ts +++ b/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 = 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 = 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) } From 619205811b31783f942574c68bebf4e3171e5d0a Mon Sep 17 00:00:00 2001 From: Bart Schuijt Date: Thu, 28 May 2026 09:47:53 +0200 Subject: [PATCH 2/2] Refactor: address PR feedback on direct-neighbors filter - Build reverse parent->children adjacency map once per lineage change so directNeighbors stays O(parents + children) per mainNode switch instead of scanning the full graph each time. - Hoist SettingsControl itemClass to module scope; rename props to withOnlyDirect / onWithOnlyDirectChange for consistency with the rest of the codebase. - Extend lineage_settings.spec.ts with a Playwright test that toggles "Only Direct Neighbors" and asserts the visible node count drops and is restored on toggle-off. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Bart Schuijt --- .../extension/tests/lineage_settings.spec.ts | 108 ++++++++++++++---- .../src/components/graph/ModelLineage.tsx | 4 +- .../src/components/graph/SettingsControl.tsx | 23 ++-- vscode/react/src/components/graph/context.tsx | 30 +++-- 4 files changed, 121 insertions(+), 44 deletions(-) diff --git a/vscode/extension/tests/lineage_settings.spec.ts b/vscode/extension/tests/lineage_settings.spec.ts index c3237f13dc..c2bd4a1e17 100644 --- 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 { + 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 index 1ff9110f7e..df8745140c 100644 --- a/vscode/react/src/components/graph/ModelLineage.tsx +++ b/vscode/react/src/components/graph/ModelLineage.tsx @@ -408,8 +408,8 @@ function ModelColumnLineage(): JSX.Element { void - onlyDirect: boolean - onOnlyDirectChange: (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, - onlyDirect, - onOnlyDirectChange, + withOnlyDirect, + onWithOnlyDirectChange, }: 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 ( onOnlyDirectChange(!onlyDirect)} + onClick={() => onWithOnlyDirectChange(!withOnlyDirect)} > Only Direct Neighbors - {onlyDirect && ( + {withOnlyDirect && (