maestro/web/src/components/structured/useStructuredSession.ts

// The `WS /api/structured/:id` half of the pane (#246). Every connect — fresh
// or reconnect — resends the whole transcript rather than resuming a cursor
// (docs/STRUCTURED-VIEW.md), so a snapshot replaces state wholesale instead of
// merging into what came before.

import { useCallback, useEffect, useMemo, useRef, useState } from "react";

import { structuredWsUrl } from "../../routes.js";
import { INITIAL_RECONNECT_MS, MAX_RECONNECT_MS } from "../../reconnect.js";
import type {
  ReducedState,
  StructuredEvent,
  TranscriptDelta,
  TranscriptRow,
  Uuid,
} from "../../types.js";

export type StructuredStatus = "connecting" | "open" | "reconnecting" | "closed" | "error";

/** Injectable so tests can drive the socket without a daemon — the same seam
 *  structured.ts used. */
export type StructuredSocketFactory = (sessionId: Uuid) => WebSocket;

const idleState = (): ReducedState => ({
  turn: { responding: false, thinking: false, active_tool_calls: [] },
});

function applyDelta(rows: TranscriptRow[], delta: TranscriptDelta): TranscriptRow[] {
  if (delta.type === "append") {
    // A patch can arrive for a row this connection has never seen (its
    // snapshot predates it); treating an unknown id as an append keeps the
    // transcript complete rather than dropping the row.
    const at = rows.findIndex((row) => row.id === delta.row.id);
    return at === -1 ? [...rows, delta.row] : rows.map((row, i) => (i === at ? delta.row : row));
  }
  if (delta.type === "patch") {
    const at = rows.findIndex((row) => row.id === delta.row.id);
    return at === -1 ? [...rows, delta.row] : rows.map((row, i) => (i === at ? delta.row : row));
  }
  return rows.filter((row) => row.id !== delta.id);
}

export interface StructuredSession {
  rows: TranscriptRow[];
  state: ReducedState;
  status: StructuredStatus;
  /** A daemon-supplied close reason (code ≥ 4000), when there is one. */
  detail: string | null;
  /** The number of sockets this pane has opened, newest last. A submission
   *  that straddles a reconnect uses it to notice its acceptance may predate
   *  the current snapshot. */
  socketGeneration: number;
  /** Drop the socket and reconnect, so the next snapshot cannot predate a
   *  POST that has already been accepted. */
  reconnect(): void;
}

export interface UseStructuredSessionOptions {
  shouldReconnect?: (code: number) => boolean;
  onClosed?: () => void;
  socketFactory?: StructuredSocketFactory;
}

/** Application-defined close codes (daemon rejection) and the clean-close code
 *  — retrying either would only loop. */
function defaultShouldReconnect(code: number): boolean {
  return code < 4000 && code !== 1000;
}

export function useStructuredSession(
  sessionId: Uuid,
  {
    shouldReconnect = defaultShouldReconnect,
    onClosed,
    socketFactory = (id) => new WebSocket(structuredWsUrl(id)),
  }: UseStructuredSessionOptions = {},
): StructuredSession {
  const [rows, setRows] = useState<TranscriptRow[]>([]);
  const [state, setState] = useState<ReducedState>(idleState);
  const [status, setStatus] = useState<StructuredStatus>("connecting");
  const [detail, setDetail] = useState<string | null>(null);
  const [socketGeneration, setSocketGeneration] = useState(0);

  // Read through refs so a changed callback never tears the socket down.
  const shouldReconnectRef = useRef(shouldReconnect);
  const onClosedRef = useRef(onClosed);
  const factoryRef = useRef(socketFactory);
  shouldReconnectRef.current = shouldReconnect;
  onClosedRef.current = onClosed;
  factoryRef.current = socketFactory;

  const reconnectRef = useRef<() => void>(() => {});

  useEffect(() => {
    let disposed = false;
    let socket: WebSocket | null = null;
    let reconnectMs = INITIAL_RECONNECT_MS;
    let timer: number | null = null;

    const closeSocket = () => {
      const existing = socket;
      socket = null;
      if (!existing) return;
      existing.onopen = null;
      existing.onmessage = null;
      existing.onerror = null;
      existing.onclose = null;
      try {
        existing.close();
      } catch {
        // ignore — an already-closing socket has nothing to close
      }
    };

    const cancelTimer = () => {
      if (timer === null) return;
      window.clearTimeout(timer);
      timer = null;
    };

    function connect(): void {
      if (disposed) return;
      cancelTimer();
      closeSocket();
      const current = factoryRef.current(sessionId);
      socket = current;
      setSocketGeneration((n) => n + 1);

      current.onopen = () => {
        if (disposed || socket !== current) return;
        reconnectMs = INITIAL_RECONNECT_MS;
        setDetail(null);
        setStatus("open");
      };

      current.onmessage = (ev) => {
        if (disposed || socket !== current) return;
        if (typeof ev.data !== "string") return;
        let msg: StructuredEvent;
        try {
          msg = JSON.parse(ev.data) as StructuredEvent;
        } catch {
          return; // not a structured event we understand
        }
        switch (msg.type) {
          case "transcript_snapshot":
            setRows([...msg.snapshot.rows]);
            break;
          case "reduced_state":
            setState(msg.state);
            break;
          case "projection_update":
            if (msg.transcript) {
              const delta = msg.transcript;
              setRows((previous) => applyDelta(previous, delta));
            }
            setState(msg.state);
            break;
        }
      };

      current.onerror = () => {
        if (disposed || socket !== current) return;
        setStatus("error");
      };

      current.onclose = (ev) => {
        if (disposed || socket !== current) return;
        socket = null;
        const reason = ev.reason?.trim();
        if (ev.code >= 4000 && reason) {
          setDetail(reason);
          setStatus("error");
          return;
        }
        if (shouldReconnectRef.current(ev.code)) {
          const delay = Math.min(reconnectMs, MAX_RECONNECT_MS);
          setStatus("reconnecting");
          timer = window.setTimeout(() => {
            timer = null;
            if (disposed) return;
            reconnectMs = Math.min(reconnectMs * 2, MAX_RECONNECT_MS);
            connect();
          }, delay);
          return;
        }
        setDetail(String(ev.code));
        setStatus("closed");
        onClosedRef.current?.();
      };
    }

    reconnectRef.current = () => {
      if (disposed) return;
      setStatus("reconnecting");
      connect();
    };

    setStatus("connecting");
    connect();

    return () => {
      disposed = true;
      reconnectRef.current = () => {};
      cancelTimer();
      closeSocket();
    };
  }, [sessionId]);

  const reconnect = useCallback(() => reconnectRef.current(), []);

  return useMemo(
    () => ({ rows, state, status, detail, socketGeneration, reconnect }),
    [rows, state, status, detail, socketGeneration, reconnect],
  );
}