Wire up LSP data source for visual designer (Migration Step 4)

src/vscode-bicep-ui/apps/visual-designer/src/App.tsx

@@ -1,29 +1,43 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
+import type { WebviewMessageChannel } from "@vscode-bicep-ui/messaging";
 import type { ComponentType } from "react";
 import type { NodeKind } from "./features/graph-engine/atoms";
+import type { DeploymentGraphPayload } from "./messages";
 
-import { PanZoomProvider } from "@vscode-bicep-ui/components";
-import { getDefaultStore, useSetAtom } from "jotai";
-import { useEffect } from "react";
+import { PanZoomProvider, useGetPanZoomDimensions, usePanZoomControl } from "@vscode-bicep-ui/components";
+import {
+  useWebviewMessageChannel,
+  useWebviewNotification,
+  WebviewMessageChannelProvider,
+} from "@vscode-bicep-ui/messaging";
+import { getDefaultStore, useAtomValue } from "jotai";
+import { lazy, Suspense, useCallback, useEffect, useLayoutEffect, useRef, useState } from "react";
 import { styled, ThemeProvider } from "styled-components";
 import { GraphControlBar } from "./features/design-view/components/GraphControlBar";
 import { ModuleDeclaration } from "./features/design-view/components/ModuleDeclaration";
 import { ResourceDeclaration } from "./features/design-view/components/ResourceDeclaration";
-import {
-  addAtomicNodeAtom,
-  addCompoundNodeAtom,
-  addEdgeAtom,
-  edgesAtom,
-  nodeConfigAtom,
-  nodesAtom,
-} from "./features/graph-engine/atoms";
+import { graphVersionAtom, nodeConfigAtom } from "./features/graph-engine/atoms";
 import { Canvas, Graph } from "./features/graph-engine/components";
-import { runLayout } from "./features/graph-engine/layout/elk-layout";
+import { applyLayout, computeFitViewTransform, computeLayout } from "./features/graph-engine/layout/elk-layout";
 import { GlobalStyle } from "./GlobalStyle";
+import { useApplyDeploymentGraph } from "./hooks/useDeploymentGraph";
+import { DEPLOYMENT_GRAPH_NOTIFICATION, READY_NOTIFICATION } from "./messages";
 import { useTheme } from "./theming/useTheme";
 
+const isDev = typeof acquireVsCodeApi === "undefined";
+
+// Lazy-load dev-only modules so they are code-split into a separate
+// chunk and never downloaded in production (where acquireVsCodeApi exists).
+const LazyDevToolbar = isDev
+  ? lazy(() => import("./dev/DevToolbar").then((m) => ({ default: m.DevToolbar })))
+  : undefined;
+
+const loadFakeMessageChannel = isDev
+  ? () => import("./dev/FakeMessageChannel").then((m) => new m.FakeMessageChannel())
+  : undefined;
+
 const store = getDefaultStore();
 const nodeConfig = store.get(nodeConfigAtom);
 
@@ -40,67 +54,142 @@ store.set(nodeConfigAtom, {
     ...nodeConfig.padding,
     top: 50,
   },
-  getContentComponent: (kind: NodeKind) =>
-    (kind === "atomic" ? ResourceDeclaration : ModuleDeclaration) as ComponentType<{ id: string; data: unknown }>,
+  getContentComponent: (kind: NodeKind, data: unknown) => {
+    if (kind === "compound") {
+      return ModuleDeclaration as ComponentType<{ id: string; data: unknown }>;
+    }
+    // An atomic node with a "path" field is a module demoted to leaf (no children).
+    const record = data as Record<string, unknown> | null;
+    if (record && "path" in record) {
+      return ModuleDeclaration as ComponentType<{ id: string; data: unknown }>;
+    }
+    return ResourceDeclaration as ComponentType<{ id: string; data: unknown }>;
+  },
 });
 
