watch: prevent stale cross-board SVG from locking the wrong fit scale

d2cli/watch.go

@@ -347,6 +347,17 @@ func (w *watcher) requestCompile() {
 	}
 }
 
+// boardStillCurrent reports whether compiledFor is still the board being
+// viewed. A compile can take long enough that the user navigates to a
+// different board before it finishes; its result must be discarded rather
+// than published, or it would overwrite the newly-selected board's cached
+// result with stale data.
+func (w *watcher) boardStillCurrent(compiledFor string) bool {
+	w.boardpathMu.Lock()
+	defer w.boardpathMu.Unlock()
+	return compiledFor == w.boardPath
+}
+
 func (w *watcher) ensureAddWatch(ctx context.Context, path string) (time.Time, error) {
 	interval := time.Millisecond * 16
 	tc := time.NewTimer(0)
@@ -449,9 +460,10 @@ func (w *watcher) compileLoop(ctx context.Context) error {
 		// localfile.Policy to reject directories and other special files.
 		fs := trackedFS{localFiles: localfile.Unrestricted()}
 		w.boardpathMu.Lock()
+		compiledBoardPath := w.boardPath
 		var boardPath []string
-		if w.boardPath != "" {
-			boardPath = strings.Split(w.boardPath, string(os.PathSeparator))
+		if compiledBoardPath != "" {
+			boardPath = strings.Split(compiledBoardPath, string(os.PathSeparator))
 		}
 		svg, _, err := compile(ctx, w.ms, &fs, w.layout, w.renderOpts, w.fontFamily, w.monoFontFamily, w.animateInterval, w.inputPath, w.outputPath, boardPath, false, w.bundle, w.forceAppendix, w.outputFormat, w.asciiMode, true)
 		w.boardpathMu.Unlock()
@@ -470,11 +482,17 @@ func (w *watcher) compileLoop(ctx context.Context) error {
 			return err
 		}
 
-		w.broadcast(&compileResult{
-			SVG:   string(svg),
-			Scale: w.renderOpts.Scale,
-			Err:   errs,
-		})
+		// handleRoot may have already switched to a different board and cleared
+		// w.res while this compile was in flight. Only publish this result if
+		// the board it was compiled for is still the one being viewed, so we
+		// never hand a newly-connecting client the wrong board's SVG.
+		if w.boardStillCurrent(compiledBoardPath) {
+			w.broadcast(&compileResult{
+				SVG:   string(svg),
+				Scale: w.renderOpts.Scale,
+				Err:   errs,
+			})
+		}
 
 		if firstCompile {
 			firstCompile = false
@@ -611,6 +629,15 @@ func (w *watcher) handleRoot(hw http.ResponseWriter, r *http.Request) {
 	}
 	w.boardpathMu.Unlock()
 	if recompile {
+		// Discard the previous board's cached result so a websocket client
+		// connecting for this navigation can't receive a stale SVG (with the
+		// wrong dimensions) before the new board finishes compiling. Without
+		// this, the client's watch.js locks in a fit-to-window scale computed
+		// from the old board and never recomputes it once the correct SVG
+		// arrives, leaving the diagram mis-scaled.
+		w.resMu.Lock()
+		w.res = nil
+		w.resMu.Unlock()
 		w.requestCompile()
 	}
 }

d2cli/watch_test.go

@@ -2,6 +2,8 @@ package d2cli
 
 import (
 	"errors"
+	"net/http"
+	"net/http/httptest"
 	"os"
 	"path/filepath"
 	"testing"
@@ -65,3 +67,67 @@ func TestTrackedFSHonorsLocalFilePolicy(t *testing.T) {
 		t.Fatalf("trusted tracked opens = %q, want [%q]", trusted.opened, outside)
 	}
 }
+
+// TestHandleRootInvalidatesStaleResultOnBoardNavigation guards against a
+// regression where navigating between boards in watch mode (e.g. descending
+// into a layer, then returning to the parent board) could momentarily hand a
+// newly-connecting websocket client the previous board's cached SVG. The
+// frontend locks in its fit-to-window scale from the first message it
+// receives, so a stale cross-board SVG left the diagram mis-scaled even
+// after the correct one arrived.
+func TestHandleRootInvalidatesStaleResultOnBoardNavigation(t *testing.T) {
+	w := &watcher{compileCh: make(chan struct{}, 1)}
+	w.boardPath = "old"
+	w.res = &compileResult{SVG: "<svg>old board</svg>"}
+
+	req := httptest.NewRequest(http.MethodGet, "/new", nil)
+	w.handleRoot(httptest.NewRecorder(), req)
+
+	if w.boardPath != "new" {
+		t.Fatalf("boardPath = %q, want %q", w.boardPath, "new")
+	}
+	if got := w.getRes(); got != nil {
+		t.Fatalf("getRes() = %+v, want nil after navigating to a different board", got)
+	}
+}
+
+// TestHandleRootKeepsResultWhenBoardUnchanged ensures the invalidation above
+// doesn't also drop the cached result on a plain reload of the same board,
+// which would cause the diagram to flicker blank while it recompiles.
+func TestHandleRootKeepsResultWhenBoardUnchanged(t *testing.T) {
+	w := &watcher{compileCh: make(chan struct{}, 1)}
+	w.boardPath = "same"
+	res := &compileResult{SVG: "<svg>same board</svg>"}
+	w.res = res
+
+	req := httptest.NewRequest(http.MethodGet, "/same", nil)
+	w.handleRoot(httptest.NewRecorder(), req)
+
+	if got := w.getRes(); got != res {
+		t.Fatalf("getRes() = %+v, want unchanged cached result %+v", got, res)
+	}
+}
+
+// TestBoardStillCurrentGuardsInFlightCompile covers the remaining half of the
+// stale-result race: a compile for the board being navigated away from can
+// still be running when handleRoot clears w.res for the new board. Without
+// this guard, compileLoop would broadcast that in-flight compile's result
+// after the clear, repopulating w.res with the wrong board's SVG.
+func TestBoardStillCurrentGuardsInFlightCompile(t *testing.T) {
+	w := &watcher{}
+	w.boardPath = "top"
+
+	if !w.boardStillCurrent("top") {
+		t.Fatal("boardStillCurrent(\"top\") = false, want true before any navigation")
+	}
+
+	// Simulate the user navigating away while a compile for "top" is in flight.
+	w.boardPath = "top.nested"
+
+	if w.boardStillCurrent("top") {
+		t.Fatal("boardStillCurrent(\"top\") = true after navigating to \"top.nested\", want false")
+	}
+	if !w.boardStillCurrent("top.nested") {
+		t.Fatal("boardStillCurrent(\"top.nested\") = false, want true for the newly-selected board")
+	}
+}