chore: switch to vite for server dev

2025-01-17 08:33:22 +00:00 · 2025-01-17 08:33:22 +00:00 · c9603ba38a
commit c9603ba38a
parent 377641265b
16 changed files with 444 additions and 265 deletions
--- a/package.json
+++ b/package.json
@ -4,8 +4,8 @@
 	"sideEffects": false,
 	"type": "module",
 	"scripts": {
-		"build": "react-router build && vite build --config server/vite.config.ts",
-		"dev": "node server/dev.mjs",
+		"build": "react-router build && vite build -c server/vite.config.ts",
+		"dev": "vite-node -w -c server/vite.config.ts server/entry.ts",
 		"start": "node build/headplane/server.js",
 		"typecheck": "tsc"
 	},
@ -40,6 +40,7 @@
 		"tailwindcss-react-aria-components": "^1.2.0",
 		"undici": "^7.2.0",
 		"usehooks-ts": "^3.1.0",
+		"vite-node": "^3.0.1",
 		"ws": "^8.18.0",
 		"yaml": "^2.7.0",
 		"zod": "^3.24.1"
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@ -103,6 +103,9 @@ importers:
      usehooks-ts:
        specifier: ^3.1.0
        version: 3.1.0(react@19.0.0)
+      vite-node:
+        specifier: ^3.0.1
+        version: 3.0.1(@types/node@22.10.1)(jiti@1.21.7)(tsx@4.19.2)(yaml@2.7.0)
      ws:
        specifier: ^8.18.0
        version: 8.18.0
@ -1791,6 +1794,9 @@ packages:
  es-module-lexer@1.5.4:
    resolution: {integrity: sha512-MVNK56NiMrOwitFB7cqDwq0CQutbw+0BvLshJSse0MUNU+y1FC3bUS/AQg7oUng+/wKrrki7JfmwtVHkVfPLlw==}

+  es-module-lexer@1.6.0:
+    resolution: {integrity: sha512-qqnD1yMU6tk/jnaMosogGySTZP8YtUgAffA9nMN+E/rjxcfRQ6IEk7IiozUjgxKoFHBGjTLnrHB/YC45r/59EQ==}
+
  esbuild@0.23.1:
    resolution: {integrity: sha512-VVNz/9Sa0bs5SELtn3f7qhJCDPCF5oMEl5cO9/SSinpE9hbPVvxbd572HH5AKiP7WD8INO53GgfDDhRjkylHEg==}
    engines: {node: '>=18'}
@ -2181,6 +2187,9 @@ packages:
  pathe@1.1.2:
    resolution: {integrity: sha512-whLdWMYL2TwI08hn8/ZqAbrVemu0LNaNNJZX73O6qaIdCTfXutsLhMkjdENX0qhsQ9uIimo4/aQOmXkoon2nDQ==}

+  pathe@2.0.1:
+    resolution: {integrity: sha512-6jpjMpOth5S9ITVu5clZ7NOgHNsv5vRQdheL9ztp2vZmM6fRbLvyua1tiBIL4lk8SAe3ARzeXEly6siXCjDHDw==}
+
  peek-stream@1.1.3:
    resolution: {integrity: sha512-FhJ+YbOSBb9/rIl2ZeE/QHEsWn7PqNYt8ARAY3kIgNGOk13g9FGyIY6JIl/xB/3TFRVoTv5as0l11weORrTekA==}

@ -2684,6 +2693,11 @@ packages:
    engines: {node: ^18.0.0 || ^20.0.0 || >=22.0.0}
    hasBin: true

+  vite-node@3.0.1:
+    resolution: {integrity: sha512-PoH9mCNsSZQXl3gdymM5IE4WR0k0WbnFd89nAyyDvltF2jVGdFcI8vpB1PBdKTcjAR7kkYiHSlIO68X/UT8Q1A==}
+    engines: {node: ^18.0.0 || ^20.0.0 || >=22.0.0}
+    hasBin: true
+
  vite-plugin-babel@1.3.0:
    resolution: {integrity: sha512-C5WKX0UwvQKH8WD2GiyWUjI62UBfLbfUhiLexnIm4asLdENX5ymrRipFlBnGeVxoOaYgTL5dh5KW6YDGpWsR8A==}
    peerDependencies:
