OpenTerminal/src/main/pty.ts

import { t } from '../shared/i18n'
import { spawn, type IPty } from '@lydell/node-pty'
import { randomUUID } from 'crypto'
import { homedir } from 'os'
import { StringDecoder } from 'string_decoder'
import { Ipc, type PtyCreateOptions, type PtyCreateResult } from '../shared/ipc'
import type { SessionOpenOptions, HostKeyPromptEvent, SshConnection } from '../shared/connections'
import { broadcast } from './broadcast'
import { connectSsh, type SshSessionHandle } from './ssh'
import type { HostKeyCheckResult } from './knownHosts'
import { configureSysinfo, registerSysinfoClient, forceStopPolling } from './sysinfo'
import { registerSftpClientProvider, closeSftp } from './sftp'
import { attachZmodem, detachZmodem, feedZmodem, isZmodemActive } from './zmodem'
import { pickAdapter } from './shellIntegration'
import { loadSettings } from './settingsStore'
import { statSync } from 'fs'

// Re-export so the session-layer loopback bundle (tests/*-e2e.mjs) can drive the
// M3 polling engine without importing src/main/sysinfo.ts separately.
export { startPolling, stopPolling, forceStopPolling } from './sysinfo'

/**
 * M5 command-store / log-service hooks (injected once by ipc.ts). Mirrors the
 * sysinfo injection pattern so pty.ts routes PTY data into the session log
 * without importing the store directly. `log` runs on every onData chunk (the
 * logger itself decides whether a session is actively logging), `stop` runs on
 * session close / kill to finalize the log file.
 */
type DataLogger = (id: string, data: string) => void
let dataLogger: DataLogger | undefined
let logStopper: (id: string) => void | undefined

export function registerLogHooks(log: DataLogger, stop: (id: string) => void): void {
  dataLogger = log
  logStopper = stop
}

function safeLog(id: string, data: string): void {
  try {
    dataLogger?.(id, data)
  } catch {
    // never crash the event loop
  }
}

function safeStopLog(id: string): void {
  try {
    logStopper?.(id)
  } catch {
    // never crash the event loop
  }
}

/**
 * Session routing table. A session is either a local pty or an established ssh
 * shell; the generic PTY_* (data plane) channels address both. PTY_DATA /
 * PTY_EXIT broadcasts are identical for both kinds.
 */
type Session = { kind: 'local'; pty: IPty; owner?: number } | { kind: 'ssh'; ssh: SshSessionHandle; owner?: number }

/** The live ssh2 client behind an ssh session, or undefined. */
export function getSshClient(id: string): SshSessionHandle['client'] | undefined {
  const session = sessions.get(id)
  return session?.kind === 'ssh' ? session.ssh.client : undefined
}

const sessions = new Map<string, Session>()

/**
 * Per-session replay of recent output (capped). Late subscribers — a terminal
 * view mounting after the shell already printed its banner, or a panel rebound
 * by template apply — read this instead of losing the head of the stream.
 */
const REPLAY_CAP = 64 * 1024
const replayBuffers = new Map<string, string>()

/**
 * Per-session UTF-8 decoder for the ssh data plane. ssh2 hands out raw TCP
 * chunks, so a multi-byte character split across a chunk boundary must be
 * held over by the decoder instead of decoding each chunk alone (which turns
 * the split char into U+FFFD — systematic for CJK output).
 */
const sshDecoders = new Map<string, StringDecoder>()

function appendReplay(id: string, data: string): void {
  const prev = replayBuffers.get(id) ?? ''
  const combined = prev + data
  let next = combined.length > REPLAY_CAP ? combined.slice(combined.length - REPLAY_CAP) : combined
  if (next !== combined) {
    // The cut may have landed inside an escape sequence or a surrogate pair.
    // Resync at the next ESC (which starts a complete sequence) and drop a
    // leading lone low surrogate so xterm neither swallows later output nor
    // renders a replacement char.
    const esc = next.indexOf('\x1b')
    if (esc > 0 && esc < 64) next = next.slice(esc)
    const code = next.charCodeAt(0)
    if (code >= 0xdc00 && code <= 0xdfff) next = next.slice(1)
  }
  replayBuffers.set(id, next)
}

