From c84e9ca4a8ae256a365c788ac44748d6a25a6867 Mon Sep 17 00:00:00 2001
From: Aarnav Tale <aarnavtale@icloud.com>
Date: Tue, 4 Nov 2025 23:06:56 -0500
Subject: [PATCH] feat: update more primitives to use the new api

---
 app/layouts/dashboard.tsx                     | 17 +++--
 app/layouts/shell.tsx                         | 31 ++++-----
 app/routes/machines/machine.tsx               |  8 +--
 app/routes/machines/overview.tsx              | 42 ++++++-------
 app/routes/ssh/console.tsx                    | 63 ++++++-------------
 app/routes/util/healthz.ts                    | 10 +--
 app/server/db/pruner.ts                       | 15 ++---
 app/server/headscale/api/endpoints/index.ts   | 31 +++++++++
 .../headscale/api/endpoints/pre-auth-keys.ts  |  8 +--
 9 files changed, 109 insertions(+), 116 deletions(-)

diff --git a/app/layouts/dashboard.tsx b/app/layouts/dashboard.tsx
index 82bb441..eecd2d2 100644
--- a/app/layouts/dashboard.tsx
+++ b/app/layouts/dashboard.tsx
@@ -1,24 +1,21 @@
-import { type LoaderFunctionArgs, Outlet, redirect } from 'react-router';
+import { Outlet, redirect } from 'react-router';
 import { ErrorPopup } from '~/components/Error';
-import type { LoadContext } from '~/server';
 import { pruneEphemeralNodes } from '~/server/db/pruner';
 import ResponseError from '~/server/headscale/api/response-error';
 import log from '~/utils/log';
+import type { Route } from './+types/dashboard';
 