@ -4825,6 +4839,8 @@ snapshots:

  es-module-lexer@1.5.4: {}

+  es-module-lexer@1.6.0: {}
+
  esbuild@0.23.1:
    optionalDependencies:
      '@esbuild/aix-ppc64': 0.23.1
@ -5218,6 +5234,8 @@ snapshots:

  pathe@1.1.2: {}

+  pathe@2.0.1: {}
+
  peek-stream@1.1.3:
    dependencies:
      buffer-from: 1.1.2
@ -5831,6 +5849,27 @@ snapshots:
      - tsx
      - yaml

+  vite-node@3.0.1(@types/node@22.10.1)(jiti@1.21.7)(tsx@4.19.2)(yaml@2.7.0):
+    dependencies:
+      cac: 6.7.14
+      debug: 4.4.0
+      es-module-lexer: 1.6.0
+      pathe: 2.0.1
+      vite: 6.0.6(@types/node@22.10.1)(jiti@1.21.7)(tsx@4.19.2)(yaml@2.7.0)
+    transitivePeerDependencies:
+      - '@types/node'
+      - jiti
+      - less
+      - lightningcss
+      - sass
+      - sass-embedded
+      - stylus
+      - sugarss
+      - supports-color
+      - terser
+      - tsx
+      - yaml
+
  vite-plugin-babel@1.3.0(@babel/core@7.26.0)(vite@6.0.6(@types/node@22.10.1)(jiti@1.21.7)(tsx@4.19.2)(yaml@2.7.0)):
    dependencies:
      '@babel/core': 7.26.0
--- a/react-router.config.ts
+++ b/react-router.config.ts
@ -2,4 +2,5 @@ import type { Config } from '@react-router/dev/config';

 export default {
 	basename: '/admin/',
+	ssr: true,
 } satisfies Config;
