stack/apps/api/src/terminal/terminal.service.ts

/**
 * TerminalService
 *
 * Manages PTY (pseudo-terminal) sessions for workspace users.
 * Spawns real shell processes via node-pty, streams I/O to connected sockets,
 * and enforces per-workspace session limits.
 *
 * Session lifecycle:
 * - createSession: spawn a new PTY, wire onData/onExit, return sessionId
 * - writeToSession: send input data to PTY stdin
 * - resizeSession: resize PTY dimensions (cols x rows)
 * - closeSession: kill PTY process, emit terminal:exit, cleanup
 * - closeWorkspaceSessions: kill all sessions for a workspace (on disconnect)
 */

import { Injectable, Logger, OnModuleInit } from "@nestjs/common";
import type { IPty } from "node-pty";
import type { Socket } from "socket.io";
import { randomUUID } from "node:crypto";

// Lazy-loaded in onModuleInit via dynamic import() to prevent crash
// if the native binary is missing. node-pty requires a compiled .node
// binary which may not be available in all Docker environments.
interface NodePtyModule {
  spawn: (file: string, args: string[], options: Record<string, unknown>) => IPty;
}
let pty: NodePtyModule | null = null;

/** Maximum concurrent PTY sessions per workspace */
export const MAX_SESSIONS_PER_WORKSPACE = parseInt(
  process.env.TERMINAL_MAX_SESSIONS_PER_WORKSPACE ?? "10",
  10
);

/** Default PTY dimensions */
const DEFAULT_COLS = 80;
const DEFAULT_ROWS = 24;

export interface TerminalSession {
  sessionId: string;
  workspaceId: string;
  pty: IPty;
  name?: string;
  createdAt: Date;
}

export interface CreateSessionOptions {
  name?: string;
  cols?: number;
  rows?: number;
  cwd?: string;
  workspaceId: string;
  socketId: string;
}

export interface SessionCreatedResult {
  sessionId: string;
  name?: string;
  cols: number;
  rows: number;
}

@Injectable()
export class TerminalService implements OnModuleInit {
  private readonly logger = new Logger(TerminalService.name);

  /**
   * Map of sessionId -> TerminalSession
   */
  private readonly sessions = new Map<string, TerminalSession>();

  /**
   * Map of workspaceId -> Set<sessionId> for fast per-workspace lookups
   */
  private readonly workspaceSessions = new Map<string, Set<string>>();

  async onModuleInit(): Promise<void> {
    if (!pty) {
      try {
        pty = await import("node-pty");
        this.logger.log("node-pty loaded successfully — terminal sessions available");
      } catch {
        this.logger.warn(
          "node-pty native module not available — terminal sessions will be disabled. " +
            "Install build tools (python3, make, g++) and rebuild node-pty to enable."
        );
      }
    }
  }

  /**
   * Create a new PTY session for the given workspace and socket.
   * Wires PTY onData -> emit terminal:output and onExit -> emit terminal:exit.
   *
   * @throws Error if workspace session limit is exceeded or node-pty is unavailable
   */
  createSession(socket: Socket, options: CreateSessionOptions): SessionCreatedResult {
    if (!pty) {
      throw new Error("Terminal sessions are unavailable: node-pty native module failed to load");
    }
    const { workspaceId, name, cwd, socketId } = options;
    const cols = options.cols ?? DEFAULT_COLS;
    const rows = options.rows ?? DEFAULT_ROWS;

    // Enforce per-workspace session limit
    const workspaceSessionIds = this.workspaceSessions.get(workspaceId) ?? new Set<string>();
    if (workspaceSessionIds.size >= MAX_SESSIONS_PER_WORKSPACE) {
      throw new Error(
        `Workspace ${workspaceId} has reached the maximum of ${String(MAX_SESSIONS_PER_WORKSPACE)} concurrent terminal sessions`
      );
    }

    const sessionId = randomUUID();
    const shell = process.env.SHELL ?? "/bin/bash";

    this.logger.log(
      `Spawning PTY session ${sessionId} for workspace ${workspaceId} (socket: ${socketId}, shell: ${shell}, ${String(cols)}x${String(rows)})`
    );

    const ptyProcess = pty.spawn(shell, [], {
      name: "xterm-256color",
      cols,
      rows,
      cwd: cwd ?? process.cwd(),
      env: process.env as Record<string, string>,
    });

    const session: TerminalSession = {
      sessionId,
      workspaceId,
      pty: ptyProcess,
      ...(name !== undefined ? { name } : {}),
      createdAt: new Date(),
    };

    this.sessions.set(sessionId, session);

    // Track by workspace
    if (!this.workspaceSessions.has(workspaceId)) {
      this.workspaceSessions.set(workspaceId, new Set());
    }
    const wsSet = this.workspaceSessions.get(workspaceId);
    if (wsSet) {
      wsSet.add(sessionId);
    }

    // Wire PTY stdout/stderr -> terminal:output
    ptyProcess.onData((data: string) => {
      socket.emit("terminal:output", { sessionId, data });
    });

    // Wire PTY exit -> terminal:exit, cleanup
    ptyProcess.onExit(({ exitCode, signal }) => {
      this.logger.log(
        `PTY session ${sessionId} exited (exitCode: ${String(exitCode)}, signal: ${String(signal ?? "none")})`
      );
      socket.emit("terminal:exit", { sessionId, exitCode, signal });
      this.cleanupSession(sessionId, workspaceId);
    });

    return { sessionId, ...(name !== undefined ? { name } : {}), cols, rows };
  }