-export async function loader({
-	request,
-	context,
-	...rest
-}: LoaderFunctionArgs<LoadContext>) {
-	const healthy = await context.client.healthcheck();
+export async function loader({ request, context, ...rest }: Route.LoaderArgs) {
 	const session = await context.sessions.auth(request);
+	const api = context.hsApi.getRuntimeClient(session.api_key);
 	await pruneEphemeralNodes({ context, request, ...rest });
+	const healthy = await api.isHealthy();
 
 	// We shouldn't session invalidate if Headscale is down
-	// TODO: Notify in the logs or the UI that OIDC auth key is wrong if enabled
+	// TODO: Notify in the logs or the UI whether or not the OIDC auth key is wrong if enabled
 	if (healthy) {
 		try {
-			await context.client.get('v1/apikey', session.api_key);
+			await api.getApiKeys();
 		} catch (error) {
 			if (error instanceof ResponseError) {
 				log.debug('api', 'API Key validation failed %o', error);
diff --git a/app/layouts/shell.tsx b/app/layouts/shell.tsx
index 0b2d247..6a6c67b 100644
--- a/app/layouts/shell.tsx
+++ b/app/layouts/shell.tsx
@@ -1,26 +1,18 @@
 import { eq } from 'drizzle-orm';
 import { CircleCheckIcon } from 'lucide-react';
-import {
-	LoaderFunctionArgs,
-	Outlet,
-	redirect,
-	useLoaderData,
-} from 'react-router';
+import { Outlet, redirect } from 'react-router';
 import Button from '~/components/Button';
 import Card from '~/components/Card';
 import Footer from '~/components/Footer';
 import Header from '~/components/Header';
-import type { LoadContext } from '~/server';
 import { users } from '~/server/db/schema';
 import { Capabilities } from '~/server/web/roles';
 import toast from '~/utils/toast';
+import { Route } from './+types/shell';
 
 // This loads the bare minimum for the application to function
 // So we know that if context fails to load then well, oops?
-export async function loader({
-	request,
-	context,
-}: LoaderFunctionArgs<LoadContext>) {
+export async function loader({ request, context }: Route.LoaderArgs) {
 	try {
 		const session = await context.sessions.auth(request);
 		if (
@@ -39,6 +31,7 @@ export async function loader({
 			}
 		}
 
+		const api = context.hsApi.getRuntimeClient(session.api_key);
 		const check = await context.sessions.check(request, Capabilities.ui_access);
 		return {
 			config: context.hs.c,
@@ -62,7 +55,7 @@ export async function loader({
 				),
 			},
 			onboarding: request.url.endsWith('/onboarding'),
-			healthy: await context.client.healthcheck(),
+			healthy: await api.isHealthy(),
 		};
 	} catch {
 		return redirect('/login', {
@@ -73,14 +66,12 @@ export async function loader({
 	}
 }
 
-export default function Shell() {
-	const data = useLoaderData<typeof loader>();
-
+export default function Shell({ loaderData }: Route.ComponentProps) {
 	return (
 		<>
-			<Header {...data} />
+			<Header {...loaderData} />
 			{/* Always show the outlet if we are onboarding */}
-			{(data.onboarding ? true : data.uiAccess) ? (
+			{(loaderData.onboarding ? true : loaderData.uiAccess) ? (
 				<Outlet />
 			) : (
 				<Card className="mx-auto w-fit mt-24">
@@ -96,13 +87,13 @@ export default function Shell() {
 						className="flex text-md font-mono"
 						onPress={async () => {
 							await navigator.clipboard.writeText(
-								`tailscale up --login-server=${data.url}`,
+								`tailscale up --login-server=${loaderData.url}`,
 							);
 
 							toast('Copied to clipboard');
 						}}
 					>
-						tailscale up --login-server={data.url}
+						tailscale up --login-server={loaderData.url}
 					</Button>
 					<p className="text-xs mt-1 opacity-50 text-center">
 						Click this button to copy the command.
@@ -113,7 +104,7 @@ export default function Shell() {
 					</p>
 				</Card>
 			)}
-			<Footer {...data} />
+			<Footer {...loaderData} />
 		</>
 	);
 }
diff --git a/app/routes/machines/machine.tsx b/app/routes/machines/machine.tsx
index 888dd74..c054ed5 100644
--- a/app/routes/machines/machine.tsx
+++ b/app/routes/machines/machine.tsx
@@ -1,6 +1,6 @@
 import { CheckCircle, CircleSlash, Info, UserCircle } from 'lucide-react';
 import { useMemo, useState } from 'react';
-import { data, Link as RemixLink, useLoaderData } from 'react-router';
+import { data, Link as RemixLink } from 'react-router';
 import Attribute from '~/components/Attribute';
 import Button from '~/components/Button';
 import Card from '~/components/Card';
@@ -58,9 +58,9 @@ export async function loader({ request, params, context }: Route.LoaderArgs) {
 
 export const action = machineAction;
 
-export default function Page() {
-	const { node, tags, magic, users, agent, stats } =
-		useLoaderData<typeof loader>();
+export default function Page({
+	loaderData: { node, tags, users, magic, agent, stats },
+}: Route.ComponentProps) {
 	const [showRouting, setShowRouting] = useState(false);
 
 	const uiTags = useMemo(() => {
diff --git a/app/routes/machines/overview.tsx b/app/routes/machines/overview.tsx
index 0348e31..cf99c1c 100644
--- a/app/routes/machines/overview.tsx
+++ b/app/routes/machines/overview.tsx
@@ -1,5 +1,4 @@
 import { Info } from 'lucide-react';
-import { useLoaderData } from 'react-router';
 import Code from '~/components/Code';
 import Link from '~/components/Link';
 import Tooltip from '~/components/Tooltip';
@@ -67,9 +66,7 @@ export async function loader({ request, context }: Route.LoaderArgs) {
 
 export const action = machineAction;
 
-export default function Page() {
-	const data = useLoaderData<typeof loader>();
-
+export default function Page({ loaderData }: Route.ComponentProps) {
 	return (
 		<>
 			<div className="flex justify-between items-center mb-6">
@@ -86,10 +83,10 @@ export default function Page() {
 					</p>
 				</div>
 				<NewMachine
-					disabledKeys={data.preAuth ? [] : ['pre-auth']}
-					isDisabled={!data.writable}
-					server={data.publicServer ?? data.server}
-					users={data.users}
+					disabledKeys={loaderData.preAuth ? [] : ['pre-auth']}
+					isDisabled={!loaderData.writable}
+					server={loaderData.publicServer ?? loaderData.server}
+					users={loaderData.users}
 				/>
 			</div>
 			<table className="table-auto w-full rounded-lg">
@@ -99,7 +96,7 @@ export default function Page() {
 						<th className="pb-2 w-1/4">
 							<div className="flex items-center gap-x-1">
 								<p className="uppercase text-xs font-bold">Addresses</p>
-								{data.magic ? (
+								{loaderData.magic ? (
 									<Tooltip>
 										<Info className="w-4 h-4" />
 										<Tooltip.Body className="font-normal">
@@ -107,7 +104,7 @@ export default function Page() {
 											their name and also at{' '}
 											<Code>
 												[name].
-												{data.magic}
+												{loaderData.magic}
 											</Code>
 										</Tooltip.Body>
 									</Tooltip>
@@ -115,7 +112,7 @@ export default function Page() {
 							</div>
 						</th>
 						{/* We only want to show the version column if there are agents */}
-						{data.agent !== undefined ? (
+						{loaderData.agent !== undefined ? (
 							<th className="uppercase text-xs font-bold pb-2">Version</th>
 						) : undefined}
 						<th className="uppercase text-xs font-bold pb-2">Last Seen</th>
@@ -127,18 +124,21 @@ export default function Page() {
 						'border-t border-headplane-100 dark:border-headplane-800',
 					)}
 				>
-					{data.populatedNodes.map((machine) => (
+					{loaderData.populatedNodes.map((node) => (
 						<MachineRow
-							isAgent={data.agent ? data.agent === machine.nodeKey : undefined}
-							isDisabled={
-								data.writable
-									? false // If the user has write permissions, they can edit all machines
-									: machine.user.providerId?.split('/').pop() !== data.subject
+							isAgent={
+								loaderData.agent ? loaderData.agent === node.nodeKey : undefined
 							}
-							key={machine.id}
-							magic={data.magic}
-							node={machine}
-							users={data.users}
+							isDisabled={
+								loaderData.writable
+									? false // If the user has write permissions, they can edit all machines
+									: node.user.providerId?.split('/').pop() !==
+										loaderData.subject
+							}
+							key={node.id}
+							magic={loaderData.magic}
+							node={node}
+							users={loaderData.users}
 						/>
 					))}
 				</tbody>
diff --git a/app/routes/ssh/console.tsx b/app/routes/ssh/console.tsx
index 8dd6fba..f5f030a 100644
--- a/app/routes/ssh/console.tsx
+++ b/app/routes/ssh/console.tsx
@@ -4,20 +4,11 @@ import { faker } from '@faker-js/faker';
 import { eq } from 'drizzle-orm';
 import { Loader2 } from 'lucide-react';
 import { useEffect, useState } from 'react';
-import {
-	ActionFunctionArgs,
-	data,
-	LinksFunction,
-	LoaderFunctionArgs,
-	ShouldRevalidateFunction,
-	useLoaderData,
-	useSubmit,
-} from 'react-router';
+import { data, type ShouldRevalidateFunction, useSubmit } from 'react-router';
 import { ExternalScriptsHandle } from 'remix-utils/external-scripts';
-import { LoadContext } from '~/server';
 import { EphemeralNodeInsert, ephemeralNodes } from '~/server/db/schema';
-import { Machine, PreAuthKey, User } from '~/types';
 import { useLiveData } from '~/utils/live-data';
+import type { Route } from './+types/console';
 import UserPrompt from './user-prompt';
 import XTerm from './xterm.client';
 
@@ -25,10 +16,7 @@ export const shouldRevalidate: ShouldRevalidateFunction = () => {
 	return false;
 };
 
-export async function loader({
-	request,
-	context,
-}: LoaderFunctionArgs<LoadContext>) {
+export async function loader({ request, context }: Route.LoaderArgs) {
 	const origin = new URL(request.url).origin;
 	const assets = ['/wasm_exec.js', '/hp_ssh.wasm'];
 	const missing: string[] = [];
@@ -53,10 +41,9 @@ export async function loader({
 	if (session.user.subject === 'unknown-non-oauth') {
 		throw data('Only OAuth users are allowed to use WebSSH', 403);
 	}
-	const { users } = await context.client.get<{ users: User[] }>(
-		'v1/user',
-		session.api_key,
-	);
+
+	const api = context.hsApi.getRuntimeClient(session.api_key);
+	const users = await api.getUsers();
 
 	// MARK: This assumes that a user has authenticated with Headscale first
 	// Since the only way to enforce permissions via ACLs is to generate a
@@ -77,15 +64,12 @@ export async function loader({
 		);
 	}
 
-	const { preAuthKey } = await context.client.post<{ preAuthKey: PreAuthKey }>(
-		'v1/preauthkey',
-		session.api_key,
-		{
-			user: lookup.id,
-			reusable: false,
-			ephemeral: true,
-			expiration: new Date(Date.now() + 60 * 1000).toISOString(), // 1 minute
-		},
+	const preAuthKey = await api.createPreAuthKey(
+		lookup.id,
+		true, // ephemeral
+		false, // reusable
+		new Date(Date.now() + 60 * 1000), // expiration: 1 minute
+		null, // aclTags
 	);
 
 	// TODO: Enable config to enforce generate_authkeys capability
@@ -130,13 +114,8 @@ export async function loader({
 	// 	);
 	// }
 
-	const { nodes } = await context.client.get<{ nodes: Machine[] }>(
-		'v1/node',
-		session.api_key,
-	);
-
-	// node.name is the hostname, given_name is the set name
-	const lookupNode = nodes.find((n) => n.name === hostname);
+	const nodes = await api.getNodes();
+	const lookupNode = nodes.find((n) => n.givenName === hostname);
 	if (!lookupNode) {
 		throw data(`Node with hostname ${hostname} not found`, 404);
 	}
@@ -182,11 +161,8 @@ function generateHostname(username: string) {
 	return `ssh-${adjective}-${noun}-${username}`;
 }
 
-export async function action({
-	request,
-	context,
-}: ActionFunctionArgs<LoadContext>) {
-	const _session = await context.sessions.auth(request);
+export async function action({ request, context }: Route.ActionArgs) {
+	await context.sessions.auth(request);
 	if (!context.agents?.agentID()) {
 		throw data(
 			'WebSSH is only available with the Headplane agent integration',
@@ -214,7 +190,7 @@ export async function action({
 		.where(eq(ephemeralNodes.auth_key, authKey));
 }
 
-export const links: LinksFunction = () => [
+export const links: Route.LinksFunction = () => [
 	{
 		rel: 'preload',
 		href: '/hp_ssh.wasm',
@@ -234,13 +210,14 @@ export const handle: ExternalScriptsHandle = {
 	],
 };
 
-export default function Page() {
+export default function Page({
+	loaderData: { ipnDetails, sshDetails },
+}: Route.ComponentProps) {
 	const submit = useSubmit();
 	const { pause } = useLiveData();
 
 	const [ipn, setIpn] = useState<TsWasmNet | null>(null);
 	const [nodeKey, setNodeKey] = useState<string | null>(null);
-	const { ipnDetails, sshDetails } = useLoaderData<typeof loader>();
 
 	useEffect(() => {
 		if (!ipnDetails) {
diff --git a/app/routes/util/healthz.ts b/app/routes/util/healthz.ts
index 5db210d..6e95493 100644
--- a/app/routes/util/healthz.ts
+++ b/app/routes/util/healthz.ts
@@ -1,8 +1,10 @@
-import { LoaderFunctionArgs } from 'react-router';
-import type { LoadContext } from '~/server';
+import type { Route } from './+types/healthz';
+
+export async function loader({ context }: Route.LoaderArgs) {
+	// Use a fake API key for healthcheck
+	const api = context.hsApi.getRuntimeClient('fake-api-key');
+	const healthy = await api.isHealthy();
 
-export async function loader({ context }: LoaderFunctionArgs<LoadContext>) {
-	const healthy = await context.client.healthcheck();
 	return new Response(JSON.stringify({ status: healthy ? 'OK' : 'ERROR' }), {
 		status: healthy ? 200 : 500,
 		headers: {
diff --git a/app/server/db/pruner.ts b/app/server/db/pruner.ts
index 9491957..b686576 100644
--- a/app/server/db/pruner.ts
+++ b/app/server/db/pruner.ts
@@ -1,14 +1,12 @@
 import { eq, isNotNull } from 'drizzle-orm';
-import { LoaderFunctionArgs } from 'react-router';
-import { Machine } from '~/types';
 import log from '~/utils/log';
-import { LoadContext } from '..';
+import type { Route } from '../../layouts/+types/dashboard';
 import { ephemeralNodes } from './schema';
 
 export async function pruneEphemeralNodes({
 	context,
 	request,
-}: LoaderFunctionArgs<LoadContext>) {
+}: Route.LoaderArgs) {
 	const session = await context.sessions.auth(request);
 	const ephemerals = await context.db
 		.select()
@@ -20,11 +18,8 @@ export async function pruneEphemeralNodes({
 		return;
 	}
 
-	const { nodes } = await context.client.get<{ nodes: Machine[] }>(
-		'v1/node',
-		session.api_key,
-	);
-
+	const api = context.hsApi.getRuntimeClient(session.api_key);
+	const nodes = await api.getNodes();
 	const toPrune = nodes.filter((node) => {
 		if (node.online) {
 			return false;
@@ -42,7 +37,7 @@ export async function pruneEphemeralNodes({
 	const promises = toPrune.map((node) => {
 		return async () => {
 			log.debug('api', `Pruning node ${node.name}`);
-			await context.client.delete(`v1/node/${node.id}`, session.api_key);
+			await api.deleteNode(node.id);
 
 			await context.db
 				.delete(ephemeralNodes)
diff --git a/app/server/headscale/api/endpoints/index.ts b/app/server/headscale/api/endpoints/index.ts
index fdd4c0f..32fe88f 100644
--- a/app/server/headscale/api/endpoints/index.ts
+++ b/app/server/headscale/api/endpoints/index.ts
@@ -1,5 +1,6 @@
 import {
 	composeEndpoints,
+	defineApiEndpoints,
 	type ExtractApiEndpoints,
 	type UnionToIntersection,
 } from '../factory';
@@ -10,12 +11,42 @@ import policyEndpoints from './policy';
 import preAuthKeyEndpoints from './pre-auth-keys';
 import userEndpoints from './users';
 
+interface HealthcheckEndpoint {
+	/**
+	 * Checks if the Headscale instance is healthy.
+	 *
+	 * @returns A boolean indicating if the instance is healthy.
+	 */
+	isHealthy(): Promise<boolean>;
+}
+
+const healthcheckEndpoint = defineApiEndpoints<HealthcheckEndpoint>(
+	(client, apiKey) => ({
+		isHealthy: async () => {
+			try {
+				const res = await client.rawFetch('/health', {
+					method: 'GET',
+					headers: {
+						// This doesn't really matter
+						Authorization: `Bearer ${apiKey}`,
+					},
+				});
+
+				return res.statusCode === 200;
+			} catch {
+				return false;
+			}
+		},
+	}),
+);
+
 /**
  * A constant list of all endpoint groups.
  * Add new endpoint groups here.
  */
 export const endpointSets = [
 	apiKeyEndpoints,
+	healthcheckEndpoint,
 	nodeEndpoints,
 	policyEndpoints,
 	preAuthKeyEndpoints,
diff --git a/app/server/headscale/api/endpoints/pre-auth-keys.ts b/app/server/headscale/api/endpoints/pre-auth-keys.ts
index 8c71e28..25c1658 100644
--- a/app/server/headscale/api/endpoints/pre-auth-keys.ts
+++ b/app/server/headscale/api/endpoints/pre-auth-keys.ts
@@ -15,7 +15,7 @@ export interface PreAuthKeyEndpoints {
 	 *
 	 * @param user The user to create the pre-authentication key for.
 	 * @param ephemeral Whether the key is ephemeral.
-	 * @param uses The number of uses for the key.
+	 * @param reusable Whether the key is reusable.
 	 * @param expiration The expiration date of the key, or `null` for no expiration.
 	 * @param aclTags An array of ACL tags to associate with the key, or `null` for none.
 	 * @returns A `PreAuthKey` object representing the newly created pre-authentication key.
@@ -23,7 +23,7 @@ export interface PreAuthKeyEndpoints {
 	createPreAuthKey(
 		user: string,
 		ephemeral: boolean,
-		uses: number,
+		reusable: boolean,
 		expiration: Date | null,
 		aclTags: string[] | null,
 	): Promise<PreAuthKey>;
@@ -46,13 +46,13 @@ export default defineApiEndpoints<PreAuthKeyEndpoints>((client, apiKey) => ({
 		return preAuthKeys;
 	},
 
-	createPreAuthKey: async (user, ephemeral, uses, expiration, aclTags) => {
+	createPreAuthKey: async (user, ephemeral, reusable, expiration, aclTags) => {
 		const { preAuthKey } = await client.apiFetch<{
 			preAuthKey: PreAuthKey;
 		}>('POST', 'v1/preauthkey', apiKey, {
 			user,
 			ephemeral,
-			uses,
+			reusable,
 			expiration: expiration ? expiration.toISOString() : null,
 			aclTags,
 		});