maestro/web/src/stores/connectionStore.ts

import { create } from "zustand";

import { createAttachConnection } from "../attach.js";
import type { AttachInfo } from "../paneTitle.js";
import type { PaneConnection } from "../pane.js";
import { createStructuredConnection } from "../structured.js";
import { setSessionFocus } from "../api.js";
import { sessionName, type Session, type SessionView, type Uuid } from "../types.js";
import type { CatalogStoreApi } from "./catalogStore.js";

/**
 * How many *parked* (off-screen but still-connected) sessions to keep warm
 * beyond whatever is currently visible. Visible panes are never counted or
 * evicted — only this background buffer is bounded, so the browser doesn't hold
 * an unbounded fan of open sockets. Switching back to a recently-viewed session
 * is instant while it's still parked; older ones are evicted (WS closed) once
 * the buffer overflows.
 */
const MAX_PARKED = 6;

/** A session's view defaults to `"Terminal"` wherever it isn't carried —
 *  matching the wire default (`SessionView`'s `#[default]`, types.ts). Only
 *  meaningful once the session itself is known; an altogether-unknown
 *  session (no `AttachInfo`/`Session` at all) is a distinct, un-resolved
 *  state callers must check for separately, not "Terminal". */
function resolveView(view: SessionView | undefined): SessionView {
  return view ?? "Terminal";
}

export interface PoolOptions {
  /** Fired when the focused session changes (drives sidebar highlight + HUD). */
  onFocusChange?: (id: Uuid | null) => void;
  /** Fired when a view's explicit disconnect control is used; the embedder
   *  decides how to update the layout/route (then typically calls `close`). */
  onRequestClose?: (id: Uuid) => void;
  /** Fired when a connection's socket closes on its own (session ended/dropped),
   *  not via an intentional close/eviction. The embedder closes the pane. */
  onClosed?: (id: Uuid) => void;
  /** Forwarded to each connection: should a dropped socket be retried? */
  shouldReconnect?: (id: Uuid, code: number) => boolean;
  /** Threaded through to a structured connection so its rail can read live
   *  worktree/git-status data (#140). Terminal connections don't need it. */
  store?: CatalogStoreApi;
}

export interface ConnectionState {
  /** Sessions currently mounted into a visible pane. */
  mounted: Set<Uuid>;
  /** Every live (mounted or parked) connection, oldest first — LRU order. */
  cacheOrder: Uuid[];
  /** The focused session, or null when nothing is focused (welcome pane). */
  focusedId: Uuid | null;
}

export interface ConnectionActions {
  isMounted(id: Uuid): boolean;
  /**
   * Ensure `id` is live and shown inside `pane`, reusing its connection if one
   * exists (refreshing LRU recency). Does not touch other mounted panes.
   */
  mount(id: Uuid, pane: HTMLElement, info?: AttachInfo, opts?: { focus?: boolean }): void;
  /** Park `id` back into the hidden cache (still connected) and stop showing it. */
  unmount(id: Uuid): void;
  /**
   * Reconcile the visible set to exactly `entries` (in render order): park any
   * mounted session not listed, mount/refresh each listed one in its pane.
   * `focus: false` shows the panes without taking the keyboard (preview).
   */
  setVisible(entries: { id: Uuid; pane: HTMLElement; info?: AttachInfo }[], opts?: { focus?: boolean }): void;
  /** Also tells the daemon which pane is focused (docs/SIDEBAR.md's
   *  `AwaitingInput` auto-clear) — the sidebar highlight never calls this,
   *  only a pane actually becoming the visible view. Best-effort: a failed
   *  report only delays the glow clearing. */
  setFocus(id: Uuid | null): void;
  /** Compat shim: show `id` alone in `pane` and focus it; park everything else. */
  activate(id: Uuid, pane: HTMLElement, info?: AttachInfo): void;
  /** Compat shim: park all visible panes and clear focus (e.g. welcome pane). */
  parkActive(): void;
  /** Put the keyboard back into the focused session's terminal. No-op when
   *  nothing is focused (welcome pane) or its connection is gone. */
  focusTerminal(): void;
  /** Type raw bytes into the focused session's PTY. Returns false when there is
   *  nothing focused, or its connection is gone — the leader's literal Ctrl-K
   *  is a no-op over the welcome pane rather than an error
   *  (docs/COMMAND-BAR.md). */
  sendToFocused(data: string): boolean;
  /** Fully disconnect + forget a session. */
  close(id: Uuid): void;
  dispose(): void;
}

export type ConnectionStore = ConnectionState & ConnectionActions;

