kjol/go/webbundler/hmr_server.go

//go:build dev

package webbundler

// The development HMR server: it serves the SPA source tree as unbundled native
// ES modules (transformed on the fly), tracks the module import graph, watches
// the filesystem, and pushes hot-update / reload messages to the browser over a
// WebSocket. Editing a .tsx component swaps it in place (solid-refresh); editing
// a non-boundary module bubbles up to a full page reload.
//
// This is the reason webbundler grew a `dev` build tag: the whole HMR
// subsystem (this file, hmr_ws.go, hmr_vendor.go, hmr_client.go, hmr_watch.go,
// and the solid-refresh bits of jsx.go) compiles only under `-tags dev`, so the
// production server and the plain `cmd/bundle` build carry none of it.

import (
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"

	esbuild "github.com/evanw/esbuild/pkg/api"
)

// devURLPrefix roots the served source tree: GET /@src/<rel> serves the
// transformed module frontend/src/<rel>.
const (
	srcURLPrefix   = "/@src/" // transformed app source modules
	kitURLPrefix   = "/@kit/" // transformed shared-kit modules (kjol web tree)
	assetURLPrefix = "/@url/" // `?url` shim: a module whose default export is the file URL
	fsURLPrefix    = "/@fs/"  // raw file passthrough (the URL the shim points at)
)

type devServer struct {
	hub        *hub
	frontend   string   // abs app frontend/
	srcRoot    string   // abs app frontend/src
	kitRoot    string   // abs kjol web (@ui/@kjol root); == frontend in single-tree
	genRoot    string   // abs app generated dir (@appgen)
	vendorDirs []string // abs vendor roots (kjol runtime first, then app vendor)
	output     string   // abs wwwroot/

	graph *moduleGraph

	importMapJSON     string
	vendorEntrypoints map[string]string

	vendorMu    sync.Mutex
	vendorCache map[string][]byte // /@vendor/<spec>.js -> bundled bytes

	cssTrigger chan struct{} // coalesced CSS rebuild requests (buffered, size 1)
}

// StartDevHMR registers the dev routes on mux, launches the filesystem watcher,
// and returns the import map (JSON object body) the SPA shell must inline so bare
// specifiers resolve to the single vendored copies. Called only from cmd/server's
// `dev`-tagged shim.
func StartDevHMR(mux *http.ServeMux, cfg Config) (importMap string, err error) {
	Configure(cfg)
	d, err := newDevServer()
	if err != nil {
		return "", err
	}
	d.register(mux)
	go d.watch()

	fmt.Println("HMR dev server: serving native-ESM source from /@src/, WebSocket at /@hmr/ws")
	return d.importMapJSON, nil
}

// newDevServer constructs the dev server (abs paths, vendor manifest, import map)
// without registering routes or starting the watcher — the split lets tests drive
// the handlers directly.
func newDevServer() (*devServer, error) {
	frontendAbs, err := filepath.Abs(frontendDir)
	if err != nil {
		return nil, err
	}
	outputAbs, err := filepath.Abs(outputDir)
	if err != nil {
		return nil, err
	}
	// esbuild reports realpaths (symlinks resolved), so the roots must be too or
	// the containment/URL math breaks on symlinked trees (e.g. macOS /var→/private/var).
	frontendAbs = evalSymlinks(frontendAbs)
	outputAbs = evalSymlinks(outputAbs)

	absSym := func(p string) string { a, _ := filepath.Abs(p); return evalSymlinks(a) }
	var vdirs []string
	for _, v := range vendorDirs() {
		vdirs = append(vdirs, absSym(v))
	}
	d := &devServer{
		hub:         newHub(),
		frontend:    frontendAbs,
		srcRoot:     filepath.Join(frontendAbs, "src"),
		kitRoot:     absSym(webRoot()),
		genRoot:     absSym(genTSDir),
		vendorDirs:  vdirs,
		output:      outputAbs,
		graph:       newModuleGraph(),
		vendorCache: map[string][]byte{},
		cssTrigger:  make(chan struct{}, 1),
	}

	eps, err := loadVendorManifest(d.vendorDirs)
	if err != nil {
		return nil, fmt.Errorf("loading vendor manifest: %w", err)
	}
	d.vendorEntrypoints = eps
	d.importMapJSON = d.buildImportMap()
	return d, nil
}