/** Recent output of a session ('' when unknown). */
export function getSessionReplay(id: string): string {
  return replayBuffers.get(id) ?? ''
}

/**
 * Dependencies injected once by ipc.ts (configureSessionRuntime) so pty.ts
 * stays free of store / known-hosts imports and the ssh service can remain
 * decoupled from Electron.
 */
export interface SessionRuntimeDeps {
  /** resolve a bookmark by id (throws a Chinese message when missing) */
  getConnection(connectionId: string): SshConnection
  /** decrypt a stored secret for a connection */
  getSecret(conn: SshConnection, field: 'password' | 'keyContent' | 'passphrase'): string | undefined
  /** mark a bookmark as recently connected */
  touch(connectionId: string): void
  /** host key pinning: check against the store, record an accepted key */
  knownHosts: {
    check(host: string, port: number, key: Buffer): HostKeyCheckResult
    accept(host: string, port: number, key: Buffer, fingerprint: string): void
  }
  /** ask the renderer to decide an unknown / changed host key */
  promptHostKey(prompt: HostKeyPromptEvent): void
  broadcast(channel: string, ...args: unknown[]): void
}

let runtimeDeps: SessionRuntimeDeps | undefined

/** Wire the real stores + renderer prompt. Called once during app startup. */
export function configureSessionRuntime(deps: SessionRuntimeDeps): void {
  runtimeDeps = deps
  configureSysinfo({
    broadcast: (channel, ...args) => deps.broadcast(channel, ...args)
  })
  registerSysinfoClient((id) => {
    const session = sessions.get(id)
    if (session?.kind === 'ssh') return session.ssh.client
    return undefined
  })
  registerSftpClientProvider((id) => {
    const session = sessions.get(id)
    if (session?.kind === 'ssh') return session.ssh.client
    return undefined
  })
}

function existingDir(candidate: string | undefined): string | null {
  if (!candidate || !candidate.trim()) return null
  try {
    return statSync(candidate).isDirectory() ? candidate : null
  } catch {
    return null
  }
}

function defaultShell(): string {
  switch (process.platform) {
    case 'win32':
      return 'powershell.exe'
    case 'darwin':
      return process.env.SHELL || '/bin/zsh'
    default:
      return process.env.SHELL || '/bin/bash'
  }
}

export function createPty(opts: PtyCreateOptions = {}, owner?: number): PtyCreateResult {
  const id = randomUUID()
  const shell = opts.shell ?? defaultShell()
  // A remembered directory can be gone by the next launch (renamed, unmounted,
  // another machine's path after syncing settings) — fall back to home rather
  // than failing the spawn.
  const cwd = existingDir(opts.cwd) ?? homedir()
  // Advertise color capability + terminal identity so TUIs (e.g. kimi CLI)
  // use their full-color theme instead of the degraded white/amber fallback.
  const env = {
    ...process.env,
    TERM: 'xterm-256color',
    COLORTERM: 'truecolor',
    TERM_PROGRAM: 'OpenTerminal',
    ...opts.env
  } as Record<string, string>
  // NO_COLOR (no-color.org) leaks in from whatever launched the app — an agent
  // CLI, an IDE, a CI shell. Every library that honors it (chalk and friends)
  // then prints uncolored output, which reads as "this terminal is broken".
  // The pty we hand out does truecolor, so the opt-out is dropped. Same for the
  // explicit FORCE_COLOR=0 form of the same request.
  delete env.NO_COLOR
  if (env.FORCE_COLOR === '0') delete env.FORCE_COLOR
  // Electron's own plumbing must not leak into the user's shell either:
  // ELECTRON_RUN_AS_NODE would turn this app's own exe into plain node, and
  // NODE_OPTIONS would be applied to every node process started from here.
  // OT_UPDATE_TOKEN is the updater credential — the updater is the only thing
  // that has a reason to see it, not every shell the user opens.
  delete env.ELECTRON_RUN_AS_NODE
  delete env.NODE_OPTIONS
  delete env.OT_UPDATE_TOKEN

  // Shell integration (opt-in): let the shell announce its own cwd over OSC 7 so
  // the remembered path is exact instead of inferred from typed `cd` commands.
  // Every platform/shell difference lives in the adapter table, not here.
  const args = [...(opts.shellArgs ?? [])]
  if (opts.shellArgs === undefined && loadSettings().system.shellIntegration) {
    const adapter = pickAdapter(shell)
    if (adapter) {
      args.push(...adapter.args(shell))
      Object.assign(env, adapter.env?.() ?? {})
    }
  }

  const pty = spawn(shell, args, { name: 'xterm-256color', cols: 80, rows: 24, cwd, env })
  sessions.set(id, { kind: 'local', pty, owner })
  replayBuffers.set(id, '')

  pty.onData((data) => {
    try {
      appendReplay(id, data)
      safeLog(id, data)
      broadcast(Ipc.PTY_DATA, { id, data })
    } catch {
      // never crash the event loop
    }
  })

  pty.onExit(({ exitCode }) => {
    try {
      sessions.delete(id)
      // The session is gone: drop its replay buffer too (killPty was the only
      // path that did, so naturally-exiting shells leaked up to 64KB each).
      replayBuffers.delete(id)
      safeStopLog(id)
      broadcast(Ipc.PTY_EXIT, { id, exitCode })
    } catch {
      // never crash the event loop
    }
  })

  return { id, shell, cwd }
}

