refactor(server): replace AppContext undefined sentinels with Feature<T>

Co-authored-by: Amp <amp@ampcode.com>
Amp-Thread-ID: https://ampcode.com/threads/T-019e5550-4435-7118-8393-cdcc97042178

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<void> {
+  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.

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<ReturnType<typeof createAppContext>>;
 
@@ -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> | 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<OidcService> {
+  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<ReturnType<typeof createDbClient>>,
+): Promise<Feature<AgentManager>> {
+  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);
+}

app/server/feature.ts

@@ -0,0 +1,20 @@
+// MARK: Feature<T>
+//
+// 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<T> = { state: "enabled"; value: T } | { state: "disabled"; reason: string };
+
+export const enabled = <T>(value: T): Feature<T> => ({ state: "enabled", value });
+
+export const disabled = (reason: string): Feature<never> => ({
+  state: "disabled",
+  reason,
+});

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,
 });