From 7a557de0b17e9c36654afa7ec4ff9d372a31af8c Mon Sep 17 00:00:00 2001
From: Rupesh <rupeshkuneed@gmail.com>
Date: Sat, 28 Mar 2026 11:18:19 +0530
Subject: [PATCH] feat: add interactive DAG WebView panel (GSoC 2026 Goal #1)

Adds a full Visual DAG Viewer to the Metaflow VS Code extension, directly
addressing Goal #1 of the GSoC 2026 VS Code Extension project.

Changes:
- dag_parser.py: zero-dependency Python AST parser that finds all FlowSpec
  subclasses, extracts @step methods and self.next() edges, outputs BFS-ordered
  JSON topology. Uses the ast module only - no runtime execution of user code.
- extension.js: new showDagView command that spawns dag_parser.py as a child
  process, creates a WebviewPanel beside the editor, sends DAG JSON via
  postMessage, registers a file-save watcher for live refresh, and handles
  node-click messages to navigate the editor cursor to the clicked step def.
- package.json: registers 'Metaflow: Show DAG' command with Ctrl+Alt+D
  keybinding active on Python files; bumps version to 0.1.0.
- media/dag_webview.html: self-contained D3.js (v7) DAG visualization with
  BFS-level layout, bezier curved edges, color-coded nodes (green=start/end,
  orange=join, blue=step), and click-to-navigate postMessage back to extension.

Tested with linear and branching flows in the Extension Development Host.
---
 dag_parser.py          | 136 ++++++++++++++++
 extension.js           | 158 +++++++++++++++++--
 media/dag_webview.html | 344 +++++++++++++++++++++++++++++++++++++++++
 package.json           |  16 +-
 4 files changed, 638 insertions(+), 16 deletions(-)
 create mode 100644 dag_parser.py
 create mode 100644 media/dag_webview.html

diff --git a/dag_parser.py b/dag_parser.py
new file mode 100644
index 0000000..e0c98da
--- /dev/null
+++ b/dag_parser.py
@@ -0,0 +1,136 @@
+"""
+Metaflow DAG Parser
+-------------------
+Parses a Metaflow flow file using the AST module (zero dependencies)
+and outputs the DAG topology as JSON:
+
+  {
+    "steps": ["start", "process", "end"],
+    "edges": [["start", "process"], ["process", "end"]],
+    "error": null
+  }
+
+Usage:
+  python dag_parser.py <path_to_flow.py>
+"""
+
+import ast
+import json
+import sys
+
+
+def parse_next_args(call_node):
+    """Extract step names from a self.next(...) call node."""
+    targets = []
+    for arg in call_node.args:
+        # self.next(self.some_step)  -> Attribute node
+        if isinstance(arg, ast.Attribute):
+            targets.append(arg.attr)
+    return targets
+
+
+def parse_flow(filepath):
+    try:
+        with open(filepath, "r", encoding="utf-8") as f:
+            source = f.read()
+    except Exception as e:
+        return {"steps": [], "edges": [], "error": str(e)}
+
+    try:
+        tree = ast.parse(source, filename=filepath)
+    except SyntaxError as e:
+        return {"steps": [], "edges": [], "error": f"SyntaxError: {e}"}
+
+    steps = []
+    edges = []
+
+    for node in ast.walk(tree):
+        # Find classes that inherit from FlowSpec
+        if not isinstance(node, ast.ClassDef):
+            continue
+
+        is_flow = any(
+            (isinstance(b, ast.Name) and b.id == "FlowSpec")
+            or (isinstance(b, ast.Attribute) and b.attr == "FlowSpec")
+            for b in node.bases
+        )
+        if not is_flow:
+            continue
+
+        # Collect all @step methods
+        step_methods = set()
+        for item in node.body:
+            if not isinstance(item, ast.FunctionDef):
+                continue
+            has_step_decorator = any(
+                (isinstance(d, ast.Name) and d.id == "step")
+                or (isinstance(d, ast.Attribute) and d.attr == "step")
+                for d in item.decorator_list
+            )
+            if has_step_decorator:
+                step_methods.add(item.name)
+
+        steps = list(step_methods)
+
+        # Extract edges from self.next(...) calls inside each step method
+        for item in node.body:
+            if not isinstance(item, ast.FunctionDef):
+                continue
+            if item.name not in step_methods:
+                continue
+
+            for subnode in ast.walk(item):
+                if not isinstance(subnode, ast.Call):
+                    continue
+                func = subnode.func
+                # Match self.next(...)
+                if (
+                    isinstance(func, ast.Attribute)
+                    and func.attr == "next"
+                    and isinstance(func.value, ast.Name)
+                    and func.value.id == "self"
+                ):
+                    for target in parse_next_args(subnode):
+                        edges.append([item.name, target])
+
+        # Only process the first FlowSpec subclass found
+        break
+
+    # Topological sort: BFS from "start" for display ordering
+    ordered = []
+    if "start" in steps:
+        visited = set()
+        queue = ["start"]
+        adj = {}
+        for s in steps:
+            adj[s] = []
+        for frm, to in edges:
+            if frm in adj:
+                adj[frm].append(to)
+
+        while queue:
+            node_name = queue.pop(0)
+            if node_name in visited:
+                continue
+            visited.add(node_name)
+            ordered.append(node_name)
+            for neighbor in adj.get(node_name, []):
+                if neighbor not in visited:
+                    queue.append(neighbor)
+
+        # Add any steps not reachable from start
+        for s in steps:
+            if s not in visited:
+                ordered.append(s)
+        steps = ordered
+
+    return {"steps": steps, "edges": edges, "error": None}
+
+
+if __name__ == "__main__":
+    if len(sys.argv) < 2:
+        print(json.dumps({"steps": [], "edges": [], "error": "No file path provided"}))
+        sys.exit(1)
+
+    result = parse_flow(sys.argv[1])
+    print(json.dumps(result))
diff --git a/extension.js b/extension.js
index deb44f5..4d097f8 100644
--- a/extension.js
+++ b/extension.js
@@ -1,5 +1,7 @@
 const vscode = require('vscode');
 const path = require('path');
+const cp = require('child_process');
+const fs = require('fs');
 
 let sharedTerminal = null;
 
@@ -27,7 +29,7 @@ async function runPythonCommand(scriptName) {
   }
 
   if (!funcName) {
-    vscode.window.showErrorMessage("No enclosing Python function found.");
+    vscode.window.showErrorMessage('No enclosing Python function found.');
     return;
   }
 
@@ -35,10 +37,10 @@ async function runPythonCommand(scriptName) {
   const fileDir = path.dirname(filePath);
 
   let command = '';
-  if (scriptName == 'spin_func')
-    command = `python ${filePath} spin ${funcName}`;
+  if (scriptName === 'spin_func')
+    command = `python "${filePath}" spin ${funcName}`;
   else
-    command = `python ${filePath} run`;
+    command = `python "${filePath}" run`;
 
   // Reuse or create a single shared terminal
   if (!sharedTerminal || sharedTerminal.exitStatus !== undefined) {
@@ -48,14 +50,139 @@ async function runPythonCommand(scriptName) {
   sharedTerminal.show();
   sharedTerminal.sendText(`cd "${fileDir}"`);
   sharedTerminal.sendText(command);
+}
 
-  /*
-  vscode.window.showInformationMessage(
-    `${scriptName.toUpperCase()}: ${funcName} from ${path.basename(filePath)}`
-  );
-  */
+// ---------------------------------------------------------------------------
+// DAG WebView
+// ---------------------------------------------------------------------------
+
+let dagPanel = null;
+let dagWatcher = null;
+
+/**
+ * Runs dag_parser.py for the given file and returns { steps, edges, error }.
+ */
+function parseDag(filePath, extensionPath) {
+  return new Promise((resolve) => {
+    const parserScript = path.join(extensionPath, 'dag_parser.py');
+    const proc = cp.spawn('python', [parserScript, filePath]);
+
+    let stdout = '';
+    let stderr = '';
+
+    proc.stdout.on('data', (d) => { stdout += d.toString(); });
+    proc.stderr.on('data', (d) => { stderr += d.toString(); });
+
+    proc.on('close', () => {
+      try {
+        resolve(JSON.parse(stdout));
+      } catch (_) {
+        resolve({ steps: [], edges: [], error: stderr || 'Failed to parse DAG output.' });
+      }
+    });
+
+    proc.on('error', (err) => {
+      resolve({ steps: [], edges: [], error: `Could not start Python: ${err.message}` });
+    });
+  });
+}
+
+/**
+ * Loads dag_webview.html from the media folder.
+ */
+function getWebviewContent(extensionUri) {
+  const htmlPath = path.join(extensionUri.fsPath, 'media', 'dag_webview.html');
+  return fs.readFileSync(htmlPath, 'utf8');
 }
 
+/**
+ * Navigates the active editor cursor to the def line of the given step.
+ */
+async function navigateToStep(stepName) {
+  const editor = vscode.window.activeTextEditor;
+  if (!editor) return;
+
+  const lines = editor.document.getText().split('\n');
+  const re = new RegExp(`^\\s*(?:async\\s+)?def\\s+${stepName}\\s*\\(`);
+
+  for (let i = 0; i < lines.length; i++) {
+    if (re.test(lines[i])) {
+      const pos = new vscode.Position(i, 0);
+      editor.selection = new vscode.Selection(pos, pos);
+      editor.revealRange(new vscode.Range(pos, pos), vscode.TextEditorRevealType.InCenter);
+      return;
+    }
+  }
+  vscode.window.showWarningMessage(`Step "${stepName}" not found in the current file.`);
+}
+
+/**
+ * Command: Open (or refresh) the Metaflow DAG WebView panel.
+ */
+async function showDagView(context) {
+  const editor = vscode.window.activeTextEditor;
+  if (!editor || editor.document.languageId !== 'python') {
+    vscode.window.showErrorMessage('Please open a Metaflow Python flow file first.');
+    return;
+  }
+
+  const filePath = editor.document.fileName;
+
+  // Create the panel if it doesn't exist, otherwise reveal it
+  if (!dagPanel) {
+    dagPanel = vscode.window.createWebviewPanel(
+      'metaflowDag',
+      'Metaflow DAG',
+      vscode.ViewColumn.Beside,
+      {
+        enableScripts: true,
+        localResourceRoots: [vscode.Uri.joinPath(context.extensionUri, 'media')],
+        retainContextWhenHidden: true,
+      }
+    );
+
+    dagPanel.webview.html = getWebviewContent(context.extensionUri);
+
+    // Node click → navigate to step in editor
+    dagPanel.webview.onDidReceiveMessage(
+      (msg) => {
+        if (msg.command === 'navigateToStep') {
+          navigateToStep(msg.step);
+        }
+      },
+      undefined,
+      context.subscriptions
+    );
+
+    dagPanel.onDidDispose(() => {
+      dagPanel = null;
+      if (dagWatcher) { dagWatcher.dispose(); dagWatcher = null; }
+    }, null, context.subscriptions);
+  } else {
+    dagPanel.reveal(vscode.ViewColumn.Beside);
+  }
+
+  // Initial render
+  const dagData = await parseDag(filePath, context.extensionPath);
+  dagData.fileName = path.basename(filePath);
+  dagPanel.webview.postMessage({ command: 'renderDag', data: dagData });
+
+  // Live-refresh on every save of this file
+  if (dagWatcher) { dagWatcher.dispose(); }
+  dagWatcher = vscode.workspace.onDidSaveTextDocument(async (savedDoc) => {
+    if (savedDoc.fileName === filePath && dagPanel) {
+      const refreshed = await parseDag(savedDoc.fileName, context.extensionPath);
+      refreshed.fileName = path.basename(savedDoc.fileName);
+      dagPanel.webview.postMessage({ command: 'renderDag', data: refreshed });
+    }
+  });
+  context.subscriptions.push(dagWatcher);
+}
+
+// ---------------------------------------------------------------------------
+// Extension lifecycle
+// ---------------------------------------------------------------------------
+
 function activate(context) {
   const runCmd = vscode.commands.registerCommand(
     'extension.runPythonFunction',
@@ -67,14 +194,17 @@ function activate(context) {
     () => runPythonCommand('spin_func')
   );
 
-  context.subscriptions.push(runCmd, spinCmd);
+  const dagCmd = vscode.commands.registerCommand(
+    'extension.showDagView',
+    () => showDagView(context)
+  );
+
+  context.subscriptions.push(runCmd, spinCmd, dagCmd);
 }
 
 function deactivate() {
-  if (sharedTerminal) {
-    sharedTerminal.dispose();
-    sharedTerminal = null;
-  }
+  if (sharedTerminal) { sharedTerminal.dispose(); sharedTerminal = null; }
+  if (dagPanel) { dagPanel.dispose(); dagPanel = null; }
 }
 
 module.exports = { activate, deactivate };
diff --git a/media/dag_webview.html b/media/dag_webview.html
new file mode 100644
index 0000000..0ce299d
--- /dev/null
+++ b/media/dag_webview.html
@@ -0,0 +1,344 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+  <meta charset="UTF-8" />
+  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
+  <title>Metaflow DAG</title>
+  <script src="https://d3js.org/d3.v7.min.js"></script>
+  <style>
+    * { box-sizing: border-box; margin: 0; padding: 0; }
+
+    body {
+      background: #1e1e2e;
+      color: #cdd6f4;
+      font-family: 'Segoe UI', system-ui, sans-serif;
+      height: 100vh;
+      display: flex;
+      flex-direction: column;
+      overflow: hidden;
+    }
+
+    #header {
+      padding: 10px 16px;
+      background: #181825;
+      border-bottom: 1px solid #313244;
+      display: flex;
+      align-items: center;
+      gap: 10px;
+      flex-shrink: 0;
+    }
+
+    #header .logo {
+      font-size: 18px;
+      font-weight: 700;
+      color: #cba6f7;
+      letter-spacing: -0.5px;
+    }
+
+    #header .filename {
+      font-size: 13px;
+      color: #6c7086;
+      font-weight: 400;
+    }
+
+    #legend {
+      display: flex;
+      gap: 14px;
+      margin-left: auto;
+      font-size: 11px;
+      color: #6c7086;
+      align-items: center;
+    }
+
+    .legend-dot {
+      width: 10px; height: 10px;
+      border-radius: 50%;
+      display: inline-block;
+      margin-right: 4px;
+      vertical-align: middle;
+    }
+
+    #hint {
+      font-size: 11px;
+      color: #45475a;
+      padding: 5px 16px;
+      background: #181825;
+      border-bottom: 1px solid #313244;
+      flex-shrink: 0;
+    }
+
+    #svg-container {
+      flex: 1;
+      overflow: hidden;
+    }
+
+    svg {
+      width: 100%;
+      height: 100%;
+    }
+
+    .node-group { cursor: pointer; }
+
+    .node-group:hover rect {
+      filter: brightness(1.25);
+    }
+
+    .node-label {
+      font-size: 13px;
+      font-weight: 600;
+      fill: white;
+      pointer-events: none;
+      dominant-baseline: middle;
+      text-anchor: middle;
+    }
+
+    .edge-path {
+      fill: none;
+      stroke: #585b70;
+      stroke-width: 2;
+      marker-end: url(#arrowhead);
+    }
+
+    #error-msg {
+      display: none;
+      padding: 24px;
+      color: #f38ba8;
+      font-size: 13px;
+      text-align: center;
+    }
+
+    #empty-msg {
+      display: none;
+      padding: 24px;
+      color: #6c7086;
+      font-size: 13px;
+      text-align: center;
+    }
+  </style>
+</head>
+<body>
+  <div id="header">
+    <span class="logo">⚡ Metaflow DAG</span>
+    <span class="filename" id="filename-label">–</span>
+    <div id="legend">
+      <span><span class="legend-dot" style="background:#a6e3a1;"></span>start / end</span>
+      <span><span class="legend-dot" style="background:#89b4fa;"></span>step</span>
+      <span><span class="legend-dot" style="background:#fab387;"></span>join</span>
+    </div>
+  </div>
+  <div id="hint">Click a node to jump to that step in the editor</div>
+  <div id="error-msg"></div>
+  <div id="empty-msg">No Metaflow FlowSpec found in this file.</div>
+  <div id="svg-container">
+    <svg id="dag-svg">
+      <defs>
+        <marker id="arrowhead" viewBox="0 -5 10 10"
+                refX="30" refY="0"
+                markerWidth="6" markerHeight="6"
+                orient="auto">
+          <path d="M0,-5L10,0L0,5" fill="#585b70" />
+        </marker>
+      </defs>
+      <g id="dag-root"></g>
+    </svg>
+  </div>
+
+  <script>
+    const vscode = acquireVsCodeApi();
+
+    const NODE_W = 130;
+    const NODE_H = 44;
+    const LEVEL_SEP = 110;  // vertical distance between BFS levels
+    const NODE_SEP = 160;   // horizontal distance between sibling nodes
+
+    // Color by role
+    function nodeColor(name, steps) {
+      if (name === 'start' || name === 'end') return '#a6e3a1';
+      // Detect join: has 2+ incoming edges
+      const incomingCount = {};
+      steps.edges && steps.edges.forEach(([, to]) => {
+        incomingCount[to] = (incomingCount[to] || 0) + 1;
+      });
+      if (incomingCount[name] >= 2) return '#fab387';
+      return '#89b4fa';
+    }
+
+    function nodeTextColor(name, steps) {
+      return '#1e1e2e';
+    }
+
+    /**
+     * Compute (x, y) positions using BFS level assignment.
+     * Returns Map<stepName, {x, y, level}>
+     */
+    function computeLayout(steps, edges) {
+      if (!steps || steps.length === 0) return new Map();
+
+      // Build adjacency
+      const adj = {};
+      const inDeg = {};
+      steps.forEach(s => { adj[s] = []; inDeg[s] = 0; });
+      edges.forEach(([f, t]) => {
+        if (adj[f]) adj[f].push(t);
+        if (t in inDeg) inDeg[t]++;
+      });
+
+      // BFS level assignment
+      const level = {};
+      const roots = steps.filter(s => inDeg[s] === 0);
+      if (roots.length === 0) roots.push(steps[0]); // fallback
+
+      const queue = [...roots];
+      roots.forEach(r => level[r] = 0);
+
+      while (queue.length) {
+        const cur = queue.shift();
+        (adj[cur] || []).forEach(nbr => {
+          const newLevel = level[cur] + 1;
+          if (!(nbr in level) || level[nbr] < newLevel) {
+            level[nbr] = newLevel;
+            queue.push(nbr);
+          }
+        });
+      }
+
+      // Assign missing levels
+      let maxLevel = 0;
+      steps.forEach(s => {
+        if (!(s in level)) level[s] = 0;
+        if (level[s] > maxLevel) maxLevel = level[s];
+      });
+
+      // Group by level
+      const byLevel = {};
+      steps.forEach(s => {
+        const lv = level[s];
+        if (!byLevel[lv]) byLevel[lv] = [];
+        byLevel[lv].push(s);
+      });
+
+      const positions = new Map();
+      Object.entries(byLevel).forEach(([lv, nodes]) => {
+        const totalW = (nodes.length - 1) * NODE_SEP;
+        nodes.forEach((name, idx) => {
+          const x = idx * NODE_SEP - totalW / 2;
+          const y = parseInt(lv) * LEVEL_SEP;
+          positions.set(name, { x, y, level: parseInt(lv) });
+        });
+      });
+
+      return positions;
+    }
+
+    function render(data) {
+      const errorEl = document.getElementById('error-msg');
+      const emptyEl = document.getElementById('empty-msg');
+      const svgContainer = document.getElementById('svg-container');
+      const fileLabel = document.getElementById('filename-label');
+
+      fileLabel.textContent = data.fileName || '';
+
+      if (data.error) {
+        errorEl.style.display = 'block';
+        errorEl.textContent = '⚠ ' + data.error;
+        emptyEl.style.display = 'none';
+        svgContainer.style.display = 'none';
+        return;
+      }
+
+      errorEl.style.display = 'none';
+
+      if (!data.steps || data.steps.length === 0) {
+        emptyEl.style.display = 'block';
+        svgContainer.style.display = 'none';
+        return;
+      }
+
+      emptyEl.style.display = 'none';
+      svgContainer.style.display = 'block';
+
+      const svg = d3.select('#dag-svg');
+      const root = d3.select('#dag-root');
+      root.selectAll('*').remove();
+
+      const positions = computeLayout(data.steps, data.edges || []);
+
+      // Shift so everything is centred at (0,0) then translate to centre of SVG
+      const xs = [...positions.values()].map(p => p.x);
+      const ys = [...positions.values()].map(p => p.y);
+      const minX = Math.min(...xs) - NODE_W / 2 - 20;
+      const maxX = Math.max(...xs) + NODE_W / 2 + 20;
+      const minY = Math.min(...ys) - NODE_H / 2 - 20;
+      const maxY = Math.max(...ys) + NODE_H / 2 + 20;
+      const graphW = maxX - minX;
+      const graphH = maxY - minY;
+
+      const svgEl = document.getElementById('dag-svg');
+      const vw = svgEl.clientWidth;
+      const vh = svgEl.clientHeight;
+
+      const scale = Math.min(1.0, Math.min(vw / graphW, vh / graphH) * 0.9);
+      const tx = vw / 2 - (minX + graphW / 2) * scale;
+      const ty = vh / 2 - (minY + graphH / 2) * scale;
+
+      root.attr('transform', `translate(${tx},${ty}) scale(${scale})`);
+
+      // Draw edges
+      (data.edges || []).forEach(([from, to]) => {
+        const fp = positions.get(from);
+        const tp = positions.get(to);
+        if (!fp || !tp) return;
+
+        // Bezier curve: start from bottom-center of from, arrive at top-center of to
+        const x1 = fp.x, y1 = fp.y + NODE_H / 2;
+        const x2 = tp.x, y2 = tp.y - NODE_H / 2;
+        const cy1 = y1 + LEVEL_SEP * 0.4;
+        const cy2 = y2 - LEVEL_SEP * 0.4;
+
+        root.append('path')
+          .attr('class', 'edge-path')
+          .attr('d', `M${x1},${y1} C${x1},${cy1} ${x2},${cy2} ${x2},${y2}`);
+      });
+
+      // Draw nodes
+      data.steps.forEach(name => {
+        const pos = positions.get(name);
+        if (!pos) return;
+
+        const color = nodeColor(name, data);
+        const g = root.append('g')
+          .attr('class', 'node-group')
+          .attr('transform', `translate(${pos.x - NODE_W / 2},${pos.y - NODE_H / 2})`)
+          .on('click', () => {
+            vscode.postMessage({ command: 'navigateToStep', step: name });
+          });
+
+        // Pill/rounded rect
+        g.append('rect')
+          .attr('width', NODE_W)
+          .attr('height', NODE_H)
+          .attr('rx', 10)
+          .attr('ry', 10)
+          .attr('fill', color)
+          .attr('stroke', '#313244')
+          .attr('stroke-width', 2);
+
+        g.append('text')
+          .attr('class', 'node-label')
+          .attr('x', NODE_W / 2)
+          .attr('y', NODE_H / 2)
+          .attr('fill', nodeTextColor(name, data))
+          .text(name);
+      });
+    }
+
+    // Listen for messages from the extension host
+    window.addEventListener('message', (event) => {
+      const msg = event.data;
+      if (msg.command === 'renderDag') {
+        render(msg.data);
+      }
+    });
+  </script>
+</body>
+</html>
diff --git a/package.json b/package.json
index 31a7c2f..ab9391d 100644
--- a/package.json
+++ b/package.json
@@ -1,12 +1,14 @@
 {
   "name": "metaflow-dev",
   "displayName": "Run and Spin Metaflow",
-  "version": "0.0.7",
+  "version": "0.1.0",
   "publisher": "local",
   "engines": { "vscode": "^1.80.0" },
   "activationEvents": [
     "onCommand:extension.runPythonFunction",
-    "onCommand:extension.spinPythonFunction"
+    "onCommand:extension.spinPythonFunction",
+    "onCommand:extension.showDagView",
+    "onLanguage:python"
   ],
   "repository": {
     "type": "git",
@@ -22,6 +24,11 @@
       {
         "command": "extension.spinPythonFunction",
         "title": "Spin the current step"
+      },
+      {
+        "command": "extension.showDagView",
+        "title": "Metaflow: Show DAG",
+        "icon": "$(type-hierarchy)"
       }
     ],
     "keybindings": [
@@ -34,6 +41,11 @@
         "command": "extension.spinPythonFunction",
         "key": "ctrl+alt+s",
         "when": "editorLangId == python"
+      },
+      {
+        "command": "extension.showDagView",
+        "key": "ctrl+alt+d",
+        "when": "editorLangId == python"
       }
     ]
   }