/**
 * Open a session. `local` reuses createPty; `ssh` goes through the ssh service
 * and resolves only once the shell stream is ready to stream data.
 */
export async function openSession(opts: SessionOpenOptions, owner?: number): Promise<{ id: string }> {
  if (opts.kind === 'local') {
    return { id: createPty(undefined, owner).id }
  }

  const deps = runtimeDeps
  if (!deps) {
    throw new Error(t('main.pty.serviceNotReady'))
  }
  if (!opts.connectionId) {
    throw new Error(t('main.pty.missingConnectionId'))
  }

  const conn = deps.getConnection(opts.connectionId)
  const handle = await connectSsh(conn, opts.secretOverride, {
    connections: {
      getSecret: (c, field) => deps.getSecret(c, field),
      touch: (id: string) => {
        // Successful connect: record lastConnectedAt on the bookmark.
        try {
          deps.touch(id)
        } catch {
          // store write failure must not break the session
        }
      }
    },
    knownHosts: deps.knownHosts,
    broadcast: deps.broadcast,
    promptHostKey: deps.promptHostKey
  })
  sessions.set(handle.id, { kind: 'ssh', ssh: handle, owner })
  sshDecoders.set(handle.id, new StringDecoder('utf8'))

  // ZMODEM (M6): attach the in-process engine to the raw ssh stream. Only ssh
  // sessions are supported -- node-pty/Windows ConPTY hands out UTF-8 strings
  // only and would corrupt binary frames. The sentry inspects every data chunk:
  // non-ZMODEM traffic is routed back through toTerminal below; once a ZMODEM
  // header is spotted the transfer suppresses the terminal data plane (and
  // writePty drops user keystrokes while isZmodemActive is true).
  attachZmodem(handle.id, {
    broadcast: (channel, ...args) => deps.broadcast(channel, ...args),
    toTerminal: (id, data) => {
      const text = sshDecoders.get(id)?.write(data) ?? data.toString('utf8')
      if (!text) return // the chunk was entirely a partial multi-byte char
      appendReplay(id, text)
      safeLog(id, text)
      deps.broadcast(Ipc.PTY_DATA, { id, data: text })
    },
    writeStream: (id, data) => {
      const s = sessions.get(id)
      if (s?.kind === 'ssh') {
        try {
          s.ssh.stream.write(data)
        } catch {
          // stream may already be dead
        }
      }
    }
  })

  handle.stream.on('data', (data: Buffer) => {
    // Route through the ZMODEM sentry. It forwards non-ZMODEM bytes to
    // deps.toTerminal and absorbs the transfer untouched; while a transfer is
    // active the engine suppresses the normal data plane entirely.
    feedZmodem(handle.id, data)
  })

  handle.stream.on('exit', (code: number | undefined) => {
    // ssh2 reports the remote command's real exit status here, before 'close'.
    if (typeof code === 'number') handle.exitCode = code
  })

  // One teardown for both terminal events. ssh2 emits 'close' after an 'error'
  // (and killSession closes the stream directly), so the path must be
  // idempotent: the flag guarantees PTY_EXIT is broadcast exactly once and the
  // polling / SFTP / ZMODEM / log cleanups run at most once per session.
  let sshClosed = false
  const teardownSshStream = (exitCode: number): void => {
    if (sshClosed) return
    sshClosed = true
    try {
      // Session is gone: no shared poll may survive any remaining panel refs.
      forceStopPolling(handle.id)
      closeSftp(handle.id)
      detachZmodem(handle.id)
      safeStopLog(handle.id)
      sessions.delete(handle.id)
      replayBuffers.delete(handle.id)
      sshDecoders.delete(handle.id)
      deps.broadcast(Ipc.PTY_EXIT, { id: handle.id, exitCode })
    } catch {
      // never crash the event loop
    }
  }

  handle.stream.on('close', () => {
    teardownSshStream(handle.exitCode)
  })

  handle.stream.on('error', (err: Error) => {
    // 'close' follows an 'error', but rely on the idempotent teardown instead
    // of waiting for it: a dead stream must release its session slot at once.
    try {
      console.warn(`[pty] ssh stream error (${handle.id}): ${err.message}`)
      if (handle.exitCode === 0) handle.exitCode = 1
      teardownSshStream(handle.exitCode)
    } catch {
      // never crash the event loop
    }
  })

  return { id: handle.id }
}