func (d *devServer) register(mux *http.ServeMux) {
	mux.HandleFunc(srcURLPrefix, d.serveModule)
	mux.HandleFunc(kitURLPrefix, d.serveKit)
	mux.HandleFunc(vendorURLPrefix, d.serveVendor)
	mux.HandleFunc(assetURLPrefix, d.serveAssetURL)
	mux.HandleFunc(fsURLPrefix, d.serveFS)
	mux.HandleFunc("/@hmr/client", d.serveClient)
	mux.HandleFunc("/@hmr/ws", d.hub.ServeWS)
}

// serveModule transforms and serves one source module as native ESM.
func (d *devServer) serveModule(w http.ResponseWriter, r *http.Request) {
	rel := strings.TrimPrefix(r.URL.Path, srcURLPrefix)
	abs := filepath.Join(d.srcRoot, filepath.FromSlash(rel))
	// Contain the request to the source tree.
	if !within(d.srcRoot, abs) {
		http.NotFound(w, r)
		return
	}
	if _, statErr := os.Stat(abs); statErr != nil {
		http.NotFound(w, r)
		return
	}

	code, err := d.transformModule(abs)
	w.Header().Set("Content-Type", "application/javascript; charset=utf-8")
	w.Header().Set("Cache-Control", "no-cache")
	if err != nil {
		// Surface the failure as a browser overlay rather than a hard 500 that
		// would break the module graph; the next save re-runs the transform.
		fmt.Fprintf(os.Stderr, "HMR transform %s: %v\n", rel, err)
		w.Write([]byte(errorModule(rel, err)))
		return
	}
	w.Write(code)
}

// serveKit transforms and serves one shared-kit module (kjol web tree) as ESM.
func (d *devServer) serveKit(w http.ResponseWriter, r *http.Request) {
	rel := strings.TrimPrefix(r.URL.Path, kitURLPrefix)
	abs := filepath.Join(d.kitRoot, filepath.FromSlash(rel))
	if !within(d.kitRoot, abs) {
		http.NotFound(w, r)
		return
	}
	if _, statErr := os.Stat(abs); statErr != nil {
		http.NotFound(w, r)
		return
	}
	code, err := d.transformModule(abs)
	w.Header().Set("Content-Type", "application/javascript; charset=utf-8")
	w.Header().Set("Cache-Control", "no-cache")
	if err != nil {
		fmt.Fprintf(os.Stderr, "HMR transform %s: %v\n", rel, err)
		w.Write([]byte(errorModule(rel, err)))
		return
	}
	w.Write(code)
}

// moduleURLBase maps an absolute module file to its dev URL (no version query):
// app sources → /@src/, shared-kit modules (kjol web) → /@kit/. "" if outside both.
func (d *devServer) moduleURLBase(abs string) string {
	if within(d.srcRoot, abs) {
		return srcURLPrefix + filepath.ToSlash(mustRel(d.srcRoot, abs))
	}
	if within(d.kitRoot, abs) {
		return kitURLPrefix + filepath.ToSlash(mustRel(d.kitRoot, abs))
	}
	return ""
}

// moduleID is a stable identifier for solid-refresh's HMR registry (survives
// recompiles), tree-aware so app and kit modules never collide.
func (d *devServer) moduleID(abs string) string {
	if within(d.srcRoot, abs) {
		return filepath.ToSlash(mustRel(d.srcRoot, abs))
	}
	if within(d.kitRoot, abs) {
		return "@kit/" + filepath.ToSlash(mustRel(d.kitRoot, abs))
	}
	return filepath.ToSlash(abs)
}