  /**
   * Write input data to a PTY session's stdin.
   *
   * @throws Error if session not found
   */
  writeToSession(sessionId: string, data: string): void {
    const session = this.sessions.get(sessionId);
    if (!session) {
      throw new Error(`Terminal session ${sessionId} not found`);
    }
    session.pty.write(data);
  }

  /**
   * Resize a PTY session's terminal dimensions.
   *
   * @throws Error if session not found
   */
  resizeSession(sessionId: string, cols: number, rows: number): void {
    const session = this.sessions.get(sessionId);
    if (!session) {
      throw new Error(`Terminal session ${sessionId} not found`);
    }
    session.pty.resize(cols, rows);
    this.logger.debug(`Resized PTY session ${sessionId} to ${String(cols)}x${String(rows)}`);
  }

  /**
   * Kill and clean up a specific PTY session.
   * Returns true if the session existed, false if it was already gone.
   */
  closeSession(sessionId: string): boolean {
    const session = this.sessions.get(sessionId);
    if (!session) {
      return false;
    }

    this.logger.log(`Closing PTY session ${sessionId} for workspace ${session.workspaceId}`);

    try {
      session.pty.kill();
    } catch (error) {
      this.logger.warn(
        `Error killing PTY session ${sessionId}: ${error instanceof Error ? error.message : String(error)}`
      );
    }

    this.cleanupSession(sessionId, session.workspaceId);
    return true;
  }

  /**
   * Close all PTY sessions for a workspace (called on client disconnect).
   */
  closeWorkspaceSessions(workspaceId: string): void {
    const sessionIds = this.workspaceSessions.get(workspaceId);
    if (!sessionIds || sessionIds.size === 0) {
      return;
    }

    this.logger.log(
      `Closing ${String(sessionIds.size)} PTY session(s) for workspace ${workspaceId} (disconnect)`
    );

    // Copy to array to avoid mutation during iteration
    const ids = Array.from(sessionIds);
    for (const sessionId of ids) {
      const session = this.sessions.get(sessionId);
      if (session) {
        try {
          session.pty.kill();
        } catch (error) {
          this.logger.warn(
            `Error killing PTY session ${sessionId} on disconnect: ${error instanceof Error ? error.message : String(error)}`
          );
        }
        this.cleanupSession(sessionId, workspaceId);
      }
    }
  }

  /**
   * Get the number of active sessions for a workspace.
   */
  getWorkspaceSessionCount(workspaceId: string): number {
    return this.workspaceSessions.get(workspaceId)?.size ?? 0;
  }

  /**
   * Check if a session belongs to a given workspace.
   * Used for access control in the gateway.
   */
  sessionBelongsToWorkspace(sessionId: string, workspaceId: string): boolean {
    const session = this.sessions.get(sessionId);
    return session?.workspaceId === workspaceId;
  }

  /**
   * Internal cleanup: remove session from tracking maps.
   * Does NOT kill the PTY (caller is responsible).
   */
  private cleanupSession(sessionId: string, workspaceId: string): void {
    this.sessions.delete(sessionId);

    const workspaceSessionIds = this.workspaceSessions.get(workspaceId);
    if (workspaceSessionIds) {
      workspaceSessionIds.delete(sessionId);
      if (workspaceSessionIds.size === 0) {
        this.workspaceSessions.delete(workspaceId);
      }
    }
  }
}