platform/plugins/offscale/sidecar/companion.ts

import { getActiveHeadscaleCredentials, type HeadscaleServerCredentials } from '../db/queries';
import { badRequest, methodNotAllowed, notFound, type OfficerContext } from './routes';

// The Officer Companion API — a small service deployed NEXT TO each Headscale server that answers what the
// admin API structurally cannot: is the container up, what do its logs say, and start/stop/restart it.
// Contract: COMMS/HEADSCALE_COMPANION_API.md.
//
// Three facts shape everything here.
//
//  1. It lives at `${server.url}/officer-api` and authenticates with the SAME admin API key we already
//     store, validated locally against Headscale's own key store — so auth keeps working while Headscale
//     is down, which is exactly when `/restart` matters. Nothing new to register, and the key still never
//     leaves this sidecar.
//
//  2. It is OPTIONAL and per-server. Of the four servers registered here today, one has it deployed. So
//     "no companion" is a normal state, not an error: every route below answers 200 with
//     `{available: false, reason}` rather than failing, and the UI degrades to what the admin API can do.
//     Distinguishing the two 502s is the whole trick — nginx returns HTML when the companion is down,
//     the companion returns JSON when a docker op fails. Branch on whether the body parses.
//
//  3. `GET /health` is ALWAYS 200, at every verdict. Never key anything off its HTTP status; read
//     `verdict`. That inversion is deliberate on their side and is preserved on ours.

/**
 * Every route answers `{available: true, ...}` or `{available: false, reason}` at HTTP 200. Not having a
 * companion is a state to render, not a request that failed — the admin API on the same domain is
 * independent and may still be working, so this must not surface as an error the UI swallows.
 */
export const unavailable = (reason: string) => ({ available: false as const, reason });

const DEFAULT_TIMEOUT_MS = 20_000;

type CompanionCall = { path: string; method?: string; body?: unknown; timeoutMs?: number; signal?: AbortSignal };

/**
 * One request to a server's companion. Returns the raw Response, or a reason string when the companion
 * itself could not be reached — the caller decides how to present that, because for this feature
 * "unreachable" is information rather than a failure.
 */
export async function callCompanion(
  creds: HeadscaleServerCredentials,
  { path, method = 'GET', body, timeoutMs = DEFAULT_TIMEOUT_MS, signal }: CompanionCall,
): Promise<Response | string> {
  let res: Response;
  try {
    res = await fetch(`${creds.url}/officer-api${path}`, {
      method,
      headers: {
        authorization: `Bearer ${creds.apiKey}`,
        accept: 'application/json',
        ...(body === undefined ? {} : { 'content-type': 'application/json' }),
      },
      body: body === undefined ? undefined : JSON.stringify(body),
      signal: signal ?? AbortSignal.timeout(timeoutMs),
    });
  } catch (err) {
    if (err instanceof Error && err.name === 'TimeoutError') return 'the companion timed out';
    // A TLS failure or DNS miss on the server's own domain: the whole host is unreachable, not just this.
    return 'could not reach the server';
  }

  if (res.status === 401) return 'the companion rejected the stored API key';

  // Both 404 and 502 are ambiguous, and the same test settles both: a JSON body means the companion
  // answered (no such container / the docker op failed) and that answer belongs to the caller; a
  // non-JSON body means we never reached it — nginx's own 502 page, or a route that isn't there at all.
  const isJson = (res.headers.get('content-type') ?? '').includes('json');
  if (res.status === 404 && !isJson) return 'this server has no companion at /officer-api';
  if (res.status === 502 && !isJson) return 'the companion is not deployed on this server';
  if (res.status >= 500 && !isJson) return `the companion returned ${res.status}`;
  return res;
}

/** Parse a companion JSON body, or a reason when it isn't JSON after all. */
export async function readBody(res: Response): Promise<Record<string, unknown> | string> {
  const text = await res.text().catch(() => '');
  if (!text) return 'the companion returned an empty body';
  try {
    const parsed = JSON.parse(text) as unknown;
    if (!parsed || typeof parsed !== 'object') return 'the companion returned an unexpected body';
    return parsed as Record<string, unknown>;
  } catch {
    return 'the companion returned a non-JSON body';
  }
}

/** The active server's credentials, or a 409 the UI already knows how to render. */
export async function activeCreds(userId: number): Promise<HeadscaleServerCredentials | Response> {
  const creds = await getActiveHeadscaleCredentials(userId);
  if (!creds) {
    return Response.json({ error: 'no active Headscale server', code: 'no_active_server' }, { status: 409 });
  }
  return creds;
}