// resolveAliasSpec resolves an @ui/@kjol/@appgen import to an absolute file, or ""
// if spec is not one of those aliases (e.g. a vendored @scope/pkg like @solidjs/router).
func (d *devServer) resolveAliasSpec(spec string) string {
	roots := aliasRoots()
	slash := strings.IndexByte(spec, '/')
	if slash < 0 {
		return ""
	}
	prefix, rest := spec[:slash], spec[slash+1:]
	root, ok := roots[prefix]
	if !ok {
		return ""
	}
	if p, ok := resolveAlias(root, rest); ok {
		return evalSymlinks(p)
	}
	return ""
}

// transformModule runs esbuild over a single entry file with all imports marked
// external (so only this file is transformed) and rewritten to dev URLs, then
// prepends the import.meta.hot bootstrap. .tsx/.jsx are Solid+refresh compiled
// via the goja pipeline first.
func (d *devServer) transformModule(abs string) ([]byte, error) {
	result := esbuild.Build(esbuild.BuildOptions{
		EntryPoints:    []string{abs},
		Bundle:         true,
		Write:          false,
		Format:         esbuild.FormatESModule,
		Platform:       esbuild.PlatformBrowser,
		Target:         esbuild.ES2022,
		Sourcemap:      esbuild.SourceMapInline,
		SourcesContent: esbuild.SourcesContentInclude,
		LogLevel:       esbuild.LogLevelSilent,
		// Same compile-time env define as the production bundle, so env.ts reads
		// the baked value under HMR too (see esbuildDefine).
		Define: esbuildDefine(),
		// Order matters: the Solid/JSX loader and the export shim handle the entry
		// file's contents; the external-rewrite resolver must be able to see every
		// import, so it runs last (its OnResolve filter is `.*`).
		Plugins: []esbuild.Plugin{
			d.solidRefreshLoadPlugin(),
			defaultExportShimPlugin(),
			d.externalRewritePlugin(),
		},
	})
	if len(result.Errors) > 0 {
		msgs := esbuild.FormatMessages(result.Errors, esbuild.FormatMessagesOptions{})
		return nil, fmt.Errorf("%s", strings.Join(msgs, "\n"))
	}
	if len(result.OutputFiles) == 0 {
		return nil, fmt.Errorf("no output")
	}

	selfURL := d.moduleURLBase(abs)
	prelude := "import { createHotContext as __createHotContext } from \"/@hmr/client\";\n" +
		"import.meta.hot = __createHotContext(" + jsString(selfURL) + ");\n"
	return append([]byte(prelude), result.OutputFiles[0].Contents...), nil
}

// solidRefreshLoadPlugin Solid-compiles .tsx/.jsx via the goja pipeline WITH
// solid-refresh instrumentation, passing the src-relative path as the stable id
// solid-refresh keys its HMR registry on (so it survives across recompiles).
func (d *devServer) solidRefreshLoadPlugin() esbuild.Plugin {
	return esbuild.Plugin{
		Name: "dev-solid-refresh",
		Setup: func(b esbuild.PluginBuild) {
			b.OnLoad(esbuild.OnLoadOptions{Filter: `\.(tsx|jsx)$`}, func(a esbuild.OnLoadArgs) (esbuild.OnLoadResult, error) {
				data, err := os.ReadFile(a.Path)
				if err != nil {
					return esbuild.OnLoadResult{}, err
				}
				id := d.moduleID(a.Path)
				out, err := CompileDev(string(data), id)
				if err != nil {
					return esbuild.OnLoadResult{}, err
				}
				// Load as TS, not JS: esbuild only re-emits (and thus rewrites this
				// module's external import paths to our /@src/ URLs) when it has to
				// transpile. A JS loader leaves the already-JS CompileDev output
				// untouched, leaking the raw ./x.js specifiers to the browser. esbuild
				// chains solid-refresh's inline sourcemap into its own, so the browser
				// still debugs against the original .tsx.
				loader := esbuild.LoaderTS
				dir := filepath.Dir(a.Path)
				return esbuild.OnLoadResult{Contents: &out, Loader: loader, ResolveDir: dir}, nil
			})
		},
	}
}

