stack/apps/web/src/hooks/useTerminalSessions.ts

/**
 * useTerminalSessions hook
 *
 * Manages multiple PTY terminal sessions over a single WebSocket connection
 * to the /terminal namespace. Supports creating, closing, renaming, and switching
 * between sessions, with per-session output callback multiplexing.
 *
 * Protocol (from terminal.gateway.ts):
 * 1. Connect with auth token in handshake
 * 2. Emit terminal:create { name?, cols?, rows? } → receive terminal:created { sessionId, name, cols, rows }
 * 3. Emit terminal:input { sessionId, data } to send keystrokes
 * 4. Receive terminal:output { sessionId, data } for stdout/stderr
 * 5. Emit terminal:resize { sessionId, cols, rows } on window resize
 * 6. Emit terminal:close { sessionId } to terminate the PTY
 * 7. Receive terminal:exit { sessionId, exitCode, signal } on PTY exit
 * 8. Receive terminal:error { message } on errors
 */

import { useEffect, useRef, useState, useCallback } from "react";
import type { Socket } from "socket.io-client";
import { io } from "socket.io-client";
import { API_BASE_URL } from "@/lib/config";

// ==========================================
// Types
// ==========================================

export type SessionStatus = "active" | "exited";

export interface SessionInfo {
  /** Session identifier returned by the server */
  sessionId: string;
  /** Human-readable tab label */
  name: string;
  /** Whether the PTY process is still running */
  status: SessionStatus;
  /** Exit code, populated when status === 'exited' */
  exitCode?: number;
}

export interface CreateSessionOptions {
  /** Optional label for the new session */
  name?: string;
  /** Terminal columns */
  cols?: number;
  /** Terminal rows */
  rows?: number;
  /** Working directory */
  cwd?: string;
}

export interface UseTerminalSessionsOptions {
  /** Authentication token for WebSocket handshake */
  token: string;
}

export interface UseTerminalSessionsReturn {
  /** Map of sessionId → SessionInfo */
  sessions: Map<string, SessionInfo>;
  /** Currently active (visible) session id, or null if none */
  activeSessionId: string | null;
  /** Whether the WebSocket is connected */
  isConnected: boolean;
  /** Connection error message, if any */
  connectionError: string | null;
  /** Create a new PTY session */
  createSession: (options?: CreateSessionOptions) => void;
  /** Close an existing PTY session */
  closeSession: (sessionId: string) => void;
  /** Rename a session (local label only, not persisted to server) */
  renameSession: (sessionId: string, name: string) => void;
  /** Switch the visible session */
  setActiveSession: (sessionId: string) => void;
  /** Send keyboard input to a session */
  sendInput: (sessionId: string, data: string) => void;
  /** Notify the server of a terminal resize */
  resize: (sessionId: string, cols: number, rows: number) => void;
  /**
   * Register a callback that receives output data for a specific session.
   * Returns an unsubscribe function — call it during cleanup.
   */
  registerOutputCallback: (sessionId: string, cb: (data: string) => void) => () => void;
}

// ==========================================
// Payload shapes matching terminal.dto.ts
// ==========================================

interface TerminalCreatedPayload {
  sessionId: string;
  name: string;
  cols: number;
  rows: number;
}

interface TerminalOutputPayload {
  sessionId: string;
  data: string;
}

interface TerminalExitPayload {
  sessionId: string;
  exitCode: number;
  signal?: string;
}

interface TerminalErrorPayload {
  message: string;
}

// ==========================================
// Security validation
// ==========================================

function validateWebSocketSecurity(url: string): void {
  const isProduction = process.env.NODE_ENV === "production";
  const isSecure = url.startsWith("https://") || url.startsWith("wss://");

  if (isProduction && !isSecure) {
    console.warn(
      "[Security Warning] Terminal WebSocket using insecure protocol (ws://). " +
        "Authentication tokens may be exposed. Use wss:// in production."
    );
  }
}

// ==========================================
// Hook
// ==========================================

/**
 * Hook for managing multiple PTY terminal sessions over a single WebSocket connection.
 *
 * @param options - Configuration including auth token
 * @returns Multi-session terminal state and control functions
 */
