offscale/web/useHeadscaleCompanion.ts

import { useEffect, useRef, useState } from 'react';
import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query';
import { useClient, getHeaders } from 'hooks/useClient';
import type { CompanionAction, CompanionActionResult, CompanionHealthResult, CompanionLogsResult } from './shared';

// Client for the active server's Officer Companion. Everything here goes through the headscale sidecar,
// because the companion authenticates with the Headscale admin key — which is encrypted in Postgres and
// decryptable only there. The browser never sees it and never talks to the companion directly.

const BASE = '/offscale/_officer/companion';
const HEALTH_KEY = ['headscale', 'companion', 'health'] as const;

/**
 * The container's health, polled.
 *
 * Polling is the point rather than a convenience: this section is what you have open while waiting for a
 * restart to take, so it has to move on its own. 10s is fast enough to watch a container come back and slow
 * enough that a degraded server isn't being hammered while it struggles.
 */
export function useCompanionHealth() {
  const { get } = useClient();
  return useQuery({
    queryKey: HEALTH_KEY,
    queryFn: () => get<CompanionHealthResult>(`${BASE}/health`),
    refetchInterval: 10_000,
    // A verdict from ten seconds ago is stale by definition; never show one on remount without refetching.
    staleTime: 0,
  });
}

/** A snapshot of the last N lines. The live tail is a separate thing — see useCompanionLogStream. */
export function useCompanionLogs(tail: number, enabled: boolean) {
  const { get } = useClient();
  return useQuery({
    queryKey: ['headscale', 'companion', 'logs', tail],
    queryFn: () => get<CompanionLogsResult>(`${BASE}/logs?tail=${tail}`),
    enabled,
    staleTime: 0,
  });
}

/**
 * Restart / stop / start the Headscale container.
 *
 * Every one of these drops every node's control-plane connection, so nothing here retries and nothing here
 * fires without the owner having confirmed. The health query is invalidated on settle — including on
 * failure, where "did it happen anyway?" is exactly the question.
 */
export function useCompanionAction() {
  const { post } = useClient();
  const qc = useQueryClient();
  return useMutation({
    mutationFn: (action: CompanionAction) => post<CompanionActionResult>(`${BASE}/${action}`),
    onSettled: () => qc.invalidateQueries({ queryKey: HEALTH_KEY }),
  });
}

/** How many lines the viewer keeps. Beyond this the browser, not the server, becomes the bottleneck. */
const MAX_LINES = 5000;

export type LogStream = {
  lines: string[];
  /** Set when the stream ended badly — the companion's own `event: error` frame, or a dropped connection. */
  error: string | null;
  /** True between opening the request and the stream ending, however it ends. */
  live: boolean;
  clear: () => void;
};

/**
 * The live log tail, over `fetch()` rather than `EventSource`.
 *
 * `EventSource` cannot send an `Authorization` header and every hop of this chain needs one — Officer's own
 * bearer token to reach the platform, and the Headscale admin key from there on. So the SSE framing is
 * parsed by hand: split on blank lines, read `data:` and `event:`. It is a small parser and it only has to
 * handle what the companion emits (one line per frame, an `event: error` frame before a fatal close).
 *
 * `follow` changing off aborts mid-stream; the abort is deliberately not reported as an error, since it is
 * the owner switching the toggle rather than anything going wrong.
 */
export function useCompanionLogStream(follow: boolean, tail: number): LogStream {
  const [lines, setLines] = useState<string[]>([]);
  const [error, setError] = useState<string | null>(null);
  const [live, setLive] = useState(false);
  // Batched through a ref: a busy container emits faster than React can render, and one setState per line
  // would spend the whole frame budget on log output.
  const pending = useRef<string[]>([]);

  useEffect(() => {
    if (!follow) return;

    const controller = new AbortController();
    let flushTimer: ReturnType<typeof setInterval> | null = null;
    setError(null);
    setLive(true);

    const flush = () => {
      if (pending.current.length === 0) return;
      const batch = pending.current;
      pending.current = [];
      setLines((prev) => {
        const next = prev.concat(batch);
        return next.length > MAX_LINES ? next.slice(next.length - MAX_LINES) : next;
      });
    };

    const run = async () => {
      try {
        const res = await fetch(`/api${BASE}/logs/stream?tail=${tail}`, {
          headers: getHeaders(),
          signal: controller.signal,
        });
        if (!res.ok || !res.body) {
          setError(`the log stream returned ${res.status}`);
          return;
        }

        flushTimer = setInterval(flush, 200);
        const reader = res.body.getReader();
        const decoder = new TextDecoder();
        let buffer = '';

        for (;;) {
          const { done, value } = await reader.read();
          if (done) break;
          buffer += decoder.decode(value, { stream: true });

          // Frames are separated by a blank line; anything after the last one is a partial frame and waits.
          const frames = buffer.split('\n\n');
          buffer = frames.pop() ?? '';

          for (const frame of frames) {
            let event = 'message';
            const data: string[] = [];
            for (const rawLine of frame.split('\n')) {
              if (rawLine.startsWith('event:')) event = rawLine.slice(6).trim();
              else if (rawLine.startsWith('data:')) data.push(rawLine.slice(5).replace(/^ /, ''));
            }
            if (data.length === 0) continue;
            const text = data.join('\n');
            if (event === 'error') setError(text);
            else pending.current.push(text);
          }
        }
      } catch (err) {
        if ((err as Error | null)?.name !== 'AbortError') setError('the log stream disconnected');
      } finally {
        if (flushTimer) clearInterval(flushTimer);
        flush();
        // An aborted stream is already being torn down by the effect that replaced this one; letting it set
        // state here would flash "not live" onto a stream that is about to reopen.
        if (!controller.signal.aborted) setLive(false);
      }
    };

    void run();
    return () => {
      controller.abort();
      if (flushTimer) clearInterval(flushTimer);
      setLive(false);
    };
  }, [follow, tail]);

  return { lines, error, live, clear: () => setLines([]) };
}