@assistkick/create 1.6.0 → 1.8.0

package/templates/assistkick-product-system/packages/backend/src/services/pty_session_manager.ts

@@ -1,63 +1,61 @@
 /**
  * PTY Session Manager — manages named, project-scoped terminal sessions.
- * Implements dec_029 (Command Prefix Whitelist) and dec_030 (Server-Side PTY Session Persistence).
+ * Sessions are persisted in the database and survive server restarts.
+ * PTY processes are ephemeral — spawned on demand when a user connects.
  *
- * Each session has a unique ID, a human-readable name (project name + creation timestamp),
- * and is permanently bound to a project. Sessions persist across WebSocket disconnects.
+ * New sessions launch claude with `--session-id <uuid>` to set a known ID.
+ * Resumed sessions (after restart) launch claude with `--resume <uuid>`.
  */
 
 import type { IPty } from 'node-pty';
-import { randomBytes } from 'node:crypto';
+import { randomBytes, randomUUID } from 'node:crypto';
+import { desc, eq } from 'drizzle-orm';
+import { terminalSessions } from '@assistkick/shared/db/schema.js';
 
 const OUTPUT_BUFFER_MAX = 50_000;
-
-const ALLOWED_COMMAND_PREFIXES = [
-  'claude',
-] as const;
-
-const PROMPT = '\x1b[32m$\x1b[0m ';
-const WELCOME_MSG = `\x1b[90mRestricted terminal — allowed commands: ${ALLOWED_COMMAND_PREFIXES.join(', ')}\x1b[0m\r\n`;
+const IDLE_TIMEOUT_MS = 1_800_000; // 30 minutes
 
 interface PtySessionManagerDeps {
   spawn: (shell: string, args: string[], options: Record<string, unknown>) => IPty;
   log: (tag: string, ...args: unknown[]) => void;
   projectRoot: string;
+  getDb: () => any;
 }
 
