From 0cfb20f6bf09e23c93abba9e1a54f46fc21deadd Mon Sep 17 00:00:00 2001
From: Alexander Wang <alex@terrastruct.com>
Date: Mon, 5 Aug 2024 17:40:00 -0700
Subject: [PATCH] d2cli: fix watch mode back links

---
 d2cli/main.go             |  8 +++---
 d2cli/watch.go            |  4 +++
 e2etests-cli/main_test.go | 53 +++++++++++++++++++++++++++++++++++++++
 3 files changed, 61 insertions(+), 4 deletions(-)

diff --git a/d2cli/main.go b/d2cli/main.go
index 39ed820d0..8f9e022a9 100644
--- a/d2cli/main.go
+++ b/d2cli/main.go
@@ -480,13 +480,13 @@ func compile(ctx context.Context, ms *xmain.State, plugins []d2plugin.Plugin, fs
 	}, time.Second*5)
 	defer cancel()
 
-	diagram, g, err := d2lib.Compile(ctx, string(input), opts, &renderOpts)
+	rootDiagram, g, err := d2lib.Compile(ctx, string(input), opts, &renderOpts)
 	if err != nil {
 		return nil, false, err
 	}
 	cancel()
 
-	diagram = diagram.GetBoard(boardPath)
+	diagram := rootDiagram.GetBoard(boardPath)
 	if diagram == nil {
 		return nil, false, fmt.Errorf(`render target "%s" not found`, strings.Join(boardPath, "."))
 	}
@@ -589,11 +589,11 @@ func compile(ctx context.Context, ms *xmain.State, plugins []d2plugin.Plugin, fs
 		compileDur := time.Since(start)
 		if animateInterval <= 0 {
 			// Rename all the "root.layers.x" to the paths that the boards get output to
-			linkToOutput, err := resolveLinks("root", outputPath, diagram)
+			linkToOutput, err := resolveLinks("root", outputPath, rootDiagram)
 			if err != nil {
 				return nil, false, err
 			}
-			err = relink("root", diagram, linkToOutput)
+			err = relink("root", rootDiagram, linkToOutput)
 			if err != nil {
 				return nil, false, err
 			}
diff --git a/d2cli/watch.go b/d2cli/watch.go
index 4a934d9c6..1dfbb5b9e 100644
--- a/d2cli/watch.go
+++ b/d2cli/watch.go
@@ -520,6 +520,10 @@ func (w *watcher) handleRoot(hw http.ResponseWriter, r *http.Request) {
 	if idx := strings.LastIndexByte(boardPath, '.'); idx != -1 {
 		boardPath = boardPath[:idx]
 	}
+	// if path is "/index", we just want "/"
+	if boardPath == "index" {
+		boardPath = ""
+	}
 	recompile := false
 	if boardPath != w.boardPath {
 		w.boardPath = boardPath
diff --git a/e2etests-cli/main_test.go b/e2etests-cli/main_test.go
index 5a2f925d7..6182232f7 100644
--- a/e2etests-cli/main_test.go
+++ b/e2etests-cli/main_test.go
@@ -986,6 +986,59 @@ layers: {
 				assert.Success(t, err)
 			},
 		},
+		{
+			name:   "watch-underscore-link",
+			serial: true,
+			run: func(t *testing.T, ctx context.Context, dir string, env *xos.Env) {
+				writeFile(t, dir, "index.d2", `
+bobby
+
+layers: {
+    cream: {
+			back.link: _
+    }
+}`)
+				stderr := &stderrWrapper{}
+				tms := testMain(dir, env, "--watch", "--browser=0", "index.d2")
+				tms.Stderr = stderr
+
+				tms.Start(t, ctx)
+				defer func() {
+					// Manually close, since watcher is daemon
+					err := tms.Signal(ctx, os.Interrupt)
+					assert.Success(t, err)
+				}()
+
+				// Wait for watch server to spin up and listen
+				urlRE := regexp.MustCompile(`127.0.0.1:([0-9]+)`)
+				watchURL, err := waitLogs(ctx, stderr, urlRE)
+				assert.Success(t, err)
+
+				stderr.Reset()
+
+				// Start a client
+				c, _, err := websocket.Dial(ctx, fmt.Sprintf("ws://%s/watch", watchURL), nil)
+				assert.Success(t, err)
+				defer c.CloseNow()
+
+				_, _, err = c.Read(ctx)
+				assert.Success(t, err)
+
+				err = getWatchPage(ctx, t, fmt.Sprintf("http://%s/%s", watchURL, "cream"))
+				assert.Success(t, err)
+
+				// Get the link
+				_, msg, err := c.Read(ctx)
+				aRE := regexp.MustCompile(`href=\\"([^\"]*)\\"`)
+				match := aRE.FindSubmatch(msg)
+				assert.Equal(t, 2, len(match))
+				assert.Equal(t, "index.svg", string(match[1]))
+
+				successRE := regexp.MustCompile(`broadcasting update to 1 client`)
+				_, err = waitLogs(ctx, stderr, successRE)
+				assert.Success(t, err)
+			},
+		},
 		{
 			name:   "watch-bad-link",
 			serial: true,