/** * Agent 的 SSH 执行层 * * 与终端 socket 的区别:这里走 exec 通道而不是 PTY。 * PTY 里混杂 ANSI 转义、提示符、用户正在编辑的半行命令,喂给模型全是噪声, * 且无法可靠判断命令何时结束、退出码是多少。exec 通道能拿到干净的 * stdout / stderr / exitCode。 * * 连接复用:一个 turn 内模型可能连续执行十几条命令,每条都重连太慢。 * 按 hostId 池化,空闲超时后回收。 */ import ssh2Module from 'ssh2' import { getConnectionOptions, handleProxyAndJumpHostConnection } from '../socket/terminal.js' import { HostListDB } from '../utils/db-class.js' const { Client: SSHClient } = ssh2Module const hostListDB = new HostListDB().getInstance() // 空闲连接保留时长 const IDLE_TTL_MS = 5 * 60 * 1000 // 单条命令默认超时 export const DEFAULT_TIMEOUT_MS = 60 * 1000 export const MAX_TIMEOUT_MS = 30 * 60 * 1000 // 单条命令最多收集的输出字节数,防止 `cat 大文件` 打爆内存 const MAX_COLLECT_BYTES = 2 * 1024 * 1024 /** hostId -> { client, jumpClients, refCount, idleTimer, connecting } */ const pool = new Map() function clearIdleTimer(entry) { if (entry.idleTimer) { clearTimeout(entry.idleTimer) entry.idleTimer = null } } function scheduleIdleClose(hostId) { const entry = pool.get(hostId) if (!entry) return clearIdleTimer(entry) entry.idleTimer = setTimeout(() => { if (entry.refCount > 0) return closeEntry(hostId, entry) }, IDLE_TTL_MS) } function closeEntry(hostId, entry) { pool.delete(hostId) clearIdleTimer(entry) try { entry.client?.end() } catch { // 忽略关闭异常 } entry.jumpClients?.forEach((client) => { try { client?.end() } catch { // 忽略 } }) } /** 建立一条新的 SSH 连接(含代理与跳板机) */ async function connect(hostId) { const hostInfo = await hostListDB.findOneAsync({ _id: hostId }) if (!hostInfo) throw new Error(`未找到主机: ${ hostId }`) const { authInfo } = await getConnectionOptions(hostId) let jumpClients = [] const proxyResult = await handleProxyAndJumpHostConnection({ hostInfo, targetConnectionOptions: authInfo, socket: null, logPrefix: 'AI Agent ' }) jumpClients = proxyResult?.jumpSshClients || [] const client = await new Promise((resolve, reject) => { const sshClient = new SSHClient() const onError = (err) => { jumpClients.forEach((item) => item?.end()) reject(new Error(`SSH 连接失败: ${ err.message }`)) } sshClient .on('ready', () => resolve(sshClient)) .on('error', onError) .on('keyboard-interactive', (name, instructions, lang, prompts, finish) => { finish([authInfo[hostInfo.authType]]) }) .connect({ tryKeyboard: true, ...authInfo }) }) // 连接被对端关闭时从池里摘掉,避免下次拿到死连接 client.on('close', () => { const entry = pool.get(hostId) if (entry?.client === client) closeEntry(hostId, entry) }) return { client, jumpClients, hostInfo } } /** * 取得一条可用连接。调用方必须在用完后调用 release()。 */ async function acquire(hostId) { let entry = pool.get(hostId) if (entry?.connecting) { await entry.connecting entry = pool.get(hostId) } if (!entry || !entry.client) { const placeholder = { client: null, jumpClients: [], refCount: 0, idleTimer: null, connecting: null } placeholder.connecting = connect(hostId) .then(({ client, jumpClients, hostInfo }) => { placeholder.client = client placeholder.jumpClients = jumpClients placeholder.hostInfo = hostInfo placeholder.connecting = null return placeholder }) .catch((error) => { pool.delete(hostId) throw error }) pool.set(hostId, placeholder) await placeholder.connecting entry = pool.get(hostId) } if (!entry?.client) throw new Error(`SSH 连接不可用: ${ hostId }`) entry.refCount += 1 clearIdleTimer(entry) return { client: entry.client, hostInfo: entry.hostInfo, release: () => { entry.refCount = Math.max(0, entry.refCount - 1) if (entry.refCount === 0) scheduleIdleClose(hostId) } } } /** * 主动断开某台主机的 agent 连接。 * * 必须尊重 refCount:连接是按 hostId 池化的,同一台主机可能正被另一个 * 会话用着。直接 close 会把别人正在跑的命令掐断。还有人在用时改为 * 重排空闲回收,等最后一个使用者释放后自然关闭。 */ export function disconnect(hostId) { const entry = pool.get(hostId) if (!entry) return if (entry.refCount > 0) { scheduleIdleClose(hostId) return } closeEntry(hostId, entry) } /** 断开全部 agent 连接(进程退出或用户结束会话时调用) */ export function disconnectAll() { for (const [hostId, entry] of pool.entries()) closeEntry(hostId, entry) } /** * 包装命令:注入非交互环境变量,可选切换工作目录。 * * 审计日志同时记录原始命令与包装后的命令,避免"用户批准的是 A、 * 实际执行的是 B"这种说不清的情况。 */ /** * 单引号转义。 * * ⚠️ 不要用 JSON.stringify 代替:它产出的是双引号字符串,而双引号里 * `$(...)`、反引号、`$VAR` 仍会被 shell 展开。单引号才是字面量, * 内部的单引号用 '\'' 的方式断开拼接。 */ export function shellQuote(value) { return `'${ String(value).replace(/'/g, '\'\\\'\'') }'` } export function wrapCommand(command, { cwd, nonInteractive = true } = {}) { const parts = [] if (nonInteractive) parts.push('export DEBIAN_FRONTEND=noninteractive') if (cwd) parts.push(`cd ${ shellQuote(cwd) }`) parts.push(command) return parts.join(' && ') } /** * 在指定主机上执行一条命令。 * * @param {string} hostId * @param {string} command * @param {object} [options] * @param {number} [options.timeoutMs] * @param {string} [options.cwd] * @param {AbortSignal} [options.signal] 用户点击停止时中断 * @returns {Promise<{stdout,stderr,exitCode,signal,timedOut,truncated,durationMs,executed}>} */ export async function execCommand(hostId, command, options = {}) { const timeoutMs = Math.min(Math.max(options.timeoutMs || DEFAULT_TIMEOUT_MS, 1000), MAX_TIMEOUT_MS) const executed = wrapCommand(command, options) const startedAt = Date.now() const { client, release } = await acquire(hostId) try { return await new Promise((resolve, reject) => { let stdout = '' let stderr = '' let collected = 0 let truncated = false let settled = false let timedOut = false let stream = null let timer = null let onAbort = null const cleanup = () => { if (timer) clearTimeout(timer) if (onAbort && options.signal) options.signal.removeEventListener('abort', onAbort) } const finish = (result) => { if (settled) return settled = true cleanup() resolve({ ...result, stdout, stderr, truncated, timedOut, durationMs: Date.now() - startedAt, executed }) } const fail = (error) => { if (settled) return settled = true cleanup() reject(error) } const append = (target, chunk) => { if (collected >= MAX_COLLECT_BYTES) { truncated = true return target } const text = chunk.toString('utf8') collected += Buffer.byteLength(text) if (collected > MAX_COLLECT_BYTES) truncated = true return target + text } const kill = () => { try { stream?.signal('KILL') stream?.close() } catch { // 通道可能已关闭 } } timer = setTimeout(() => { timedOut = true kill() finish({ exitCode: null, signal: 'TIMEOUT' }) }, timeoutMs) if (options.signal) { if (options.signal.aborted) return fail(new Error('已取消')) onAbort = () => { kill() finish({ exitCode: null, signal: 'ABORTED' }) } options.signal.addEventListener('abort', onAbort, { once: true }) } client.exec(executed, (err, execStream) => { if (err) return fail(new Error(`命令执行失败: ${ err.message }`)) stream = execStream stream .on('close', (code, signalName) => { finish({ exitCode: typeof code === 'number' ? code : null, signal: signalName || null }) }) .on('data', (chunk) => { stdout = append(stdout, chunk) }) .on('error', (streamErr) => fail(new Error(`命令流错误: ${ streamErr.message }`))) stream.stderr.on('data', (chunk) => { stderr = append(stderr, chunk) }) }) }) } finally { release() } } /** 打开一个 SFTP 会话,回调结束后自动释放 */ export async function withSftp(hostId, handler) { const { client, release } = await acquire(hostId) let sftp = null try { sftp = await new Promise((resolve, reject) => { client.sftp((err, session) => (err ? reject(new Error(`SFTP 打开失败: ${ err.message }`)) : resolve(session))) }) return await handler(sftp) } finally { try { sftp?.end() } catch { // 忽略 } release() } }