--- a/server/dev-handler.ts
+++ b/server/dev-handler.ts
@ -0,0 +1,7 @@
+import { createRequestHandler } from 'react-router'
+
+export default createRequestHandler(
+	// @ts-expect-error: React Router Vite plugin
+	() => import('virtual:react-router/server-build'),
+	'development'
+);
--- a/server/dev.mjs
+++ b/server/dev.mjs
@ -1,33 +0,0 @@
-// This is a polyglot entrypoint for Headplane when running in development
-// It does some silly magic to load the vite config, set some globals that
-// are required to function, and create the vite development server.
-import { createServer } from 'vite';
-import { env, exit } from 'node:process';
-import { log } from './utils.mjs';
-
-log('DEVX', 'INFO', 'This script is only intended for development');
-env.NODE_ENV = 'development';
-
-// The production entrypoint uses a global called "PREFIX" to determine
-// what route the application is being served at and a global called "BUILD"
-// to determine the Remix handler. We need to set these globals here so that
-// the development server can function correctly and override the production
-// values.
-
-log('DEVX', 'INFO', 'Creating Vite Development Server');
-const server = await createServer({
-	server: {
-		middlewareMode: true,
-	},
-});
-
-// This entrypoint is defined in the documentation to load the server
-const build = await server.ssrLoadModule('virtual:react-router/server-build');
-
-// We already handle this logic in the Vite configuration
-global.PREFIX = server.config.base.slice(0, -1);
-global.BUILD = build;
-global.MODE = 'development';
-global.MIDDLEWARE = server.middlewares;
-
-await import('./prod.mjs');
--- a/server/dev.ts
+++ b/server/dev.ts
@ -0,0 +1,57 @@
+import log from '~server/log';
+import { createServer, type ViteDevServer } from 'vite';
+import { type createRequestHandler } from 'react-router';
+
+let server: ViteDevServer | undefined;
+export async function loadDevtools() {
+	log.info('DEVX', 'Starting Vite Development server')
+	process.env.NODE_ENV = 'development';
+
+	// This is loading the ROOT vite.config.ts
+	server = await createServer({
+		server: {
+			middlewareMode: true,
+		}
+	});
+
+	// We can't just do ssrLoadModule for virtual:react-router/server-build
+	// because for hot reload to work server side it needs to be imported
+	// using builtin import in its own file.
+	const handler = await server.ssrLoadModule('./server/dev-handler.ts');
+	return {
+		server,
+		handler: handler.default,
+	};
+}
+
+export async function stacksafeTry(
+	devtools: {
+		server: ViteDevServer,
+		handler: any, // import() is dynamic
+	},
+	req: Request,
+	context: unknown
+) {
+	try {
+		const result = await devtools.handler(req, context);
+		return result;
+	} catch (error) {
+		log.error('DEVX', 'Error in request handler', error);
+		if (typeof error === 'object' && error instanceof Error) {
+			console.log('got error');
+			devtools.server.ssrFixStacktrace(error);
+		}
+
+		throw error;
+	}
+}
+
+if (import.meta.hot) {
+	import.meta.hot.on('vite:beforeFullReload', () => {
+		server?.close();
+	});
+
+	import.meta.hot.dispose(() => {
+		server?.close();
+	});
+}
--- a/server/entry.ts
+++ b/server/entry.ts
@ -0,0 +1,45 @@
+import { createServer } from 'node:http';
+import { listener } from '~server/listener';
+import { initWebsocket } from '~server/ws';
+import { access, constants } from 'node:fs/promises';
+import log from '~server/log';
+
+log.info('SRVX', 'Running Node.js %s', process.versions.node);
+
+try {
+	await access('./node_modules/react-router', constants.F_OK | constants.R_OK);
+	log.info('SRVX', 'Found dependencies');
+} catch (error) {
+	log.error('SRVX', 'No dependencies found. Please run `npm install`');
+	console.error(error);
+	process.exit(1);
+}
+
+const server = createServer(listener);
+const port = process.env.PORT || 3000;
+const host = process.env.HOST || '0.0.0.0';
+
+const ws = initWebsocket();
+if (ws) {
+	server.on('upgrade', (req, socket, head) => {
+		ws.handleUpgrade(req, socket, head, (ws) => {
+			ws.emit('connection', ws, req);
+		});
+	});
+}
+
+server.listen(port, host, () => {
+	log.info('SRVX', 'Running on %s:%s', host, port);
+});
+
+if (import.meta.hot) {
+	import.meta.hot.on('vite:beforeFullReload', () => {
+		server.close();
+	});
+
+	import.meta.hot.dispose(() => {
+		server.close();
+	});
+}
+
+// export const app = listener;
--- a/server/listener.ts
+++ b/server/listener.ts
@ -0,0 +1,137 @@
+import { type RequestListener } from 'node:http';
+import { resolve, join } from 'node:path'
+import { createServer } from 'vite'
+import { createRequestHandler } from 'react-router'
+import { access, constants } from 'node:fs/promises';
+import { createReadStream, existsSync, statSync } from 'node:fs';
+import {
+	createReadableStreamFromReadable,
+	writeReadableStreamToWritable,
+} from '@react-router/node';
+import mime from 'mime/lite'
+
+import { loadDevtools, stacksafeTry } from '~server/dev';
+import { appContext } from '~server/ws';
+import prodBuild from '~server/prod-handler';
+
+declare global {
+	// Prefix is a build-time constant
+	const __hp_prefix: string;
+}
+
+const devtools = import.meta.env.DEV
+	? await loadDevtools()
+	: undefined;
+
+const prodHandler = import.meta.env.PROD
+	? await prodBuild()
+	: undefined;
+
+const buildPath = process.env.BUILD_PATH ?? './build';
+const baseDir = resolve(join(buildPath, 'client'));
+
+export const listener: RequestListener = async (req, res) => {
+	const url = new URL(`http://${req.headers.host}${req.url}`);
+
+	// build:strip
+	if (devtools) {
+		await new Promise((resolve) => {
+			devtools.server.middlewares(req, res, resolve);
+		});
+	}
+
+	if (!url.pathname.startsWith(__hp_prefix)) {
+		res.writeHead(404);
+		res.end();
+		return;
+	}
+
+	// We need to handle an issue where say we are navigating to $PREFIX
+	// but Remix does not handle it without the trailing slash. This is
+	// because Remix uses the URL constructor to parse the URL and it
+	// will remove the trailing slash. We need to redirect to the correct
+	// URL so that Remix can handle it correctly.
+	if (url.pathname === __hp_prefix) {
+		res.writeHead(301, {
+			Location: `${__hp_prefix}/`,
+		});
+		res.end();
+		return;
+	}
+
+	// Before we pass any requests to our Remix handler we need to check
+	// if we can handle a raw file request. This is important for the
+	// Remix loader to work correctly.
+	//
+	// To optimize this, we send them as readable streams in the node
+	// response and we also set headers for aggressive caching.
+	if (url.pathname.startsWith(`${__hp_prefix}/assets/`)) {
+		const filePath = join(baseDir, url.pathname.slice(__hp_prefix.length));
+		const exists = existsSync(filePath);
+		const stats = exists ? statSync(filePath) : null;
+
+		if (exists && stats && stats.isFile()) {
+			// Build assets are cache-bust friendly so we can cache them heavily
+			res.setHeader('Cache-Control', 'public, max-age=31536000, immutable');
+		}
+
+		// Send the file as a readable stream
+		const fileStream = createReadStream(filePath);
+		const type = mime.getType(filePath);
+
+		res.writeHead(200, {
+			'Content-Type': type || 'application/octet-stream',
+		});
+
+		fileStream.pipe(res);
+		return;
+	}
+
+	// Handling the request
+	const controller = new AbortController();
+	res.on('close', () => controller.abort());
+
+	const headers = new Headers();
+	for (const [key, value] of Object.entries(req.headers)) {
+		if (!value) continue;
+		if (Array.isArray(value)) {
+			for (const v of value) {
+				headers.append(key, v);
+			}
+
+			continue;
+		}
+
+		headers.append(key, value);
+	}
+
+	const frameworkReq = new Request(url.href, {
+		headers,
+		method: req.method,
+		signal: controller.signal,
+
+		// If we have a body, we set the duplex and load it
+		...(req.method !== 'GET' && req.method !== 'HEAD'
+			? {
+				body: createReadableStreamFromReadable(req),
+				duplex: 'half',
+			} : {}),
+	});
+
+	const response = devtools
+		? await stacksafeTry(devtools, frameworkReq, appContext())
+		: await prodHandler(frameworkReq, appContext());
+
+	res.statusCode = response.status;
+	res.statusMessage = response.statusText;
+
+	for (const [key, value] of response.headers.entries()) {
+		res.appendHeader(key, value);
+	}
+
+	if (response.body) {
+		await writeReadableStreamToWritable(response.body, res);
+	}
+
+	res.end();
+}
--- a/server/log.ts
+++ b/server/log.ts
@ -0,0 +1,29 @@
+export default {
+	info: (category: string, message: string, ...args: unknown[]) => {
+		defaultLog('INFO', category, message, ...args);
+	},
+
+	warn: (category: string, message: string, ...args: unknown[]) => {
+		defaultLog('WARN', category, message, ...args);
+	},
+
+	error: (category: string, message: string, ...args: unknown[]) => {
+		defaultLog('ERRO', category, message, ...args);
+	},
+
+	debug: (category: string, message: string, ...args: unknown[]) => {
+		if (process.env.DEBUG === 'true') {
+			defaultLog('DEBG', category, message, ...args);
+		}
+	},
+};
+
+function defaultLog(
+	level: string,
+	category: string,
+	message: string,
+	...args: unknown[]
+) {
+	const date = new Date().toISOString();
+	console.log(`${date} (${level}) [${category}] ${message}`, ...args);
+}
--- a/server/prod-handler.ts
+++ b/server/prod-handler.ts
@ -0,0 +1,23 @@
+import { createRequestHandler } from 'react-router'
+import { access, constants } from 'node:fs/promises';
+import { join, resolve } from 'node:path';
+import log from '~server/log';
+
+export default async function() {
+	const buildPath = process.env.BUILD_PATH ?? './build';
+	const server = resolve(join(buildPath, 'server'));
+
+	try {
+		await access(server, constants.F_OK | constants.R_OK);
+		log.info('SRVX', 'Using build directory %s', resolve(buildPath));
+	} catch (error) {
+		log.error('SRVX', 'No build found. Please refer to the documentation');
+		log.error('SRVX', 'https://github.com/tale/headplane/blob/main/docs/integration/Native.md');
+		console.error(error);
+		process.exit(1);
+	}
+
+	// @vite-ignore
+	const build = await import(resolve(join(server, 'index.js')));
+	return createRequestHandler(build, 'production');
+}
--- a/server/prod.mjs
+++ b/server/prod.mjs
@ -1,170 +0,0 @@
-// This is a polyglot entrypoint for Headplane when running in production
-// It doesn't use any dependencies aside from @remix-run/node and mime
-// During build we bundle the used dependencies into the file so that
-// we can only need this file and a Node.js installation to run the server.
-// PREFIX is defined globally, see vite.config.ts
-
-import { access, constants } from 'node:fs/promises';
-import { createReadStream, existsSync, statSync } from 'node:fs';
-import { createServer } from 'node:http';
-import { join, resolve } from 'node:path';
-import { env } from 'node:process';
-import { log } from './utils.mjs';
-import { getWss, registerWss } from './ws.mjs';
-import {
-	createReadableStreamFromReadable,
-	writeReadableStreamToWritable,
-} from '@react-router/node';
-
-log('SRVX', 'INFO', `Running with Node.js ${process.versions.node}`);
-
-try {
-	await access('./node_modules/react-router', constants.F_OK | constants.R_OK);
-	log('SRVX', 'INFO', 'Found node_modules dependencies');
-} catch (error) {
-	log('SRVX', 'ERROR', 'No node_modules found. Please run `pnpm install`');
-	log('SRVX', 'ERROR', error);
-	process.exit(1);
-}
-
-const { createRequestHandler } = await import('react-router');
-const { default: mime } = await import('mime');
-
-const port = env.PORT || 3000;
-const host = env.HOST || '0.0.0.0';
-const buildPath = env.BUILD_PATH || './build';
-const baseDir = resolve(join(buildPath, 'client'));
-
-if (!global.BUILD) {
-	try {
-		await access(join(buildPath, 'server'), constants.F_OK | constants.R_OK);
-		log('SRVX', 'INFO', 'Found build directory');
-	} catch (error) {
-		const date = new Date().toISOString();
-		log('SRVX', 'ERROR', 'No build found. Please run `pnpm build`');
-		log('SRVX', 'ERROR', error);
-		process.exit(1);
-	}
-
-	// Because this is a dynamic import without an easily discernable path
-	// we gain the "deoptimization" we want so that Vite doesn't bundle this
-	const build = await import(resolve(join(buildPath, 'server', 'index.js')));
-	global.BUILD = build;
-	global.MODE = 'production';
-}
-
-const handler = createRequestHandler(global.BUILD, global.MODE);
-const http = createServer(async (req, res) => {
-	const url = new URL(`http://${req.headers.host}${req.url}`);
-
-	if (global.MIDDLEWARE) {
-		await new Promise((resolve) => {
-			global.MIDDLEWARE(req, res, resolve);
-		});
-	}
-
-	if (!url.pathname.startsWith(PREFIX)) {
-		res.writeHead(404);
-		res.end();
-		return;
-	}
-
-	// We need to handle an issue where say we are navigating to $PREFIX
-	// but Remix does not handle it without the trailing slash. This is
-	// because Remix uses the URL constructor to parse the URL and it
-	// will remove the trailing slash. We need to redirect to the correct
-	// URL so that Remix can handle it correctly.
-	if (url.pathname === PREFIX) {
-		res.writeHead(302, {
-			Location: `${PREFIX}/`,
-		});
-		res.end();
-		return;
-	}
-
-	// Before we pass any requests to our Remix handler we need to check
-	// if we can handle a raw file request. This is important for the
-	// Remix loader to work correctly.
-	//
-	// To optimize this, we send them as readable streams in the node
-	// response and we also set headers for aggressive caching.
-	if (url.pathname.startsWith(`${PREFIX}/assets/`)) {
-		const filePath = join(baseDir, url.pathname.replace(PREFIX, ''));
-		const exists = existsSync(filePath);
-		const stats = statSync(filePath);
-
-		if (exists && stats.isFile()) {
-			// Build assets are cache-bust friendly so we can cache them heavily
-			if (req.url.startsWith('/build')) {
-				res.setHeader('Cache-Control', 'public, max-age=31536000, immutable');
-			}
-
-			// Send the file as a readable stream
-			const fileStream = createReadStream(filePath);
-			const type = mime.getType(filePath);
-
-			res.setHeader('Content-Length', stats.size);
-			res.setHeader('Content-Type', type);
-			fileStream.pipe(res);
-			return;
-		}
-	}
-
-	// Handling the request
-	const controller = new AbortController();
-	res.on('close', () => controller.abort());
-
-	const headers = new Headers();
-	for (const [key, value] of Object.entries(req.headers)) {
-		if (!value) continue;
-
-		if (Array.isArray(value)) {
-			for (const v of value) {
-				headers.append(key, v);
-			}
-
-			continue;
-		}
-
-		headers.append(key, value);
-	}
-
-	const remixReq = new Request(url.href, {
-		headers,
-		method: req.method,
-		signal: controller.signal,
-
-		// If we have a body we set a duplex and we load the body
-		...(req.method !== 'GET' && req.method !== 'HEAD'
-			? {
-					body: createReadableStreamFromReadable(req),
-					duplex: 'half',
-				}
-			: {}),
-	});
-
-	// Pass our request to the Remix handler and get a response
-	const response = await handler(remixReq, {
-		ws: getWss(),
-	});
-
-	// Handle our response and reply
-	res.statusCode = response.status;
-	res.statusMessage = response.statusText;
-
-	for (const [key, value] of response.headers.entries()) {
-		res.appendHeader(key, value);
-	}
-
-	if (response.body) {
-		await writeReadableStreamToWritable(response.body, res);
-		return;
-	}
-
-	res.end();
-});
-
-registerWss(http);
-http.listen(port, host, () => {
-	log('SRVX', 'INFO', `Running on ${host}:${port}`);
-});
--- a/server/utils.mjs
+++ b/server/utils.mjs
@ -1,4 +0,0 @@
-export function log(topic, level, message) {
-	const date = new Date().toISOString();
-	console.log(`${date} (${level}) [${topic}] ${message}`);
-}
--- a/server/vite.config.ts
+++ b/server/vite.config.ts
@ -1,33 +1,50 @@
 import { defineConfig } from 'vite';
