From d4eee702e928eb6e4e38336f5a67fb711b46c05e Mon Sep 17 00:00:00 2001 From: Aarnav Tale Date: Sat, 23 May 2026 16:32:22 -0400 Subject: [PATCH] refactor(server): replace AppContext undefined sentinels with Feature Co-authored-by: Amp Amp-Thread-ID: https://ampcode.com/threads/T-019e5550-4435-7118-8393-cdcc97042178 --- app/layout/app.tsx | 5 +- app/routes/acls/acl-action.ts | 3 +- app/routes/acls/acl-loader.ts | 3 +- app/routes/auth/login/page.tsx | 9 +- app/routes/auth/logout.ts | 13 +- app/routes/auth/oidc-callback.ts | 8 +- app/routes/auth/oidc-start.ts | 6 +- app/routes/home.tsx | 8 +- app/routes/machines/machine-actions.ts | 3 +- app/routes/machines/machine.tsx | 10 +- app/routes/machines/overview.tsx | 9 +- app/routes/settings/agent.tsx | 18 +-- app/routes/settings/auth-keys/actions.ts | 4 +- app/routes/settings/auth-keys/overview.tsx | 4 +- app/routes/settings/overview.tsx | 3 +- app/routes/ssh/page.tsx | 7 +- app/routes/users/overview.tsx | 3 +- app/routes/users/user-actions.ts | 3 +- app/routes/util/live.ts | 4 +- app/server/app.ts | 11 +- app/server/context.ts | 173 +++++++++++++++------ app/server/feature.ts | 20 +++ app/server/main.ts | 3 +- runtime/http.ts | 24 ++- runtime/vite-plugin.ts | 23 ++- 25 files changed, 253 insertions(+), 124 deletions(-) create mode 100644 app/server/feature.ts diff --git a/app/layout/app.tsx b/app/layout/app.tsx index 0c764ef..632f41b 100644 --- a/app/layout/app.tsx +++ b/app/layout/app.tsx @@ -31,10 +31,7 @@ export const shouldRevalidate: ShouldRevalidateFunction = ({ export async function loader({ request, context }: Route.LoaderArgs) { try { - const principal = await context.auth.require(request); - - const apiKey = context.auth.getHeadscaleApiKey(principal); - const api = context.hsApi.getRuntimeClient(apiKey); + const { principal, api } = await context.apiForRequest(request); const user = principal.kind === "oidc" diff --git a/app/routes/acls/acl-action.ts b/app/routes/acls/acl-action.ts index 5716e8f..a7e7288 100644 --- a/app/routes/acls/acl-action.ts +++ b/app/routes/acls/acl-action.ts @@ -26,8 +26,7 @@ export async function aclAction({ request, context }: Route.ActionArgs) { }); } - const apiKey = context.auth.getHeadscaleApiKey(principal); - const api = context.hsApi.getRuntimeClient(apiKey); + const { api } = await context.apiForRequest(request); try { const { policy, updatedAt } = await api.setPolicy(policyData); return data({ diff --git a/app/routes/acls/acl-loader.ts b/app/routes/acls/acl-loader.ts index 24dc3b1..6b4382c 100644 --- a/app/routes/acls/acl-loader.ts +++ b/app/routes/acls/acl-loader.ts @@ -29,8 +29,7 @@ export async function aclLoader({ request, context }: Route.LoaderArgs) { }; // Try to load the ACL policy from the API. - const apiKey = context.auth.getHeadscaleApiKey(principal); - const api = context.hsApi.getRuntimeClient(apiKey); + const { api } = await context.apiForRequest(request); try { const { policy, updatedAt } = await api.getPolicy(); flags.writable = updatedAt !== null; diff --git a/app/routes/auth/login/page.tsx b/app/routes/auth/login/page.tsx index 166ebdc..bd9924d 100644 --- a/app/routes/auth/login/page.tsx +++ b/app/routes/auth/login/page.tsx @@ -24,7 +24,7 @@ export async function loader({ request, context }: Route.LoaderArgs) { const qp = new URL(request.url).searchParams; const urlState = qp.get("s") ?? undefined; - const oidcService = context.oidc?.service; + const oidcService = context.oidc.state === "enabled" ? context.oidc.value : undefined; const oidcStatus = oidcService ? await oidcService.discover().then( (r) => (r.ok ? oidcService.status() : oidcService.status()), @@ -32,7 +32,12 @@ export async function loader({ request, context }: Route.LoaderArgs) { ) : undefined; - if (context.oidc?.disableApiKeyLogin && oidcStatus?.state === "ready" && urlState !== "logout") { + if ( + oidcService && + context.config.oidc?.disable_api_key_login && + oidcStatus?.state === "ready" && + urlState !== "logout" + ) { return redirect("/oidc/start"); } diff --git a/app/routes/auth/logout.ts b/app/routes/auth/logout.ts index 5510f5e..8151bb3 100644 --- a/app/routes/auth/logout.ts +++ b/app/routes/auth/logout.ts @@ -23,15 +23,20 @@ export async function action({ request, context }: ActionFunctionArgs node.nodeKey)); + const agents = context.agents.state === "enabled" ? context.agents.value : undefined; + const stats = await agents?.lookup(nodes.map((node) => node.nodeKey)); const populatedNodes = mapNodes(nodes, stats); const supportsNodeOwnerChange = !context.hsApi.clientHelpers.isAtleast("0.28.0"); - const agentSync = context.agents?.lastSync(); + const agentSync = agents?.lastSync(); return { agent: agentSync ? { syncedAt: agentSync.syncedAt?.toISOString() ?? null, nodeCount: agentSync.nodeCount, - nodeKey: context.agents?.agentNodeKey(), + nodeKey: agents?.agentNodeKey(), } : undefined, headscaleUserId: principal.kind === "oidc" ? principal.user.headscaleUserId : undefined, diff --git a/app/routes/settings/agent.tsx b/app/routes/settings/agent.tsx index 55b54de..9fac44e 100644 --- a/app/routes/settings/agent.tsx +++ b/app/routes/settings/agent.tsx @@ -11,13 +11,13 @@ import { formatTimeDelta } from "~/utils/time"; import type { Route } from "./+types/agent"; export async function loader({ request, context }: Route.LoaderArgs) { - const principal = await context.auth.require(request); + await context.auth.require(request); - if (!context.agents) { - return { enabled: false as const }; + if (context.agents.state !== "enabled") { + return { enabled: false as const, reason: context.agents.reason }; } - const sync = context.agents.lastSync(); + const sync = context.agents.value.lastSync(); return { enabled: true as const, syncedAt: sync.syncedAt?.toISOString() ?? null, @@ -29,12 +29,12 @@ export async function loader({ request, context }: Route.LoaderArgs) { export async function action({ request, context }: Route.ActionArgs) { await context.auth.require(request); - if (!context.agents) { - return { success: false, error: "Agent is not enabled" }; + if (context.agents.state !== "enabled") { + return { success: false, error: context.agents.reason }; } - await context.agents.triggerSync(); - const sync = context.agents.lastSync(); + await context.agents.value.triggerSync(); + const sync = context.agents.value.lastSync(); return { success: !sync.error, error: sync.error }; } @@ -47,7 +47,7 @@ export default function Page({ loaderData }: Route.ComponentProps) {
Headplane Agent - The Headplane Agent is not enabled. To learn how to set up the agent, visit the{" "} + {loaderData.reason}. To learn how to set up the agent, visit the{" "} documentation diff --git a/app/routes/settings/auth-keys/actions.ts b/app/routes/settings/auth-keys/actions.ts index 35ec59c..918386f 100644 --- a/app/routes/settings/auth-keys/actions.ts +++ b/app/routes/settings/auth-keys/actions.ts @@ -7,9 +7,7 @@ import type { PreAuthKey } from "~/types"; import type { Route } from "./+types/overview"; export async function authKeysAction({ request, context }: Route.ActionArgs) { - const principal = await context.auth.require(request); - const apiKey = context.auth.getHeadscaleApiKey(principal); - const api = context.hsApi.getRuntimeClient(apiKey); + const { principal, api } = await context.apiForRequest(request); const canGenerateAny = context.auth.can(principal, Capabilities.generate_authkeys); const canGenerateOwn = context.auth.can(principal, Capabilities.generate_own_authkeys); diff --git a/app/routes/settings/auth-keys/overview.tsx b/app/routes/settings/auth-keys/overview.tsx index a75f555..bc2ac23 100644 --- a/app/routes/settings/auth-keys/overview.tsx +++ b/app/routes/settings/auth-keys/overview.tsx @@ -19,9 +19,7 @@ import AuthKeyRow from "./auth-key-row"; import AddAuthKey from "./dialogs/add-auth-key"; export async function loader({ request, context }: Route.LoaderArgs) { - const principal = await context.auth.require(request); - const apiKey = context.auth.getHeadscaleApiKey(principal); - const api = context.hsApi.getRuntimeClient(apiKey); + const { principal, api } = await context.apiForRequest(request); const usersSnap = await context.hsLive.get(usersResource, api); const users = usersSnap.data; diff --git a/app/routes/settings/overview.tsx b/app/routes/settings/overview.tsx index 8ded5c0..e90324a 100644 --- a/app/routes/settings/overview.tsx +++ b/app/routes/settings/overview.tsx @@ -8,7 +8,8 @@ import type { Route } from "./+types/overview"; export async function loader({ context }: Route.LoaderArgs) { return { config: context.hs.writable(), - isOidcEnabled: context.oidc?.service.status().state === "ready", + isOidcEnabled: + context.oidc.state === "enabled" && context.oidc.value.status().state === "ready", }; } diff --git a/app/routes/ssh/page.tsx b/app/routes/ssh/page.tsx index 8d31673..c0e051e 100644 --- a/app/routes/ssh/page.tsx +++ b/app/routes/ssh/page.tsx @@ -37,18 +37,15 @@ export async function loader({ request, params, context }: Route.LoaderArgs) { throw data(sshErrors.wasm_missing, 405); } - if (context.agents == null) { + if (context.agents.state !== "enabled") { throw data(sshErrors.agent_required, 400); } - const principal = await context.auth.require(request); + const { principal, api } = await context.apiForRequest(request); if (principal.kind === "api_key") { throw data(sshErrors.oidc_required, 403); } - const apiKey = context.auth.getHeadscaleApiKey(principal); - const api = context.hsApi.getRuntimeClient(apiKey); - const hostname = params.id; const username = new URL(request.url).searchParams.get("user") || undefined; diff --git a/app/routes/users/overview.tsx b/app/routes/users/overview.tsx index 25b2d79..f2591d5 100644 --- a/app/routes/users/overview.tsx +++ b/app/routes/users/overview.tsx @@ -54,8 +54,7 @@ export async function loader({ request, context }: Route.LoaderArgs) { let apiError: string | undefined; try { - const apiKey = context.auth.getHeadscaleApiKey(principal); - const api = context.hsApi.getRuntimeClient(apiKey); + const { api } = await context.apiForRequest(request); const [nodesSnap, usersSnap] = await Promise.all([ context.hsLive.get(nodesResource, api), context.hsLive.get(usersResource, api), diff --git a/app/routes/users/user-actions.ts b/app/routes/users/user-actions.ts index 33c1a21..23a3aaa 100644 --- a/app/routes/users/user-actions.ts +++ b/app/routes/users/user-actions.ts @@ -24,8 +24,7 @@ export async function userAction({ request, context }: Route.ActionArgs) { }); } - const apiKey = context.auth.getHeadscaleApiKey(principal); - const api = context.hsApi.getRuntimeClient(apiKey); + const { api } = await context.apiForRequest(request); switch (action) { case "create_user": { const name = formData.get("username")?.toString(); diff --git a/app/routes/util/live.ts b/app/routes/util/live.ts index 9d7a38c..67af56e 100644 --- a/app/routes/util/live.ts +++ b/app/routes/util/live.ts @@ -4,9 +4,7 @@ import log from "~/utils/log"; import type { Route } from "./+types/live"; export async function loader({ request, context }: Route.LoaderArgs) { - const principal = await context.auth.require(request); - const apiKey = context.auth.getHeadscaleApiKey(principal); - const api = context.hsApi.getRuntimeClient(apiKey); + const { api } = await context.apiForRequest(request); // Ensure resources are loaded before streaming await Promise.all([ diff --git a/app/server/app.ts b/app/server/app.ts index cb65050..aaee400 100644 --- a/app/server/app.ts +++ b/app/server/app.ts @@ -36,10 +36,19 @@ try { } const ctx = await createAppContext(config); -ctx.auth.start(); +ctx.startServices(); export { config }; +/** + * Disposes the per-process context. Invoked by the production + * supervisor on SIGTERM/SIGINT and by the dev Vite plugin on HMR + * reload. + */ +export async function dispose(): Promise { + await ctx.dispose(); +} + // TODO: `getLoadContext` is the right place to handle reverse proxy // translation — better than doing it in the OIDC client because it // applies to all requests, not just OIDC ones. diff --git a/app/server/context.ts b/app/server/context.ts index 7f524c9..41dc27e 100644 --- a/app/server/context.ts +++ b/app/server/context.ts @@ -5,12 +5,14 @@ import log from "~/utils/log"; import type { HeadplaneConfig } from "./config/config-schema"; import { loadIntegration } from "./config/integration"; import { createDbClient } from "./db/client.server"; +import { disabled, enabled, type Feature } from "./feature"; import { createHeadscaleInterface } from "./headscale/api"; +import type { RuntimeApiClient } from "./headscale/api/endpoints"; import { loadHeadscaleConfig } from "./headscale/config-loader"; import { createLiveStore, nodesResource, usersResource } from "./headscale/live-store"; -import { createAgentManager } from "./hp-agent"; -import { createOidcService } from "./oidc/provider"; -import { createAuthService } from "./web/auth"; +import { type AgentManager, createAgentManager } from "./hp-agent"; +import { createOidcService, type OidcService } from "./oidc/provider"; +import { createAuthService, type Principal } from "./web/auth"; export type AppContext = Awaited>; @@ -29,18 +31,12 @@ export async function createAppContext(config: HeadplaneConfig) { // falling back to the deprecated oidc.headscale_api_key for compatibility. const headscaleApiKey = config.headscale.api_key ?? config.oidc?.headscale_api_key; - let agents; - if (headscaleApiKey) { - agents = await createAgentManager( - config.integration?.agent, - config.headscale.url, - hsApi.getRuntimeClient(headscaleApiKey), - hsApi.clientHelpers.isAtleast("0.28.0"), - db, - ); - } else if (config.integration?.agent?.enabled) { - log.warn("agent", "Agent is enabled but no headscale.api_key is configured"); - } + const agents = await buildAgents( + config, + hsApi.clientHelpers.isAtleast("0.28.0"), + headscaleApiKey ? hsApi.getRuntimeClient(headscaleApiKey) : undefined, + db, + ); const auth = createAuthService({ secret: config.server.cookie_secret, @@ -54,34 +50,43 @@ export async function createAppContext(config: HeadplaneConfig) { }, }); - const oidc = - config.oidc && config.oidc.enabled !== false && headscaleApiKey - ? { - service: createOidcService({ - issuer: config.oidc.issuer, - clientId: config.oidc.client_id, - clientSecret: config.oidc.client_secret, - baseUrl: config.server.base_url ?? "", - authorizationEndpoint: config.oidc.authorization_endpoint, - tokenEndpoint: config.oidc.token_endpoint, - userinfoEndpoint: config.oidc.userinfo_endpoint, - endSessionEndpoint: config.oidc.end_session_endpoint, - tokenEndpointAuthMethod: - config.oidc.token_endpoint_auth_method === "client_secret_jwt" - ? undefined - : config.oidc.token_endpoint_auth_method, - usePkce: config.oidc.use_pkce, - scope: config.oidc.scope, - subjectClaims: config.oidc.subject_claims, - allowWeakRsaKeys: config.oidc.allow_weak_rsa_keys, - extraParams: config.oidc.extra_params, - profilePictureSource: config.oidc.profile_picture_source, - postLogoutRedirectUri: config.oidc.post_logout_redirect_uri, - }), - disableApiKeyLogin: config.oidc.disable_api_key_login, - useEndSession: config.oidc.use_end_session, - } - : undefined; + const oidc = buildOidc(config, headscaleApiKey); + + const hsLive = createLiveStore([nodesResource, usersResource]); + const hs = await loadHeadscaleConfig( + config.headscale.config_path, + config.headscale.config_strict, + config.headscale.dns_records_path, + ); + const integration = await loadIntegration(config.integration); + + // Disposers run in reverse-registration order on shutdown. + const disposers: Array<() => Promise | void> = [() => auth.stop(), () => hsLive.dispose()]; + if (agents.state === "enabled") { + disposers.push(() => agents.value.dispose()); + } + + async function apiForRequest( + request: Request, + ): Promise<{ principal: Principal; api: RuntimeApiClient }> { + const principal = await auth.require(request); + const apiKey = auth.getHeadscaleApiKey(principal); + return { principal, api: hsApi.getRuntimeClient(apiKey) }; + } + + function startServices() { + auth.start(); + } + + async function dispose() { + for (const d of [...disposers].reverse()) { + try { + await d(); + } catch (error) { + log.warn("server", "Error during shutdown: %s", String(error)); + } + } + } return { config, @@ -91,12 +96,80 @@ export async function createAppContext(config: HeadplaneConfig) { agents, auth, oidc, - hsLive: createLiveStore([nodesResource, usersResource]), - hs: await loadHeadscaleConfig( - config.headscale.config_path, - config.headscale.config_strict, - config.headscale.dns_records_path, - ), - integration: await loadIntegration(config.integration), + hsLive, + hs, + integration, + apiForRequest, + startServices, + dispose, }; } + +function buildOidc( + config: HeadplaneConfig, + headscaleApiKey: string | undefined, +): Feature { + if (!config.oidc) { + return disabled("OIDC is not configured"); + } + if (config.oidc.enabled === false) { + return disabled("OIDC is disabled in the configuration"); + } + if (!headscaleApiKey) { + return disabled("OIDC requires headscale.api_key to be configured"); + } + + return enabled( + createOidcService({ + issuer: config.oidc.issuer, + clientId: config.oidc.client_id, + clientSecret: config.oidc.client_secret, + baseUrl: config.server.base_url ?? "", + authorizationEndpoint: config.oidc.authorization_endpoint, + tokenEndpoint: config.oidc.token_endpoint, + userinfoEndpoint: config.oidc.userinfo_endpoint, + endSessionEndpoint: config.oidc.end_session_endpoint, + tokenEndpointAuthMethod: + config.oidc.token_endpoint_auth_method === "client_secret_jwt" + ? undefined + : config.oidc.token_endpoint_auth_method, + usePkce: config.oidc.use_pkce, + scope: config.oidc.scope, + subjectClaims: config.oidc.subject_claims, + allowWeakRsaKeys: config.oidc.allow_weak_rsa_keys, + extraParams: config.oidc.extra_params, + profilePictureSource: config.oidc.profile_picture_source, + postLogoutRedirectUri: config.oidc.post_logout_redirect_uri, + }), + ); +} + +async function buildAgents( + config: HeadplaneConfig, + supportsTagOnlyKeys: boolean, + apiClient: RuntimeApiClient | undefined, + db: Awaited>, +): Promise> { + const agentConfig = config.integration?.agent; + if (!agentConfig?.enabled) { + return disabled("Agent is not enabled in the configuration"); + } + if (!apiClient) { + return disabled("Agent requires headscale.api_key to be configured"); + } + if (!supportsTagOnlyKeys) { + return disabled("Agent requires Headscale 0.28 or newer"); + } + + const manager = await createAgentManager( + agentConfig, + config.headscale.url, + apiClient, + supportsTagOnlyKeys, + db, + ); + if (!manager) { + return disabled("Agent failed to initialize (see logs)"); + } + return enabled(manager); +} diff --git a/app/server/feature.ts b/app/server/feature.ts new file mode 100644 index 0000000..b51a6f5 --- /dev/null +++ b/app/server/feature.ts @@ -0,0 +1,20 @@ +// MARK: Feature +// +// A two-state tagged union used in place of `T | undefined` for +// optional features on the AppContext. Carries a human-readable +// reason when the feature is disabled so loaders can surface it (or +// choose to ignore it). +// +// This is deliberately *not* a Result/Either. It does not represent +// async readiness or retryable failure — it represents "is this +// feature wired up at all." Services that have their own runtime +// status (e.g. OIDC discovery) keep that on the service itself. + +export type Feature = { state: "enabled"; value: T } | { state: "disabled"; reason: string }; + +export const enabled = (value: T): Feature => ({ state: "enabled", value }); + +export const disabled = (reason: string): Feature => ({ + state: "disabled", + reason, +}); diff --git a/app/server/main.ts b/app/server/main.ts index 2babfef..542cf9f 100644 --- a/app/server/main.ts +++ b/app/server/main.ts @@ -12,7 +12,7 @@ import { dirname, resolve } from "node:path"; import { fileURLToPath } from "node:url"; import { composeListener, startHttpServer } from "../../runtime/http"; -import requestListener, { config } from "./app"; +import requestListener, { config, dispose } from "./app"; // `import.meta.url` resolves to `build/server/index.js`; the built // client lives next to it at `build/client/`. @@ -27,4 +27,5 @@ startHttpServer({ immutableAssets: true, requestListener, }), + onShutdown: dispose, }); diff --git a/runtime/http.ts b/runtime/http.ts index 875de5f..d96d7cb 100644 --- a/runtime/http.ts +++ b/runtime/http.ts @@ -172,6 +172,13 @@ export interface StartOptions { listener: RequestListener; tls?: HttpsServerOptions; logger?: Logger; + /** + * Optional async hook invoked on SIGINT/SIGTERM after the HTTP + * server stops accepting new connections but before the process + * exits. Use this to dispose long-lived resources (timers, + * subprocesses, DB handles). + */ + onShutdown?: () => Promise | void; } /** @@ -189,15 +196,24 @@ export function startHttpServer(opts: StartOptions): Server { log.info("Listening on %s://%s:%s", proto, opts.host, opts.port); }); - const shutdown = (signal: string) => { + const shutdown = async (signal: string) => { log.info("Received %s, shutting down...", signal); - server.close(() => process.exit(0)); + server.close(async () => { + if (opts.onShutdown) { + try { + await opts.onShutdown(); + } catch (err) { + log.error("Error during shutdown hook: %o", err); + } + } + process.exit(0); + }); // Force exit if connections don't drain in time. setTimeout(() => process.exit(0), 5_000).unref(); }; - process.once("SIGINT", () => shutdown("SIGINT")); - process.once("SIGTERM", () => shutdown("SIGTERM")); + process.once("SIGINT", () => void shutdown("SIGINT")); + process.once("SIGTERM", () => void shutdown("SIGTERM")); return server; } diff --git a/runtime/vite-plugin.ts b/runtime/vite-plugin.ts index 0f6ea29..022ad4d 100644 --- a/runtime/vite-plugin.ts +++ b/runtime/vite-plugin.ts @@ -18,6 +18,11 @@ export interface DevServerOptions { publicDir: string; } +interface AppModule { + default: RequestListener; + dispose?: () => Promise | void; +} + export function headplaneDevServer(options: DevServerOptions): Plugin { return { name: "headplane:dev-server", @@ -30,6 +35,12 @@ export function headplaneDevServer(options: DevServerOptions): Plugin { res.end("Server entry not loaded yet"); }; + // Track the last-loaded module identity so we can call its + // `dispose` when Vite swaps in a new one on HMR. Without this + // the LiveStore and other long-lived services leak intervals + // and subprocesses on every reload. + let currentModule: AppModule | null = null; + const composed = composeListener({ basename: options.basename, staticRoot: options.publicDir, @@ -43,9 +54,15 @@ export function headplaneDevServer(options: DevServerOptions): Plugin { return () => { server.middlewares.use(async (req, res, next) => { try { - const mod = (await server.ssrLoadModule(options.entry)) as { - default: RequestListener; - }; + const mod = (await server.ssrLoadModule(options.entry)) as AppModule; + if (currentModule && currentModule !== mod) { + try { + await currentModule.dispose?.(); + } catch (err) { + server.config.logger.warn(`[headplane:dev-server] dispose failed: ${String(err)}`); + } + } + currentModule = mod; appListener = mod.default; composed(req, res); } catch (err) {