/** `GET /_officer/companion/health` — verdict, container state and, when unhappy, its own diagnosis. */
async function health(creds: HeadscaleServerCredentials): Promise<Response> {
  const res = await callCompanion(creds, { path: '/health' });
  if (typeof res === 'string') return Response.json(unavailable(res));

  const body = await readBody(res);
  if (typeof body === 'string') return Response.json(unavailable(body));
  // Passed through as-is. The companion owns this vocabulary and versions it; re-shaping it here would mean
  // a new verdict or a new likely-cause silently disappearing on the way to the screen.
  return Response.json({ available: true, health: body });
}

/** `GET /_officer/companion/logs?tail=N` — a snapshot of the last N lines. */
async function logs(creds: HeadscaleServerCredentials, url: URL): Promise<Response> {
  const tail = Number(url.searchParams.get('tail') ?? 200);
  if (!Number.isInteger(tail) || tail < 1 || tail > 2000) return badRequest('tail must be 1–2000');

  const res = await callCompanion(creds, { path: `/logs?tail=${tail}` });
  if (typeof res === 'string') return Response.json(unavailable(res));

  const body = await readBody(res);
  if (typeof body === 'string') return Response.json(unavailable(body));
  const lines = Array.isArray(body.lines) ? body.lines.filter((l): l is string => typeof l === 'string') : [];
  return Response.json({ available: true, lines });
}

/**
 * `GET /_officer/companion/logs/stream?tail=N` — the live tail, relayed frame for frame.
 *
 * The browser cannot open this itself: EventSource sends no Authorization header, and the key it would need
 * is one this sidecar exists to keep. So the stream is proxied, and the body is returned UNTOUCHED — a
 * ReadableStream that Bun cancels when the client disconnects, which aborts the upstream fetch in turn.
 * Buffering it into frames here would break that, and would also mean a log line waiting on our own flush.
 */
async function logStream(creds: HeadscaleServerCredentials, ctx: OfficerContext): Promise<Response> {
  const tail = Number(ctx.url.searchParams.get('tail') ?? 200);
  if (!Number.isInteger(tail) || tail < 1 || tail > 2000) return badRequest('tail must be 1–2000');

  // No timeout: a quiet log is the normal case and must not look like a dropped connection. The request's
  // own signal is the lifetime — when the panel closes, this closes.
  const res = await callCompanion(creds, {
    path: `/logs?tail=${tail}&follow=1`,
    signal: ctx.req.signal,
  });

  // An unavailable companion still answers in the stream's own vocabulary, so the client has one parser and
  // one place to show a problem rather than a second, JSON-shaped failure mode.
  if (typeof res === 'string') {
    return new Response(`event: error\ndata: ${res}\n\n`, {
      headers: { 'content-type': 'text/event-stream', 'cache-control': 'no-cache' },
    });
  }

  return new Response(res.body, {
    status: 200,
    headers: {
      'content-type': 'text/event-stream',
      'cache-control': 'no-cache',
      // Belt and braces through our own proxy chain, matching what the companion already sets.
      'x-accel-buffering': 'no',
    },
  });
}

const ACTIONS = new Set(['restart', 'stop', 'start']);

/**
 * `POST /_officer/companion/:action` — restart / stop / start the Headscale container.
 *
 * Every one of these drops every node's control-plane connection for the duration. That is the intended
 * "kill it" behaviour and the reason the UI asks twice; it is not something to retry automatically.
 */
async function action(creds: HeadscaleServerCredentials, name: string): Promise<Response> {
  // 60s: docker restart on a busy container is not fast, and a timeout here reads as "did it work?" — the
  // one question this feature exists to answer.
  const res = await callCompanion(creds, { path: `/${name}`, method: 'POST', timeoutMs: 60_000 });
  if (typeof res === 'string') return Response.json(unavailable(res));

  const body = await readBody(res);
  if (typeof body === 'string') return Response.json(unavailable(body));
  return Response.json({ available: true, ...body });
}

/** Dispatch `/_officer/companion/...`. Always acts on the ACTIVE server, like every other domain route. */
export async function handleCompanionRoute(ctx: OfficerContext, rest: string[]): Promise<Response> {
  const creds = await activeCreds(ctx.userId);
  if (creds instanceof Response) return creds;

  const [head, tail] = rest;

  if (head === 'health' && !tail) {
    if (ctx.req.method !== 'GET') return methodNotAllowed();
    return health(creds);
  }

  if (head === 'logs') {
    if (ctx.req.method !== 'GET') return methodNotAllowed();
    if (!tail) return logs(creds, ctx.url);
    if (tail === 'stream') return logStream(creds, ctx);
    return notFound();
  }

  if (head && ACTIONS.has(head) && !tail) {
    if (ctx.req.method !== 'POST') return methodNotAllowed();
    return action(creds, head);
  }

  return notFound();
}