// externalRewritePlugin marks every non-entry import external and rewrites its
// path to a dev URL: relative/absolute → /@src/<resolved> (recording a graph
// edge), `?url` → /@fs/<file> raw passthrough, bare → left as-is for the import
// map. esbuild emits the returned path verbatim, so we control the URL.
func (d *devServer) externalRewritePlugin() esbuild.Plugin {
	return esbuild.Plugin{
		Name: "dev-external-rewrite",
		Setup: func(b esbuild.PluginBuild) {
			b.OnResolve(esbuild.OnResolveOptions{Filter: `.*`}, func(a esbuild.OnResolveArgs) (esbuild.OnResolveResult, error) {
				if a.Kind == esbuild.ResolveEntryPoint {
					return esbuild.OnResolveResult{}, nil // let esbuild load the entry
				}
				spec := a.Path

				// `import x from "...?url"` — emit the referenced file as a raw asset.
				if strings.HasSuffix(spec, "?url") {
					real := strings.TrimSuffix(spec, "?url")
					if target := d.resolveAsset(a.ResolveDir, real); target != "" {
						return esbuild.OnResolveResult{Path: target, External: true}, nil
					}
				}

				// @ui/@kjol/@appgen → resolve into the kjol web tree (or the app's
				// generated dir) and rewrite to the /@src/ or /@kit/ URL that serves it.
				if strings.HasPrefix(spec, "@") {
					if resolved := d.resolveAliasSpec(spec); resolved != "" {
						if url := d.moduleURLBase(resolved); url != "" {
							if a.Importer != "" {
								d.graph.recordEdge(a.Importer, resolved)
							}
							if v := d.graph.versionOf(resolved); v > 0 {
								url += fmt.Sprintf("?t=%d", v)
							}
							return esbuild.OnResolveResult{Path: url, External: true}, nil
						}
					}
					// Not an @ui/@kjol/@appgen alias (e.g. @solidjs/router): fall through
					// to the bare-specifier / import-map path below.
				}

				// Bare specifier → import map (single vendored copy).
				if !strings.HasPrefix(spec, ".") && !filepath.IsAbs(spec) {
					return esbuild.OnResolveResult{Path: spec, External: true}, nil
				}

				// Relative/absolute → resolve to the real source file (app OR kit tree)
				// and rewrite to its /@src/ or /@kit/ URL, recording the edge for HMR.
				resolved := resolveSource(a.ResolveDir, spec)
				url := ""
				if resolved != "" {
					url = d.moduleURLBase(resolved)
				}
				if url == "" {
					// Unknown import: leave it and let the browser 404 loudly.
					return esbuild.OnResolveResult{Path: spec, External: true}, nil
				}
				if a.Importer != "" {
					d.graph.recordEdge(a.Importer, resolved)
				}
				if v := d.graph.versionOf(resolved); v > 0 {
					url += fmt.Sprintf("?t=%d", v)
				}
				return esbuild.OnResolveResult{Path: url, External: true}, nil
			})
		},
	}
}

// resolveAsset resolves a `?url` target (relative to importer, or a vendored
// subpath) to a /@url/ shim-module path rooted at the frontend dir. The shim's
// default export is the raw /@fs/ URL — matching esbuild's file-loader `?url`
// semantics, where `import u from "x?url"` binds u to the asset's URL string, not
// the module's exports.
func (d *devServer) resolveAsset(resolveDir, spec string) string {
	var abs string
	if strings.HasPrefix(spec, ".") || filepath.IsAbs(spec) {
		abs = filepath.Join(resolveDir, spec)
	} else {
		for _, vd := range d.vendorDirs { // vendored subpath
			if cand := filepath.Join(vd, filepath.FromSlash(spec)); fileExists(cand) {
				abs = cand
				break
			}
		}
	}
	if abs == "" || !within(d.frontend, abs) || !fileExists(abs) {
		return ""
	}
	return assetURLPrefix + filepath.ToSlash(mustRel(d.frontend, abs))
}

