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/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({

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;

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

app/routes/auth/logout.ts

@@ -23,15 +23,20 @@ export async function action({ request, context }: ActionFunctionArgs<AppContext
   // ended. Disabled by default because the post_logout_redirect_uri must be
   // pre-registered on the IdP — turning this on without registering it would
   // strand users on the IdP's error page.
-  if (principal?.kind === "oidc" && context.oidc?.useEndSession && context.oidc.service) {
-    const status = context.oidc.service.status();
+  if (
+    principal?.kind === "oidc" &&
+    context.oidc.state === "enabled" &&
+    context.config.oidc?.use_end_session
+  ) {
+    const service = context.oidc.value;
+    const status = service.status();
     if (status.state !== "ready") {
       // Trigger discovery if it hasn't happened yet so we can find the
       // end_session_endpoint without forcing a re-login.
-      await context.oidc.service.discover();
+      await service.discover();
     }
 
-    const endSessionUrl = context.oidc.service.buildEndSessionUrl(principal.idToken);
+    const endSessionUrl = service.buildEndSessionUrl(principal.idToken);
     if (endSessionUrl) {
       url = endSessionUrl;
     }

app/routes/auth/oidc-callback.ts

@@ -7,10 +7,10 @@ import { createOidcStateCookie } from "~/utils/oidc-state";
 import type { Route } from "./+types/oidc-callback";
 
 export async function loader({ request, context }: Route.LoaderArgs) {
-  const service = context.oidc?.service;
-  if (!service) {
-    throw data("OIDC is not enabled or misconfigured", { status: 501 });
+  if (context.oidc.state !== "enabled") {
+    throw data(`OIDC is unavailable: ${context.oidc.reason}`, { status: 501 });
   }
+  const service = context.oidc.value;
 
   const url = new URL(request.url);
   if (url.searchParams.toString().length === 0) {
@@ -68,7 +68,7 @@ export async function loader({ request, context }: Route.LoaderArgs) {
 
   // Only persist the id_token when RP-initiated logout is enabled — otherwise
   // we'd be storing a credential we never use.
-  const idToken = context.oidc?.useEndSession ? identity.idToken : undefined;
+  const idToken = context.config.oidc?.use_end_session ? identity.idToken : undefined;
 
   return redirect("/", {
     headers: {

app/routes/auth/oidc-start.ts

@@ -10,10 +10,10 @@ export async function loader({ request, context }: Route.LoaderArgs) {
     return redirect("/");
   } catch {}
 
-  const service = context.oidc?.service;
-  if (!service) {
-    throw data("OIDC is not enabled or misconfigured", { status: 501 });
+  if (context.oidc.state !== "enabled") {
+    throw data(`OIDC is unavailable: ${context.oidc.reason}`, { status: 501 });
   }
+  const service = context.oidc.value;
 
   const result = await service.startFlow();
   if (!result.ok) {

app/routes/home.tsx

@@ -24,12 +24,11 @@ export async function loader({ request, context }: Route.LoaderArgs) {
   // Unclaimed users they can pick from before anything else.
   let unlinked = false;
   if (
-    typeof context.oidc === "object" &&
+    context.oidc.state === "enabled" &&
     principal.kind === "oidc" &&
     !principal.user.headscaleUserId
   ) {
-    const apiKey = context.auth.getHeadscaleApiKey(principal);
-    const api = context.hsApi.getRuntimeClient(apiKey);
+    const { api } = await context.apiForRequest(request);
 
     let headscaleUsers: { id: string; name: string }[] = [];
     try {
@@ -64,8 +63,7 @@ export async function loader({ request, context }: Route.LoaderArgs) {
   }
 
   // No UI access — show the download/connect page
-  const apiKey = context.auth.getHeadscaleApiKey(principal);
-  const api = context.hsApi.getRuntimeClient(apiKey);
+  const { api } = await context.apiForRequest(request);
 
   let linkedUserName: string | undefined;
   if (principal.kind === "oidc" && principal.user.headscaleUserId) {

app/routes/machines/machine-actions.ts

@@ -7,10 +7,9 @@ import { Capabilities } from "~/server/web/roles";
 import type { Route } from "./+types/machine";
 
 export async function machineAction({ request, context }: Route.ActionArgs) {
-  const principal = await context.auth.require(request);
+  const { principal, api } = await context.apiForRequest(request);
 
   const formData = await request.formData();
-  const api = context.hsApi.getRuntimeClient(context.auth.getHeadscaleApiKey(principal));
 
   const action = formData.get("action_id")?.toString();
   if (!action) {

app/routes/machines/machine.tsx

@@ -22,7 +22,6 @@ import Routes from "./dialogs/routes";
 import { machineAction } from "./machine-actions";
 
 export async function loader({ request, params, context }: Route.LoaderArgs) {
-  const principal = await context.auth.require(request);
   if (!params.id) {
     throw new Error("No machine ID provided");
   }
@@ -38,7 +37,7 @@ export async function loader({ request, params, context }: Route.LoaderArgs) {
     }
   }
 
-  const api = context.hsApi.getRuntimeClient(context.auth.getHeadscaleApiKey(principal));
+  const { api } = await context.apiForRequest(request);
   const [nodesSnap, usersSnap] = await Promise.all([
     context.hsLive.get(nodesResource, api),
     context.hsLive.get(usersResource, api),
@@ -50,18 +49,19 @@ export async function loader({ request, params, context }: Route.LoaderArgs) {
     throw data(null, { status: 404 });
   }
 
-  const lookup = await context.agents?.lookup([node.nodeKey]);
+  const agents = context.agents.state === "enabled" ? context.agents.value : undefined;
+  const lookup = await agents?.lookup([node.nodeKey]);
   const [enhancedNode] = mapNodes([node], lookup);
   const tags = [...node.tags].toSorted();
   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,
     existingTags: sortNodeTags(nodes),

app/routes/machines/overview.tsx

@@ -30,7 +30,7 @@ export async function loader({ request, context }: Route.LoaderArgs) {
 
   const writablePermission = context.auth.can(principal, Capabilities.write_machines);
 
-  const api = context.hsApi.getRuntimeClient(context.auth.getHeadscaleApiKey(principal));
+  const { api } = await context.apiForRequest(request);
   const [nodesSnap, usersSnap] = await Promise.all([
     context.hsLive.get(nodesResource, api),
     context.hsLive.get(usersResource, api),
@@ -45,17 +45,18 @@ export async function loader({ request, context }: Route.LoaderArgs) {
     }
   }
 
-  const stats = await context.agents?.lookup(nodes.map((node) => 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,

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) {
       <div className="flex max-w-(--breakpoint-lg) flex-col gap-8">
         <Title>Headplane Agent</Title>
         <Notice title="Agent Not Enabled">
-          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{" "}
           <Link external styled to="https://headplane.dev/docs/agent">
             documentation
           </Link>

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

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;

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

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;
 

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

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();

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([