export function writePty(id: string, data: string): void {
  // While a ZMODEM transfer is active the stream carries binary frames; user
  // keystrokes must be dropped (they would corrupt the transfer).
  if (isZmodemActive(id)) return
  const session = sessions.get(id)
  if (!session) return
  try {
    if (session.kind === 'local') session.pty.write(data)
    else session.ssh.stream.write(data)
  } catch {
    // session may already be dead
  }
}

export function resizePty(id: string, cols: number, rows: number): void {
  const session = sessions.get(id)
  if (!session) return
  try {
    if (session.kind === 'local') session.pty.resize(cols, rows)
    // ssh2 Channel.setWindow(rows, cols, height, width)
    else session.ssh.stream.setWindow(rows, cols, 0, 0)
  } catch {
    // session may already be dead
  }
}

function killSession(id: string): void {
  // Forced: the session is being destroyed, so the shared poll must stop even
  // if split panels still hold references.
  forceStopPolling(id)
  closeSftp(id)
  detachZmodem(id)
  safeStopLog(id)
  replayBuffers.delete(id)
  sshDecoders.delete(id)
  const session = sessions.get(id)
  if (!session) return
  if (session.kind === 'ssh') {
    try {
      session.ssh.stream.close()
    } catch {
      // best effort
    }
    try {
      session.ssh.client.end()
    } catch {
      // best effort
    }
  } else {
    try {
      session.pty.kill()
    } catch {
      // best effort
    }
  }
  sessions.delete(id)
}

export function killPty(id: string): void {
  killSession(id)
}

/** Kill every session owned by a renderer that crashed or was destroyed
 *  without running its own cleanup — normal close/reload kills its own
 *  sessions via beforeunload before we ever get here. */
export function killPtysByOwner(owner: number): void {
  for (const id of [...sessions.keys()]) {
    if (sessions.get(id)?.owner !== owner) continue
    killSession(id)
  }
}

export function killAllPtys(): void {
  for (const id of sessions.keys()) {
    forceStopPolling(id)
    closeSftp(id)
    detachZmodem(id)
    safeStopLog(id)
    const session = sessions.get(id)
    if (!session) continue
    if (session.kind === 'ssh') {
      try {
        session.ssh.stream.close()
      } catch {
        // best effort
      }
      try {
        session.ssh.client.end()
      } catch {
        // best effort
      }
    } else {
      try {
        session.pty.kill()
      } catch {
        // best effort
      }
    }
  }
  sessions.clear()
  sshDecoders.clear()
}