-interface PtySession {
-  id: string;
-  name: string;
-  projectId: string;
-  projectName: string;
-  pty: IPty | null;
-  state: 'idle' | 'running';
-  cols: number;
-  rows: number;
+/** In-memory state for a live PTY process. */
+interface LivePty {
+  pty: IPty;
   outputBuffer: string;
-  inputBuffer: string;
   listeners: Set<(data: string) => void>;
-  createdAt: Date;
+  cols: number;
+  rows: number;
 }
 
 export interface SessionInfo {
   id: string;
+  claudeSessionId: string;
   name: string;
   projectId: string;
   projectName: string;
-  state: 'idle' | 'running';
+  state: 'suspended' | 'running';
   createdAt: string;
+  lastUsedAt: string;
 }
 
 export class PtySessionManager {
-  private readonly sessions = new Map<string, PtySession>();
+  /** Live PTY processes — keyed by session ID. Only exists while PTY is running. */
+  private readonly live = new Map<string, LivePty>();
+  private readonly disconnectTimers = new Map<string, NodeJS.Timeout>();
   private readonly spawn: PtySessionManagerDeps['spawn'];
   private readonly log: PtySessionManagerDeps['log'];
   private readonly projectRoot: string;
+  private readonly getDb: PtySessionManagerDeps['getDb'];
 
-  constructor({ spawn, log, projectRoot }: PtySessionManagerDeps) {
+  constructor({ spawn, log, projectRoot, getDb }: PtySessionManagerDeps) {
     this.spawn = spawn;
     this.log = log;
     this.projectRoot = projectRoot;
+    this.getDb = getDb;
   }
 
   generateId = (): string => {
@@ -74,215 +72,256 @@ export class PtySessionManager {
     return `${projectName} - ${dd}/${mm}/${yy} - ${hh}:${min}`;
   };
 
-  validateCommand = (input: string): { valid: boolean; error?: string } => {
-    const trimmed = input.trim();
-    if (!trimmed) {
-      return { valid: false, error: 'Empty command' };
-    }
-
-    const command = trimmed.split(/\s+/)[0];
-    const isAllowed = ALLOWED_COMMAND_PREFIXES.some(prefix => command === prefix);
-
-    if (!isAllowed) {
-      return {
-        valid: false,
-        error: `Command "${command}" is not allowed. Allowed commands: ${ALLOWED_COMMAND_PREFIXES.join(', ')}`,
-      };
-    }
-
-    return { valid: true };
-  };
-
-  listSessions = (): SessionInfo[] => {
-    return Array.from(this.sessions.values()).map(s => ({
-      id: s.id,
-      name: s.name,
-      projectId: s.projectId,
-      projectName: s.projectName,
-      state: s.state,
-      createdAt: s.createdAt.toISOString(),
+  /** List all sessions from the database, annotated with live/suspended state. */
+  listSessions = async (): Promise<SessionInfo[]> => {
+    const db = this.getDb();
+    const rows = await db.select().from(terminalSessions).orderBy(desc(terminalSessions.createdAt));
+    return rows.map((r: any) => ({
+      id: r.id,
+      claudeSessionId: r.claudeSessionId,
+      name: r.name,
+      projectId: r.projectId,
+      projectName: r.projectName,
+      state: this.live.has(r.id) ? 'running' as const : 'suspended' as const,
+      createdAt: r.createdAt,
+      lastUsedAt: r.lastUsedAt,
     }));
   };
 
-  getSession = (sessionId: string): PtySession | undefined => {
-    return this.sessions.get(sessionId);
+  /** Check whether a session exists (in DB). */
+  getSession = async (sessionId: string): Promise<SessionInfo | undefined> => {
+    const db = this.getDb();
+    const rows = await db.select().from(terminalSessions).where(eq(terminalSessions.id, sessionId));
+    if (rows.length === 0) return undefined;
+    const r = rows[0];
+    return {
+      id: r.id,
+      claudeSessionId: r.claudeSessionId,
+      name: r.name,
+      projectId: r.projectId,
+      projectName: r.projectName,
+      state: this.live.has(r.id) ? 'running' as const : 'suspended' as const,
+      createdAt: r.createdAt,
+      lastUsedAt: r.lastUsedAt,
+    };
   };
 
-  createSession = (projectId: string, projectName: string, cols: number, rows: number): PtySession => {
+  /** Create a new session — persists to DB and spawns claude with --session-id. */
+  createSession = async (projectId: string, projectName: string, cols: number, rows: number): Promise<SessionInfo> => {
     const id = this.generateId();
+    const claudeSessionId = randomUUID();
     const createdAt = new Date();
     const name = this.formatSessionName(projectName, createdAt);
+    const now = createdAt.toISOString();
 
-    const session: PtySession = {
+    // Persist to DB
+    const db = this.getDb();
+    await db.insert(terminalSessions).values({
       id,
-      name,
+      claudeSessionId,
       projectId,
       projectName,
-      pty: null,
-      state: 'idle',
-      cols,
-      rows,
-      outputBuffer: '',
-      inputBuffer: '',
-      listeners: new Set(),
-      createdAt,
-    };
+      name,
+      createdAt: now,
+      lastUsedAt: now,
+    });
 
-    this.sessions.set(id, session);
-    this.log('PTY', `Created session "${name}" (${id}) for project ${projectId}`);
+    this.log('PTY', `Created session "${name}" (${id}) for project ${projectId}, claude session ${claudeSessionId}`);
 
-    // Show welcome message then auto-launch claude with project context
-    this.emitOutput(session, WELCOME_MSG);
-    this.autoLaunchClaude(session);
+    // Spawn claude with --session-id
+    this.spawnClaude(id, claudeSessionId, projectId, cols, rows, false);
 
-    return session;
+    return {
+      id,
+      claudeSessionId,
+      name,
+      projectId,
+      projectName,
+      state: 'running',
+      createdAt: now,
+      lastUsedAt: now,
+    };
+  };
+
+  /**
+   * Ensure the PTY is running for a session. If suspended (no PTY), resume claude.
+   * Returns true if PTY is (now) running; false if session not found.
+   */
+  ensureRunning = async (sessionId: string, cols: number, rows: number): Promise<boolean> => {
+    if (this.live.has(sessionId)) return true;
+
+    const session = await this.getSession(sessionId);
+    if (!session) return false;
+
+    // Update last used timestamp
+    const db = this.getDb();
+    await db.update(terminalSessions)
+      .set({ lastUsedAt: new Date().toISOString() })
+      .where(eq(terminalSessions.id, sessionId));
+
+    // Resume claude
+    this.spawnClaude(sessionId, session.claudeSessionId, session.projectId, cols, rows, true);
+    return true;
   };
 
   addListener = (sessionId: string, listener: (data: string) => void): void => {
-    const session = this.sessions.get(sessionId);
-    if (session) {
-      session.listeners.add(listener);
+    const entry = this.live.get(sessionId);
+    if (entry) {
+      const pendingTimer = this.disconnectTimers.get(sessionId);
+      if (pendingTimer) {
+        clearTimeout(pendingTimer);
+        this.disconnectTimers.delete(sessionId);
+        this.log('PTY', `Cancelled idle timer for session ${sessionId} — client reconnected`);
+      }
+      entry.listeners.add(listener);
     }
   };
 
   removeListener = (sessionId: string, listener: (data: string) => void): void => {
-    const session = this.sessions.get(sessionId);
-    if (session) {
-      session.listeners.delete(listener);
+    const entry = this.live.get(sessionId);
+    if (entry) {
+      entry.listeners.delete(listener);
+      if (entry.listeners.size === 0 && !this.disconnectTimers.has(sessionId)) {
+        this.log('PTY', `All listeners disconnected from session ${sessionId} — starting ${IDLE_TIMEOUT_MS / 60_000}min idle timer`);
+        const timer = setTimeout(() => {
+          this.disconnectTimers.delete(sessionId);
+          const current = this.live.get(sessionId);
+          if (current && current.listeners.size === 0) {
+            this.log('PTY', `Idle timer fired for session ${sessionId} — killing orphan PTY`);
+            current.pty.kill();
+            this.live.delete(sessionId);
+          }
+        }, IDLE_TIMEOUT_MS);
+        this.disconnectTimers.set(sessionId, timer);
+      }
     }
   };
 
   writeToSession = (sessionId: string, data: string): void => {
-    const session = this.sessions.get(sessionId);
-    if (!session) return;
-
-    if (session.state === 'running' && session.pty) {
-      // Forward raw input to the running command
-      session.pty.write(data);
-    } else if (session.state === 'idle') {
-      // Handle command-line input
-      this.handleIdleInput(session, data);
+    const entry = this.live.get(sessionId);
+    if (entry) {
+      entry.pty.write(data);
     }
   };
 
   resizeSession = (sessionId: string, cols: number, rows: number): void => {
-    const session = this.sessions.get(sessionId);
-    if (session) {
-      session.cols = cols;
-      session.rows = rows;
-      if (session.pty) {
-        session.pty.resize(cols, rows);
-      }
+    const entry = this.live.get(sessionId);
+    if (entry) {
+      entry.cols = cols;
+      entry.rows = rows;
+      entry.pty.resize(cols, rows);
     }
   };
 
-  destroySession = (sessionId: string): void => {
-    const session = this.sessions.get(sessionId);
-    if (session) {
-      if (session.pty) {
-        session.pty.kill();
-      }
-      this.sessions.delete(sessionId);
-      this.log('PTY', `Destroyed session "${session.name}" (${sessionId})`);
+  /** Destroy a session — kills PTY and removes from DB permanently. */
+  destroySession = async (sessionId: string): Promise<void> => {
+    // Clear any pending idle timer
+    const pendingTimer = this.disconnectTimers.get(sessionId);
+    if (pendingTimer) {
+      clearTimeout(pendingTimer);
+      this.disconnectTimers.delete(sessionId);
     }
+
+    // Kill live PTY if running
+    const entry = this.live.get(sessionId);
+    if (entry) {
+      entry.pty.kill();
+      this.live.delete(sessionId);
+    }
+
+    // Remove from DB
+    const db = this.getDb();
+    await db.delete(terminalSessions).where(eq(terminalSessions.id, sessionId));
+    this.log('PTY', `Destroyed session ${sessionId}`);
   };
 
-  destroyAll = (): void => {
-    for (const [sessionId] of this.sessions) {
-      this.destroySession(sessionId);
+  /** Kill all live PTY processes (on server shutdown). Does NOT remove DB records. */
+  destroyAllPty = (): void => {
+    for (const timer of this.disconnectTimers.values()) {
+      clearTimeout(timer);
+    }
+    this.disconnectTimers.clear();
+    for (const [id, entry] of this.live) {
+      entry.pty.kill();
+      this.log('PTY', `Killed PTY for session ${id} (shutdown)`);
     }
+    this.live.clear();
   };
 
-  getBufferedOutput = (sessionId: string): string => {
-    const session = this.sessions.get(sessionId);
-    return session?.outputBuffer ?? '';
+  getDisconnectTimerCount = (): number => {
+    return this.disconnectTimers.size;
   };
 
-  private handleIdleInput = (session: PtySession, data: string): void => {
-    for (const char of data) {
-      if (char === '\r' || char === '\n') {
-        // Enter pressed — validate and execute
-        this.emitOutput(session, '\r\n');
-        const command = session.inputBuffer.trim();
-        session.inputBuffer = '';
-
-        if (!command) {
-          this.emitOutput(session, PROMPT);
-          return;
-        }
-
-        const validation = this.validateCommand(command);
-        if (!validation.valid) {
-          this.emitOutput(session, `\x1b[31m${validation.error}\x1b[0m\r\n`);
-          this.emitOutput(session, PROMPT);
-          return;
-        }
-
-        this.spawnCommand(session, command);
-      } else if (char === '\x7f' || char === '\b') {
-        // Backspace
-        if (session.inputBuffer.length > 0) {
-          session.inputBuffer = session.inputBuffer.slice(0, -1);
-          this.emitOutput(session, '\b \b');
-        }
-      } else if (char === '\x03') {
-        // Ctrl+C — clear input
-        session.inputBuffer = '';
-        this.emitOutput(session, '^C\r\n');
-        this.emitOutput(session, PROMPT);
-      } else if (char >= ' ') {
-        // Printable character — echo and buffer
-        session.inputBuffer += char;
-        this.emitOutput(session, char);
-      }
-    }
+  hasDisconnectTimer = (sessionId: string): boolean => {
+    return this.disconnectTimers.has(sessionId);
   };
 
-  private autoLaunchClaude = (session: PtySession): void => {
-    const projectContext = `We are working on project-id ${session.projectId}`;
-    const args = ['--dangerously-skip-permissions', '--append-system-prompt', projectContext];
+  getBufferedOutput = (sessionId: string): string => {
+    return this.live.get(sessionId)?.outputBuffer ?? '';
+  };
 
-    this.log('PTY', `Auto-launching claude for session ${session.id} (project ${session.projectId})`);
+  isLive = (sessionId: string): boolean => {
+    return this.live.has(sessionId);
+  };
+
+  // --- Private helpers ---
+
+  private spawnClaude = (
+    sessionId: string,
+    claudeSessionId: string,
+    projectId: string,
+    cols: number,
+    rows: number,
+    resume: boolean,
+  ): void => {
+    const projectContext = `We are working on project-id ${projectId}`;
+    const args: string[] = ['--dangerously-skip-permissions', '--append-system-prompt', projectContext];
+
+    if (resume) {
+      args.push('--resume', claudeSessionId);
+      this.log('PTY', `Resuming claude session ${claudeSessionId} for ${sessionId}`);
+    } else {
+      args.push('--session-id', claudeSessionId);
+      this.log('PTY', `Launching new claude session ${claudeSessionId} for ${sessionId}`);
+    }
 
     let spawnedPty: IPty;
     try {
       spawnedPty = this.spawn('claude', args, {
         name: 'xterm-256color',
-        cols: session.cols,
-        rows: session.rows,
+        cols,
+        rows,
         cwd: this.projectRoot,
         env: this.buildEnv(),
       });
     } catch (err) {
       const msg = err instanceof Error ? err.message : String(err);
-      this.emitOutput(session, `\x1b[31mFailed to auto-launch claude: ${msg}\x1b[0m\r\n`);
-      this.log('PTY', `Auto-launch failed for session ${session.id}: ${msg}`);
-      this.emitOutput(session, PROMPT);
+      this.log('PTY', `Failed to spawn claude for session ${sessionId}: ${msg}`);
       return;
     }
 
-    session.pty = spawnedPty;
-    session.state = 'running';
+    const entry: LivePty = {
+      pty: spawnedPty,
+      outputBuffer: '',
+      listeners: new Set(),
+      cols,
+      rows,
+    };
+    this.live.set(sessionId, entry);
 
     spawnedPty.onData((data: string) => {
-      this.emitOutput(session, data);
+      this.emitOutput(entry, data);
     });
 
     spawnedPty.onExit(({ exitCode }) => {
-      this.log('PTY', `Auto-launched claude exited with code ${exitCode} for session ${session.id}`);
-      session.pty = null;
-      session.state = 'idle';
-      this.emitOutput(session, '\r\n');
-      this.emitOutput(session, PROMPT);
+      this.log('PTY', `Claude exited with code ${exitCode} for session ${sessionId}`);
+      this.live.delete(sessionId);
     });
   };
 
   private buildEnv = (): Record<string, string> => {
     const env = { ...process.env } as Record<string, string>;
     const home = env.HOME || '/root';
-    // Ensure common user-local binary directories are on PATH so that
-    // tools installed under the user profile (e.g. claude via npm -g)
-    // are found even when the server wasn't started from a login shell.
     const extraPaths = [
       `${home}/.local/bin`,
       `${home}/.npm-global/bin`,
@@ -298,56 +337,12 @@ export class PtySessionManager {
     return env;
   };
 
-  private spawnCommand = (session: PtySession, command: string): void => {
-    const parts = command.trim().split(/\s+/);
-    const cmd = parts[0];
-    const args = parts.slice(1);
-
-    this.log('PTY', `Executing validated command "${command}" for session ${session.id}`);
-
-    let spawnedPty: IPty;
-    try {
-      spawnedPty = this.spawn(cmd, args, {
-        name: 'xterm-256color',
-        cols: session.cols,
-        rows: session.rows,
-        cwd: this.projectRoot,
-        env: this.buildEnv(),
-      });
-    } catch (err) {
-      const msg = err instanceof Error ? err.message : String(err);
-      const env = this.buildEnv();
-      this.emitOutput(session, `\x1b[31mFailed to execute: ${msg}\x1b[0m\r\n`);
-      this.emitOutput(session, `\x1b[90mPATH: ${env.PATH}\x1b[0m\r\n`);
-      this.emitOutput(session, `\x1b[90mHOME: ${env.HOME || '(unset)'}\x1b[0m\r\n`);
-      this.emitOutput(session, `\x1b[90mSHELL: ${env.SHELL || '(unset)'}\x1b[0m\r\n`);
-      this.log('PTY', `Spawn failed for "${command}": ${msg} | PATH=${env.PATH}`);
-      this.emitOutput(session, PROMPT);
-      return;
-    }
-
-    session.pty = spawnedPty;
-    session.state = 'running';
-
-    spawnedPty.onData((data: string) => {
-      this.emitOutput(session, data);
-    });
-
-    spawnedPty.onExit(({ exitCode }) => {
-      this.log('PTY', `Command exited with code ${exitCode} for session ${session.id}`);
-      session.pty = null;
-      session.state = 'idle';
-      this.emitOutput(session, '\r\n');
-      this.emitOutput(session, PROMPT);
-    });
-  };
-
-  private emitOutput = (session: PtySession, data: string): void => {
-    session.outputBuffer += data;
-    if (session.outputBuffer.length > OUTPUT_BUFFER_MAX) {
-      session.outputBuffer = session.outputBuffer.slice(-OUTPUT_BUFFER_MAX);
+  private emitOutput = (entry: LivePty, data: string): void => {
+    entry.outputBuffer += data;
+    if (entry.outputBuffer.length > OUTPUT_BUFFER_MAX) {
+      entry.outputBuffer = entry.outputBuffer.slice(-OUTPUT_BUFFER_MAX);
     }
-    for (const listener of session.listeners) {
+    for (const listener of entry.listeners) {
       listener(data);
     }
   };