wind-agent-cli 1.37.2

package/server/modules/websocket/services/shell-websocket.service.ts

@@ -0,0 +1,776 @@
+import fs from 'node:fs';
+import crypto from 'node:crypto';
+import os from 'node:os';
+import path from 'node:path';
+
+import pty, { type IPty } from 'node-pty';
+import { WebSocket, type RawData } from 'ws';
+
+import { parseIncomingJsonObject } from '@/shared/utils.js';
+import { terminalReplayStore } from '@/modules/websocket/services/terminal-replay-store.js';
+
+type ShellIncomingMessage = {
+  type?: string;
+  data?: string;
+  cols?: number;
+  rows?: number;
+  projectPath?: string;
+  sessionId?: string;
+  hasSession?: boolean;
+  provider?: string;
+  initialCommand?: string;
+  isPlainShell?: boolean;
+  forceRestart?: boolean;
+  shellMode?: string;
+};
+
+type PtySessionEntry = {
+  pty: IPty;
+  ws: WebSocket | null;
+  timeoutId: NodeJS.Timeout | null;
+  projectPath: string;
+  sessionId: string | null;
+};
+
+const ptySessionsMap = new Map<string, PtySessionEntry>();
+const PTY_SESSION_TIMEOUT = 30 * 60 * 1000;
+const SHELL_URL_PARSE_BUFFER_LIMIT = 32768;
+const ANSI_ESCAPE_SEQUENCE_REGEX = /\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~]|\][^\x07]*(?:\x07|\x1B\\))/g;
+const TRAILING_URL_PUNCTUATION_REGEX = /[)\]}>.,;:!?]+$/;
+const TERMINAL_REPLAY_TRUNCATED_MESSAGE =
+  '\x1b[33m[Earlier terminal output was omitted to limit memory usage]\x1b[0m\r\n';
+
+function stripAnsiSequences(value: string): string {
+  return value.replace(ANSI_ESCAPE_SEQUENCE_REGEX, '');
+}
+
+function normalizeDetectedUrl(url: string): string | null {
+  const cleanedUrl = url.trim().replace(TRAILING_URL_PUNCTUATION_REGEX, '');
+  if (!cleanedUrl) {
+    return null;
+  }
+
+  try {
+    const parsedUrl = new URL(cleanedUrl);
+    if (parsedUrl.protocol !== 'http:' && parsedUrl.protocol !== 'https:') {
+      return null;
+    }
+    return parsedUrl.toString();
+  } catch {
+    return null;
+  }
+}
+
+function extractUrlsFromText(value: string): string[] {
+  const directMatches = value.match(/https?:\/\/[^\s<>"'`\\\x1b\x07]+/gi) ?? [];
+
+  // Terminal width can split a URL across lines, so valid URL characters on
+  // immediately following lines are joined before the URL is validated.
+  const wrappedMatches: string[] = [];
+  const urlContinuationPattern = /^[A-Za-z0-9\-._~:/?#\[\]@!$&'()*+,;=%]+$/;
+  const lines = value.split(/\r?\n/);
+  for (let lineIndex = 0; lineIndex < lines.length; lineIndex += 1) {
+    const line = lines[lineIndex].trim();
+    const startMatch = line.match(/https?:\/\/[^\s<>"'`\\\x1b\x07]+/i);
+    if (!startMatch) {
+      continue;
+    }
+
+    let combinedUrl = startMatch[0];
+    let continuationIndex = lineIndex + 1;
+    while (continuationIndex < lines.length) {
+      const continuation = lines[continuationIndex].trim();
+      if (!continuation || !urlContinuationPattern.test(continuation)) {
+        break;
+      }
+      combinedUrl += continuation;
+      continuationIndex += 1;
+    }
+
+    wrappedMatches.push(combinedUrl);
+  }
+
+  return Array.from(new Set([...directMatches, ...wrappedMatches]));
+}
+
+function shouldAutoOpenUrlFromOutput(value: string): boolean {
+  const normalizedOutput = value.toLowerCase();
+  return (
+    normalizedOutput.includes("browser didn't open") ||
+    normalizedOutput.includes('open this url') ||
+    normalizedOutput.includes('continue in your browser') ||
+    normalizedOutput.includes('press enter to open') ||
+    normalizedOutput.includes('open_url:')
+  );
+}
+
+type ShellWebSocketDependencies = {
+  resolveProviderSessionId: (
+    sessionId: string,
+    provider: string,
+  ) => string | null | undefined;
+  spawnPty?: typeof pty.spawn;
+  platform?: () => NodeJS.Platform;
+  environment?: NodeJS.ProcessEnv;
+};
+
+/**
+ * Reads a string field from untyped payloads and falls back when absent.
+ */
+function readString(value: unknown, fallback = ''): string {
+  return typeof value === 'string' ? value : fallback;
+}
+
+/**
+ * Reads a boolean field from untyped payloads and falls back when absent.
+ */
+function readBoolean(value: unknown, fallback = false): boolean {
+  return typeof value === 'boolean' ? value : fallback;
+}
+
+/**
+ * Reads a finite number field from untyped payloads and falls back when absent.
+ */
+function readNumber(value: unknown, fallback: number): number {
+  return typeof value === 'number' && Number.isFinite(value) ? value : fallback;
+}
+
+/**
+ * Parses incoming websocket shell messages and keeps processing safe when
+ * malformed payloads are received.
+ */
+function parseShellMessage(rawMessage: RawData): ShellIncomingMessage | null {
+  const payload = parseIncomingJsonObject(rawMessage);
+  if (!payload) {
+    return null;
+  }
+
+  return payload as ShellIncomingMessage;
+}
+
+const SAFE_SESSION_ID_PATTERN = /^[a-zA-Z0-9_.\-:]+$/;
+const BASE64_PATTERN = /^(?:[A-Za-z0-9+/]{4})*(?:[A-Za-z0-9+/]{2}==|[A-Za-z0-9+/]{3}=)?$/;
+const MAX_TERMINAL_BINARY_INPUT_BYTES = 1024;
+const HERDR_REDRAW_DELAY_MS = 16;
+const HERDR_PANE_ENVIRONMENT_KEYS = [
+  'HERDR_ENV',
+  'HERDR_PANE_ID',
+  'HERDR_TAB_ID',
+  'HERDR_WORKSPACE_ID',
+] as const;
+
+function resolveResumeSessionId(
+  message: ShellIncomingMessage,
+  dependencies: ShellWebSocketDependencies
+): string {
+  const hasSession = readBoolean(message.hasSession);
+  const sessionId = readString(message.sessionId);
+  const provider = readString(message.provider, 'claude');
+
+  if (!hasSession || !sessionId) {
+    return '';
+  }
+
+  let resumeSessionId: string | null | undefined;
+  try {
+    resumeSessionId = dependencies.resolveProviderSessionId(sessionId, provider);
+  } catch (error) {
+    console.error('Failed to resolve provider session ID:', error);
+    resumeSessionId = undefined;
+  }
+
+  const resolvedSessionId = resumeSessionId === undefined ? sessionId : resumeSessionId;
+  if (!resolvedSessionId || !SAFE_SESSION_ID_PATTERN.test(resolvedSessionId)) {
+    return '';
+  }
+
+  return resolvedSessionId;
+}
+
+/**
+ * Resolves provider command line for plain shell and agent-backed shell modes.
+ */
+function buildShellCommand(
+  message: ShellIncomingMessage,
+  dependencies: ShellWebSocketDependencies
+): string {
+  const hasSession = readBoolean(message.hasSession);
+  const initialCommand = readString(message.initialCommand);
+  const provider = readString(message.provider, 'claude');
+  const resumeSessionId = resolveResumeSessionId(message, dependencies);
+  const platform = dependencies.platform?.() ?? os.platform();
+  const isPlainShell =
+    readBoolean(message.isPlainShell) ||
+    (!!initialCommand && !hasSession) ||
+    provider === 'plain-shell';
+
+  if (isPlainShell) {
+    return initialCommand;
+  }
+
+  if (provider === 'cursor') {
+    if (resumeSessionId) {
+      return `cursor-agent --resume="${resumeSessionId}"`;
+    }
+    return 'cursor-agent';
+  }
+
+  if (provider === 'codex') {
+    const codexCommand = platform === 'win32' ? 'codex.cmd' : 'codex';
+    if (resumeSessionId) {
+      if (platform === 'win32') {
+        return `${codexCommand} resume "${resumeSessionId}"; if ($LASTEXITCODE -ne 0) { ${codexCommand} }`;
+      }
+      return `${codexCommand} resume "${resumeSessionId}" || ${codexCommand}`;
+    }
+    return codexCommand;
+  }
+
+  if (provider === 'opencode') {
+    if (resumeSessionId) {
+      return `opencode --session "${resumeSessionId}"`;
+    }
+    return initialCommand || 'opencode';
+  }
+
+  const command = initialCommand || 'claude';
+  if (resumeSessionId) {
+    if (platform === 'win32') {
+      return `claude --resume "${resumeSessionId}"; if ($LASTEXITCODE -ne 0) { claude }`;
+    }
+    return `claude --resume "${resumeSessionId}" || claude`;
+  }
+  return command;
+}
+
+function readEnvValue(env: NodeJS.ProcessEnv, key: string): string | undefined {
+  const resolvedKey = Object.keys(env).find((envKey) => envKey.toLowerCase() === key.toLowerCase());
+  return resolvedKey ? env[resolvedKey] : undefined;
+}
+
+function getPathEnvKey(env: NodeJS.ProcessEnv): string {
+  return Object.keys(env).find((key) => key.toLowerCase() === 'path') || 'PATH';
+}
+
+function prioritizeUserNpmGlobalBin(env: NodeJS.ProcessEnv): { key: string; value: string | undefined } {
+  const pathKey = getPathEnvKey(env);
+  const currentPath = env[pathKey];
+  if (!currentPath) {
+    return { key: pathKey, value: currentPath };
+  }
+
+  const delimiter = path.delimiter;
+  const pathEntries = currentPath.split(delimiter).filter(Boolean);
+  const npmPrefix = readEnvValue(env, 'npm_config_prefix');
+  const appData = readEnvValue(env, 'APPDATA');
+  const candidates = [
+    npmPrefix || '',
+    npmPrefix ? path.join(npmPrefix, 'bin') : '',
+    appData ? path.join(appData, 'npm') : '',
+    path.join(os.homedir(), 'AppData', 'Roaming', 'npm'),
+    path.join(os.homedir(), '.npm-global', 'bin'),
+  ].filter(Boolean);
+
+  const normalizedPathEntries = pathEntries.map((entry) => os.platform() === 'win32' ? entry.toLowerCase() : entry);
+  const preferredEntries = candidates.filter((candidate, index) => {
+    const normalizedCandidate = os.platform() === 'win32' ? candidate.toLowerCase() : candidate;
+    return (
+      candidates.indexOf(candidate) === index &&
+      normalizedPathEntries.includes(normalizedCandidate)
+    );
+  });
+
+  if (preferredEntries.length === 0) {
+    return { key: pathKey, value: currentPath };
+  }
+
+  const normalizedPreferredEntries = preferredEntries.map((entry) =>
+    os.platform() === 'win32' ? entry.toLowerCase() : entry
+  );
+
+  const value = [
+    ...preferredEntries,
+    ...pathEntries.filter((entry) => {
+      const normalizedEntry = os.platform() === 'win32' ? entry.toLowerCase() : entry;
+      return !normalizedPreferredEntries.includes(normalizedEntry);
+    }),
+  ].join(delimiter);
+
+  return { key: pathKey, value };
+}
+
+function buildShellEnvironment(
+  message: ShellIncomingMessage,
+  dependencies: ShellWebSocketDependencies,
+): NodeJS.ProcessEnv {
+  const baseEnvironment = dependencies.environment ?? process.env;
+  const prioritizedPath = prioritizeUserNpmGlobalBin(baseEnvironment);
+  const environment: NodeJS.ProcessEnv = {
+    ...baseEnvironment,
+    [prioritizedPath.key]: prioritizedPath.value,
+    TERM: 'xterm-256color',
+    COLORTERM: 'truecolor',
+    FORCE_COLOR: '3',
+  };
+
+  if (readString(message.shellMode) === 'herdr') {
+    // The web terminal is a new Herdr client, not a pane nested inside the
+    // Herdr pane that may have launched CloudCLI. Keep HERDR_SOCKET_PATH so it
+    // attaches to the running session, but remove inherited pane identity.
+    HERDR_PANE_ENVIRONMENT_KEYS.forEach((key) => delete environment[key]);
+  }
+
+  return environment;
+}
+
+function decodeTerminalBinaryInput(value: unknown): Buffer | null {
+  const encoded = readString(value);
+  if (!encoded || !BASE64_PATTERN.test(encoded)) {
+    return null;
+  }
+
+  const decoded = Buffer.from(encoded, 'base64');
+  if (decoded.length > MAX_TERMINAL_BINARY_INPUT_BYTES) {
+    return null;
+  }
+
+  return decoded;
+}
+
+function redrawHerdrViewport(
+  sessionKey: string,
+  session: PtySessionEntry,
+  ws: WebSocket,
+  cols: number,
+  rows: number,
+): void {
+  const redrawCols = cols > 2 ? cols - 1 : cols + 1;
+  session.pty.resize(redrawCols, rows);
+  setTimeout(() => {
+    const activeSession = ptySessionsMap.get(sessionKey);
+    if (activeSession === session && activeSession.ws === ws) {
+      activeSession.pty.resize(cols, rows);
+    }
+  }, HERDR_REDRAW_DELAY_MS);
+}
+
+function buildPtySessionKey({
+  projectPath,
+  sessionId,
+  provider,
+  shellMode,
+  initialCommand,
+}: {
+  projectPath: string;
+  sessionId: string | null;
+  provider: string;
+  shellMode: string;
+  initialCommand: string;
+}): string {
+  const identity = JSON.stringify({
+    projectPath: process.platform === 'win32' ? projectPath.toLowerCase() : projectPath,
+    sessionId: sessionId ?? 'default',
+    provider,
+    shellMode,
+    commandHash: crypto.createHash('sha256').update(initialCommand).digest('hex'),
+  });
+  return `pty:${crypto.createHash('sha256').update(identity).digest('hex')}`;
+}
+
+/**
+ * Used by this module's websocket gateway to connect the standalone Shell UI
+ * to a retained PTY while keeping process lifecycle ownership on the server.
+ */
+export function handleShellConnection(
+  ws: WebSocket,
+  dependencies: ShellWebSocketDependencies
+): void {
+  console.log('[INFO] Shell websocket connected');
+
+  let shellProcess: IPty | null = null;
+  let ptySessionKey: string | null = null;
+  let isHerdrConnection = false;
+  let urlDetectionBuffer = '';
+  const announcedAuthUrls = new Set<string>();
+
+  ws.on('message', async (rawMessage) => {
+    try {
+      const data = parseShellMessage(rawMessage);
+      if (!data?.type) {
+        throw new Error('Invalid websocket payload');
+      }
+
+      if (data.type === 'init') {
+        const projectPath = readString(data.projectPath, process.cwd());
+        const sessionId = readString(data.sessionId) || null;
+        const hasSession = readBoolean(data.hasSession);
+        const provider = readString(data.provider, 'claude');
+        const initialCommand = readString(data.initialCommand);
+        const forceRestart = readBoolean(data.forceRestart);
+        const isHerdrMode = readString(data.shellMode) === 'herdr';
+        isHerdrConnection = isHerdrMode;
+        const isPlainShell =
+          readBoolean(data.isPlainShell) ||
+          (!!initialCommand && !hasSession) ||
+          provider === 'plain-shell';
+
+        urlDetectionBuffer = '';
+        announcedAuthUrls.clear();
+
+        const isLoginCommand =
+          !!initialCommand &&
+          (initialCommand.includes('setup-token') ||
+            initialCommand.includes('cursor-agent login') ||
+            initialCommand.includes('auth login'));
+
+        const resolvedProjectPath = path.resolve(projectPath);
+        try {
+          const stats = fs.statSync(resolvedProjectPath);
+          if (!stats.isDirectory()) {
+            throw new Error('Not a directory');
+          }
+        } catch {
+          ws.send(JSON.stringify({ type: 'error', message: 'Invalid project path' }));
+          return;
+        }
+
+        if (sessionId && !SAFE_SESSION_ID_PATTERN.test(sessionId)) {
+          ws.send(JSON.stringify({ type: 'error', message: 'Invalid session ID' }));
+          return;
+        }
+
+        ptySessionKey = buildPtySessionKey({
+          projectPath: resolvedProjectPath,
+          sessionId,
+          provider,
+          shellMode: readString(data.shellMode),
+          initialCommand: isPlainShell ? initialCommand : '',
+        });
+
+        if (isLoginCommand || forceRestart) {
+          const oldSession = ptySessionsMap.get(ptySessionKey);
+          if (oldSession) {
+            if (oldSession.timeoutId) {
+              clearTimeout(oldSession.timeoutId);
+            }
+            terminalReplayStore.delete(ptySessionKey);
+            oldSession.pty.kill();
+            ptySessionsMap.delete(ptySessionKey);
+          }
+        }
+
+        const existingSession =
+          isLoginCommand || forceRestart ? null : ptySessionsMap.get(ptySessionKey);
+        if (existingSession) {
+          shellProcess = existingSession.pty;
+          if (existingSession.timeoutId) {
+            clearTimeout(existingSession.timeoutId);
+            existingSession.timeoutId = null;
+          }
+
+          existingSession.ws = ws;
+          const reconnectCols = readNumber(data.cols, 80);
+          const reconnectRows = readNumber(data.rows, 24);
+
+          if (isHerdrMode) {
+            // Replaying a full-screen TUI's ANSI history reconstructs stale
+            // frames from earlier sizes. Ask Herdr for a fresh frame instead.
+            terminalReplayStore.clear(ptySessionKey);
+            redrawHerdrViewport(
+              ptySessionKey,
+              existingSession,
+              ws,
+              reconnectCols,
+              reconnectRows,
+            );
+            return;
+          }
+
+          ws.send(
+            JSON.stringify({
+              type: 'output',
+              data: '\x1b[36m[Reconnected to existing session]\x1b[0m\r\n',
+            })
+          );
+
+          const replay = terminalReplayStore.snapshot(ptySessionKey);
+          if (replay.truncated) {
+            ws.send(JSON.stringify({ type: 'output', data: TERMINAL_REPLAY_TRUNCATED_MESSAGE }));
+          }
+          if (replay.chunks.length > 0) {
+            replay.chunks.forEach((bufferedData) => {
+              ws.send(
+                JSON.stringify({
+                  type: 'output',
+                  data: bufferedData,
+                })
+              );
+            });
+          }
+
+          existingSession.pty.resize(reconnectCols, reconnectRows);
+          return;
+        }
+
+        const shellCommand = buildShellCommand(data, dependencies);
+        const resumeSessionId = resolveResumeSessionId(data, dependencies);
+        const platform = dependencies.platform?.() ?? os.platform();
+        // Herdr is already an interactive executable. Launching it through
+        // PowerShell can briefly expose a console window before ConPTY attaches.
+        const startsHerdrDirectly = isHerdrMode && platform === 'win32';
+        const shell = startsHerdrDirectly
+          ? 'herdr.exe'
+          : platform === 'win32'
+            ? 'powershell.exe'
+            : 'bash';
+        const shellArgs = startsHerdrDirectly
+          ? []
+          : platform === 'win32'
+            ? ['-Command', shellCommand]
+            : ['-c', shellCommand];
+        const termCols = readNumber(data.cols, 80);
+        const termRows = readNumber(data.rows, 24);
+        const shellEnvironment = buildShellEnvironment(data, dependencies);
+
+        shellProcess = (dependencies.spawnPty ?? pty.spawn)(shell, shellArgs, {
+          name: 'xterm-256color',
+          cols: termCols,
+          rows: termRows,
+          cwd: resolvedProjectPath,
+          env: shellEnvironment,
+        });
+
+        ptySessionsMap.set(ptySessionKey, {
+          pty: shellProcess,
+          ws,
+          timeoutId: null,
+          projectPath,
+          sessionId,
+        });
+
+        shellProcess.onData((chunk) => {
+          if (!ptySessionKey) {
+            return;
+          }
+
+          const session = ptySessionsMap.get(ptySessionKey);
+          if (!session) {
+            return;
+          }
+
+          if (!isHerdrMode) {
+            terminalReplayStore.append(ptySessionKey, chunk);
+          }
+
+          if (session.ws && session.ws.readyState === WebSocket.OPEN) {
+            let outputData = chunk;
+            const cleanChunk = stripAnsiSequences(chunk);
+            urlDetectionBuffer = `${urlDetectionBuffer}${cleanChunk}`.slice(-SHELL_URL_PARSE_BUFFER_LIMIT);
+
+            outputData = outputData.replace(
+              /OPEN_URL:\s*(https?:\/\/[^\s\x1b\x07]+)/g,
+              '[INFO] Opening in browser: $1'
+            );
+
+            const emitAuthUrl = (detectedUrl: string, autoOpen = false) => {
+              const normalizedUrl = normalizeDetectedUrl(detectedUrl);
+              if (!normalizedUrl) {
+                return;
+              }
+
+              const isNewUrl = !announcedAuthUrls.has(normalizedUrl);
+              if (isNewUrl) {
+                announcedAuthUrls.add(normalizedUrl);
+                session.ws?.send(
+                  JSON.stringify({
+                    type: 'auth_url',
+                    url: normalizedUrl,
+                    autoOpen,
+                  })
+                );
+              }
+            };
+
+            const normalizedDetectedUrls = extractUrlsFromText(urlDetectionBuffer)
+              .map((url) => normalizeDetectedUrl(url))
+              .filter((url): url is string => Boolean(url));
+
+            const dedupedDetectedUrls = Array.from(new Set(normalizedDetectedUrls)).filter(
+              (url, _, urls) =>
+                !urls.some((otherUrl) => otherUrl !== url && otherUrl.startsWith(url))
+            );
+
+            dedupedDetectedUrls.forEach((url) => emitAuthUrl(url, false));
+
+            if (
+              shouldAutoOpenUrlFromOutput(cleanChunk) &&
+              dedupedDetectedUrls.length > 0
+            ) {
+              const bestUrl = dedupedDetectedUrls.reduce((longest, current) =>
+                current.length > longest.length ? current : longest
+              );
+              emitAuthUrl(bestUrl, true);
+            }
+
+            session.ws.send(
+              JSON.stringify({
+                type: 'output',
+                data: outputData,
+              })
+            );
+          }
+        });
+
+        shellProcess.onExit((exitCode) => {
+          if (!ptySessionKey) {
+            return;
+          }
+
+          const session = ptySessionsMap.get(ptySessionKey);
+          if (session && session.pty !== shellProcess) {
+            return;
+          }
+
+          if (session && session.ws && session.ws.readyState === WebSocket.OPEN) {
+            session.ws.send(
+              JSON.stringify({
+                type: 'output',
+                data: `\r\n\x1b[33mProcess exited with code ${exitCode.exitCode}${
+                  exitCode.signal != null ? ` (${exitCode.signal})` : ''
+                }\x1b[0m\r\n`,
+              })
+            );
+          }
+
+          if (session?.timeoutId) {
+            clearTimeout(session.timeoutId);
+          }
+
+          terminalReplayStore.delete(ptySessionKey);
+          ptySessionsMap.delete(ptySessionKey);
+          shellProcess = null;
+        });
+
+        let welcomeMsg = `\x1b[36mStarting terminal in: ${projectPath}\x1b[0m\r\n`;
+        if (!isPlainShell) {
+          const providerName =
+            provider === 'cursor'
+              ? 'Cursor'
+              : provider === 'codex'
+                ? 'Codex'
+                : provider === 'opencode'
+                    ? 'OpenCode'
+                  : 'Claude';
+          welcomeMsg = hasSession && resumeSessionId
+            ? `\x1b[36mResuming ${providerName} session ${resumeSessionId} in: ${projectPath}\x1b[0m\r\n`
+            : `\x1b[36mStarting new ${providerName} session in: ${projectPath}\x1b[0m\r\n`;
+        }
+
+        ws.send(
+          JSON.stringify({
+            type: 'output',
+            data: welcomeMsg,
+          })
+        );
+        return;
+      }
+
+      if (data.type === 'input') {
+        const activeSession = ptySessionKey ? ptySessionsMap.get(ptySessionKey) : null;
+        if (shellProcess && activeSession?.pty === shellProcess && activeSession.ws === ws) {
+          shellProcess.write(readString(data.data));
+        }
+        return;
+      }
+
+      if (data.type === 'input_binary') {
+        const activeSession = ptySessionKey ? ptySessionsMap.get(ptySessionKey) : null;
+        const binaryInput = decodeTerminalBinaryInput(data.data);
+        if (
+          binaryInput &&
+          shellProcess &&
+          activeSession?.pty === shellProcess &&
+          activeSession.ws === ws
+        ) {
+          shellProcess.write(binaryInput);
+        }
+        return;
+      }
+
+      if (data.type === 'resize') {
+        const activeSession = ptySessionKey ? ptySessionsMap.get(ptySessionKey) : null;
+        if (shellProcess && activeSession?.pty === shellProcess && activeSession.ws === ws) {
+          shellProcess.resize(readNumber(data.cols, 80), readNumber(data.rows, 24));
+        }
+        return;
+      }
+
+      if (data.type === 'viewport_claim') {
+        const activeSession = ptySessionKey ? ptySessionsMap.get(ptySessionKey) : null;
+        if (
+          isHerdrConnection &&
+          ptySessionKey &&
+          shellProcess &&
+          activeSession?.pty === shellProcess &&
+          activeSession.ws === ws
+        ) {
+          redrawHerdrViewport(
+            ptySessionKey,
+            activeSession,
+            ws,
+            readNumber(data.cols, 80),
+            readNumber(data.rows, 24),
+          );
+        }
+      }
+    } catch (error) {
+      const message = error instanceof Error ? error.message : String(error);
+      console.error('[ERROR] Shell WebSocket error:', message);
+      if (ws.readyState === WebSocket.OPEN) {
+        ws.send(
+          JSON.stringify({
+            type: 'output',
+            data: `\r\n\x1b[31mError: ${message}\x1b[0m\r\n`,
+          })
+        );
+      }
+    }
+  });
+
+  ws.on('close', () => {
+    if (!ptySessionKey) {
+      return;
+    }
+
+    const session = ptySessionsMap.get(ptySessionKey);
+    if (!session) {
+      return;
+    }
+
+    // Mobile networks can deliver an old socket's close after its replacement
+    // has attached. Only the socket that currently owns the PTY may detach it.
+    if (session.ws !== ws) {
+      return;
+    }
+
+    session.ws = null;
+    if (session.timeoutId) {
+      clearTimeout(session.timeoutId);
+    }
+    session.timeoutId = setTimeout(() => {
+      // A reconnect may win just as this timer becomes runnable. Re-check the
+      // active socket so a queued cleanup can never kill a reattached PTY.
+      if (ptySessionsMap.get(ptySessionKey as string) !== session || session.ws !== null) {
+        return;
+      }
+
+      session.pty.kill();
+      terminalReplayStore.delete(ptySessionKey as string);
+      ptySessionsMap.delete(ptySessionKey as string);
+    }, PTY_SESSION_TIMEOUT);
+  });
+
+  ws.on('error', (error) => {
+    console.error('[ERROR] Shell WebSocket error:', error);
+  });
+}