-export function App() {
-  const setNodesAtom = useSetAtom(nodesAtom);
-  const setEdgesAtom = useSetAtom(edgesAtom);
-  const addAtomicNode = useSetAtom(addAtomicNodeAtom);
-  const addCompoundNode = useSetAtom(addCompoundNodeAtom);
-  const addEdge = useSetAtom(addEdgeAtom);
+/**
+ * Inner component that lives inside PanZoomProvider so it can
+ * access both the messaging channel and the pan-zoom controls.
+ */
+function GraphContainer() {
+  const applyGraph = useApplyDeploymentGraph();
+  const messageChannel = useWebviewMessageChannel();
+  const getPanZoomDimensions = useGetPanZoomDimensions();
+  const { transform } = usePanZoomControl();
+  const graphVersion = useAtomValue(graphVersionAtom);
 
+  // Send READY notification on mount
   useEffect(() => {
-    addAtomicNode(
-      "A",
-      { x: 200, y: 200 },
-      { symbolicName: "foobar", resourceType: "Microsoft.Compute/virtualMachines" },
-    );
-    addAtomicNode("B", { x: 500, y: 200 }, { symbolicName: "bar", resourceType: "Foo" });
-    addAtomicNode("C", { x: 800, y: 500 }, { symbolicName: "someRandomStorage", resourceType: "Foo" });
-    addAtomicNode("D", { x: 1200, y: 700 }, { symbolicName: "Tricep", resourceType: "Foo" });
-    addCompoundNode("E", ["A", "C"], {
-      symbolicName: "myMod",
-      path: "modules/foooooooooooooooooooooooooooooooooooooobar.bicep",
+    messageChannel.sendNotification({
+      method: READY_NOTIFICATION,
     });
+  }, [messageChannel]);
 
-    addEdge("A->B", "A", "B");
-    addEdge("E->D", "E", "D");
-    addEdge("C->B", "C", "B");
+  // Listen for deployment graph updates
+  useWebviewNotification(
+    DEPLOYMENT_GRAPH_NOTIFICATION,
+    useCallback(
+      (params: unknown) => {
+        const payload = params as DeploymentGraphPayload;
+        applyGraph(payload.deploymentGraph);
+      },
+      [applyGraph],
+    ),
+  );
 
-    // Wait for DOM measurement (two frames) then run auto-layout
-    const frame1 = requestAnimationFrame(() => {
-      const frame2 = requestAnimationFrame(() => {
-        void runLayout(store);
-      });
+  // Run ELK layout after the DOM has been updated with the new graph.
+  // useLayoutEffect fires synchronously after React commits DOM changes,
+  // which is the reliable moment to measure and lay out.
+  //
+  // Guard against overlapping layouts: if a newer graph arrives while a
+  // previous layout is still in flight, the stale layout's completion
+  // is ignored (via the `cancelled` flag set in the cleanup function).
+  useLayoutEffect(() => {
+    if (graphVersion === 0) {
+      return;
+    }
 
-      cleanup = () => cancelAnimationFrame(frame2);
-    });
+    let cancelled = false;
+    // Always pass viewport dimensions so computeLayout can center
+    // the graph in the viewport on every layout.
+    const viewport = getPanZoomDimensions();
+    void computeLayout(store, viewport).then(async (result) => {
+      if (cancelled) {
+        return;
+      }
 
-    let cleanup: (() => void) | undefined;
+      // Compute and apply the fit-view transform immediately from the
+      // ELK result (final positions are known), so the viewport adjusts
+      // before the spring animations start.
+      const { width, height } = viewport;
+      const { translateX, translateY, scale } = computeFitViewTransform(result, width, height);
+      transform(translateX, translateY, scale);
+
+      await applyLayout(store, result, /* animate */ true);
+    });
 
     return () => {
-      cancelAnimationFrame(frame1);
-      cleanup?.();
-      setEdgesAtom([]);
-      setNodesAtom({});
+      cancelled = true;
     };
-  }, [addCompoundNode, addAtomicNode, addEdge, setNodesAtom, setEdgesAtom]);
+  }, [graphVersion, getPanZoomDimensions, transform]);
 
+  return (
+    <>
+      <$ControlBarContainer>
+        <GraphControlBar />
+      </$ControlBarContainer>
+      <Canvas>
+        <Graph />
+      </Canvas>
+    </>
+  );
+}
+
+const $AppContainer = styled.div`
+  flex: 1 1 auto;
+  position: relative;
+  overflow: hidden;
+`;
+
+function VisualDesignerApp() {
   const theme = useTheme();
 
   return (
     <ThemeProvider theme={theme}>
       <GlobalStyle />
-      <PanZoomProvider>
-        <$ControlBarContainer>
-          <GraphControlBar />
-        </$ControlBarContainer>
-        <Canvas>
-          <Graph />
-        </Canvas>
-      </PanZoomProvider>
+      <$AppContainer>
+        <PanZoomProvider>
+          <GraphContainer />
+        </PanZoomProvider>
+      </$AppContainer>
     </ThemeProvider>
   );
 }
+
+export function App() {
+  const [fakeChannel, setFakeChannel] = useState<WebviewMessageChannel | undefined>(undefined);
+  const fakeChannelRef = useRef<unknown>(undefined);
+
+  useEffect(() => {
+    if (!loadFakeMessageChannel || fakeChannelRef.current) return;
+    fakeChannelRef.current = true; // prevent double-loading in StrictMode
+    void loadFakeMessageChannel().then((ch) => {
+      fakeChannelRef.current = ch;
+      setFakeChannel(ch as unknown as WebviewMessageChannel);
+    });
+  }, []);
+
+  // In production, render immediately. In dev, wait for the fake channel to load.
+  if (isDev && !fakeChannel) return null;
+
+  return (
+    <WebviewMessageChannelProvider messageChannel={fakeChannel as unknown as WebviewMessageChannel}>
+      {isDev && LazyDevToolbar && (
+        <Suspense>
+          <LazyDevToolbar channel={fakeChannelRef.current as never} />
+        </Suspense>
+      )}
+      <VisualDesignerApp />
+    </WebviewMessageChannelProvider>
+  );
+}

src/vscode-bicep-ui/apps/visual-designer/src/GlobalStyle.ts

@@ -9,14 +9,18 @@ export const GlobalStyle = createGlobalStyle`
     margin: 0;
     padding: 0;
     display: flex;
+    flex-direction: column;
     overflow: hidden;
     font-family: var(--vscode-font-family, sans-serif);
     background-color: ${({ theme }) => theme.canvas.background};
     color: ${({ theme }) => theme.text.primary};
   }
 
   #root {
+    position: relative;
     flex: 1 1 auto;
     overflow: hidden;
+    display: flex;
+    flex-direction: column;
   }
 `;

src/vscode-bicep-ui/apps/visual-designer/src/dev/DevToolbar.tsx

@@ -0,0 +1,98 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+import { styled } from "styled-components";
+import { FakeMessageChannel, GRAPH_MUTATIONS, SAMPLE_GRAPHS } from "./FakeMessageChannel";
+
+interface DevToolbarProps {
+  channel: FakeMessageChannel;
+}
+
+const $Toolbar = styled.div`
+  flex: 0 0 auto;
+  display: flex;
+  flex-wrap: wrap;
+  align-items: center;
+  gap: 6px 8px;
+  padding: 8px 12px;
+  background: rgba(30, 30, 30, 0.92);
+  border-bottom: 1px solid rgba(255, 255, 255, 0.12);
+  font-family: system-ui, sans-serif;
+  font-size: 12px;
+  color: #ccc;
+`;
+
+const $Label = styled.span`
+  font-weight: 600;
+  color: #e0a030;
+  margin-right: 4px;
+  user-select: none;
+`;
+
+const $Separator = styled.div`
+  width: 1px;
+  height: 20px;
+  background: rgba(255, 255, 255, 0.2);
+  margin: 0 4px;
+`;
+
+const $SectionLabel = styled.span`
+  font-size: 11px;
+  color: #999;
+  user-select: none;
+`;
+
+const $Button = styled.button`
+  padding: 4px 10px;
+  border: 1px solid rgba(255, 255, 255, 0.2);
+  border-radius: 4px;
+  background: rgba(255, 255, 255, 0.08);
+  color: #ddd;
+  font-size: 12px;
+  cursor: pointer;
+  white-space: nowrap;
+
+  &:hover {
+    background: rgba(255, 255, 255, 0.16);
+    border-color: rgba(255, 255, 255, 0.35);
+  }
+
+  &:active {
+    background: rgba(255, 255, 255, 0.24);
+  }
+`;
+
+/**
+ * A floating toolbar rendered only in dev mode (`npm run dev`).
+ * Each button pushes a different sample graph through the
+ * {@link FakeMessageChannel}, simulating the extension host
+ * sending `deploymentGraph` notifications.
+ */
+export function DevToolbar({ channel }: DevToolbarProps) {
+  const applyMutation = (
+    apply: (graph: import("../messages").DeploymentGraph) => import("../messages").DeploymentGraph,
+  ) => {
+    const current = channel.getCurrentGraph();
+    if (!current) return;
+    channel.pushGraph(apply(current));
+  };
+
+  return (
+    <$Toolbar>
+      <$Label>DEV</$Label>
+      <$SectionLabel>Graphs</$SectionLabel>
+      {Object.entries(SAMPLE_GRAPHS).map(([name, graph]) => (
+        <$Button key={name} onClick={() => channel.pushGraph(graph)}>
+          {name}
+        </$Button>
+      ))}
+      <$Separator />
+      <$SectionLabel>Mutations</$SectionLabel>
+      {GRAPH_MUTATIONS.map((mutation) => (
+        <$Button key={mutation.label} title={mutation.description} onClick={() => applyMutation(mutation.apply)}>
+          {mutation.label}
+        </$Button>
+      ))}
+    </$Toolbar>
+  );
+}