Enable HMR (#260)

* feat: enable HMR in `createSnowpack`

* feat: enable Snowpack's HMR

* chore: add changeset

* chore: remove unused file

* chore: add changeset

8 files changed, 113 insertions, 68 deletions

diff --git a/.changeset/brave-panthers-heal.md b/.changeset/brave-panthers-heal.md
new file mode 100644
index 000000000..59f12df50
--- /dev/null
+++ b/.changeset/brave-panthers-heal.md
@@ -0,0 +1,5 @@
+---
+'astro': minor
+---
+
+Enable Snowpack's [built-in HMR support](https://www.snowpack.dev/concepts/hot-module-replacement) to enable seamless live updates while editing.
diff --git a/.changeset/khaki-avocados-lie.md b/.changeset/khaki-avocados-lie.md
new file mode 100644
index 000000000..3a47406f5
--- /dev/null
+++ b/.changeset/khaki-avocados-lie.md
@@ -0,0 +1,5 @@
+---
+'astro': minor
+---
+
+Enabled Snowpack's built-in HMR engine for Astro pages
diff --git a/packages/astro/snowpack-plugin.cjs b/packages/astro/snowpack-plugin.cjs
index f701e8b70..40dbf628a 100644
--- a/packages/astro/snowpack-plugin.cjs
+++ b/packages/astro/snowpack-plugin.cjs
@@ -3,7 +3,7 @@ const { readFile } = require('fs').promises;
 // Snowpack plugins must be CommonJS :(
 const transformPromise = import('./dist/compiler/index.js');
 
-module.exports = function (snowpackConfig, { resolvePackageUrl, renderers, astroConfig } = {}) {
+module.exports = (snowpackConfig, { resolvePackageUrl, hmrPort, renderers, astroConfig } = {}) => {
   return {
     name: 'snowpack-astro',
     resolve: {
@@ -48,6 +48,7 @@ ${contents}`;
       const contents = await readFile(filePath, 'utf-8');
       const compileOptions = {
         astroConfig,
+        hmrPort,
         resolvePackageUrl,
         renderers,
       };
diff --git a/packages/astro/src/@types/compiler.ts b/packages/astro/src/@types/compiler.ts
index 10703c523..6435aea80 100644
--- a/packages/astro/src/@types/compiler.ts
+++ b/packages/astro/src/@types/compiler.ts
@@ -5,6 +5,7 @@ export interface CompileOptions {
   logging: LogOptions;
   resolvePackageUrl: (p: string) => Promise<string>;
   astroConfig: AstroConfig;
+  hmrPort?: number;
   mode: RuntimeMode;
   tailwindConfig?: string;
 }
diff --git a/packages/astro/src/compiler/transform/head.ts b/packages/astro/src/compiler/transform/head.ts
new file mode 100644
index 000000000..58b14c451
--- /dev/null
+++ b/packages/astro/src/compiler/transform/head.ts
@@ -0,0 +1,86 @@
+import type { Transformer, TransformOptions } from '../../@types/transformer';
+import type { TemplateNode } from 'astro-parser';
+
+/** If there are hydrated components, inject styles for [data-astro-root] and [data-astro-children] */
+export default function (opts: TransformOptions): Transformer {
+  let head: TemplateNode;
+  let hasComponents = false;
+  let isHmrEnabled = typeof opts.compileOptions.hmrPort !== 'undefined';
+
+  return {
+    visitors: {
+      html: {
+        InlineComponent: {
+          enter(node, parent) {
+            const [name, kind] = node.name.split(':');
+            if (kind && !hasComponents) {
+              hasComponents = true;
+            }
+          },
+        },
+        Element: {
+          enter(node) {
+            if (!hasComponents) return;
+            switch (node.name) {
+              case 'head': {
+                head = node;
+                return;
+              }
+              default:
+                return;
+            }
+          },
+        },
+      },
+    },
+    async finalize() {
+      if (!head) return;
+
+      const children = [];
+      if (hasComponents) {
+        children.push({
+          type: 'Element',
+          name: 'style',
+          attributes: [{ name: 'type', type: 'Attribute', value: [{ type: 'Text', raw: 'text/css', data: 'text/css' }] }],
+          start: 0,
+          end: 0,
+          children: [
+            {
+              start: 0,
+              end: 0,
+              type: 'Text',
+              data: 'astro-root, astro-fragment { display: contents; }',
+              raw: 'astro-root, astro-fragment { display: contents; }',
+            },
+          ],
+        });
+      }
+
+      if (isHmrEnabled) {
+        const { hmrPort } = opts.compileOptions;
+        children.push({
+              type: 'Element',
+              name: 'script',
+              attributes: [],
+              children: [
+                { type: 'Text', data: `window.HMR_WEBSOCKET_URL = 'ws://localhost:${hmrPort}'`, start: 0, end: 0 }
+              ],
+              start: 0,
+              end: 0
+            }, {
+              type: 'Element',
+              name: 'script',
+              attributes: [
+                { type: 'Attribute', name: 'type', value: [{ type: 'Text', data: 'module', start: 0, end: 0 }], start: 0, end: 0 },
+                { type: 'Attribute', name: 'src', value: [{ type: 'Text', data: '/_snowpack/hmr-client.js', start: 0, end: 0 }], start: 0, end: 0 },
+              ],
+              children: [],
+              start: 0,
+              end: 0
+            })
+      }
+      head.children = head.children ?? [];
+      head.children.push(...children);
+    },
+  };
+}
diff --git a/packages/astro/src/compiler/transform/hydration.ts b/packages/astro/src/compiler/transform/hydration.ts
deleted file mode 100644
index 8a6b2700a..000000000
--- a/packages/astro/src/compiler/transform/hydration.ts
+++ /dev/null
@@ -1,62 +0,0 @@
-import { Transformer } from '../../@types/transformer';
-import type { TemplateNode } from 'astro-parser';
-
-/** If there are hydrated components, inject styles for [data-astro-root] and [data-astro-children] */
-export default function (): Transformer {
-  let head: TemplateNode;
-  let body: TemplateNode;
-  let hasComponents = false;
-
-  return {
-    visitors: {
-      html: {
-        InlineComponent: {
-          enter(node, parent) {
-            const [name, kind] = node.name.split(':');
-            if (kind && !hasComponents) {
-              hasComponents = true;
-            }
-          },
-        },
-        Element: {
-          enter(node) {
-            if (!hasComponents) return;
-            switch (node.name) {
-              case 'head': {
-                head = node;
-                return;
-              }
-              case 'body': {
-                body = node;
-                return;
-              }
-              default:
-                return;
-            }
-          },
-        },
-      },
-    },
-    async finalize() {
-      if (!(head && hasComponents)) return;
-
-      const style: TemplateNode = {
-        type: 'Element',
-        name: 'style',
-        attributes: [{ name: 'type', type: 'Attribute', value: [{ type: 'Text', raw: 'text/css', data: 'text/css' }] }],
-        start: 0,
-        end: 0,
-        children: [
-          {
-            start: 0,
-            end: 0,
-            type: 'Text',
-            data: 'astro-root, astro-fragment { display: contents; }',
-            raw: 'astro-root, astro-fragment { display: contents; }',
-          },
-        ],
-      };
-      head.children = [...(head.children ?? []), style];
-    },
-  };
-}
diff --git a/packages/astro/src/compiler/transform/index.ts b/packages/astro/src/compiler/transform/index.ts
index d622846d9..9df050ee5 100644
--- a/packages/astro/src/compiler/transform/index.ts
+++ b/packages/astro/src/compiler/transform/index.ts
@@ -8,7 +8,7 @@ import transformStyles from './styles.js';
 import transformDoctype from './doctype.js';
 import transformModuleScripts from './module-scripts.js';
 import transformCodeBlocks from './prism.js';
-import transformHydration from './hydration.js';
+import transformHead from './head.js';
 
 interface VisitorCollection {
   enter: Map<string, VisitorFn[]>;
@@ -85,7 +85,7 @@ export async function transform(ast: Ast, opts: TransformOptions) {
   const cssVisitors = createVisitorCollection();
   const finalizers: Array<() => Promise<void>> = [];
 
-  const optimizers = [transformHydration(), transformStyles(opts), transformDoctype(opts), transformModuleScripts(opts), transformCodeBlocks(ast.module)];
+  const optimizers = [transformHead(opts), transformStyles(opts), transformDoctype(opts), transformModuleScripts(opts), transformCodeBlocks(ast.module)];
 
   for (const optimizer of optimizers) {
     collectVisitors(optimizer, htmlVisitors, cssVisitors, finalizers);
diff --git a/packages/astro/src/runtime.ts b/packages/astro/src/runtime.ts
index c665fe1e9..1acf1685e 100644
--- a/packages/astro/src/runtime.ts
+++ b/packages/astro/src/runtime.ts
@@ -264,26 +264,31 @@ interface CreateSnowpackOptions {
   env: Record<string, any>;
   mode: RuntimeMode;
   resolvePackageUrl?: (pkgName: string) => Promise<string>;
+  target: 'frontend'|'backend';
 }
 
-const defaultRenderers = ['@astro-renderer/vue', '@astro-renderer/svelte', '@astro-renderer/react', '@astro-renderer/preact'];
+const DEFAULT_HMR_PORT = 12321;
+const DEFAULT_RENDERERS = ['@astro-renderer/vue', '@astro-renderer/svelte', '@astro-renderer/react', '@astro-renderer/preact'];
 
 /** Create a new Snowpack instance to power Astro */
 async function createSnowpack(astroConfig: AstroConfig, options: CreateSnowpackOptions) {
-  const { projectRoot, astroRoot, renderers = defaultRenderers } = astroConfig;
-  const { env, mode, resolvePackageUrl } = options;
+  const { projectRoot, astroRoot, renderers = DEFAULT_RENDERERS } = astroConfig;
+  const { env, mode, resolvePackageUrl, target } = options;
 
   const internalPath = new URL('./frontend/', import.meta.url);
   const resolveDependency = (dep: string) => resolve.sync(dep, { basedir: fileURLToPath(projectRoot) });
+  const isHmrEnabled = mode === 'development' && target === 'backend';
 
   let snowpack: SnowpackDevServer;
   let astroPluginOptions: {
     resolvePackageUrl?: (s: string) => Promise<string>;
     renderers?: { name: string; client: string; server: string }[];
     astroConfig: AstroConfig;
+    hmrPort?: number
   } = {
     astroConfig,
     resolvePackageUrl,
+    hmrPort: isHmrEnabled ? DEFAULT_HMR_PORT : undefined
   };
 
   const mountOptions = {
@@ -360,6 +365,8 @@ async function createSnowpack(astroConfig: AstroConfig, options: CreateSnowpackO
       open: 'none',
       output: 'stream',
       port: 0,
+      hmr: isHmrEnabled,
+      hmrPort: isHmrEnabled ? DEFAULT_HMR_PORT : undefined,
       tailwindConfig: astroConfig.devOptions.tailwindConfig,
     },
     buildOptions: {
@@ -394,6 +401,7 @@ export async function createRuntime(astroConfig: AstroConfig, { mode, logging }:
     snowpackRuntime: backendSnowpackRuntime,
     snowpackConfig: backendSnowpackConfig,
   } = await createSnowpack(astroConfig, {
+    target: 'backend',
     env: {
       astro: true,
     },
@@ -408,6 +416,7 @@ export async function createRuntime(astroConfig: AstroConfig, { mode, logging }:
     snowpackRuntime: frontendSnowpackRuntime,
     snowpackConfig: frontendSnowpackConfig,
   } = await createSnowpack(astroConfig, {
+    target: 'frontend',
     env: {
       astro: false,
     },