+import { VitePluginNode } from 'vite-plugin-node';
+import tsconfigPaths from 'vite-tsconfig-paths';

 const prefix = process.env.__INTERNAL_PREFIX || '/admin';
 if (prefix.endsWith('/')) {
 	throw new Error('Prefix must not end with a slash');
 }

-export default defineConfig(() => {
-	return {
-		build: {
-			minify: false,
-			target: 'esnext',
-			rollupOptions: {
-				input: './server/prod.mjs',
-				output: {
-					entryFileNames: 'server.js',
-					dir: 'build/headplane',
-					banner: '#!/usr/bin/env node\n',
-				},
-				external: (id) => id.startsWith('node:') || id === 'ws',
+export default defineConfig({
+	define: {
+		__hp_prefix: JSON.stringify(prefix),
+	},
+	resolve: {
+		preserveSymlinks: true,
+		alias: {
+			stream: 'node:stream',
+			crypto: 'node:crypto'
+		},
+	},
+	plugins: [tsconfigPaths()],
+	build: {
+		minify: false,
+		target: 'esnext',
+		rollupOptions: {
+			input: './server/entry.ts',
+			treeshake: {
+				moduleSideEffects: false,
 			},
-		},
-		define: {
-			PREFIX: JSON.stringify(prefix),
-		},
-		resolve: {
-			alias: {
-				stream: 'node:stream',
-				crypto: 'node:crypto',
+			output: {
+				entryFileNames: 'server.js',
+				dir: 'build/headplane',
+				banner: '#!/usr/bin/env node\n',
 			},
+			// external: (id) => id.startsWith('node:') || id === 'ws',
+			external: (id) => {
+				// Resolve happens before side-effects are removed
+				// ie. vite import because of viteDevServer
+				if (/node_modules/.test(id)) {
+					return true;
+				}
+
+				return id.startsWith('node:')
+					|| id === 'ws'
+					|| id === 'mime/lite'
+					|| id === '@react-router/node';
+			}
 		},
-	};
-});
+	}
+})
--- a/server/ws.mjs
+++ b/server/ws.mjs
@ -1,28 +0,0 @@
-// The Websocket server is wholly responsible for ingesting messages from
-// Headplane agent instances (hopefully not more than 1 is running lol)
-import { WebSocketServer } from 'ws';
-import { log } from './utils.mjs';
-
-const wss = new WebSocketServer({ noServer: true });
-wss.on('connection', (ws, req) => {
-	// On connection the agent will send its NodeID via Headers
-	// We store this for later use to validate and show on the UI
-	const nodeID = req.headers['x-headplane-ts-node-id'];
-	if (!nodeID) {
-		ws.close(1008, 'ERR_NO_HP_TS_NODE_ID');
-		return;
-	}
-});
-
-export async function registerWss(server) {
-	log('SRVX', 'INFO', 'Registering Websocket Server');
-	server.on('upgrade', (request, socket, head) => {
-		wss.handleUpgrade(request, socket, head, (ws) => {
-			wss.emit('connection', ws, request);
-		});
-	});
-}
-
-export function getWss() {
-	return wss;
-}
--- a/server/ws.ts
+++ b/server/ws.ts
@ -0,0 +1,57 @@
+import WebSocket, { WebSocketServer } from 'ws'
+import log from '~server/log'
+
+const server = new WebSocketServer({ noServer: true });
+export function initWebsocket() {
+	const key = process.env.LOCAL_AGENT_AUTHKEY;
+	if (!key) {
+		return;
+	}
+
+	log.info('CACH', 'Initializing agent WebSocket');
+	server.on('connection', (ws, req) => {
+		const auth = req.headers['authorization'];
+		if (auth !== `Bearer ${key}`) {
+			log.warn('CACH', 'Invalid agent WebSocket connection');
+			ws.close(1008, 'ERR_INVALID_AUTH');
+			return;
+		}
+
+
+		const nodeID = req.headers['x-headplane-ts-node-id'];
+		if (!nodeID) {
+			log.warn('CACH', 'Invalid agent WebSocket connection');
+			ws.close(1008, 'ERR_INVALID_NODE_ID');
+			return;
+		}
+
+		const pinger = setInterval(() => {
+			if (ws.readyState !== WebSocket.OPEN) {
+				clearInterval(pinger);
+				return;
+			}
+
+			ws.ping();
+		}, 30000);
+
+		ws.on('close', () => {
+			clearInterval(pinger);
+		});
+
+		ws.on('error', (error) => {
+			clearInterval(pinger);
+			log.error('CACH', 'Closing agent WebSocket connection');
+			log.error('CACH', 'Agent WebSocket error: %s', error);
+			ws.close(1011, 'ERR_INTERNAL_ERROR');
+		})
+	});
+
+	return server;
+}
+
+export function appContext() {
+	return {
+		ws: server,
+		wsAuthKey: process.env.LOCAL_AGENT_AUTHKEY,
+	}
+}
--- a/tsconfig.json
+++ b/tsconfig.json
@ -23,10 +23,11 @@
 		"forceConsistentCasingInFileNames": true,
 		"baseUrl": ".",
 		"paths": {
-	  "~/*": ["./app/*"]
-	},
+			"~/*": ["./app/*"],
+			"~server/*": ["./server/*"]
+		},

-	// Vite takes care of building everything, not tsc.
-	"noEmit": true
-  }
+		// Vite takes care of building everything, not tsc.
+		"noEmit": true
+	}
 }