/**
 * Builds an isolated connection store — owns the set of live attach
 * connections, decoupled from the router/UI. The returned value is both a
 * React hook and a plain object (`getState`/`subscribe`), same shape as
 * {@link createCatalogStore}.
 *
 * Unlike a single-active pool, several connections can be **mounted**
 * (visible) at once — one per leaf of the active tab's layout tree. The
 * caller reconciles the visible set with {@link ConnectionActions.setVisible}.
 * Sessions no longer shown are **parked** in a hidden cache (still streaming)
 * and bounded by an LRU. A connection only dies on eviction, an explicit
 * {@link ConnectionActions.close}, or {@link ConnectionActions.dispose}.
 */
export function createConnectionStore(opts: PoolOptions = {}) {
  /** All live connections. Insertion order doubles as LRU: oldest first. */
  const cache = new Map<Uuid, PaneConnection>();
  /** The view each cached connection was actually created for — a deep link
   *  or a restored layout can `mount()` a session before the catalog
   *  snapshot names its view, defaulting it to `"Terminal"`; this is what
   *  lets a later reconciliation notice the resolved view disagrees (#140
   *  review). */
  const connViews = new Map<Uuid, SessionView>();
  // Off-screen home for parked, still-connected terminals.
  const hidden = document.createElement("div");
  hidden.className = "conn-cache";
  hidden.setAttribute("aria-hidden", "true");
  document.body.appendChild(hidden);

  let unsubscribeStore: (() => void) | undefined;

  const store = create<ConnectionStore>((set, get) => {
    function createConn(id: Uuid, info?: AttachInfo): PaneConnection {
      const view = resolveView(info?.view);
      connViews.set(id, view);
      const connOpts = {
        onRequestClose: (sid: Uuid) => opts.onRequestClose?.(sid),
        onClosed: (sid: Uuid) => opts.onClosed?.(sid),
        shouldReconnect: (sid: Uuid, code: number) => opts.shouldReconnect?.(sid, code) ?? false,
      };
      // structured.ts's own `store` param still expects the pre-Zustand
      // `CatalogStore` class (events.ts) and doesn't use it yet (#140) — left
      // unpassed here until that rail wiring is ported too.
      return view === "Structured"
        ? createStructuredConnection(id, connOpts, info)
        : createAttachConnection(id, connOpts, info);
    }

    /** Dispose whatever is cached for `id` and cache a freshly created
     *  connection in its place. Does not touch the DOM, `mounted`, or LRU
     *  order — callers place the new element themselves. */
    function replace(id: Uuid, info?: AttachInfo): PaneConnection {
      cache.get(id)?.dispose();
      const fresh = createConn(id, info);
      cache.set(id, fresh);
      return fresh;
    }

    function insertOrder(id: Uuid): void {
      set((s) => ({ cacheOrder: [...s.cacheOrder, id] }));
    }

    function touchOrder(id: Uuid): void {
      set((s) => ({ cacheOrder: [...s.cacheOrder.filter((x) => x !== id), id] }));
    }

    function removeOrder(id: Uuid): void {
      set((s) => ({ cacheOrder: s.cacheOrder.filter((x) => x !== id) }));
    }

    function setMounted(id: Uuid, isMounted: boolean): void {
      set((s) => {
        if (s.mounted.has(id) === isMounted) return {};
        const next = new Set(s.mounted);
        if (isMounted) next.add(id);
        else next.delete(id);
        return { mounted: next };
      });
    }

    function setFocus(id: Uuid | null): void {
      const previous = get().focusedId;
      if (id === previous) return;
      set({ focusedId: id });
      if (previous !== null) void setSessionFocus(previous, false).catch(() => {});
      if (id !== null) void setSessionFocus(id, true).catch(() => {});
      opts.onFocusChange?.(id);
    }

    function close(id: Uuid): void {
      const conn = cache.get(id);
      if (!conn) return;
      conn.dispose();
      cache.delete(id);
      connViews.delete(id);
      removeOrder(id);
      setMounted(id, false);
      if (get().focusedId === id) setFocus(null);
    }

    /** Drop least-recently-viewed *parked* connections beyond the buffer cap. */
    function evict(): void {
      while (cache.size - get().mounted.size > MAX_PARKED) {
        const mounted = get().mounted;
        let victim: Uuid | null = null;
        for (const id of cache.keys()) {
          if (!mounted.has(id)) {
            victim = id;
            break;
          }
        }
        if (victim === null) break;
        close(victim);
      }
    }

    function mount(id: Uuid, pane: HTMLElement, info?: AttachInfo, mountOpts?: { focus?: boolean }): void {
      let conn = cache.get(id);
      if (conn) {
        // `info === undefined` means "session not known yet", not "Terminal" —
        // never downgrade a resolved connection just because a later caller
        // couldn't resolve the session. Only a *resolved* view that disagrees
        // with what this connection was built for forces a replace.
        if (info !== undefined && resolveView(info.view) !== connViews.get(id)) {
          conn = replace(id, info);
        } else {
          touchOrder(id);
          if (info) conn.updateInfo(info);
        }
      } else {
        conn = createConn(id, info);
        cache.set(id, conn);
        insertOrder(id);
      }
      if (conn.el.parentElement !== pane) {
        pane.replaceChildren(conn.el);
        conn.onMount(mountOpts);
      } else if (!get().mounted.has(id)) {
        conn.onMount(mountOpts);
      }
      setMounted(id, true);
      evict();
    }

    function unmount(id: Uuid): void {
      if (!get().mounted.has(id)) return;
      setMounted(id, false);
      const conn = cache.get(id);
      if (conn) {
        hidden.appendChild(conn.el); // stays alive, just off-screen
        conn.onUnmount();
      }
    }

    function setVisible(
      entries: { id: Uuid; pane: HTMLElement; info?: AttachInfo }[],
      visOpts?: { focus?: boolean },
    ): void {
      const next = new Set(entries.map((e) => e.id));
      for (const id of [...get().mounted]) {
        if (!next.has(id)) unmount(id);
      }
      for (const e of entries) mount(e.id, e.pane, e.info, visOpts);
    }

    /** Proactively fix any cached connection whose view no longer matches the
     *  catalog — the only path that repairs a pane nothing will `mount()`
     *  again on its own (parked, or mounted but not currently being
     *  re-rendered by anything else). A session absent from `sessions` (not
     *  yet known, or removed) is left alone — that's the embedder's own
     *  close/eviction concern, not a view mismatch. */
    function reconcileViews(sessions: Map<Uuid, Session>): void {
      for (const id of [...cache.keys()]) {
        const session = sessions.get(id);
        if (session === undefined || resolveView(session.view) === connViews.get(id)) continue;
        const existing = cache.get(id)!;
        const parent = existing.el.parentElement;
        const wasMounted = get().mounted.has(id);
        // Capture before anything below touches the DOM: dispose()/el.remove()
        // detach `existing.el`, and once a focused descendant is detached the
        // browser moves `document.activeElement` to `<body>` — too late to
        // tell whether the keyboard was actually in this pane.
        const hadDomFocus = existing.el.contains(document.activeElement);
        const info: AttachInfo = { name: sessionName(session), command: session.command, view: session.view };
        const fresh = replace(id, info);
        existing.el.remove(); // don't rely on dispose() to clear the DOM
        if (parent && parent !== hidden) {
          parent.replaceChildren(fresh.el);
        } else {
          // Parked: `hidden` holds every parked connection's element as
          // siblings — append, never replaceChildren, or they'd all vanish.
          hidden.appendChild(fresh.el);
        }
        if (wasMounted) {
          // `focusedId` names the session the pool was last told to focus —
          // but a same-tab sidebar preview moves the keyboard into the
          // sidebar without clearing it, so `focusedId` alone can point at a
          // pane the keyboard already left. Only re-focus when the DOM
          // agrees it never did.
          fresh.onMount({ focus: get().focusedId === id && hadDomFocus });
        } else {
          fresh.onUnmount();
        }
      }
    }

    // A session `mount()`ed before its catalog entry exists gets a
    // `"Terminal"` connection by default; once the snapshot actually names
    // the view, fix any connection that guessed wrong — nothing else
    // re-`mount()`s an already-open pane just because the catalog updated.
    unsubscribeStore = opts.store?.subscribe((state) => reconcileViews(state.sessions));

    return {
      mounted: new Set<Uuid>(),
      cacheOrder: [],
      focusedId: null,

      isMounted(id) {
        return get().mounted.has(id);
      },

      mount,
      unmount,
      setVisible,
      setFocus,

      activate(id, pane, info) {
        setVisible([{ id, pane, info }]);
        setFocus(id);
      },

      parkActive() {
        for (const id of [...get().mounted]) unmount(id);
        setFocus(null);
      },

      focusTerminal() {
        const id = get().focusedId;
        if (id === null) return;
        cache.get(id)?.focusTerminal();
      },

      sendToFocused(data) {
        const id = get().focusedId;
        if (id === null) return false;
        const conn = cache.get(id);
        if (!conn) return false;
        conn.sendInput(data);
        return true;
      },

      close,

      dispose() {
        unsubscribeStore?.();
        for (const conn of cache.values()) conn.dispose();
        cache.clear();
        connViews.clear();
        set({ mounted: new Set(), cacheOrder: [], focusedId: null });
        hidden.remove();
      },
    };
  });

  return store;
}

export type ConnectionStoreApi = ReturnType<typeof createConnectionStore>;