diff --git a/Termind/App/AppModel.swift b/Termind/App/AppModel.swift index a4c0b1ae..1f2cf84a 100644 --- a/Termind/App/AppModel.swift +++ b/Termind/App/AppModel.swift @@ -96,8 +96,12 @@ final class AppModel { connections = connectionCoordinator agentRunRepository = dependencies.agentRunRepository ?? InMemoryAgentRunRepository() agentHosts = AgentHostResolver(store: hosts) + let agentHostResolver = agentHosts agentConnections = AgentConnectionPool( - hosts: agentHosts, coordinator: connectionCoordinator + configuration: { try await agentHostResolver.configuration(for: $0) }, + routeFingerprint: { try await agentHostResolver.routeFingerprint(for: $0) }, + openExec: { try await connectionCoordinator.connectExec(using: $0) }, + openSFTP: { try await connectionCoordinator.connectSFTP(using: $0) } ) hostRepository = dependencies.hostRepository hostLibrary = HostLibraryModel(repository: dependencies.hostRepository) @@ -121,8 +125,12 @@ final class AppModel { self?.temporaryAgentConfiguration(for: hostID) } ) - hosts.willDeleteHost = { [weak forwarding] _, ruleIDs in + hosts.willDeleteHost = { [weak forwarding, agentConnections] hostID, ruleIDs in forwarding?.stop(ruleIDs) + // Not a `HostDeletionOutbox` step like the others: a live socket + // cannot outlive the process, so there is nothing to replay after a + // crash and nothing to make durable. + Task { await agentConnections.disconnect(hostID: hostID) } } hosts.willDeletePortForwardings = { [weak forwarding] ruleIDs in forwarding?.stop(ruleIDs) @@ -252,10 +260,7 @@ final class AppModel { let runtime = AgentRuntime( model: model, tools: tools, approval: approval, repository: agentRunRepository, - userInteraction: WeakAgentUserInteraction(conversation), - runCleanup: { [agentConnections] runID in - await agentConnections.cleanup(runID: runID) - } + userInteraction: WeakAgentUserInteraction(conversation) ) // No second Shell Integration gate here. `Pane.remoteSystemInfo` already // withholds the probe when the setting is off, and `context` is the same @@ -401,6 +406,7 @@ final class AppModel { for window in liveWindows { window.closeAll() } sftpTransfers.cancelAll() forwarding.stopAll() + Task { [agentConnections] in await agentConnections.closeAll() } outboxRunner?.stop() } diff --git a/Termind/Application/Agent/AgentBuiltInTools.swift b/Termind/Application/Agent/AgentBuiltInTools.swift index 02fa4455..0cdfd542 100644 --- a/Termind/Application/Agent/AgentBuiltInTools.swift +++ b/Termind/Application/Agent/AgentBuiltInTools.swift @@ -56,6 +56,40 @@ final class AgentHostResolver: Sendable { throw AgentToolError.invalidHost } + /// A stable description of *which machine, reached how, and as whom* — the + /// only question a pooled connection has to keep answering the same way. + /// + /// Built from records the store already holds in memory, with deliberately + /// no Keychain read: this runs before every reuse, and resolving credentials + /// here would put private-key material on the hot path of every command to + /// no purpose. A rotated credential on an unchanged identity is therefore + /// not a mismatch — the live session is still the same box and the same + /// login, and the new credential applies to the next connect. + func routeFingerprint(for id: UUID) throws -> String { + guard let host = store.host(id) else { + // A pane-backed target has no Vault record to describe, and its + // configuration is rebuilt from the live pane on every connect. Tie + // the fingerprint to that pane still existing, so a closed tab + // invalidates the Agent's connection to it. + guard temporaryTargets().contains(where: { $0.id == id }) else { + throw AgentToolError.invalidHost + } + return "temporary:\(id.uuidString)" + } + return store.route(for: host).map { node in + let identity = store.identity(for: node) + let proxy = node.proxy.map { + "\($0.kind.rawValue)://\($0.username)@\($0.hostname):\($0.port)" + } + let keys = store.knownHosts(for: node).map(\.fingerprint).sorted() + return [ + node.id.uuidString, node.host, String(node.port), + identity?.id.uuidString ?? "", identity?.username ?? "", + proxy ?? "", keys.joined(separator: ","), + ].joined(separator: "\u{1F}") + }.joined(separator: "\u{1E}") + } + private static func stateName(_ status: ConnectionStatus) -> String { switch status { case .disconnected: "disconnected" @@ -67,34 +101,238 @@ final class AgentHostResolver: Sendable { } } +/// Keeps one authenticated SSH-exec session and one SFTP session per host alive +/// across tool calls, runs, and conversations. +/// +/// Before this existed, `ssh_exec` paid a full TCP connect, handshake, host-key +/// check, and signature for *every command* — and a jump-host route rebuilt +/// every hop each time. Keying by host rather than by run is what lets a +/// follow-up message reuse the connection; the idle sweep, not the end of a run, +/// is what gives it back. +/// +/// Exec and SFTP stay on separate sessions so one long-running command cannot +/// block a directory listing. That is at most two connections per host, against +/// one per command before. actor AgentConnectionPool { - private struct Key: Hashable { var runID: UUID; var hostID: UUID } - private let hosts: AgentHostResolver - private let coordinator: RemoteConnectionCoordinator - private var sftpClients: [Key: any SFTPClientProtocol] = [:] + nonisolated enum Kind: String, Hashable, Sendable, CaseIterable { case exec, sftp } - init(hosts: AgentHostResolver, coordinator: RemoteConnectionCoordinator) { - self.hosts = hosts - self.coordinator = coordinator + typealias ConfigurationResolving = + @Sendable (UUID) async throws -> SSHConnectionConfiguration + typealias FingerprintResolving = @Sendable (UUID) async throws -> String + typealias ExecOpening = + @Sendable (SSHConnectionConfiguration) async throws -> any SSHExecSessionProtocol + typealias SFTPOpening = + @Sendable (SSHConnectionConfiguration) async throws -> any SFTPClientProtocol + + /// Long enough that a user reading output and asking a follow-up keeps the + /// connection, short enough that walking away does not leave an + /// authenticated session open on a server indefinitely. + static let defaultIdleTimeout = Duration.seconds(180) + + private enum Client: Sendable { + case exec(any SSHExecSessionProtocol) + case sftp(any SFTPClientProtocol) + + func disconnect() async { + switch self { + case .exec(let session): await session.disconnect() + case .sftp(let client): await client.disconnect() + } + } } - func configuration(for hostID: UUID) async throws -> SSHConnectionConfiguration { - try await hosts.configuration(for: hostID) + private struct Key: Hashable { var hostID: UUID; var kind: Kind } + + private struct Entry { + var client: Client + var fingerprint: String + var lastUsedAt: ContinuousClock.Instant } - func sftp(runID: UUID, hostID: UUID) async throws -> any SFTPClientProtocol { - let key = Key(runID: runID, hostID: hostID) - if let client = sftpClients[key] { return client } - let configuration = try await hosts.configuration(for: hostID) - let client = try await coordinator.connectSFTP(using: configuration) - sftpClients[key] = client + private struct InFlight { + var fingerprint: String + var task: Task + } + + private let configuration: ConfigurationResolving + private let routeFingerprint: FingerprintResolving + private let openExec: ExecOpening + private let openSFTP: SFTPOpening + private let idleTimeout: Duration + private let sweepInterval: Duration? + private var entries: [Key: Entry] = [:] + private var pending: [Key: InFlight] = [:] + private var sweeper: Task? + + init( + configuration: @escaping ConfigurationResolving, + routeFingerprint: @escaping FingerprintResolving, + openExec: @escaping ExecOpening, + openSFTP: @escaping SFTPOpening, + idleTimeout: Duration = AgentConnectionPool.defaultIdleTimeout, + sweepInterval: Duration? = .seconds(30) + ) { + self.configuration = configuration + self.routeFingerprint = routeFingerprint + self.openExec = openExec + self.openSFTP = openSFTP + self.idleTimeout = idleTimeout + self.sweepInterval = sweepInterval + } + + /// Runs one command on the host's pooled session, opening one if needed. + /// + /// Retries exactly once, and only on `channelUnavailable` — the failure that + /// is thrown while the command is provably not running. A pooled session + /// that died quietly between commands is invisible to the caller; anything + /// that failed *after* the command started is surfaced, because re-running + /// half of an `rm` is worse than a failed tool call. + func execute( + _ command: String, workingDirectory: String?, stdin: Data?, + timeout: TimeInterval, outputLimit: Int, hostID: UUID + ) async throws -> SSHExecResult { + do { + return try await exec(hostID: hostID).execute( + command, workingDirectory: workingDirectory, stdin: stdin, + timeout: timeout, outputLimit: outputLimit + ) + } catch SSHExecSessionError.channelUnavailable { + await disconnect(hostID: hostID, kind: .exec) + return try await exec(hostID: hostID).execute( + command, workingDirectory: workingDirectory, stdin: stdin, + timeout: timeout, outputLimit: outputLimit + ) + } + } + + func exec(hostID: UUID) async throws -> any SSHExecSessionProtocol { + guard case .exec(let session) = try await resolve(hostID: hostID, kind: .exec) + else { throw AgentToolError.invalidHost } + return session + } + + func sftp(hostID: UUID) async throws -> any SFTPClientProtocol { + guard case .sftp(let client) = try await resolve(hostID: hostID, kind: .sftp) + else { throw AgentToolError.invalidHost } return client } - func cleanup(runID: UUID) async { - let keys = sftpClients.keys.filter { $0.runID == runID } - let clients = keys.compactMap { sftpClients.removeValue(forKey: $0) } - for client in clients { await client.disconnect() } + /// Reports whether anything was actually open, which is what lets callers + /// treat closing an already-closed host as a no-op rather than a failure. + @discardableResult + func disconnect(hostID: UUID, kind: Kind) async -> Bool { + guard let entry = entries.removeValue(forKey: Key(hostID: hostID, kind: kind)) + else { return false } + stopSweepingIfEmpty() + await entry.client.disconnect() + return true + } + + @discardableResult + func disconnect(hostID: UUID) async -> Bool { + var closed = false + for kind in Kind.allCases { + closed = await disconnect(hostID: hostID, kind: kind) || closed + } + return closed + } + + func closeAll() async { + let closing = entries.values.map(\.client) + entries.removeAll() + stopSweepingIfEmpty() + for client in closing { await client.disconnect() } + } + + /// Closes everything untouched for longer than the idle timeout and reports + /// how many connections remain. Driven by a background task in the app and + /// called directly, with a synthetic instant, by tests. + @discardableResult + func sweepIdle(now: ContinuousClock.Instant) async -> Int { + let expired = entries.filter { now - $0.value.lastUsedAt >= idleTimeout } + for key in expired.keys { entries[key] = nil } + stopSweepingIfEmpty() + for entry in expired.values { await entry.client.disconnect() } + return entries.count + } + + func isConnected(hostID: UUID, kind: Kind) -> Bool { + entries[Key(hostID: hostID, kind: kind)] != nil + } + + var openConnectionCount: Int { entries.count } + + private func resolve(hostID: UUID, kind: Kind) async throws -> Client { + let fingerprint: String + do { + fingerprint = try await routeFingerprint(hostID) + } catch { + // The host is gone — deleted from the Vault, or a temporary target + // whose pane closed. Do not keep a session to it alive. + await disconnect(hostID: hostID) + throw error + } + let key = Key(hostID: hostID, kind: kind) + if let entry = entries.removeValue(forKey: key) { + guard entry.fingerprint == fingerprint else { + // Nothing waits on a stale socket. Its worker also frees the + // libssh2 session in `deinit`, so dropping the reference alone + // would be safe; this just makes it prompt. + Task { await entry.client.disconnect() } + return try await open(key: key, fingerprint: fingerprint) + } + entries[key] = Entry( + client: entry.client, fingerprint: fingerprint, lastUsedAt: .now + ) + return entry.client + } + return try await open(key: key, fingerprint: fingerprint) + } + + private func open(key: Key, fingerprint: String) async throws -> Client { + // Without this join, two tool calls that both miss the cache would each + // build a connection and one of them would never be pooled — or closed. + if let inFlight = pending[key], inFlight.fingerprint == fingerprint { + let client = try await inFlight.task.value + commit(client, key: key, fingerprint: fingerprint) + return client + } + let task = Task { [configuration, openExec, openSFTP] () async throws -> Client in + let resolved = try await configuration(key.hostID) + switch key.kind { + case .exec: return .exec(try await openExec(resolved)) + case .sftp: return .sftp(try await openSFTP(resolved)) + } + } + pending[key] = InFlight(fingerprint: fingerprint, task: task) + defer { if pending[key]?.fingerprint == fingerprint { pending[key] = nil } } + // Committed on every path that receives a client, not only on the path + // that started the task: whoever started it may have been cancelled + // while a second caller was still waiting on the same connection. + let client = try await task.value + commit(client, key: key, fingerprint: fingerprint) + return client + } + + private func commit(_ client: Client, key: Key, fingerprint: String) { + entries[key] = Entry(client: client, fingerprint: fingerprint, lastUsedAt: .now) + startSweeping() + } + + private func startSweeping() { + guard sweeper == nil, let sweepInterval else { return } + sweeper = Task { [weak self, sweepInterval] in + while !Task.isCancelled { + do { try await Task.sleep(for: sweepInterval) } catch { return } + guard let self, await self.sweepIdle(now: .now) > 0 else { return } + } + } + } + + private func stopSweepingIfEmpty() { + guard entries.isEmpty, pending.isEmpty else { return } + sweeper?.cancel() + sweeper = nil } } @@ -175,11 +413,10 @@ nonisolated enum AgentBuiltInTools { ) input = Data(secret.utf8) } - let configuration = try await pool.configuration(for: hostID) - let output = try await SSHSessionActor.execute( + let output = try await pool.execute( command, workingDirectory: arguments.optionalString("cwd"), stdin: input, timeout: TimeInterval(arguments.optionalInt("timeout_seconds") ?? 60), - outputLimit: 256 * 1_024, using: configuration + outputLimit: 256 * 1_024, hostID: hostID ) return result(invocation, .object([ "stdout": .string(output.stdout), "stderr": .string(output.stderr), @@ -397,9 +634,9 @@ nonisolated enum AgentBuiltInTools { return AnyAgentTool( descriptor: tool, preflight: hostPathPreflight(safety: .locallyReadOnly), - execute: { invocation, context in + execute: { invocation, _ in let arguments = try Arguments(invocation) - let client = try await pool.sftp(runID: context.runID, hostID: arguments.uuid("host_id")) + let client = try await pool.sftp(hostID: arguments.uuid("host_id")) return try await operation(client, normalizeRemotePath(arguments.string("path")), invocation) } ) @@ -424,23 +661,32 @@ nonisolated enum AgentBuiltInTools { return AnyAgentTool( descriptor: tool, preflight: hostPathPreflight(safety: .requiresAuthorization), - execute: { invocation, context in + execute: { invocation, _ in let arguments = try Arguments(invocation) - let client = try await pool.sftp(runID: context.runID, hostID: arguments.uuid("host_id")) + let client = try await pool.sftp(hostID: arguments.uuid("host_id")) return try await operation(client, normalizeRemotePath(arguments.string("path")), arguments, invocation) } ) } + /// Resolves `host_id` before approval so the tool card names the host the + /// action is aimed at rather than an opaque argument blob. + private static func hostPreflight( + safety: AgentToolDescriptor.Safety + ) -> @Sendable (AgentToolInvocation) async throws -> AgentToolPreflight { + { invocation in + var checked = invocation + checked.targetHostID = try Arguments(invocation).uuid("host_id") + return AgentToolPreflight(invocation: checked, safety: safety) + } + } + private static func hostPathPreflight( safety: AgentToolDescriptor.Safety ) -> @Sendable (AgentToolInvocation) async throws -> AgentToolPreflight { { invocation in - let arguments = try Arguments(invocation) - _ = try normalizeRemotePath(arguments.string("path")) - var checked = invocation - checked.targetHostID = try arguments.uuid("host_id") - return AgentToolPreflight(invocation: checked, safety: safety) + _ = try normalizeRemotePath(Arguments(invocation).string("path")) + return try await hostPreflight(safety: safety)(invocation) } } diff --git a/Termind/Application/Agent/AgentRuntime.swift b/Termind/Application/Agent/AgentRuntime.swift index 633ffb9c..3ae97b45 100644 --- a/Termind/Application/Agent/AgentRuntime.swift +++ b/Termind/Application/Agent/AgentRuntime.swift @@ -174,7 +174,6 @@ actor AgentRuntime { private let approval: any AgentApprovalHandling private let repository: any AgentRunPersisting private let userInteraction: any AgentUserInteractionHandling - private let runCleanup: @Sendable (UUID) async -> Void let secretBroker: SecretBroker private var tools: [String: AnyAgentTool] private var runTask: Task? @@ -191,7 +190,6 @@ actor AgentRuntime { repository: any AgentRunPersisting = InMemoryAgentRunRepository(), secretBroker: SecretBroker = SecretBroker(), userInteraction: any AgentUserInteractionHandling = UnavailableAgentUserInteraction(), - runCleanup: @escaping @Sendable (UUID) async -> Void = { _ in }, maxTurns: Int = 24, maxToolCalls: Int = 64 ) { @@ -207,7 +205,6 @@ actor AgentRuntime { self.repository = repository self.secretBroker = secretBroker self.userInteraction = userInteraction - self.runCleanup = runCleanup self.maxTurns = maxTurns self.maxToolCalls = maxToolCalls } @@ -542,10 +539,10 @@ actor AgentRuntime { } private func finish(runID: UUID) { - Task { - await secretBroker.discardSecrets(for: runID) - await runCleanup(runID) - } + // Secrets, and nothing else. SSH and SFTP connections used to be closed + // here too; they are now pooled per host and given back by an idle + // sweep, so a run ending is no longer the event that releases them. + Task { await secretBroker.discardSecrets(for: runID) } activeRunID = nil runTask = nil continuation?.finish() diff --git a/Termind/Application/Connections/RemoteConnectionCoordinator.swift b/Termind/Application/Connections/RemoteConnectionCoordinator.swift index 45ff9acb..c28406a8 100644 --- a/Termind/Application/Connections/RemoteConnectionCoordinator.swift +++ b/Termind/Application/Connections/RemoteConnectionCoordinator.swift @@ -16,6 +16,23 @@ nonisolated final class RemoteConnectionCoordinator: Sendable { func makeSFTPClient() -> any SFTPClientProtocol { SFTPClient() } + func makeExecSession() -> any SSHExecSessionProtocol { SSHExecSession() } + + /// Mirrors `connectSFTP`, including its guarantee that a session which fails + /// to come up is torn down here rather than handed back half-open. + func connectExec( + using configuration: SSHConnectionConfiguration + ) async throws -> any SSHExecSessionProtocol { + let session = makeExecSession() + do { + try await session.connect(using: configuration) + return session + } catch { + await session.disconnect() + throw error + } + } + func connectSFTP( using configuration: SSHConnectionConfiguration, progress: SSHRouteProgressHandler? = nil, diff --git a/Termind/Core/Models/RemoteSystemInfo.swift b/Termind/Core/Models/RemoteSystemInfo.swift index 1ac2e571..44ea4fa3 100644 --- a/Termind/Core/Models/RemoteSystemInfo.swift +++ b/Termind/Core/Models/RemoteSystemInfo.swift @@ -52,6 +52,26 @@ nonisolated struct RemoteSystemInfo: Hashable, Codable, Sendable { try container.encodeIfPresent(architecture, forKey: .architecture) } + /// Reads the `id|version|kernel_name|kernel_release|arch` line that both + /// probes print — the PTY bootstrap's token-delimited one and the Agent's + /// dedicated exec channel. Returns nil when the line is malformed or when + /// every field was rejected by `sanitized`. + /// + /// Empty fields are meaningful (a BSD has no os-release), so the split keeps + /// them rather than collapsing the run of separators. + init?(probeFields line: some StringProtocol) { + let fields = line.split( + separator: "|", maxSplits: 4, omittingEmptySubsequences: false + ) + guard fields.count == 5 else { return nil } + let values = fields.map { Self.unquoted(String($0)) } + self.init( + distributionID: values[0], distributionVersion: values[1], + kernelName: values[2], kernelRelease: values[3], architecture: values[4] + ) + if isEmpty { return nil } + } + var isEmpty: Bool { distributionID == nil && distributionVersion == nil && kernelName == nil && kernelRelease == nil && architecture == nil @@ -75,6 +95,17 @@ nonisolated struct RemoteSystemInfo: Hashable, Codable, Sendable { return lines } + /// os-release quotes values that could contain spaces. Stripping a matched + /// pair here rather than in `sanitized` keeps the ASCII-token grammar itself + /// free of exceptions: `"22.04"` becomes `22.04`, while a lone quote stays + /// and is rejected. + private static func unquoted(_ raw: String) -> String { + guard raw.count >= 2, let first = raw.first, let last = raw.last, + first == last, first == "\"" || first == "'" + else { return raw } + return String(raw.dropFirst().dropLast()) + } + private static func sanitized(_ proposed: String?) -> String? { guard let proposed else { return nil } let value = proposed.trimmingCharacters(in: .whitespacesAndNewlines) diff --git a/Termind/Core/SSH/SSHExecSession.swift b/Termind/Core/SSH/SSHExecSession.swift new file mode 100644 index 00000000..ec7c550f --- /dev/null +++ b/Termind/Core/SSH/SSHExecSession.swift @@ -0,0 +1,377 @@ +import CSSH2 +import Darwin +import Foundation + +nonisolated struct SSHExecResult: Hashable, Sendable { + var stdout: String + var stderr: String + var exitCode: Int32 + var signal: String? + var stdoutTruncated: Bool + var stderrTruncated: Bool +} + +/// The one failure that is safe to retry on a fresh connection. +/// +/// It is thrown only while the remote command is provably *not running*: the +/// session was already gone, or the channel could not be opened or started. +/// Every failure after `process_startup` succeeds is rethrown unchanged, because +/// by then the command may have partially run and a transparent retry could +/// delete something twice. +nonisolated enum SSHExecSessionError: LocalizedError, Sendable { + case channelUnavailable(any Error) + + var errorDescription: String? { + switch self { + case .channelUnavailable(let underlying): underlying.localizedDescription + } + } +} + +/// A long-lived authenticated SSH session that runs each command on its own +/// non-PTY channel. +/// +/// Splitting connect from execute is the whole point: the handshake, host-key +/// check, and signature are paid once per host instead of once per command, and +/// a jump-host route stops rebuilding every hop for every `ls`. Semantics are +/// unchanged — a command still gets a private channel and a private +/// `/bin/sh -lc`, so there is no shared shell state and no cwd carry-over. +nonisolated protocol SSHExecSessionProtocol: Sendable { + func connect(using configuration: SSHConnectionConfiguration) async throws + func execute( + _ command: String, workingDirectory: String?, stdin: Data?, + timeout: TimeInterval, outputLimit: Int + ) async throws -> SSHExecResult + func disconnect() async +} + +actor SSHExecSession: SSHExecSessionProtocol { + private nonisolated let worker = SSHExecWorker() + + func connect(using configuration: SSHConnectionConfiguration) async throws { + try await worker.connect(using: configuration) + } + + /// Actor-isolated, so commands on one session queue behind each other. That + /// is not incidental: a libssh2 session cannot be driven from two places at + /// once, and the worker's queue would serialize them anyway. + func execute( + _ command: String, workingDirectory: String?, stdin: Data?, + timeout: TimeInterval, outputLimit: Int + ) async throws -> SSHExecResult { + let cancellation = TransferCancellationToken() + return try await withTaskCancellationHandler { + try await worker.execute( + command, workingDirectory: workingDirectory, stdin: stdin, + timeout: timeout, outputLimit: outputLimit, cancellation: cancellation + ) + } onCancel: { + cancellation.cancel() + } + } + + func disconnect() async { await worker.disconnect() } +} + +private nonisolated final class SSHExecWorker: @unchecked Sendable { + /// Every libssh2 call for this session happens here, the idle keepalive + /// included. libssh2 sessions are not thread-safe, and serializing on one + /// queue is what makes reuse across commands safe at all. + private let queue = DispatchQueue(label: "com.akinokaede.termind.ssh.exec") + private var routeTransport: SSHRoutedSocket? + private var session: OpaquePointer? + private var keepaliveTimer: (any DispatchSourceTimer)? + /// Set when the transport is known to be gone. The next `execute` reports + /// `channelUnavailable` instead of calling into a session that cannot answer. + private var isDead = false + + deinit { cleanup() } + + func connect(using configuration: SSHConnectionConfiguration) async throws { + try await perform { + self.cleanup() + let (transport, session) = try SSHSessionWorker.openAuthenticatedSession( + using: configuration, timeout: 20 + ) + self.routeTransport = transport + self.session = session + self.isDead = false + self.startKeepalive(for: session) + } + } + + func disconnect() async { await performNoThrow { self.cleanup() } } + + func execute( + _ command: String, workingDirectory: String?, stdin: Data?, + timeout: TimeInterval, outputLimit: Int, cancellation: TransferCancellationToken + ) async throws -> SSHExecResult { + try await perform { + do { + return try self.executeOnQueue( + command, workingDirectory: workingDirectory, stdin: stdin, + timeout: timeout, outputLimit: outputLimit, cancellation: cancellation + ) + } catch { + // Still on the queue, and the channel defer has already run, so + // this is the right moment to give the socket back rather than + // leave a corpse in the pool for the idle sweep to find. + if self.isDead { self.cleanup() } + throw error + } + } + } + + private func executeOnQueue( + _ command: String, workingDirectory: String?, stdin: Data?, + timeout: TimeInterval, outputLimit: Int, cancellation: TransferCancellationToken + ) throws -> SSHExecResult { + guard !command.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { + throw TermindSSHError.library("The command is empty.") + } + guard let session, !isDead else { + throw SSHExecSessionError.channelUnavailable(TermindSSHError.disconnected) + } + // Blocking for the two setup calls, non-blocking for the pump below. + libssh2_session_set_blocking(session, 1) + guard let channel = libssh2_channel_open_ex( + session, "session", 7, 2 * 1_024 * 1_024, 32 * 1_024, nil, 0 + ) else { + throw notStarted(session, prefix: "Unable to open SSH exec channel") + } + // Only a command that ran to completion leaves a channel worth closing + // politely, and a session worth keeping. + var completed = false + defer { + if completed { + libssh2_session_set_blocking(session, 1) + _ = libssh2_channel_free(channel) + } else { + // Timed out, cancelled, or the transport failed: the command may + // still be running and the channel is mid-stream. Blocking here + // to close it courteously would stall every later command on this + // host, so give the whole session up instead — `cleanup` frees + // it, and freeing a session reclaims its channels with it. + isDead = true + } + } + + let effective: String + if let cwd = workingDirectory?.trimmingCharacters(in: .whitespacesAndNewlines), + !cwd.isEmpty + { + effective = "cd -- \(Self.shellQuote(cwd)) && /bin/sh -lc \(Self.shellQuote(command))" + } else { + effective = "/bin/sh -lc \(Self.shellQuote(command))" + } + let started = effective.withCString { + libssh2_channel_process_startup(channel, "exec", 4, $0, UInt32(effective.utf8.count)) + } + guard started == 0 else { + throw notStarted(session, prefix: "Unable to start remote command") + } + // Past this line the command is running. Nothing below may be retried. + libssh2_session_set_blocking(session, 0) + + var pendingInput = stdin ?? Data() + if !pendingInput.isEmpty, pendingInput.last != 0x0A { pendingInput.append(0x0A) } + var inputOffset = 0 + var sentEOF = pendingInput.isEmpty + if sentEOF { _ = libssh2_channel_send_eof(channel) } + var standardOut = Data() + var standardError = Data() + var outTruncated = false + var errorTruncated = false + var buffer = [CChar](repeating: 0, count: 16 * 1_024) + let deadline = Date().addingTimeInterval(max(1, min(timeout, 15 * 60))) + var lastKeepalive = Date() + + while libssh2_channel_eof(channel) == 0 { + if cancellation.isCancelled { throw CancellationError() } + if Date() >= deadline { throw TermindSSHError.library("The remote command timed out.") } + var progressed = false + if inputOffset < pendingInput.count { + let written = pendingInput.withUnsafeBytes { bytes in + libssh2_channel_write_ex( + channel, 0, + bytes.baseAddress!.advanced(by: inputOffset) + .assumingMemoryBound(to: CChar.self), + bytes.count - inputOffset + ) + } + if written > 0 { inputOffset += written; progressed = true } + else if written != LIBSSH2_ERROR_EAGAIN { + throw failed(session, prefix: "Unable to write command input") + } + } else if !sentEOF { + let result = libssh2_channel_send_eof(channel) + if result == 0 { sentEOF = true; progressed = true } + else if result != LIBSSH2_ERROR_EAGAIN { + throw failed(session, prefix: "Unable to finish command input") + } + } + for stream in [Int32(0), Int32(1)] { + while true { + let count = libssh2_channel_read_ex(channel, stream, &buffer, buffer.count) + if count > 0 { + progressed = true + let chunk = Data(bytes: buffer, count: count) + if stream == 0 { + Self.appendCapped( + chunk, to: &standardOut, limit: outputLimit, + truncated: &outTruncated + ) + } else { + Self.appendCapped( + chunk, to: &standardError, limit: outputLimit, + truncated: &errorTruncated + ) + } + } else if count == LIBSSH2_ERROR_EAGAIN || count == 0 { break } + else { throw failed(session, prefix: "Unable to read command output") } + } + } + if !progressed { + // The idle timer cannot fire while this queue is busy, so a + // long-running quiet command has to keep its own session alive. + if Date().timeIntervalSince(lastKeepalive) >= 30 { + lastKeepalive = Date() + var next = Int32(0) + let sent = libssh2_keepalive_send(session, &next) + if sent < 0, sent != LIBSSH2_ERROR_EAGAIN { + throw failed(session, prefix: "SSH keepalive failed") + } + } + usleep(10_000) + } + } + // Drain bytes buffered alongside EOF. + for stream in [Int32(0), Int32(1)] { + while true { + let count = libssh2_channel_read_ex(channel, stream, &buffer, buffer.count) + guard count > 0 else { break } + let chunk = Data(bytes: buffer, count: count) + if stream == 0 { + Self.appendCapped( + chunk, to: &standardOut, limit: outputLimit, truncated: &outTruncated + ) + } else { + Self.appendCapped( + chunk, to: &standardError, limit: outputLimit, truncated: &errorTruncated + ) + } + } + } + while true { + let closed = libssh2_channel_close(channel) + if closed == 0 { break } + if closed != LIBSSH2_ERROR_EAGAIN { + throw failed(session, prefix: "Unable to close SSH exec channel") + } + if cancellation.isCancelled { throw CancellationError() } + if Date() >= deadline { throw TermindSSHError.library("The remote command timed out.") } + usleep(10_000) + } + // The remote is done and the channel is closed, so the session is clean + // and can carry the next command. + completed = true + var exitSignal: UnsafeMutablePointer? + var exitSignalLength = 0 + let signalResult = libssh2_channel_get_exit_signal( + channel, &exitSignal, &exitSignalLength, nil, nil, nil, nil + ) + let signal: String? + if signalResult == 0, let exitSignal { + signal = String(decoding: UnsafeBufferPointer( + start: exitSignal, count: exitSignalLength + ).map(UInt8.init(bitPattern:)), as: UTF8.self) + libssh2_free(session, exitSignal) + } else { + signal = nil + } + return SSHExecResult( + stdout: String(decoding: standardOut, as: UTF8.self), + stderr: String(decoding: standardError, as: UTF8.self), + exitCode: Int32(libssh2_channel_get_exit_status(channel)), + signal: signal, + stdoutTruncated: outTruncated, stderrTruncated: errorTruncated + ) + } + + /// The command never started, so the caller may open a fresh connection and + /// try again exactly once. + private func notStarted(_ session: OpaquePointer, prefix: String) -> SSHExecSessionError { + SSHExecSessionError.channelUnavailable(failed(session, prefix: prefix)) + } + + /// Records whether the session is still usable, then describes the failure. + private func failed(_ session: OpaquePointer, prefix: String) -> TermindSSHError { + let transport: [Int32] = [ + LIBSSH2_ERROR_SOCKET_SEND, LIBSSH2_ERROR_SOCKET_RECV, + LIBSSH2_ERROR_SOCKET_DISCONNECT, LIBSSH2_ERROR_SOCKET_TIMEOUT, + LIBSSH2_ERROR_TIMEOUT, LIBSSH2_ERROR_CHANNEL_CLOSED, + ] + if transport.contains(libssh2_session_last_errno(session)) { isDead = true } + return SSHSessionWorker.sessionError(session, prefix: prefix) + } + + private func startKeepalive(for session: OpaquePointer) { + libssh2_keepalive_config(session, 1, 30) + let timer = DispatchSource.makeTimerSource(queue: queue) + timer.schedule(deadline: .now() + 30, repeating: 30) + timer.setEventHandler { [weak self] in + guard let self, self.session == session else { return } + var secondsToNext: Int32 = 0 + let result = libssh2_keepalive_send(session, &secondsToNext) + if result < 0, result != LIBSSH2_ERROR_EAGAIN { + self.isDead = true + self.cleanup() + } + } + keepaliveTimer = timer + timer.resume() + } + + private func cleanup() { + keepaliveTimer?.cancel() + keepaliveTimer = nil + if let session { + _ = libssh2_session_disconnect_ex( + session, SSH_DISCONNECT_BY_APPLICATION, "Termind closed the exec session", "" + ) + _ = libssh2_session_free(session) + self.session = nil + } + routeTransport?.close() + routeTransport = nil + } + + private func perform( + _ operation: @escaping @Sendable () throws -> T + ) async throws -> T { + try await withCheckedThrowingContinuation { continuation in + queue.async { + do { continuation.resume(returning: try operation()) } + catch { continuation.resume(throwing: error) } + } + } + } + + private func performNoThrow(_ operation: @escaping @Sendable () -> Void) async { + await withCheckedContinuation { continuation in + queue.async { operation(); continuation.resume() } + } + } + + private static func shellQuote(_ value: String) -> String { + "'" + value.replacingOccurrences(of: "'", with: "'\\''") + "'" + } + + private static func appendCapped( + _ chunk: Data, to destination: inout Data, limit: Int, truncated: inout Bool + ) { + let remaining = max(0, limit - destination.count) + if chunk.count > remaining { truncated = true } + if remaining > 0 { destination.append(chunk.prefix(remaining)) } + } +} diff --git a/Termind/Core/SSH/SSHSessionActor.swift b/Termind/Core/SSH/SSHSessionActor.swift index f1fbb74a..0d359d13 100644 --- a/Termind/Core/SSH/SSHSessionActor.swift +++ b/Termind/Core/SSH/SSHSessionActor.swift @@ -356,27 +356,7 @@ nonisolated struct RemoteShellIntegrationBootstrap: Sendable { guard payload.count >= 2, payload.suffix(2).elementsEqual(Data("__".utf8)) else { return nil } payload.removeLast(2) - let fields = String(decoding: payload, as: UTF8.self).split( - separator: "|", maxSplits: 4, omittingEmptySubsequences: false - ) - guard fields.count == 5 else { return nil } - let values = fields.map { normalizedOSReleaseValue(String($0)) } - let info = RemoteSystemInfo( - distributionID: values[0], distributionVersion: values[1], - kernelName: values[2], kernelRelease: values[3], architecture: values[4] - ) - return info.isEmpty ? nil : info - } - - private func normalizedOSReleaseValue(_ raw: String) -> String? { - guard !raw.isEmpty else { return nil } - if raw.count >= 2, - (raw.first == "\"" && raw.last == "\"" - || raw.first == "'" && raw.last == "'") - { - return String(raw.dropFirst().dropLast()) - } - return raw + return RemoteSystemInfo(probeFields: String(decoding: payload, as: UTF8.self)) } private mutating func scanForLateProbe(now: UInt64) -> Update { @@ -616,15 +596,6 @@ nonisolated enum SSHShellIntegrationEvent: Sendable { case draft(String) } -nonisolated struct SSHExecResult: Hashable, Sendable { - var stdout: String - var stderr: String - var exitCode: Int32 - var signal: String? - var stdoutTruncated: Bool - var stderrTruncated: Bool -} - actor SSHSessionActor { typealias OutputHandler = @Sendable (Data) -> Void typealias ExitHandler = @Sendable (Result) -> Void @@ -666,29 +637,6 @@ actor SSHSessionActor { ) } - /// Executes one command on a dedicated non-PTY channel. This never shares or - /// writes the user's interactive shell and therefore has deterministic output - /// and exit-code boundaries suitable for Agent tools. - static func execute( - _ command: String, - workingDirectory: String? = nil, - stdin: Data? = nil, - timeout: TimeInterval = 120, - outputLimit: Int = 256 * 1024, - using configuration: SSHConnectionConfiguration - ) async throws -> SSHExecResult { - let cancellation = TransferCancellationToken() - return try await withTaskCancellationHandler { - try await SSHSessionWorker.execute( - command, workingDirectory: workingDirectory, stdin: stdin, - timeout: timeout, outputLimit: outputLimit, - using: configuration, cancellation: cancellation - ) - } onCancel: { - cancellation.cancel() - } - } - func connect( hostID: Host.ID, host: String, @@ -800,207 +748,54 @@ nonisolated final class SSHSessionWorker: @unchecked Sendable { } } - static func execute( - _ command: String, - workingDirectory: String?, - stdin: Data?, - timeout: TimeInterval, - outputLimit: Int, + /// Opens a transport, handshakes, validates the host key, and authenticates — + /// the identical prologue every non-PTY entry point needs before it can open + /// a channel. On any failure it frees what it created and rethrows, so a + /// caller only owns the pair once it has been handed one. + /// + /// The returned session is left in blocking mode. Callers that want the + /// non-blocking pump switch it themselves once their channel is running. + static func openAuthenticatedSession( using configuration: SSHConnectionConfiguration, - cancellation: TransferCancellationToken - ) async throws -> SSHExecResult { + timeout: TimeInterval, + progress: SSHRouteProgressHandler? = nil + ) throws -> (transport: SSHRoutedSocket, session: OpaquePointer) { _ = initialization - return try await withCheckedThrowingContinuation { continuation in - DispatchQueue.global(qos: .userInitiated).async { - do { - continuation.resume(returning: try executeBlocking( - command, workingDirectory: workingDirectory, stdin: stdin, - timeout: timeout, outputLimit: outputLimit, - using: configuration, cancellation: cancellation - )) - } catch { continuation.resume(throwing: error) } - } - } - } - - private static func executeBlocking( - _ command: String, - workingDirectory: String?, - stdin: Data?, - timeout: TimeInterval, - outputLimit: Int, - using configuration: SSHConnectionConfiguration, - cancellation: TransferCancellationToken - ) throws -> SSHExecResult { - guard !command.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { - throw TermindSSHError.library("The command is empty.") - } var transport: SSHRoutedSocket? var session: OpaquePointer? - var channel: OpaquePointer? - defer { - if let channel { - _ = libssh2_channel_send_eof(channel) - _ = libssh2_channel_close(channel) - _ = libssh2_channel_free(channel) + do { + let opened = try SSHRoutedSocket.open( + host: configuration.host, port: configuration.port, + proxy: configuration.proxy, jumps: configuration.jumps, + timeout: timeout, progress: progress + ) + transport = opened + progress?(.connecting(configuration.hostID)) + guard let created = libssh2_session_init_ex(nil, nil, nil, nil) else { + throw TermindSSHError.library("Unable to create an SSH session.") } - if let session { - _ = libssh2_session_disconnect_ex( - session, SSH_DISCONNECT_BY_APPLICATION, - "Termind completed agent command", "" - ) - _ = libssh2_session_free(session) + session = created + libssh2_session_set_blocking(created, 1) + guard libssh2_session_handshake(created, opened.socket) == 0 else { + throw sessionError(created, prefix: "SSH handshake failed") } + let observed = try observation( + host: configuration.host, port: configuration.port, session: created + ) + try validateHostKey(observed, expected: configuration.expectedHostKeys) + progress?(.authenticating(configuration.hostID)) + try authenticate( + username: configuration.username, + credentials: configuration.credentials, + session: created + ) + progress?(.authenticated(configuration.hostID)) + return (opened, created) + } catch { + if let session { _ = libssh2_session_free(session) } transport?.close() + throw error } - - transport = try SSHRoutedSocket.open( - host: configuration.host, port: configuration.port, - proxy: configuration.proxy, jumps: configuration.jumps, timeout: 20 - ) - guard let created = libssh2_session_init_ex(nil, nil, nil, nil) else { - throw TermindSSHError.library("Unable to create an SSH session.") - } - session = created - libssh2_session_set_blocking(created, 1) - guard let transport, libssh2_session_handshake(created, transport.socket) == 0 else { - throw sessionError(created, prefix: "SSH handshake failed") - } - let observed = try observation( - host: configuration.host, port: configuration.port, session: created - ) - try validateHostKey(observed, expected: configuration.expectedHostKeys) - try authenticate( - username: configuration.username, - credentials: configuration.credentials, - session: created - ) - guard let opened = libssh2_channel_open_ex( - created, "session", 7, 2 * 1024 * 1024, 32 * 1024, nil, 0 - ) else { throw sessionError(created, prefix: "Unable to open SSH exec channel") } - channel = opened - - let effective: String - if let cwd = workingDirectory?.trimmingCharacters(in: .whitespacesAndNewlines), !cwd.isEmpty { - effective = "cd -- \(shellQuote(cwd)) && /bin/sh -lc \(shellQuote(command))" - } else { - effective = "/bin/sh -lc \(shellQuote(command))" - } - let started = effective.withCString { - libssh2_channel_process_startup(opened, "exec", 4, $0, UInt32(effective.utf8.count)) - } - guard started == 0 else { throw sessionError(created, prefix: "Unable to start remote command") } - libssh2_session_set_blocking(created, 0) - - var pendingInput = stdin ?? Data() - if !pendingInput.isEmpty, pendingInput.last != 0x0A { pendingInput.append(0x0A) } - var inputOffset = 0 - var sentEOF = pendingInput.isEmpty - if sentEOF { _ = libssh2_channel_send_eof(opened) } - var standardOut = Data() - var standardError = Data() - var outTruncated = false - var errorTruncated = false - var buffer = [CChar](repeating: 0, count: 16 * 1024) - let deadline = Date().addingTimeInterval(max(1, min(timeout, 15 * 60))) - - while libssh2_channel_eof(opened) == 0 { - if cancellation.isCancelled { throw CancellationError() } - if Date() >= deadline { throw TermindSSHError.library("The remote command timed out.") } - var progressed = false - if inputOffset < pendingInput.count { - let written = pendingInput.withUnsafeBytes { bytes in - libssh2_channel_write_ex( - opened, 0, - bytes.baseAddress!.advanced(by: inputOffset).assumingMemoryBound(to: CChar.self), - bytes.count - inputOffset - ) - } - if written > 0 { inputOffset += written; progressed = true } - else if written != LIBSSH2_ERROR_EAGAIN { - throw sessionError(created, prefix: "Unable to write command input") - } - } else if !sentEOF { - let result = libssh2_channel_send_eof(opened) - if result == 0 { sentEOF = true; progressed = true } - else if result != LIBSSH2_ERROR_EAGAIN { - throw sessionError(created, prefix: "Unable to finish command input") - } - } - for stream in [Int32(0), Int32(1)] { - while true { - let count = libssh2_channel_read_ex(opened, stream, &buffer, buffer.count) - if count > 0 { - progressed = true - let chunk = Data(bytes: buffer, count: count) - if stream == 0 { - appendCapped(chunk, to: &standardOut, limit: outputLimit, truncated: &outTruncated) - } else { - appendCapped(chunk, to: &standardError, limit: outputLimit, truncated: &errorTruncated) - } - } else if count == LIBSSH2_ERROR_EAGAIN || count == 0 { break } - else { throw sessionError(created, prefix: "Unable to read command output") } - } - } - if !progressed { usleep(10_000) } - } - // Drain bytes buffered alongside EOF. - for stream in [Int32(0), Int32(1)] { - while true { - let count = libssh2_channel_read_ex(opened, stream, &buffer, buffer.count) - guard count > 0 else { break } - let chunk = Data(bytes: buffer, count: count) - if stream == 0 { - appendCapped(chunk, to: &standardOut, limit: outputLimit, truncated: &outTruncated) - } else { - appendCapped(chunk, to: &standardError, limit: outputLimit, truncated: &errorTruncated) - } - } - } - while true { - let closed = libssh2_channel_close(opened) - if closed == 0 { break } - if closed != LIBSSH2_ERROR_EAGAIN { - throw sessionError(created, prefix: "Unable to close SSH exec channel") - } - if cancellation.isCancelled { throw CancellationError() } - if Date() >= deadline { throw TermindSSHError.library("The remote command timed out.") } - usleep(10_000) - } - var exitSignal: UnsafeMutablePointer? - var exitSignalLength = 0 - let signalResult = libssh2_channel_get_exit_signal( - opened, &exitSignal, &exitSignalLength, nil, nil, nil, nil - ) - let signal: String? - if signalResult == 0, let exitSignal { - signal = String(decoding: UnsafeBufferPointer( - start: exitSignal, count: exitSignalLength - ).map(UInt8.init(bitPattern:)), as: UTF8.self) - libssh2_free(created, exitSignal) - } else { - signal = nil - } - return SSHExecResult( - stdout: String(decoding: standardOut, as: UTF8.self), - stderr: String(decoding: standardError, as: UTF8.self), - exitCode: Int32(libssh2_channel_get_exit_status(opened)), - signal: signal, - stdoutTruncated: outTruncated, stderrTruncated: errorTruncated - ) - } - - private static func shellQuote(_ value: String) -> String { - "'" + value.replacingOccurrences(of: "'", with: "'\\''") + "'" - } - - private static func appendCapped( - _ chunk: Data, to destination: inout Data, - limit: Int, truncated: inout Bool - ) { - let remaining = max(0, limit - destination.count) - if chunk.count > remaining { truncated = true } - if remaining > 0 { destination.append(chunk.prefix(remaining)) } } static func installAuthorizedKey( @@ -1032,8 +827,9 @@ nonisolated final class SSHSessionWorker: @unchecked Sendable { private static func installAuthorizedKeyBlocking( _ publicKey: String, using configuration: SSHConnectionConfiguration ) throws { - var transport: SSHRoutedSocket? - var session: OpaquePointer? + let (transport, created) = try openAuthenticatedSession( + using: configuration, timeout: 20 + ) var channel: OpaquePointer? defer { if let channel { @@ -1041,40 +837,13 @@ nonisolated final class SSHSessionWorker: @unchecked Sendable { _ = libssh2_channel_close(channel) _ = libssh2_channel_free(channel) } - if let session { - _ = libssh2_session_disconnect_ex( - session, SSH_DISCONNECT_BY_APPLICATION, "Termind completed key export", "" - ) - _ = libssh2_session_free(session) - } - transport?.close() + _ = libssh2_session_disconnect_ex( + created, SSH_DISCONNECT_BY_APPLICATION, "Termind completed key export", "" + ) + _ = libssh2_session_free(created) + transport.close() } - transport = try SSHRoutedSocket.open( - host: configuration.host, port: configuration.port, - proxy: configuration.proxy, jumps: configuration.jumps, timeout: 20 - ) - guard let created = libssh2_session_init_ex(nil, nil, nil, nil) else { - throw TermindSSHError.library("Unable to create an SSH session.") - } - session = created - libssh2_session_set_blocking(created, 1) - guard let transport, - libssh2_session_handshake(created, transport.socket) == 0 - else { - throw sessionError(created, prefix: "SSH handshake failed") - } - - let observed = try observation( - host: configuration.host, port: configuration.port, session: created - ) - try validateHostKey(observed, expected: configuration.expectedHostKeys) - try authenticate( - username: configuration.username, - credentials: configuration.credentials, - session: created - ) - guard let opened = libssh2_channel_open_ex( created, "session", 7, 2 * 1024 * 1024, 32 * 1024, nil, 0 diff --git a/TermindTests/AgentConnectionPoolTests.swift b/TermindTests/AgentConnectionPoolTests.swift new file mode 100644 index 00000000..8e5c3876 --- /dev/null +++ b/TermindTests/AgentConnectionPoolTests.swift @@ -0,0 +1,345 @@ +import Foundation +import Testing + +@testable import Termind + +@Suite +struct AgentConnectionPoolTests { + @Test + func repeatedCommandsOnOneHostShareASingleConnection() async throws { + let backend = StubBackend() + let pool = makePool(backend) + let host = UUID() + + for index in 0..<5 { + let result = try await pool.execute( + "echo \(index)", workingDirectory: nil, stdin: nil, + timeout: 60, outputLimit: 1_024, hostID: host + ) + #expect(result.exitCode == 0) + } + + // The whole point of the pool: five commands, one handshake. + #expect(await backend.execOpens == 1) + #expect(await backend.execDisconnects == 0) + } + + @Test + func concurrentCallersShareOneConnectionInsteadOfLeakingOne() async throws { + let backend = StubBackend() + let gate = Gate() + let pool = makePool(backend, openerGate: gate) + let host = UUID() + + async let first = pool.exec(hostID: host) + // Only start the second caller once the first is parked mid-connect, so + // it is guaranteed to find work already in flight rather than racing it. + while await gate.arrivals == 0 { await Task.yield() } + async let second = pool.exec(hostID: host) + await gate.open() + + let sessions = try await [first, second] + #expect(try sessions.map(serial) == [1, 1]) + #expect(await backend.execOpens == 1) + } + + @Test + func execAndSFTPAreSeparateConnectionsThatCloseIndependently() async throws { + let backend = StubBackend() + let pool = makePool(backend) + let host = UUID() + + _ = try await pool.exec(hostID: host) + _ = try await pool.sftp(hostID: host) + #expect(await pool.openConnectionCount == 2) + + #expect(await pool.disconnect(hostID: host, kind: .exec)) + #expect(await pool.isConnected(hostID: host, kind: .sftp)) + #expect(await backend.execDisconnects == 1) + #expect(await backend.sftpDisconnects == 0) + } + + @Test + func disconnectingATwiceReportsThatNothingWasOpenTheSecondTime() async throws { + let backend = StubBackend() + let pool = makePool(backend) + let host = UUID() + + _ = try await pool.exec(hostID: host) + #expect(await pool.disconnect(hostID: host, kind: .exec)) + #expect(await pool.disconnect(hostID: host, kind: .exec) == false) + #expect(await backend.execDisconnects == 1) + } + + @Test + func idleConnectionsCloseAndTheNextCommandReconnects() async throws { + let backend = StubBackend() + let pool = makePool(backend) + let host = UUID() + + _ = try await pool.exec(hostID: host) + let remaining = await pool.sweepIdle( + now: .now.advanced(by: AgentConnectionPool.defaultIdleTimeout + .seconds(1)) + ) + + #expect(remaining == 0) + #expect(await backend.execDisconnects == 1) + #expect(await pool.isConnected(hostID: host, kind: .exec) == false) + + _ = try await pool.exec(hostID: host) + #expect(await backend.execOpens == 2) + } + + @Test + func aRecentlyUsedConnectionSurvivesTheSweep() async throws { + let backend = StubBackend() + let pool = makePool(backend) + let host = UUID() + + _ = try await pool.exec(hostID: host) + let remaining = await pool.sweepIdle(now: .now) + + #expect(remaining == 1) + #expect(await backend.execDisconnects == 0) + } + + @Test + func editingAHostsRouteRetiresTheConnectionBoundToTheOldOne() async throws { + let backend = StubBackend() + let pool = makePool(backend) + let host = UUID() + + let before = try await pool.exec(hostID: host) + await backend.setFingerprint("route-b", for: host) + let after = try await pool.exec(hostID: host) + + // Not the same machine any more, so not the same session — otherwise a + // command aimed at the edited host would run on the old one. + #expect(try serial(before) == 1) + #expect(try serial(after) == 2) + #expect(await backend.execOpens == 2) + #expect(await pool.openConnectionCount == 1) + } + + @Test + func aHostThatNoLongerResolvesHasItsConnectionClosed() async throws { + let backend = StubBackend() + let pool = makePool(backend) + let host = UUID() + + _ = try await pool.exec(hostID: host) + await backend.forget(host) + + await #expect(throws: AgentToolError.self) { _ = try await pool.exec(hostID: host) } + #expect(await backend.execDisconnects == 1) + #expect(await pool.openConnectionCount == 0) + } + + @Test + func aSessionThatDiedBetweenCommandsIsReplacedTransparently() async throws { + let backend = StubBackend() + let pool = makePool(backend) + let host = UUID() + + await backend.queueExecFailures([ + SSHExecSessionError.channelUnavailable(TermindSSHError.disconnected) + ]) + let result = try await pool.execute( + "uptime", workingDirectory: nil, stdin: nil, + timeout: 60, outputLimit: 1_024, hostID: host + ) + + #expect(result.exitCode == 0) + #expect(await backend.execOpens == 2) + #expect(await backend.execDisconnects == 1) + } + + @Test + func aCommandThatFailedAfterItStartedIsNeverRunTwice() async throws { + let backend = StubBackend() + let pool = makePool(backend) + let host = UUID() + + // Not `channelUnavailable`: the command reached the remote shell, so a + // silent retry could delete something a second time. + await backend.queueExecFailures([TermindSSHError.library("connection reset")]) + await #expect(throws: TermindSSHError.self) { + _ = try await pool.execute( + "rm -rf /srv//old", workingDirectory: nil, stdin: nil, + timeout: 60, outputLimit: 1_024, hostID: host + ) + } + + #expect(await backend.execOpens == 1) + #expect(await backend.totalExecuteCount() == 1) + } + + @Test + func closeAllReleasesEveryHost() async throws { + let backend = StubBackend() + let pool = makePool(backend) + + for _ in 0..<3 { + let host = UUID() + _ = try await pool.exec(hostID: host) + _ = try await pool.sftp(hostID: host) + } + await pool.closeAll() + + #expect(await pool.openConnectionCount == 0) + #expect(await backend.execDisconnects == 3) + #expect(await backend.sftpDisconnects == 3) + } +} + +/// Which stub connection the pool handed back. Distinguishing "reused" from +/// "reconnected" is the whole subject of these tests. +private func serial(_ session: any SSHExecSessionProtocol) throws -> Int { + try #require(session as? StubExecSession).stubSerial +} + +/// `sweepInterval: nil` keeps the background sweeper out of the tests; every +/// case here drives `sweepIdle` itself with a synthetic instant. +private func makePool( + _ backend: StubBackend, openerGate: Gate? = nil +) -> AgentConnectionPool { + AgentConnectionPool( + configuration: { hostID in + SSHConnectionConfiguration( + hostID: hostID, host: "host.example", port: 22, username: "root", + credentials: [], expectedHostKeys: [], proxy: nil, jumps: [] + ) + }, + routeFingerprint: { try await backend.fingerprint(for: $0) }, + openExec: { _ in + if let openerGate { await openerGate.wait() } + return await backend.openExec() + }, + openSFTP: { _ in await backend.openSFTP() }, + sweepInterval: nil + ) +} + +private actor Gate { + private(set) var arrivals = 0 + private var isOpen = false + private var waiting: [CheckedContinuation] = [] + + func wait() async { + arrivals += 1 + guard !isOpen else { return } + await withCheckedContinuation { waiting.append($0) } + } + + func open() { + isOpen = true + for continuation in waiting { continuation.resume() } + waiting.removeAll() + } +} + +private actor StubBackend { + private(set) var execOpens = 0 + private(set) var sftpOpens = 0 + private(set) var execDisconnects = 0 + private(set) var sftpDisconnects = 0 + private var fingerprints: [UUID: String] = [:] + private var forgotten: Set = [] + private var nextExecFailures: [any Error] = [] + private var execSessions: [StubExecSession] = [] + + func fingerprint(for hostID: UUID) throws -> String { + guard !forgotten.contains(hostID) else { throw AgentToolError.invalidHost } + return fingerprints[hostID] ?? "route-a" + } + + func setFingerprint(_ value: String, for hostID: UUID) { fingerprints[hostID] = value } + func forget(_ hostID: UUID) { forgotten.insert(hostID) } + + /// Applies to the next session opened, so a test can arrange for a pooled + /// connection to fail once and the replacement to succeed. + func queueExecFailures(_ errors: [any Error]) { nextExecFailures = errors } + + func openExec() -> StubExecSession { + execOpens += 1 + let session = StubExecSession( + serial: execOpens, backend: self, failures: nextExecFailures + ) + nextExecFailures = [] + execSessions.append(session) + return session + } + + func openSFTP() -> StubSFTPClient { + sftpOpens += 1 + return StubSFTPClient(backend: self) + } + + func noteExecDisconnect() { execDisconnects += 1 } + func noteSFTPDisconnect() { sftpDisconnects += 1 } + + func totalExecuteCount() async -> Int { + var total = 0 + for session in execSessions { total += await session.executeCount } + return total + } +} + +private actor StubExecSession: SSHExecSessionProtocol { + nonisolated let stubSerial: Int + private let backend: StubBackend + private var failures: [any Error] + private(set) var executeCount = 0 + + init(serial: Int, backend: StubBackend, failures: [any Error]) { + stubSerial = serial + self.backend = backend + self.failures = failures + } + + func connect(using configuration: SSHConnectionConfiguration) async throws {} + + func execute( + _ command: String, workingDirectory: String?, stdin: Data?, + timeout: TimeInterval, outputLimit: Int + ) async throws -> SSHExecResult { + executeCount += 1 + if !failures.isEmpty { throw failures.removeFirst() } + return SSHExecResult( + stdout: command, stderr: "", exitCode: 0, signal: nil, + stdoutTruncated: false, stderrTruncated: false + ) + } + + func disconnect() async { await backend.noteExecDisconnect() } +} + +private actor StubSFTPClient: SFTPClientProtocol { + private let backend: StubBackend + + init(backend: StubBackend) { self.backend = backend } + + func connect( + using configuration: SSHConnectionConfiguration, + progress: SSHRouteProgressHandler? + ) async throws {} + func disconnect() async { await backend.noteSFTPDisconnect() } + func homeDirectory() async throws -> String { "/root" } + func listDirectory(at path: String) async throws -> [RemoteFile] { [] } + func attributes(at path: String) async throws -> RemoteFile? { nil } + func setPermissions(at path: String, to permissions: UInt32) async throws {} + func createDirectory(at path: String) async throws {} + func rename(from source: String, to destination: String, overwrite: Bool) async throws {} + func removeFile(at path: String) async throws {} + func removeDirectory(at path: String) async throws {} + func uploadFile( + from localURL: URL, to remotePath: String, overwrite: Bool, + cancellation: TransferCancellationToken, + progress: @escaping @Sendable (SFTPProgress) -> Void + ) async throws {} + func downloadFile( + from remotePath: String, to localURL: URL, overwrite: Bool, + cancellation: TransferCancellationToken, + progress: @escaping @Sendable (SFTPProgress) -> Void + ) async throws {} +}