// serveAssetURL serves the `?url` shim module: `export default "<raw file URL>"`.
func (d *devServer) serveAssetURL(w http.ResponseWriter, r *http.Request) {
	rel := strings.TrimPrefix(r.URL.Path, assetURLPrefix)
	abs := filepath.Join(d.frontend, filepath.FromSlash(rel))
	if !within(d.frontend, abs) || !fileExists(abs) {
		http.NotFound(w, r)
		return
	}
	w.Header().Set("Content-Type", "application/javascript; charset=utf-8")
	w.Header().Set("Cache-Control", "no-cache")
	fmt.Fprintf(w, "export default %s;\n", jsString(fsURLPrefix+rel))
}

// serveFS serves a raw file under the frontend dir (the target of a `?url` shim,
// e.g. the pdfjs worker, which must load from a real URL for import.meta.url).
func (d *devServer) serveFS(w http.ResponseWriter, r *http.Request) {
	rel := strings.TrimPrefix(r.URL.Path, fsURLPrefix)
	abs := filepath.Join(d.frontend, filepath.FromSlash(rel))
	if !within(d.frontend, abs) || !fileExists(abs) {
		http.NotFound(w, r)
		return
	}
	// Pin the JS MIME so module workers (.mjs) aren't rejected for a text/plain type.
	switch strings.ToLower(filepath.Ext(abs)) {
	case ".mjs", ".js":
		w.Header().Set("Content-Type", "application/javascript; charset=utf-8")
	}
	w.Header().Set("Cache-Control", "no-cache")
	http.ServeFile(w, r, abs)
}

// ---- module graph ----------------------------------------------------------

// moduleGraph tracks importer→import edges and per-module versions so a change
// can be propagated to the nearest self-accepting boundary (a .tsx/.jsx compiled
// with solid-refresh) or, failing that, a full reload. Keys are absolute paths.
type moduleGraph struct {
	mu        sync.Mutex
	importers map[string]map[string]bool // module -> set of modules that import it
	version   map[string]int64           // module -> current version (0 = never changed)
	counter   int64
}

func newModuleGraph() *moduleGraph {
	return &moduleGraph{importers: map[string]map[string]bool{}, version: map[string]int64{}}
}

func (g *moduleGraph) recordEdge(importer, imported string) {
	g.mu.Lock()
	defer g.mu.Unlock()
	set := g.importers[imported]
	if set == nil {
		set = map[string]bool{}
		g.importers[imported] = set
	}
	set[importer] = true
}

func (g *moduleGraph) versionOf(module string) int64 {
	g.mu.Lock()
	defer g.mu.Unlock()
	return g.version[module]
}

// invalidate walks up from the changed file to the nearest HMR boundaries,
// bumping the version of every module on the path (so a re-imported boundary
// re-fetches the changed dep, not a cached copy). It returns the boundary URLs to
// re-import and the shared version, or fullReload if the change reaches a
// non-accepting root (the entry, a routes table).
func (g *moduleGraph) invalidate(changed string) (boundaries []string, version int64, fullReload bool) {
	g.mu.Lock()
	defer g.mu.Unlock()

	g.counter++
	v := g.counter
	visited := map[string]bool{}
	queue := []string{changed}
	for len(queue) > 0 {
		m := queue[0]
		queue = queue[1:]
		if visited[m] {
			continue
		}
		visited[m] = true
		g.version[m] = v

		if isHMRBoundary(m) {
			boundaries = append(boundaries, m)
			continue // don't climb past a self-accepting boundary
		}
		imps := g.importers[m]
		if len(imps) == 0 {
			fullReload = true // reached a root that can't accept -> reload
			continue
		}
		for imp := range imps {
			queue = append(queue, imp)
		}
	}
	if fullReload {
		return nil, v, true
	}
	return boundaries, v, false
}