export function useTerminalSessions(
  options: UseTerminalSessionsOptions
): UseTerminalSessionsReturn {
  const { token } = options;

  const socketRef = useRef<Socket | null>(null);
  // Per-session output callback registry; keyed by sessionId
  const outputCallbacksRef = useRef<Map<string, Set<(data: string) => void>>>(new Map());

  const [sessions, setSessions] = useState<Map<string, SessionInfo>>(new Map());
  const [activeSessionId, setActiveSessionIdState] = useState<string | null>(null);
  const [isConnected, setIsConnected] = useState(false);
  const [connectionError, setConnectionError] = useState<string | null>(null);

  // ==========================================
  // Auto-select first available session when active becomes null
  // ==========================================

  useEffect(() => {
    if (activeSessionId === null && sessions.size > 0) {
      const firstId = sessions.keys().next().value;
      if (firstId !== undefined) {
        setActiveSessionIdState(firstId);
      }
    }
  }, [activeSessionId, sessions]);

  // ==========================================
  // WebSocket connection
  // ==========================================

  useEffect(() => {
    if (!token) {
      return;
    }

    const wsUrl = API_BASE_URL;
    validateWebSocketSecurity(wsUrl);

    setConnectionError(null);

    const socket = io(`${wsUrl}/terminal`, {
      auth: { token },
      transports: ["websocket", "polling"],
    });

    socketRef.current = socket;

    const handleConnect = (): void => {
      setIsConnected(true);
      setConnectionError(null);
    };

    const handleDisconnect = (): void => {
      setIsConnected(false);
      // Sessions remain in the Map but are no longer interactive
      setSessions((prev) => {
        const next = new Map(prev);
        for (const [id, info] of next) {
          if (info.status === "active") {
            next.set(id, { ...info, status: "exited" });
          }
        }
        return next;
      });
    };

    const handleConnectError = (error: Error): void => {
      setConnectionError(error.message || "Terminal connection failed");
      setIsConnected(false);
    };

    const handleTerminalCreated = (payload: TerminalCreatedPayload): void => {
      setSessions((prev) => {
        const next = new Map(prev);
        next.set(payload.sessionId, {
          sessionId: payload.sessionId,
          name: payload.name,
          status: "active",
        });
        return next;
      });
      // Set as active session if none is currently active
      setActiveSessionIdState((prev) => prev ?? payload.sessionId);
    };

    const handleTerminalOutput = (payload: TerminalOutputPayload): void => {
      const callbacks = outputCallbacksRef.current.get(payload.sessionId);
      if (callbacks) {
        for (const cb of callbacks) {
          cb(payload.data);
        }
      }
    };

    const handleTerminalExit = (payload: TerminalExitPayload): void => {
      setSessions((prev) => {
        const next = new Map(prev);
        const session = next.get(payload.sessionId);
        if (session) {
          next.set(payload.sessionId, {
            ...session,
            status: "exited",
            exitCode: payload.exitCode,
          });
        }
        return next;
      });
    };

    const handleTerminalError = (payload: TerminalErrorPayload): void => {
      console.error("[Terminal] Error:", payload.message);
    };

    socket.on("connect", handleConnect);
    socket.on("disconnect", handleDisconnect);
    socket.on("connect_error", handleConnectError);
    socket.on("terminal:created", handleTerminalCreated);
    socket.on("terminal:output", handleTerminalOutput);
    socket.on("terminal:exit", handleTerminalExit);
    socket.on("terminal:error", handleTerminalError);

    return (): void => {
      socket.off("connect", handleConnect);
      socket.off("disconnect", handleDisconnect);
      socket.off("connect_error", handleConnectError);
      socket.off("terminal:created", handleTerminalCreated);
      socket.off("terminal:output", handleTerminalOutput);
      socket.off("terminal:exit", handleTerminalExit);
      socket.off("terminal:error", handleTerminalError);

      // Close all active sessions before disconnecting
      const currentSessions = sessions;
      for (const [id, info] of currentSessions) {
        if (info.status === "active") {
          socket.emit("terminal:close", { sessionId: id });
        }
      }

      socket.disconnect();
      socketRef.current = null;
    };
    // Intentional: token is the only dep that should trigger reconnection
  }, [token]);

  // ==========================================
  // Control functions
  // ==========================================

  const createSession = useCallback((createOptions: CreateSessionOptions = {}): void => {
    const socket = socketRef.current;
    if (!socket?.connected) {
      return;
    }

    const payload: Record<string, unknown> = {};
    if (createOptions.name !== undefined) payload.name = createOptions.name;
    if (createOptions.cols !== undefined) payload.cols = createOptions.cols;
    if (createOptions.rows !== undefined) payload.rows = createOptions.rows;
    if (createOptions.cwd !== undefined) payload.cwd = createOptions.cwd;

    socket.emit("terminal:create", payload);
  }, []);

  const closeSession = useCallback((sessionId: string): void => {
    const socket = socketRef.current;
    if (socket?.connected) {
      socket.emit("terminal:close", { sessionId });
    }

    setSessions((prev) => {
      const next = new Map(prev);
      next.delete(sessionId);
      return next;
    });

    // If closing the active session, activeSessionId becomes null
    // and the auto-select useEffect will pick the first remaining session
    setActiveSessionIdState((prev) => (prev === sessionId ? null : prev));
  }, []);

  const renameSession = useCallback((sessionId: string, name: string): void => {
    setSessions((prev) => {
      const next = new Map(prev);
      const session = next.get(sessionId);
      if (session) {
        next.set(sessionId, { ...session, name });
      }
      return next;
    });
  }, []);

  const setActiveSession = useCallback((sessionId: string): void => {
    setActiveSessionIdState(sessionId);
  }, []);

  const sendInput = useCallback((sessionId: string, data: string): void => {
    const socket = socketRef.current;
    if (!socket?.connected) {
      return;
    }
    socket.emit("terminal:input", { sessionId, data });
  }, []);

  const resize = useCallback((sessionId: string, cols: number, rows: number): void => {
    const socket = socketRef.current;
    if (!socket?.connected) {
      return;
    }
    socket.emit("terminal:resize", { sessionId, cols, rows });
  }, []);

  const registerOutputCallback = useCallback(
    (sessionId: string, cb: (data: string) => void): (() => void) => {
      const registry = outputCallbacksRef.current;
      if (!registry.has(sessionId)) {
        registry.set(sessionId, new Set());
      }
      // Safe: we just ensured the key exists
      const callbackSet = registry.get(sessionId);
      if (callbackSet) {
        callbackSet.add(cb);
      }

      return (): void => {
        registry.get(sessionId)?.delete(cb);
      };
    },
    []
  );

  return {
    sessions,
    activeSessionId,
    isConnected,
    connectionError,
    createSession,
    closeSession,
    renameSession,
    setActiveSession,
    sendInput,
    resize,
    registerOutputCallback,
  };
}