// isHMRBoundary reports whether a module can self-accept — only solid-refresh
// -instrumented JSX (.tsx/.jsx). .ts/.js (incl. the solid-js/html app pages)
// bubble to a full reload.
func isHMRBoundary(path string) bool {
	switch filepath.Ext(path) {
	case ".tsx", ".jsx":
		return true
	default:
		return false
	}
}

// ---- helpers ---------------------------------------------------------------

// resolveSource replicates esbuild's resolution of a relative import to a real
// source file, including the project's `.js`-specifier-for-a-.ts-file convention
// (see import_check.go) and extensionless / index resolution.
func resolveSource(resolveDir, spec string) string {
	base := filepath.Join(resolveDir, filepath.FromSlash(spec))
	if fileExists(base) {
		return base
	}
	// `./Foo.js` may name a sibling .ts/.tsx/.jsx (esbuild resolves it silently).
	if strings.HasSuffix(base, ".js") {
		stem := strings.TrimSuffix(base, ".js")
		for _, ext := range []string{".ts", ".tsx", ".jsx"} {
			if fileExists(stem + ext) {
				return stem + ext
			}
		}
	}
	// Extensionless specifier.
	for _, ext := range []string{".ts", ".tsx", ".jsx", ".js"} {
		if fileExists(base + ext) {
			return base + ext
		}
	}
	// Directory index.
	for _, ext := range []string{".ts", ".tsx", ".jsx", ".js"} {
		if idx := filepath.Join(base, "index"+ext); fileExists(idx) {
			return idx
		}
	}
	return ""
}

func fileExists(p string) bool {
	info, err := os.Stat(p)
	return err == nil && !info.IsDir()
}

// evalSymlinks resolves symlinks so paths compare equal to esbuild's realpaths;
// returns the input unchanged if it can't be resolved (e.g. doesn't exist yet).
func evalSymlinks(p string) string {
	if r, err := filepath.EvalSymlinks(p); err == nil {
		return r
	}
	return p
}

// within reports whether abs is inside root (after cleaning), guarding against
// `..` traversal out of the served tree.
func within(root, abs string) bool {
	rel, err := filepath.Rel(root, abs)
	if err != nil {
		return false
	}
	return rel != ".." && !strings.HasPrefix(rel, ".."+string(filepath.Separator))
}

func mustRel(root, abs string) string {
	rel, err := filepath.Rel(root, abs)
	if err != nil {
		return abs
	}
	return rel
}

// jsString renders s as a double-quoted JS string literal.
func jsString(s string) string {
	b := strings.Builder{}
	b.WriteByte('"')
	for _, r := range s {
		switch r {
		case '"':
			b.WriteString(`\"`)
		case '\\':
			b.WriteString(`\\`)
		case '\n':
			b.WriteString(`\n`)
		default:
			b.WriteRune(r)
		}
	}
	b.WriteByte('"')
	return b.String()
}

// errorModule returns an ES module that surfaces a compile/transform failure as
// a full-screen overlay in the browser (Vite-style) instead of breaking the page
// hard. It imports the HMR client's showErrorOverlay and calls it with the error
// text; because the stub evaluates synchronously as part of the importing
// boundary's module graph, the overlay appears the moment a broken module is
// (re-)imported — on first load or on a hot re-import. A later successful hot
// update clears it (see applyUpdates in hmr_client.go). rel is the src-relative
// path of the failing module, shown in the overlay header.
func errorModule(rel string, err error) string {
	payload, jerr := json.Marshal(hmrError{Message: err.Error(), File: rel})
	if jerr != nil {
		return "console.error(" + jsString("[hmr] build error:\n"+err.Error()) + ");\n"
	}
	return "import { showErrorOverlay as __hmrShowError } from \"/@hmr/client\";\n" +
		"__hmrShowError(" + string(payload) + ");\n"
}