diff --git a/Frameworks/Resources/Localizable.xcstrings b/Frameworks/Resources/Localizable.xcstrings index ae35f4e6..bbfd8269 100644 --- a/Frameworks/Resources/Localizable.xcstrings +++ b/Frameworks/Resources/Localizable.xcstrings @@ -1,6 +1,380 @@ { "sourceLanguage" : "en", "strings" : { + "The SFTP copy source grew during transfer." : { + "extractionState" : "manual", + "localizations" : { + "zh-Hans" : { + "stringUnit" : { + "state" : "translated", + "value" : "SFTP复制源在传输期间变大。" + } + }, + "zh-Hant" : { + "stringUnit" : { + "state" : "translated", + "value" : "SFTP複製來源在傳輸期間變大。" + } + } + } + }, + "Unsupported or invalid SFTP copy source: %@" : { + "extractionState" : "manual", + "localizations" : { + "zh-Hans" : { + "stringUnit" : { + "state" : "translated", + "value" : "不支持或无效的SFTP复制源:%@" + } + }, + "zh-Hant" : { + "stringUnit" : { + "state" : "translated", + "value" : "不支援或無效的SFTP複製來源:%@" + } + } + } + }, + "Invalid SFTP copy source path." : { + "extractionState" : "manual", + "localizations" : { + "zh-Hans" : { + "stringUnit" : { + "state" : "translated", + "value" : "SFTP复制源路径无效。" + } + }, + "zh-Hant" : { + "stringUnit" : { + "state" : "translated", + "value" : "SFTP複製來源路徑無效。" + } + } + } + }, + "The SFTP copy source does not exist: %@" : { + "extractionState" : "manual", + "localizations" : { + "zh-Hans" : { + "stringUnit" : { + "state" : "translated", + "value" : "SFTP复制源不存在:%@" + } + }, + "zh-Hant" : { + "stringUnit" : { + "state" : "translated", + "value" : "SFTP複製來源不存在:%@" + } + } + } + }, + "A path cannot be copied into itself." : { + "extractionState" : "manual", + "localizations" : { + "zh-Hans" : { + "stringUnit" : { + "state" : "translated", + "value" : "无法将路径复制到自身。" + } + }, + "zh-Hant" : { + "stringUnit" : { + "state" : "translated", + "value" : "無法將路徑複製到自身。" + } + } + } + }, + "The SFTP copy destination must be an existing directory." : { + "extractionState" : "manual", + "localizations" : { + "zh-Hans" : { + "stringUnit" : { + "state" : "translated", + "value" : "SFTP复制目标必须是现有目录。" + } + }, + "zh-Hant" : { + "stringUnit" : { + "state" : "translated", + "value" : "SFTP複製目標必須是現有目錄。" + } + } + } + }, + "The SFTP copy is too large to count." : { + "extractionState" : "manual", + "localizations" : { + "zh-Hans" : { + "stringUnit" : { + "state" : "translated", + "value" : "SFTP复制大小超出可统计范围。" + } + }, + "zh-Hant" : { + "stringUnit" : { + "state" : "translated", + "value" : "SFTP複製大小超出可統計範圍。" + } + } + } + }, + "The SFTP copy source changed: %@" : { + "extractionState" : "manual", + "localizations" : { + "zh-Hans" : { + "stringUnit" : { + "state" : "translated", + "value" : "SFTP复制源已更改:%@" + } + }, + "zh-Hant" : { + "stringUnit" : { + "state" : "translated", + "value" : "SFTP複製來源已變更:%@" + } + } + } + }, + "Source host" : { + "extractionState" : "manual", + "localizations" : { + "zh-Hans" : { + "stringUnit" : { + "state" : "translated", + "value" : "源主机" + } + }, + "zh-Hant" : { + "stringUnit" : { + "state" : "translated", + "value" : "來源主機" + } + } + } + }, + "Source path" : { + "extractionState" : "manual", + "localizations" : { + "zh-Hans" : { + "stringUnit" : { + "state" : "translated", + "value" : "源路径" + } + }, + "zh-Hant" : { + "stringUnit" : { + "state" : "translated", + "value" : "來源路徑" + } + } + } + }, + "Destination host" : { + "extractionState" : "manual", + "localizations" : { + "zh-Hans" : { + "stringUnit" : { + "state" : "translated", + "value" : "目标主机" + } + }, + "zh-Hant" : { + "stringUnit" : { + "state" : "translated", + "value" : "目標主機" + } + } + } + }, + "Destination path" : { + "extractionState" : "manual", + "localizations" : { + "zh-Hans" : { + "stringUnit" : { + "state" : "translated", + "value" : "目标路径" + } + }, + "zh-Hant" : { + "stringUnit" : { + "state" : "translated", + "value" : "目標路徑" + } + } + } + }, + "Actual destination" : { + "extractionState" : "manual", + "localizations" : { + "zh-Hans" : { + "stringUnit" : { + "state" : "translated", + "value" : "实际目标" + } + }, + "zh-Hant" : { + "stringUnit" : { + "state" : "translated", + "value" : "實際目標" + } + } + } + }, + "Conflict policy" : { + "extractionState" : "manual", + "localizations" : { + "zh-Hans" : { + "stringUnit" : { + "state" : "translated", + "value" : "冲突策略" + } + }, + "zh-Hant" : { + "stringUnit" : { + "state" : "translated", + "value" : "衝突策略" + } + } + } + }, + "Copied files" : { + "extractionState" : "manual", + "localizations" : { + "zh-Hans" : { + "stringUnit" : { + "state" : "translated", + "value" : "已复制文件" + } + }, + "zh-Hant" : { + "stringUnit" : { + "state" : "translated", + "value" : "已複製檔案" + } + } + } + }, + "Created directories" : { + "extractionState" : "manual", + "localizations" : { + "zh-Hans" : { + "stringUnit" : { + "state" : "translated", + "value" : "已创建目录" + } + }, + "zh-Hant" : { + "stringUnit" : { + "state" : "translated", + "value" : "已建立目錄" + } + } + } + }, + "Skipped items" : { + "extractionState" : "manual", + "localizations" : { + "zh-Hans" : { + "stringUnit" : { + "state" : "translated", + "value" : "已跳过项目" + } + }, + "zh-Hant" : { + "stringUnit" : { + "state" : "translated", + "value" : "已略過項目" + } + } + } + }, + "Completed bytes" : { + "extractionState" : "manual", + "localizations" : { + "zh-Hans" : { + "stringUnit" : { + "state" : "translated", + "value" : "已传输大小" + } + }, + "zh-Hant" : { + "stringUnit" : { + "state" : "translated", + "value" : "已傳輸大小" + } + } + } + }, + "Total bytes" : { + "extractionState" : "manual", + "localizations" : { + "zh-Hans" : { + "stringUnit" : { + "state" : "translated", + "value" : "总大小" + } + }, + "zh-Hant" : { + "stringUnit" : { + "state" : "translated", + "value" : "總大小" + } + } + } + }, + "Transfer speed" : { + "extractionState" : "manual", + "localizations" : { + "zh-Hans" : { + "stringUnit" : { + "state" : "translated", + "value" : "传输速度" + } + }, + "zh-Hant" : { + "stringUnit" : { + "state" : "translated", + "value" : "傳輸速度" + } + } + } + }, + "Failed path" : { + "extractionState" : "manual", + "localizations" : { + "zh-Hans" : { + "stringUnit" : { + "state" : "translated", + "value" : "失败路径" + } + }, + "zh-Hant" : { + "stringUnit" : { + "state" : "translated", + "value" : "失敗路徑" + } + } + } + }, + "Overwrite may delete conflicting target directories." : { + "extractionState" : "manual", + "localizations" : { + "zh-Hans" : { + "stringUnit" : { + "state" : "translated", + "value" : "覆盖可能删除冲突的目标目录。" + } + }, + "zh-Hant" : { + "stringUnit" : { + "state" : "translated", + "value" : "覆寫可能刪除衝突的目標目錄。" + } + } + } + }, "" : { "localizations" : { "zh-Hans" : { @@ -27372,4 +27746,4 @@ } }, "version" : "1.2" -} \ No newline at end of file +} diff --git a/Frameworks/TermindRemoteInfrastructure/Agent/AgentConnections.swift b/Frameworks/TermindRemoteInfrastructure/Agent/AgentConnections.swift index a1335857..312afcc0 100644 --- a/Frameworks/TermindRemoteInfrastructure/Agent/AgentConnections.swift +++ b/Frameworks/TermindRemoteInfrastructure/Agent/AgentConnections.swift @@ -781,6 +781,8 @@ public actor AgentSSHAuthenticationRegistry { private func openAuthenticated( hostID: Host.ID, authentication: AgentSSHConfigurationSession, pool: AgentConnectionPool, kind: AgentConnectionPool.Kind, + evictsPooledConnection: Bool = true, + cancellation: TransferCancellationToken? = nil, open: @Sendable (AgentSSHConfigurationSession.Attempt?) async throws -> Client ) async throws -> Client { var attempt: AgentSSHConfigurationSession.Attempt? @@ -788,6 +790,8 @@ public actor AgentSSHAuthenticationRegistry { attempt = try await authentication.configurationAttempt() } while true { + try Task.checkCancellation() + if cancellation?.isCancelled == true { throw SFTPError.cancelled } let attemptedGeneration: UInt64 if let attempt { attemptedGeneration = attempt.generation @@ -797,7 +801,8 @@ public actor AgentSSHAuthenticationRegistry { do { return try await open(attempt) } catch let error as TermindSSHError where error.shouldRequestAuthentication { - await pool.disconnect(hostID: hostID, kind: kind) + if cancellation?.isCancelled == true { throw SFTPError.cancelled } + if evictsPooledConnection { await pool.disconnect(hostID: hostID, kind: kind) } await authentication.retry(after: error, generation: attemptedGeneration) attempt = try await authentication.configurationAttempt() } catch let error as TermindSSHError where error.isLocalAuthenticationRequirement { @@ -830,6 +835,42 @@ public actor AgentSSHAuthenticationRegistry { return try await operation(session) } + /// Retry authentication only while opening; never replay the transfer operation. + public func withDedicatedSFTP( + hostID: Host.ID, invocationID: AgentToolInvocation.ID, + interaction: AgentSSHInteraction, pool: AgentConnectionPool, + cancellation: TransferCancellationToken, + operation: @escaping @Sendable (any SFTPClientProtocol) async throws -> Value + ) async throws -> Value { + try Task.checkCancellation() + guard !cancellation.isCancelled else { throw SFTPError.cancelled } + return try await withTaskCancellationHandler { + try await withSession(hostID: hostID, invocationID: invocationID, interaction: interaction) { session in + let lease = try await self.openAuthenticated( + hostID: hostID, authentication: session, pool: pool, kind: .sftp, + evictsPooledConnection: false, cancellation: cancellation + ) { attempt in + try await pool.acquireDedicatedSFTP(hostID: hostID, cancellation: cancellation) { id in + if let attempt { return attempt.request } + return SSHConnectionRequest(configuration: try await self.resolver.configuration(for: id)) + } + } + do { + try Task.checkCancellation() + guard !cancellation.isCancelled else { throw SFTPError.cancelled } + let value = try await operation(lease.client) + await pool.releaseDedicatedSFTP(lease.id) + return value + } catch { + await pool.releaseDedicatedSFTP(lease.id) + throw error + } + } + } onCancel: { + cancellation.cancel() + } + } + private func release(hostID: Host.ID, token: UUID) { guard var entry = entries[hostID], entry.token == token else { return } entry.leases -= 1 @@ -907,6 +948,47 @@ public actor AgentConnectionPool { private let sweeper: AgentIdleSweeper private var entries: [Key: Entry] = [:] private var pending: [Key: InFlight] = [:] + struct DedicatedLease: Sendable { + var id: UUID + var client: any SFTPClientProtocol + } + private struct DedicatedEntry { + var hostID: UUID + var cancellation: TransferCancellationToken + var client: (any SFTPClientProtocol)? + } + private var dedicated: [UUID: DedicatedEntry] = [:] + + func acquireDedicatedSFTP( + hostID: UUID, cancellation: TransferCancellationToken, + configuration: @escaping RequestResolving + ) async throws -> DedicatedLease { + try Task.checkCancellation() + guard !cancellation.isCancelled else { throw SFTPError.cancelled } + let id = UUID() + dedicated[id] = DedicatedEntry(hostID: hostID, cancellation: cancellation) + do { + _ = try await routeFingerprint(hostID) + let request = try await configuration(hostID) + try Task.checkCancellation() + guard !cancellation.isCancelled else { throw SFTPError.cancelled } + let client = try await openSFTP(request) + guard dedicated[id] != nil, !cancellation.isCancelled, !Task.isCancelled else { + await client.disconnect() + throw SFTPError.cancelled + } + dedicated[id]?.client = client + return DedicatedLease(id: id, client: client) + } catch { + dedicated[id] = nil + throw error + } + } + + func releaseDedicatedSFTP(_ id: UUID) async { + let entry = dedicated.removeValue(forKey: id) + await entry?.client?.disconnect() + } public init( configuration: @escaping ConfigurationResolving, @@ -1003,7 +1085,10 @@ public actor AgentConnectionPool { @discardableResult public func disconnect(hostID: UUID) async -> Bool { - var closed = false + let ids = dedicated.filter { $0.value.hostID == hostID }.map(\.key) + for id in ids { dedicated[id]?.cancellation.cancel() } + for id in ids { await releaseDedicatedSFTP(id) } + var closed = !ids.isEmpty for kind in Kind.allCases { closed = await disconnect(hostID: hostID, kind: kind) || closed } @@ -1011,6 +1096,9 @@ public actor AgentConnectionPool { } public func closeAll() async { + let ids = Array(dedicated.keys) + for id in ids { dedicated[id]?.cancellation.cancel() } + for id in ids { await releaseDedicatedSFTP(id) } let closing = entries.values.map(\.client) entries.removeAll() stopSweepingIfEmpty() @@ -1026,14 +1114,14 @@ public actor AgentConnectionPool { for key in expired.keys { entries[key] = nil } stopSweepingIfEmpty() for entry in expired.values { await entry.client.disconnect() } - return entries.count + return openConnectionCount } public func isConnected(hostID: UUID, kind: Kind) -> Bool { entries[Key(hostID: hostID, kind: kind)] != nil } - public var openConnectionCount: Int { entries.count } + public var openConnectionCount: Int { entries.count + dedicated.values.filter { $0.client != nil }.count } private func resolve( hostID: UUID, kind: Kind, configuration override: RequestResolving? diff --git a/Frameworks/TermindRemoteInfrastructure/Agent/SFTPCopyTool.swift b/Frameworks/TermindRemoteInfrastructure/Agent/SFTPCopyTool.swift new file mode 100644 index 00000000..6caf12d9 --- /dev/null +++ b/Frameworks/TermindRemoteInfrastructure/Agent/SFTPCopyTool.swift @@ -0,0 +1,207 @@ +import AgentKit +import Foundation +import TermindApplication +import TermindDomain + +nonisolated public struct SFTPCopyTool: AgentToolDefinition, AgentToolSchemaBuilding { + public static let presenter = AgentToolDetailPresenter( + id: "builtin.sftp_copy", present: present, presentArguments: presentArguments + ) + public let hosts: AgentHostResolver + public let authentication: AgentSSHAuthenticationRegistry + public let connections: AgentConnectionPool + + public init( + hosts: AgentHostResolver, authentication: AgentSSHAuthenticationRegistry, connections: AgentConnectionPool + ) { + self.hosts = hosts + self.authentication = authentication + self.connections = connections + } + + public var descriptor: AgentToolDescriptor { + Self.descriptor( + "sftp_copy", + "Copy a regular file or directory recursively between hosts using streamed SFTP. " + + "Preserves names and empty directories, not permissions or timestamps. Never deletes the source.", + properties: [ + "source_host_id": Self.string(), "source_path": Self.string(max: 4_096), + "destination_host_id": Self.string(), + "destination_path": Self.described(Self.string(max: 4_096), "An existing destination directory."), + "conflict_policy": Self.described( + Self.enumeration(["skip", "replace", "keep_both"]), + "Defaults to skip. Same-name directories merge. Replace may delete conflicting target directories." + ), + ], + required: ["source_host_id", "source_path", "destination_host_id", "destination_path"], + target: .host, approvalPolicy: .ask, concurrency: .sequential, + presentation: .init( + symbol: "doc.on.doc", activity: .semanticArgument(key: "source_path", fallback: .remoteFile), + output: .json, actionKind: .copy + ) + ) + } + + private struct Request: Sendable { + var sourceHostID: UUID + var sourcePath: String + var destinationHostID: UUID + var destinationPath: String + var policy: TransferConflictPolicy + + init(_ invocation: AgentToolInvocation) throws { + let arguments = try Arguments(invocation) + sourceHostID = try arguments.uuid("source_host_id") + destinationHostID = try arguments.uuid("destination_host_id") + sourcePath = try AgentToolSupport.normalizeRemotePath(arguments.string("source_path")) + destinationPath = try AgentToolSupport.normalizeRemotePath(arguments.string("destination_path")) + switch arguments.optionalString("conflict_policy") ?? "skip" { + case "skip": policy = .skip + case "replace": policy = .replace + case "keep_both": policy = .keepBoth + default: throw AgentToolError.invalidArguments("Unknown SFTP conflict policy.") + } + let name = (sourcePath as NSString).lastPathComponent + let target = destinationPath == "/" ? "/\(name)" : destinationPath + "/" + name + guard sourcePath != "/", sourceHostID != destinationHostID || target != sourcePath else { + throw AgentToolError.invalidPath + } + } + } + + public func preflight(_ invocation: AgentToolInvocation) async throws -> AgentToolPreflight { + let request = try Request(invocation) + let listing = await hosts.listing().hosts + guard let source = listing.first(where: { $0.id == request.sourceHostID }), + let destination = listing.first(where: { $0.id == request.destinationHostID }) + else { throw AgentToolError.invalidHost } + var checked = invocation + checked.targetHostID = request.destinationHostID + return AgentToolPreflight( + invocation: checked, approvalPolicy: .ask, + reasons: [ + "\(source.displayName) (\(source.id)): \(request.sourcePath) → " + + "\(destination.displayName) (\(destination.id)): \(request.destinationPath)" + ], + executionMetadata: [ + "source_host_name": .string(source.displayName), + "destination_host_name": .string(destination.displayName), + ] + ) + } + + public func execute( + _ invocation: AgentToolInvocation, context: AgentToolExecutionContext + ) async throws -> AgentToolResult { + let request = try Request(invocation) + let cancellation = TransferCancellationToken() + return try await withTaskCancellationHandler { + try await authentication.withDedicatedSFTP( + hostID: request.sourceHostID, invocationID: invocation.id, + interaction: context.sshInteraction, pool: connections, cancellation: cancellation + ) { source in + try await authentication.withDedicatedSFTP( + hostID: request.destinationHostID, invocationID: invocation.id, + interaction: context.sshInteraction, pool: connections, cancellation: cancellation + ) { destination in + let outcome = await SFTPRemoteCopyOperation.run( + sourcePath: request.sourcePath, destinationPath: request.destinationPath, + sameHost: request.sourceHostID == request.destinationHostID, + source: source, destination: destination, policy: request.policy, cancellation: cancellation + ) { outcome, progress in + context.report(Self.copyResult(invocation, request, context, outcome, progress: progress)) + } + return Self.copyResult(invocation, request, context, outcome) + } + } + } onCancel: { + cancellation.cancel() + } + } + + private static func copyResult( + _ invocation: AgentToolInvocation, _ request: Request, _ context: AgentToolExecutionContext, + _ outcome: SFTPRemoteCopyOperation.Outcome, progress: SFTPProgress? = nil + ) -> AgentToolResult { + var object: [String: AgentJSONValue] = [ + "source_host_id": .string(request.sourceHostID.uuidString), "source_path": .string(request.sourcePath), + "destination_host_id": .string(request.destinationHostID.uuidString), + "destination_path": .string(request.destinationPath), + "actual_destination_path": .string(outcome.actualDestinationPath), + "conflict_policy": .string(request.policy == .keepBoth ? "keep_both" : request.policy.rawValue), + "copied_files": .number(Double(outcome.copiedFiles)), + "created_directories": .number(Double(outcome.createdDirectories)), + "skipped_items": .number(Double(outcome.skippedItems)), "bytes": .number(Double(outcome.bytes)), + "completed_bytes": .number(Double(outcome.bytes + (progress?.completedBytes ?? 0))), + "total_bytes": .number(Double(outcome.totalBytes)), + "bytes_per_second": .number(Double(progress?.bytesPerSecond ?? 0)), + ] + for key in ["source_host_name", "destination_host_name"] { object[key] = context.preflightMetadata[key] } + if let error = outcome.error { object["error"] = .string(error) } + if let path = outcome.failedPath { object["failed_path"] = .string(path) } + var result = Self.result(invocation, .object(object)) + result.isError = outcome.error != nil + return result + } + + private static func fields(_ object: [String: AgentJSONValue], locale: Locale) -> [AgentToolDetail.Item] { + func field(_ label: String.LocalizationValue, _ value: String) -> AgentToolDetail.Item { + .field( + .init(label: AppLocalization.string(label, locale: locale), value: value, isMonospaced: true)) + } + var items: [AgentToolDetail.Item] = [] + let values: [(String.LocalizationValue, String)] = [ + ("Source host", "source_host_id"), ("Source path", "source_path"), + ("Destination host", "destination_host_id"), ("Destination path", "destination_path"), + ("Actual destination", "actual_destination_path"), + ] + for (label, key) in values { + if let value = object[key]?.stringValue { + let name = + key == "source_host_id" + ? object["source_host_name"]?.stringValue + : key == "destination_host_id" ? object["destination_host_name"]?.stringValue : nil + items.append(field(label, name.map { "\($0) (\(value))" } ?? value)) + } + } + let policy = object["conflict_policy"]?.stringValue ?? "skip" + let policyLabel: String.LocalizationValue = + policy == "replace" ? "Replace" : policy == "keep_both" ? "Keep Both" : "Skip" + items.append(field("Conflict policy", AppLocalization.string(policyLabel, locale: locale))) + if policy == "replace" { + items.append( + .message( + AppLocalization.string("Overwrite may delete conflicting target directories.", locale: locale), + .warning + )) + } + let counts: [(String.LocalizationValue, String)] = [ + ("Copied files", "copied_files"), ("Created directories", "created_directories"), + ("Skipped items", "skipped_items"), + ] + for (label, key) in counts { + if let value = object[key]?.integerValue { items.append(field(label, String(value))) } + } + for (label, key) in [ + (String.LocalizationValue("Completed bytes"), "completed_bytes"), ("Total bytes", "total_bytes"), + ] { + if let value = object[key]?.integerValue { + items.append(field(label, AgentToolDetailFormatting.byteCount(value))) + } + } + if let value = object["bytes_per_second"]?.integerValue { + items.append(field("Transfer speed", AgentToolDetailFormatting.byteCount(value) + "/s")) + } + if let path = object["failed_path"]?.stringValue { items.append(field("Failed path", path)) } + if let error = object["error"]?.stringValue { items.append(.message(error, .failure)) } + return items + } + + public static func present(_ input: AgentToolDetailInput) -> [AgentToolDetail.Item] { + AgentToolDetailFormatting.objectItems(input) { fields($0, locale: input.locale) } + } + + public static func presentArguments(_ input: AgentToolArgumentDetailInput) -> [AgentToolDetail.Item] { + fields(input.arguments, locale: input.locale) + } +} diff --git a/Frameworks/TermindRemoteInfrastructure/Agent/TermindAgentTools.swift b/Frameworks/TermindRemoteInfrastructure/Agent/TermindAgentTools.swift index abad2ae4..b5d20cd9 100644 --- a/Frameworks/TermindRemoteInfrastructure/Agent/TermindAgentTools.swift +++ b/Frameworks/TermindRemoteInfrastructure/Agent/TermindAgentTools.swift @@ -122,6 +122,9 @@ nonisolated public enum TermindAgentTools { .init(SFTPMoveTool.self, availableIn: [.acting]) { SFTPMoveTool(authentication: $0.authentication, connections: $0.connections) }, + .init(SFTPCopyTool.self, availableIn: [.acting]) { + SFTPCopyTool(hosts: $0.hosts, authentication: $0.authentication, connections: $0.connections) + }, .init(SFTPChmodTool.self, availableIn: [.acting]) { SFTPChmodTool(authentication: $0.authentication, connections: $0.connections) }, diff --git a/Frameworks/TermindRemoteInfrastructure/SSH/SFTPClient.swift b/Frameworks/TermindRemoteInfrastructure/SSH/SFTPClient.swift index 53d08635..d60cba2c 100644 --- a/Frameworks/TermindRemoteInfrastructure/SSH/SFTPClient.swift +++ b/Frameworks/TermindRemoteInfrastructure/SSH/SFTPClient.swift @@ -9,9 +9,31 @@ public nonisolated final class TransferCancellationToken: @unchecked Sendable { public init() {} private let lock = NSLock() private var cancelled = false + private var observers: [UUID: @Sendable () -> Void] = [:] public var isCancelled: Bool { lock.withLock { cancelled } } - public func cancel() { lock.withLock { cancelled = true } } + public func cancel() { + let callbacks = lock.withLock { + cancelled = true + let callbacks = Array(observers.values) + observers.removeAll() + return callbacks + } + for callback in callbacks { callback() } + } + + func observeCancellation(_ callback: @escaping @Sendable () -> Void) -> UUID { + let id = UUID() + let alreadyCancelled = lock.withLock { + if cancelled { return true } + observers[id] = callback + return false + } + if alreadyCancelled { callback() } + return id + } + + func removeObserver(_ id: UUID) { _ = lock.withLock { observers.removeValue(forKey: id) } } } public nonisolated struct SFTPProgress: Sendable { @@ -283,39 +305,56 @@ public actor SFTPClient: SFTPClientProtocol { throw SFTPError.operation("The destination cannot receive a streamed SFTP copy.") } let relay = try SFTPSocketRelay() - defer { relay.close() } - try await withThrowingTaskGroup(of: Void.self) { group in - group.addTask { - defer { relay.finishWriting() } - try await self.worker.downloadFile( - from: sourcePath, toSocket: relay.writer, - cancellation: cancellation - ) - } - group.addTask { - try await destination.worker.uploadFile( - fromSocket: relay.reader, to: destinationPath, - byteCount: byteCount, overwrite: overwrite, - cancellation: cancellation, progress: progress - ) - } - do { - while try await group.next() != nil {} - } catch { - cancellation.cancel() - relay.stop() - group.cancelAll() - throw error + let observer = cancellation.observeCancellation { relay.stop() } + defer { + cancellation.removeObserver(observer) + relay.close() + } + try await withTaskCancellationHandler { + try Task.checkCancellation() + try await withThrowingTaskGroup(of: Void.self) { group in + group.addTask { + do { + try await self.worker.downloadFile( + from: sourcePath, toSocket: relay.writer, + cancellation: cancellation + ) + relay.finishWriting() + } catch { + cancellation.cancel() + throw error + } + } + group.addTask { + try await destination.worker.uploadFile( + fromSocket: relay.reader, to: destinationPath, + byteCount: byteCount, overwrite: overwrite, + cancellation: cancellation, progress: progress, + sourceSucceeded: { relay.sourceSucceeded } + ) + } + do { + while try await group.next() != nil {} + } catch { + cancellation.cancel() + relay.stop() + group.cancelAll() + throw error + } } + } onCancel: { + cancellation.cancel() } } } -private nonisolated final class SFTPSocketRelay: @unchecked Sendable { +nonisolated final class SFTPSocketRelay: @unchecked Sendable { public let writer: Int32 public let reader: Int32 private let lock = NSLock() private var closed = false + private var readSucceeded = false + public var sourceSucceeded: Bool { lock.withLock { readSucceeded } } public init() throws { var pair: [Int32] = [-1, -1] @@ -336,6 +375,7 @@ private nonisolated final class SFTPSocketRelay: @unchecked Sendable { public func finishWriting() { lock.withLock { guard !closed else { return } + readSucceeded = true _ = Darwin.shutdown(writer, SHUT_WR) } } @@ -767,7 +807,8 @@ private nonisolated final class SFTPWorker: @unchecked Sendable { fromSocket socket: Int32, to remotePath: String, byteCount: Int64, overwrite: Bool, cancellation: TransferCancellationToken, - progress: @escaping @Sendable (SFTPProgress) -> Void + progress: @escaping @Sendable (SFTPProgress) -> Void, + sourceSucceeded: @escaping @Sendable () -> Bool ) async throws { try await queue.perform { let sftp = try self.requireSFTP() @@ -780,6 +821,9 @@ private nonisolated final class SFTPWorker: @unchecked Sendable { if cancellation.isCancelled { throw SFTPError.cancelled } let count = Darwin.recv(socket, &buffer, buffer.count, 0) if count > 0 { + guard Int64(count) <= byteCount - meter.completed else { + throw SFTPError.operation(String(localized: "The SFTP copy source grew during transfer.")) + } try buffer.withUnsafeBytes { bytes in guard let base = bytes.baseAddress?.assumingMemoryBound(to: CChar.self) else { return } @@ -802,6 +846,7 @@ private nonisolated final class SFTPWorker: @unchecked Sendable { throw SFTPError.operation("The SFTP copy relay was interrupted.") } } + guard !cancellation.isCancelled, sourceSucceeded() else { throw SFTPError.cancelled } guard meter.completed == byteCount else { throw SFTPError.operation( "The source changed while it was being copied " diff --git a/Frameworks/TermindRemoteInfrastructure/Transfers/SFTPRemoteCopyOperation.swift b/Frameworks/TermindRemoteInfrastructure/Transfers/SFTPRemoteCopyOperation.swift new file mode 100644 index 00000000..18e7a280 --- /dev/null +++ b/Frameworks/TermindRemoteInfrastructure/Transfers/SFTPRemoteCopyOperation.swift @@ -0,0 +1,157 @@ +import AgentKit +import Foundation +import TermindApplication +import TermindDomain + +/// An awaited copy independent of browser payloads and transfer UI lifetimes. +public nonisolated enum SFTPRemoteCopyOperation { + public struct Outcome: Sendable { + public var actualDestinationPath: String + public var copiedFiles = 0 + public var createdDirectories = 0 + public var skippedItems = 0 + public var bytes: Int64 = 0 + public var totalBytes: Int64 = 0 + public var failedPath: String? + public var error: String? + } + + public static func run( + sourcePath: String, destinationPath: String, sameHost: Bool, + source: any SFTPClientProtocol, destination: any SFTPClientProtocol, + policy: TransferConflictPolicy, cancellation: TransferCancellationToken, + report: @escaping @Sendable (Outcome, SFTPProgress?) -> Void + ) async -> Outcome { + var outcome = Outcome(actualDestinationPath: destinationPath) + var currentPath = sourcePath + func checkCancellation() throws { + try Task.checkCancellation() + guard !cancellation.isCancelled else { throw SFTPError.cancelled } + } + func contains(_ path: String, in prefix: String) -> Bool { + path == prefix || prefix == "/" || path.hasPrefix(prefix + "/") + } + func join(_ base: String, _ relative: String) -> String { + base == "/" ? "/\(relative)" : base + "/" + relative + } + do { + try checkCancellation() + guard try AgentToolSupport.normalizeRemotePath(sourcePath) == sourcePath, + try AgentToolSupport.normalizeRemotePath(destinationPath) == destinationPath + else { throw SFTPError.operation(String(localized: "Invalid SFTP copy source path.")) } + // LSTAT each ancestor so a link cannot alias a source tree into the destination. + for (path, client) in [(sourcePath, source), (destinationPath, destination)] { + var ancestor = "" + for component in path.split(separator: "/").dropLast() { + try checkCancellation() + ancestor += "/" + component + currentPath = ancestor + guard try await client.attributes(at: ancestor)?.kind == .directory else { + throw SFTPRemoteCopySourceError(path: ancestor) + } + } + } + currentPath = sourcePath + guard let root = try await source.attributes(at: sourcePath) else { + throw SFTPError.operation(String(localized: "The SFTP copy source does not exist: \(sourcePath)")) + } + guard root.path == sourcePath, sourcePath != "/", + root.name == (sourcePath as NSString).lastPathComponent + else { + throw SFTPError.operation(String(localized: "Invalid SFTP copy source path.")) + } + let rootTarget = join(destinationPath, root.name) + outcome.actualDestinationPath = rootTarget + guard + !sameHost || (rootTarget != sourcePath && (!root.isDirectory || !contains(rootTarget, in: sourcePath))) + else { throw SFTPError.operation(String(localized: "A path cannot be copied into itself.")) } + currentPath = destinationPath + guard try await destination.attributes(at: destinationPath)?.isDirectory == true else { + throw SFTPError.operation(String(localized: "The SFTP copy destination must be an existing directory.")) + } + currentPath = sourcePath + let entries = try await SFTPRemoteCopyPlanner.manifest( + files: [root], client: source, regularItemsOnly: true, cancellation: cancellation + ) + for entry in entries where !entry.isDirectory { + let (total, overflow) = outcome.totalBytes.addingReportingOverflow(entry.byteCount) + guard !overflow else { + throw SFTPError.operation(String(localized: "The SFTP copy is too large to count.")) + } + outcome.totalBytes = total + } + report(outcome, nil) + var skipped: [String] = [] + var mappings: [(String, String)] = [] + for entry in entries { + try checkCancellation() + let original = join(destinationPath, entry.relativePath) + if skipped.contains(where: { contains(original, in: $0) }) { + outcome.skippedItems += 1 + continue + } + currentPath = entry.sourcePath + guard let fresh = try await source.attributes(at: entry.sourcePath), fresh.kind == entry.kind, + entry.isDirectory || fresh.size == entry.byteCount + else { + throw SFTPError.operation(String(localized: "The SFTP copy source changed: \(entry.sourcePath)")) + } + var target = original + for mapping in mappings where contains(original, in: mapping.0) { + target = mapping.1 + original.dropFirst(mapping.0.count) + break + } + currentPath = target + try checkCancellation() + let existing = try await destination.attributes(at: target) + if entry.isDirectory { + let resolution = try await SFTPRemoteCopyPlanner.resolveDirectories( + [target], conflicts: existing != nil && existing?.isDirectory != true ? [target] : [], + policy: policy, client: destination + ) + outcome.createdDirectories += resolution.created.count + if !resolution.skipped.isEmpty { + skipped.append(original) + outcome.skippedItems += 1 + } + if let mapping = resolution.mappings.first { + // Most specific mapping wins, including remapped nested directories. + mappings.insert((original, mapping.1), at: 0) + if original == rootTarget { outcome.actualDestinationPath = mapping.1 } + } + } else { + var overwrite = false + if let existing { + guard + let resolution = try await SFTPRemoteCopyPlanner.resolveConflict( + path: target, conflictingTargetIsDirectory: existing.isDirectory, + policy: policy, client: destination + ) + else { + outcome.skippedItems += 1 + report(outcome, nil) + continue + } + target = resolution.path + overwrite = resolution.overwrite + } + currentPath = target + if original == rootTarget { outcome.actualDestinationPath = target } + let snapshot = outcome + try checkCancellation() + try await source.copyFile( + from: entry.sourcePath, to: target, on: destination, + byteCount: entry.byteCount, overwrite: overwrite, cancellation: cancellation + ) { progress in report(snapshot, progress) } + outcome.copiedFiles += 1 + outcome.bytes += entry.byteCount + } + report(outcome, nil) + } + } catch { + outcome.failedPath = (error as? SFTPRemoteCopySourceError)?.path ?? currentPath + outcome.error = error.localizedDescription + } + return outcome + } +} diff --git a/Frameworks/TermindRemoteInfrastructure/Transfers/SFTPRemoteTransferPlanner.swift b/Frameworks/TermindRemoteInfrastructure/Transfers/SFTPRemoteTransferPlanner.swift index b29152f1..4311af07 100644 --- a/Frameworks/TermindRemoteInfrastructure/Transfers/SFTPRemoteTransferPlanner.swift +++ b/Frameworks/TermindRemoteInfrastructure/Transfers/SFTPRemoteTransferPlanner.swift @@ -4,6 +4,13 @@ import TermindDomain import struct TermindDomain.Host +nonisolated struct SFTPRemoteCopySourceError: LocalizedError { + var path: String + var errorDescription: String? { + String(localized: "Unsupported or invalid SFTP copy source: \(path)") + } +} + public nonisolated struct SFTPRemoteDownloadBatch: Sendable { public struct Job: Sendable { var remotePath: String @@ -180,6 +187,37 @@ public nonisolated struct SFTPRemoteCopyBatch: Sendable { } public nonisolated enum SFTPRemoteCopyPlanner { + static func manifest( + files: [RemoteFile], client: any SFTPClientProtocol, regularItemsOnly: Bool = false, + cancellation: TransferCancellationToken? = nil + ) async throws -> [SFTPRemoteManifestEntry] { + var entries: [SFTPRemoteManifestEntry] = [] + func collect(_ file: RemoteFile, relativePath: String) async throws { + try Task.checkCancellation() + if cancellation?.isCancelled == true { throw SFTPError.cancelled } + if regularItemsOnly { + guard file.kind == .file || file.kind == .directory, + RemoteFileName.isValid(file.name), file.size >= 0 + else { throw SFTPRemoteCopySourceError(path: file.path) } + } + entries.append( + SFTPRemoteManifestEntry( + sourcePath: file.path, relativePath: relativePath, + kind: file.kind, byteCount: file.size + )) + guard file.isDirectory else { return } + for child in try await client.listDirectory(at: file.path) { + if regularItemsOnly { + let expected = file.path == "/" ? "/\(child.name)" : "\(file.path)/\(child.name)" + guard child.path == expected else { throw SFTPRemoteCopySourceError(path: child.path) } + } + try await collect(child, relativePath: relativePath + "/" + child.name) + } + } + for file in files { try await collect(file, relativePath: file.name) } + return entries + } + public static func prepare( _ payload: SFTPRemoteItemsPayload, destination: String, @@ -275,9 +313,10 @@ public nonisolated enum SFTPRemoteCopyPlanner { var conflictCount: Int } - private struct DirectoryResolution { + struct DirectoryResolution { var skipped: [String] = [] var mappings: [(String, String)] = [] + var created: [String] = [] } private static func preflight( @@ -309,12 +348,13 @@ public nonisolated enum SFTPRemoteCopyPlanner { ) } - private static func resolveDirectories( + static func resolveDirectories( _ directories: [String], conflicts: Set, policy: TransferConflictPolicy, client: any SFTPClientProtocol ) async throws -> DirectoryResolution { var result = DirectoryResolution() for original in directories { + try Task.checkCancellation() if result.skipped.contains(where: { contains(original, in: $0) }) { continue } var destination = original for mapping in result.mappings where contains(destination, in: mapping.0) { @@ -334,12 +374,13 @@ public nonisolated enum SFTPRemoteCopyPlanner { } if try await client.attributes(at: destination) == nil { try await client.createDirectory(at: destination) + result.created.append(destination) } } return result } - private static func resolveConflict( + static func resolveConflict( path: String, conflictingTargetIsDirectory: Bool, policy: TransferConflictPolicy, client: any SFTPClientProtocol ) async throws -> (path: String, overwrite: Bool)? { @@ -361,6 +402,7 @@ public nonisolated enum SFTPRemoteCopyPlanner { var index = 2 var candidate = RemotePath.keepBothCandidate(for: path, index: index) while try await client.attributes(at: candidate) != nil { + try Task.checkCancellation() index += 1 candidate = RemotePath.keepBothCandidate(for: path, index: index) } @@ -370,8 +412,13 @@ public nonisolated enum SFTPRemoteCopyPlanner { private static func deleteRecursively( _ file: RemoteFile, client: any SFTPClientProtocol ) async throws { + try Task.checkCancellation() if file.isDirectory { for child in try await client.listDirectory(at: file.path) { + let expected = file.path == "/" ? "/\(child.name)" : file.path + "/" + child.name + guard RemoteFileName.isValid(child.name), child.path == expected else { + throw SFTPRemoteCopySourceError(path: child.path) + } try await deleteRecursively(child, client: client) } try await client.removeDirectory(at: file.path) diff --git a/Frameworks/TermindRemoteInfrastructure/Transfers/SFTPTransferCoordinator.swift b/Frameworks/TermindRemoteInfrastructure/Transfers/SFTPTransferCoordinator.swift index f590edc7..79cfb67b 100644 --- a/Frameworks/TermindRemoteInfrastructure/Transfers/SFTPTransferCoordinator.swift +++ b/Frameworks/TermindRemoteInfrastructure/Transfers/SFTPTransferCoordinator.swift @@ -196,24 +196,7 @@ public final class SFTPTransferCoordinator { let lease = try await clientLease(for: source.connection) let client = lease.client do { - var entries: [SFTPRemoteManifestEntry] = [] - func collect(_ file: RemoteFile, relativePath: String) async throws { - entries.append( - SFTPRemoteManifestEntry( - sourcePath: file.path, relativePath: relativePath, - kind: file.kind, byteCount: file.size - )) - guard file.isDirectory else { return } - for child in try await client.listDirectory(at: file.path) { - try await collect( - child, - relativePath: relativePath + "/" + child.name - ) - } - } - for file in source.files { - try await collect(file, relativePath: file.name) - } + let entries = try await SFTPRemoteCopyPlanner.manifest(files: source.files, client: client) await lease.disconnect() return (source, entries) } catch { diff --git a/Memory/ARCHITECTURE.md b/Memory/ARCHITECTURE.md index dc15e570..bf288ab8 100644 --- a/Memory/ARCHITECTURE.md +++ b/Memory/ARCHITECTURE.md @@ -170,6 +170,30 @@ Infrastructure planners. Download batches retain their destination independently file jobs, so directory-only selections still create their empty directory trees. App sessions retain pending-batch presentation, cancellation, and completion handling. +`sftp_copy` is Termind-owned and acting-only in Assistant, with Ask/sequential +policy. It copies regular files and directory trees into an existing directory; +unsupported descendants fail before writes. Its awaited operation shares browser +manifest and conflict helpers, including directory merging and destructive type +replacement. Source data is never deleted. Batches are not atomic: failures retain +completed items and disclose the failing path and completed counts. +Two dedicated SFTP leases use the same authenticated connection opener as other +Agent tools, but are neither pooled nor idle-swept. Their shared cancellation token +is cancelled by host deletion, transport shutdown, or caller cancellation. Leases +remain tracked while opening and are always disconnected after use. Credential +retries apply only to acquisition, never to an operation that has started. +Socket-relayed copies require successful source EOF and an exact byte count before +the remote temporary file is committed. Cancellation shuts down the relay to unblock +socket I/O. No scratch staging, local temporary file, or model-visible payload is used. + +The macOS MCP Server exports the same `sftp_copy` definition without a scratch +workspace. Its existing Bearer token/tool switches and Full access execution remain +the authorization boundary; SSH credentials and host keys still use external +authentication. The tool is annotated as state-changing, potentially destructive, +non-idempotent, and open-world. Session termination/server shutdown cancels and awaits +active copies. Audit consumes byte progress and final counts only, never file bodies; +the summary records both endpoints and the conflict policy. The stdio bridge remains +a protocol-only forwarder. Mobile exposes Assistant tools but no local MCP listener. + ## Callback isolation Speech services and coordinator state are MainActor-isolated, but authorization diff --git a/TFI/App/MobileAppModel.swift b/TFI/App/MobileAppModel.swift index 46ea4f56..38822069 100644 --- a/TFI/App/MobileAppModel.swift +++ b/TFI/App/MobileAppModel.swift @@ -269,6 +269,7 @@ final class MobileAppModel { func deleteHost(_ id: UUID) { do { + assistant.disconnectHost(id) for workspace in workspaces.filter({ $0.descriptor.hostID == id }) { close(workspace.id) } transfers.cancelAll(for: id) forwardings.remove(library.portForwardings(for: id).map(\.id)) diff --git a/TFI/Application/Agent/MobileAssistantModel.swift b/TFI/Application/Agent/MobileAssistantModel.swift index 42328830..6edb5fa2 100644 --- a/TFI/Application/Agent/MobileAssistantModel.swift +++ b/TFI/Application/Agent/MobileAssistantModel.swift @@ -29,6 +29,10 @@ final class MobileAssistantModel { var isRunning: Bool { store.conversations.contains { $0.isStreaming } } + func disconnectHost(_ id: UUID) { + Task { await connections.disconnect(hostID: id) } + } + func makeSceneModel() -> MobileAssistantSceneModel { MobileAssistantSceneModel(runtime: self, library: AssistantLibraryModel(store: store)) } diff --git a/TFITests/MobileSFTPCopyAvailabilityTests.swift b/TFITests/MobileSFTPCopyAvailabilityTests.swift new file mode 100644 index 00000000..9755ded5 --- /dev/null +++ b/TFITests/MobileSFTPCopyAvailabilityTests.swift @@ -0,0 +1,57 @@ +import AgentKit +import Foundation +import TermindApplication +import TermindRemoteInfrastructure +import Testing + +@Suite +@MainActor +struct MobileSFTPCopyAvailabilityTests { + @Test + func copyIsActingOnlyAndDoesNotRequireAScratchWorkspace() throws { + let hosts = AgentHostResolver(store: HostLibraryModel()) + let connections = AgentConnectionPool( + configuration: { _ in throw AgentToolError.invalidHost }, routeFingerprint: { _ in "test" }, + openExec: { _ in throw AgentToolError.invalidHost }, openSFTP: { _ in throw AgentToolError.invalidHost }, + sweepInterval: nil + ) + let registry = AgentToolRegistry( + TermindAgentTools.tools( + in: .init( + hosts: hosts, metrics: EmptyMetricsReader(), connections: connections + ))) + let descriptor = try #require(registry.available(in: .acting).descriptors.first { $0.name == "sftp_copy" }) + #expect(descriptor.approvalPolicy == .ask) + #expect(descriptor.concurrency == .sequential) + #expect(!registry.available(in: .planning).descriptors.contains { $0.name == "sftp_copy" }) + #expect(registry.available(in: .reviewing).descriptors.isEmpty) + #expect(TermindAgentTools.presenter(for: descriptor)?.id == "builtin.sftp_copy") + } + + @Test(arguments: ["en", "zh-Hans", "zh-Hant"]) + func copyApprovalLabelsAndOverwriteWarningUseAppResources(_ language: String) { + let items = SFTPCopyTool.presentArguments( + .init( + arguments: [ + "source_host_id": .string("source"), "destination_host_id": .string("destination"), + "conflict_policy": .string("replace"), + ], locale: Locale(identifier: language) + )) + let label = language == "en" ? "Source host" : language == "zh-Hant" ? "來源主機" : "源主机" + #expect( + items.contains { + if case .field(let field) = $0 { return field.label == label } + return false + }) + #expect( + items.contains { + if case .message(_, .warning) = $0 { return true } + return false + }) + } +} + +@MainActor +private struct EmptyMetricsReader: MetricsSnapshotReading { + func snapshot(for hostID: UUID) -> MetricsSnapshot? { nil } +} diff --git a/TFM/Application/FeatureModels/ToolAuditModel.swift b/TFM/Application/FeatureModels/ToolAuditModel.swift index fc382545..930fa7a5 100644 --- a/TFM/Application/FeatureModels/ToolAuditModel.swift +++ b/TFM/Application/FeatureModels/ToolAuditModel.swift @@ -90,6 +90,15 @@ final class ToolAuditModel { begin(invocation, descriptor: descriptor) case .toolFinished(let invocation, let result): finish(invocation, result: result) + case .toolProgress(let invocation, let result) where invocation.call.name == "sftp_copy": + if let id = rowsByInvocation[invocation.id], var row = value(id: id), + let payload = result.untrustedPayload, + let object = try? AgentJSONValue.decode(Data(payload.utf8)).objectValue + { + row.bytesRead = object["completed_bytes"]?.integerValue + row.bytesWritten = row.bytesRead + update(row) + } default: break } @@ -132,6 +141,7 @@ final class ToolAuditModel { { row.exitCode = object["exit_code"]?.integerValue row.bytesRead = object["bytes"]?.integerValue + if invocation.call.name == "sftp_copy" { row.bytesWritten = object["bytes"]?.integerValue } if invocation.call.name == "ssh_exec", let raw = object["handle"]?.stringValue, let handle = UUID(uuidString: raw), let hostID = row.hostID @@ -193,6 +203,12 @@ final class ToolAuditModel { case "sftp_move": return "\(oneLine(arguments["path"]?.stringValue ?? "")) → " + oneLine(arguments["destination"]?.stringValue ?? "") + case "sftp_copy": + return "\(arguments["source_host_id"]?.stringValue ?? ""):" + + "\(oneLine(arguments["source_path"]?.stringValue ?? "")) → " + + "\(arguments["destination_host_id"]?.stringValue ?? ""):" + + "\(oneLine(arguments["destination_path"]?.stringValue ?? "")) " + + "[\(arguments["conflict_policy"]?.stringValue ?? "skip")]" case "sftp_chmod": return "\(oneLine(arguments["path"]?.stringValue ?? "")) " + (arguments["mode"]?.stringValue ?? "") @@ -226,6 +242,9 @@ final class ToolAuditModel { ].contains(call.name) { return .readOnly } - return call.name == "sftp_delete" ? .destructive : .modifiesState + let replacingCopy = + call.name == "sftp_copy" + && call.arguments.objectValue?["conflict_policy"]?.stringValue == "replace" + return call.name == "sftp_delete" || replacingCopy ? .destructive : .modifiesState } } diff --git a/TFM/Application/MCPServer/MCPServerTools.swift b/TFM/Application/MCPServer/MCPServerTools.swift index b4a310ae..91700cae 100644 --- a/TFM/Application/MCPServer/MCPServerTools.swift +++ b/TFM/Application/MCPServer/MCPServerTools.swift @@ -8,7 +8,7 @@ final class MCPServerTools { static let exportedNames: Set = [ "list_hosts", "get_metrics", "request_user_secret", "ssh_exec", "ssh_wait", "sftp_list", "sftp_stat", "sftp_read", "sftp_write", - "sftp_mkdir", "sftp_move", "sftp_chmod", "sftp_delete", + "sftp_mkdir", "sftp_move", "sftp_copy", "sftp_chmod", "sftp_delete", ] private let channel: AgentEventChannel @@ -16,6 +16,7 @@ final class MCPServerTools { private let externalAuthentication: ExternalAuthCoordinator private let approval: AgentApprovalBroker private var registry: AgentToolRegistry + private var copies: [UUID: [UUID: Task]] = [:] init( hosts: AgentHostResolver, metrics: any MetricsSnapshotReading, @@ -75,11 +76,26 @@ final class MCPServerTools { userIntent: "MCP client: \(clientName)" ) let call = AgentToolCall(id: callID, name: name, arguments: arguments) - return await executor.execute(executor.plan(call)) + guard name == "sftp_copy" else { return await executor.execute(executor.plan(call)) } + let id = UUID() + let task = Task { await executor.execute(executor.plan(call)) } + copies[runID, default: [:]][id] = task + defer { + copies[runID]?[id] = nil + if copies[runID]?.isEmpty == true { copies[runID] = nil } + } + return await withTaskCancellationHandler { + await task.value + } onCancel: { + task.cancel() + } } func endSession(runID: UUID, secretBroker: SecretBroker) async { + let tasks = Array(copies.removeValue(forKey: runID)?.values ?? [:].values) + for task in tasks { task.cancel() } externalAuthentication.cancel(runID: runID) + for task in tasks { _ = await task.value } history.unregister(runID: runID) await secretBroker.discardAll() } diff --git a/TFM/Infrastructure/MCP/MCPServerProtocol.swift b/TFM/Infrastructure/MCP/MCPServerProtocol.swift index e95e09ae..fe47b3fe 100644 --- a/TFM/Infrastructure/MCP/MCPServerProtocol.swift +++ b/TFM/Infrastructure/MCP/MCPServerProtocol.swift @@ -369,7 +369,7 @@ actor MCPServerProtocol { "list_hosts", "get_metrics", "request_user_secret", "ssh_wait", "sftp_list", "sftp_stat", "sftp_read", ].contains(descriptor.name) - let destructive = ["ssh_exec", "sftp_write", "sftp_move", "sftp_chmod", "sftp_delete"] + let destructive = ["ssh_exec", "sftp_write", "sftp_move", "sftp_copy", "sftp_chmod", "sftp_delete"] .contains(descriptor.name) let openWorld = !["list_hosts", "ssh_wait", "request_user_secret"] .contains(descriptor.name) diff --git a/TFMTests/AgentBuiltInToolAvailabilityTests.swift b/TFMTests/AgentBuiltInToolAvailabilityTests.swift index b731eb5f..4352dae1 100644 --- a/TFMTests/AgentBuiltInToolAvailabilityTests.swift +++ b/TFMTests/AgentBuiltInToolAvailabilityTests.swift @@ -54,6 +54,8 @@ struct AgentBuiltInToolAvailabilityTests { #expect(registry.available(in: .reviewing).descriptors.isEmpty) #expect(planning.contains("ssh_exec")) #expect(!planning.contains("sftp_write")) + #expect(!planning.contains("sftp_copy")) + #expect(acting.contains("sftp_copy")) #expect(acting.isSuperset(of: ["ssh_exec", "sftp_write"])) } diff --git a/TFMTests/AgentHostAuthenticationTests.swift b/TFMTests/AgentHostAuthenticationTests.swift index 53832e8e..db74c080 100644 --- a/TFMTests/AgentHostAuthenticationTests.swift +++ b/TFMTests/AgentHostAuthenticationTests.swift @@ -931,6 +931,71 @@ struct AgentHostAuthenticationTests { #expect(await setup.backend.sftpOpens == 2) } + @Test + func dedicatedSFTPRetriesCredentialsOnlyBeforeTheOperation() async throws { + let setup = authenticationSetup(sftpAuthenticationFailures: 1) + var prompts = 0 + let interaction = AuthenticationResponder { _ in + prompts += 1 + return .submit(.init(username: "deploy", password: "secret-\(prompts)", save: false)) + } + let files = try await setup.registry.withDedicatedSFTP( + hostID: setup.host.id, invocationID: UUID(), interaction: interaction.sshInteraction, + pool: setup.pool, cancellation: TransferCancellationToken() + ) { client in try await client.listDirectory(at: "/") } + #expect(files.isEmpty) + #expect(prompts == 2) + #expect(await setup.backend.sftpOpens == 2) + #expect(await setup.backend.sftpOperations == 1) + #expect(await setup.pool.openConnectionCount == 0) + } + + @Test + func dedicatedSFTPOperationFailureNeverRetries() async throws { + let setup = authenticationSetup(sftpOperationFailures: 1) + let interaction = AuthenticationResponder { _ in + .submit(.init(username: "deploy", password: "secret", save: false)) + } + await #expect(throws: SFTPError.self) { + _ = try await setup.registry.withDedicatedSFTP( + hostID: setup.host.id, invocationID: UUID(), interaction: interaction.sshInteraction, + pool: setup.pool, cancellation: TransferCancellationToken() + ) { client in try await client.listDirectory(at: "/") } + } + #expect(await setup.backend.sftpOpens == 1) + #expect(await setup.backend.sftpOperations == 1) + #expect(await setup.pool.openConnectionCount == 0) + } + + @Test + func sameHostDedicatedConnectionsShareOnlyTheScopedAuthenticationSession() async throws { + let setup = authenticationSetup() + var prompts = 0 + let interaction = AuthenticationResponder { _ in + prompts += 1 + return .submit(.init(username: "deploy", password: "secret", save: false)) + } + let registry = setup.registry + let pool = setup.pool + let hostID = setup.host.id + let cancellation = TransferCancellationToken() + _ = try await registry.withDedicatedSFTP( + hostID: hostID, invocationID: UUID(), interaction: interaction.sshInteraction, + pool: pool, cancellation: cancellation + ) { first in + try await registry.withDedicatedSFTP( + hostID: hostID, invocationID: UUID(), interaction: interaction.sshInteraction, + pool: pool, cancellation: cancellation + ) { second in + _ = try await first.listDirectory(at: "/") + return try await second.listDirectory(at: "/") + } + } + #expect(prompts == 1) + #expect(await setup.backend.sftpOpens == 2) + #expect(await setup.pool.openConnectionCount == 0) + } + private func authenticationSetup( sftpAuthenticationFailures: Int = 0, sftpOperationFailures: Int = 0, diff --git a/TFMTests/ArchitectureBoundaryTests.swift b/TFMTests/ArchitectureBoundaryTests.swift index 2ae17c4e..d3b9dac6 100644 --- a/TFMTests/ArchitectureBoundaryTests.swift +++ b/TFMTests/ArchitectureBoundaryTests.swift @@ -337,6 +337,17 @@ struct ArchitectureBoundaryTests { ) #expect(transport.contains("SFTPClientProtocol: RemoteFileSystem")) #expect(!transport.contains("struct RemoteFile")) + #expect(transport.contains("sourceSucceeded()")) + #expect(transport.contains("observeCancellation { relay.stop() }")) + + let copy = try String( + contentsOf: sourceRoot.appending( + path: "../Frameworks/TermindRemoteInfrastructure/Agent/SFTPCopyTool.swift"), + encoding: .utf8 + ) + #expect(copy.contains("withDedicatedSFTP")) + #expect(!copy.contains("AgentScratchWorkspace")) + #expect(!copy.contains("SFTPClient()")) let mobile = try String( contentsOf: sourceRoot.appending(path: "../TFI/Application/Transfers/MobileSFTPSessionModel.swift"), diff --git a/TFMTests/MCPServerSessionTests.swift b/TFMTests/MCPServerSessionTests.swift index 7120bff6..27ea616d 100644 --- a/TFMTests/MCPServerSessionTests.swift +++ b/TFMTests/MCPServerSessionTests.swift @@ -56,7 +56,7 @@ struct MCPServerSessionTests { let tools = try #require( try responseObject(listed)["result"]?.objectValue?["tools"]?.arrayValue ) - #expect(tools.count == 13) + #expect(tools.count == 14) let readOnly = Set([ "list_hosts", "get_metrics", "request_user_secret", "ssh_wait", "sftp_list", "sftp_stat", "sftp_read", @@ -65,6 +65,12 @@ struct MCPServerSessionTests { let name = try #require(tool.objectValue?["name"]?.stringValue) let hint = tool.objectValue?["annotations"]?.objectValue?["readOnlyHint"] #expect(hint == .bool(readOnly.contains(name)), "\(name)") + if name == "sftp_copy" { + let annotations = try #require(tool.objectValue?["annotations"]?.objectValue) + #expect(annotations["destructiveHint"] == .bool(true)) + #expect(annotations["idempotentHint"] == .bool(false)) + #expect(annotations["openWorldHint"] == .bool(true)) + } } } diff --git a/TFMTests/MCPServerToolsTests.swift b/TFMTests/MCPServerToolsTests.swift index 84a39955..1d9cd7ce 100644 --- a/TFMTests/MCPServerToolsTests.swift +++ b/TFMTests/MCPServerToolsTests.swift @@ -20,7 +20,7 @@ struct MCPServerToolsTests { let server = makeMCPServerTools() let descriptors = server.descriptors - #expect(descriptors.count == 13) + #expect(descriptors.count == 14) #expect(Set(descriptors.map(\.name)) == MCPServerTools.exportedNames) for descriptor in descriptors { #expect(descriptor.inputSchema.objectValue?["type"]?.stringValue == "object") diff --git a/TFMTests/SFTPCopyToolTests.swift b/TFMTests/SFTPCopyToolTests.swift new file mode 100644 index 00000000..f16b93a9 --- /dev/null +++ b/TFMTests/SFTPCopyToolTests.swift @@ -0,0 +1,518 @@ +import AgentKit +import Foundation +import TermindApplication +import TermindDomain +import TermindRemoteInfrastructure +import Testing + +@testable import Termind +@testable import TermindRemoteInfrastructure + +@Suite(.serialized) +@MainActor +struct SFTPCopyToolTests { + @Test(arguments: [Data(), Data([0, 255, 128, 10]), Data(count: AgentScratchWorkspace.Limits.fileBytes + 1)]) + func copiesBytesWithoutScratchOrInlineIO(_ bytes: Data) async throws { + let fixture = CopyFixture() + await fixture.source.put("/source/a.bin", bytes: bytes) + let result = try await fixture.copy() + #expect(!result.isError) + #expect(await fixture.destination.data["/target/a.bin"] == bytes) + #expect(await fixture.source.data["/source/a.bin"] == bytes) + let object = try fixture.object(result) + #expect(object["bytes"]?.integerValue == bytes.count) + #expect(object["copied_files"]?.integerValue == 1) + #expect(await fixture.backend.opens == 2) + #expect(await fixture.backend.disconnects == 2) + #expect(await fixture.pool.openConnectionCount == 0) + } + + @Test + func preservesHierarchyAndEmptyDirectories() async throws { + let fixture = CopyFixture() + await fixture.source.directory("/source/tree") + await fixture.source.directory("/source/tree/empty") + await fixture.source.directory("/source/tree/nested") + await fixture.source.put("/source/tree/nested/a", bytes: Data([1, 2])) + let result = try await fixture.copy(path: "/source/tree") + let object = try fixture.object(result) + #expect(!result.isError) + #expect(object["created_directories"]?.integerValue == 3) + #expect(await fixture.destination.nodes["/target/tree/empty"]?.isDirectory == true) + #expect(await fixture.destination.data["/target/tree/nested/a"] == Data([1, 2])) + } + + @Test(arguments: ["skip", "replace", "keep_both"]) + func resolvesSameNameFilesUsingBrowserPolicies(_ policy: String) async throws { + let fixture = CopyFixture() + await fixture.source.put("/source/a.bin", bytes: Data([1])) + await fixture.destination.put("/target/a.bin", bytes: Data([2])) + let result = try await fixture.copy(policy: policy) + #expect(!result.isError) + let object = try fixture.object(result) + #expect(object["skipped_items"]?.integerValue == (policy == "skip" ? 1 : 0)) + let target = policy == "keep_both" ? "/target/a.2.bin" : "/target/a.bin" + #expect(await fixture.destination.data[target] == Data([policy == "skip" ? 2 : 1])) + #expect(object["actual_destination_path"]?.stringValue == target) + } + + @Test(arguments: ["skip", "replace", "keep_both"]) + func resolvesFileOverDirectoryConflicts(_ policy: String) async throws { + let fixture = CopyFixture() + await fixture.source.put("/source/a.bin", bytes: Data([1])) + await fixture.destination.directory("/target/a.bin") + await fixture.destination.put("/target/a.bin/old", bytes: Data([2])) + let result = try await fixture.copy(policy: policy) + #expect(!result.isError) + if policy == "replace" { + #expect(await fixture.destination.nodes["/target/a.bin/old"] == nil) + #expect(await fixture.destination.data["/target/a.bin"] == Data([1])) + } else { + #expect(await fixture.destination.data["/target/a.bin/old"] == Data([2])) + } + } + + @Test(arguments: ["skip", "replace", "keep_both"]) + func resolvesDirectoryOverFileAndMapsDescendants(_ policy: String) async throws { + let fixture = CopyFixture() + await fixture.source.directory("/source/tree") + await fixture.source.directory("/source/tree/empty") + await fixture.source.put("/source/tree/a", bytes: Data([1])) + await fixture.destination.put("/target/tree", bytes: Data([2])) + let result = try await fixture.copy(path: "/source/tree", policy: policy) + #expect(!result.isError) + let object = try fixture.object(result) + if policy == "skip" { + #expect(object["skipped_items"]?.integerValue == 3) + #expect(await fixture.destination.data["/target/tree"] == Data([2])) + } else { + let target = policy == "keep_both" ? "/target/tree.2" : "/target/tree" + #expect(object["actual_destination_path"]?.stringValue == target) + #expect(await fixture.destination.nodes[target + "/empty"]?.isDirectory == true) + #expect(await fixture.destination.data[target + "/a"] == Data([1])) + } + } + + @Test + func mergesDirectoriesAndKeepsUnrelatedDestinationFiles() async throws { + let fixture = CopyFixture() + await fixture.source.directory("/source/tree") + await fixture.source.put("/source/tree/a", bytes: Data([1])) + await fixture.destination.directory("/target/tree") + await fixture.destination.put("/target/tree/old", bytes: Data([2])) + #expect(try await !fixture.copy(path: "/source/tree", policy: "replace").isError) + #expect(await fixture.destination.data["/target/tree/old"] == Data([2])) + } + + @Test(arguments: [RemoteFile.Kind.symbolicLink, .special, .unknown]) + func rejectsUnsupportedDescendantsBeforeWriting(_ kind: RemoteFile.Kind) async throws { + let fixture = CopyFixture() + await fixture.source.directory("/source/tree") + await fixture.source.put("/source/tree/a", bytes: Data([1])) + await fixture.source.node("/source/tree/z", kind: kind) + let result = try await fixture.copy(path: "/source/tree") + #expect(result.isError) + #expect(await fixture.destination.nodes.count == 1) + } + + @Test(arguments: ["relative", "/../escape", "/a\nb", "/a\0b", "/"]) + func rejectsInvalidSourcePathsWithoutConnecting(_ path: String) async throws { + let fixture = CopyFixture() + await #expect(throws: (any Error).self) { _ = try await fixture.copy(path: path) } + #expect(await fixture.backend.opens == 0) + } + + @Test + func sameHostUsesIndependentConnectionsAndRejectsSelfCopies() async throws { + let fixture = CopyFixture(sameHost: true) + await fixture.source.directory("/target") + await fixture.source.put("/source/a.bin", bytes: Data([1])) + #expect(try await !fixture.copy().isError) + #expect(await fixture.source.data["/target/a.bin"] == Data([1])) + #expect(await fixture.backend.opens == 2) + await #expect(throws: (any Error).self) { + _ = try await fixture.copy(destination: "/source") + } + await fixture.source.directory("/source/tree") + await fixture.source.directory("/source/tree/nested") + #expect(try await fixture.copy(path: "/source/tree", destination: "/source/tree/nested").isError) + #expect(await fixture.source.nodes["/source/tree/nested/tree"] == nil) + } + + @Test + func detectsChangedSourceAndReportsPartialCompletionWithoutRetry() async throws { + let fixture = CopyFixture() + await fixture.source.directory("/source/tree") + await fixture.source.put("/source/tree/a", bytes: Data([1])) + await fixture.source.put("/source/tree/b", bytes: Data([2])) + await fixture.source.changeOnSecondStat("/source/tree/b") + let result = try await fixture.copy(path: "/source/tree") + #expect(result.isError) + let object = try fixture.object(result) + #expect(object["copied_files"]?.integerValue == 1) + #expect(object["failed_path"]?.stringValue == "/source/tree/b") + #expect(await fixture.destination.data["/target/tree/a"] == Data([1])) + #expect(await fixture.backend.opens == 2) + } + + @Test + func copyFailureDoesNotReplayAndReleasesBothConnections() async throws { + let fixture = CopyFixture() + await fixture.source.put("/source/a.bin", bytes: Data([1])) + await fixture.source.failCopy() + #expect(try await fixture.copy().isError) + #expect(await fixture.source.copyAttempts == 1) + #expect(await fixture.backend.disconnects == 2) + } + + @Test + func sourceAndDestinationAncestorsCannotAliasThroughSymlinks() async throws { + let fixture = CopyFixture() + await fixture.source.put("/source/a.bin", bytes: Data([1])) + await fixture.source.node("/source", kind: .symbolicLink) + #expect(try await fixture.copy().isError) + #expect(await fixture.destination.nodes.count == 1) + await fixture.source.directory("/source") + await fixture.destination.node("/alias", kind: .symbolicLink) + await fixture.destination.directory("/alias/target") + #expect(try await fixture.copy(destination: "/alias/target").isError) + #expect(await fixture.destination.data.isEmpty) + } + + @Test + func activeTransfersSurviveIdleSweepAndHostDeletionCancelsThem() async throws { + let fixture = CopyFixture() + await fixture.source.put("/source/a.bin", bytes: Data([1])) + await fixture.source.blockCopy() + let task = Task { try await fixture.copy() } + try await copyWait { await fixture.source.copyAttempts == 1 } + _ = await fixture.pool.sweepIdle(now: .now.advanced(by: .seconds(1_000))) + #expect(await fixture.pool.openConnectionCount == 2) + await fixture.pool.disconnect(hostID: fixture.destinationID) + #expect(try await task.value.isError) + #expect(await fixture.pool.openConnectionCount == 0) + #expect(await fixture.destination.data["/target/a.bin"] == nil) + } + + @Test + func taskCancellationReleasesConnections() async throws { + let fixture = CopyFixture() + await fixture.source.put("/source/a.bin", bytes: Data([1])) + await fixture.source.blockCopy() + let task = Task { try await fixture.copy() } + try await copyWait { await fixture.source.copyAttempts == 1 } + task.cancel() + #expect(try await task.value.isError) + #expect(await fixture.backend.disconnects == 2) + } + + @Test + func mcpExportsSharedSchemaAndHonorsDisabledTools() async throws { + let fixture = CopyFixture() + let server = fixture.server() + let descriptor = try #require(server.descriptors.first { $0.name == "sftp_copy" }) + #expect(descriptor.inputSchema == fixture.tool.descriptor.inputSchema) + #expect(descriptor.approvalPolicy == .ask) + #expect(descriptor.concurrency == .sequential) + let disabled = fixture.server(disabled: ["sftp_copy"]) + #expect(!disabled.contains("sftp_copy")) + await #expect(throws: MCPServerToolsError.unknownTool("sftp_copy")) { + _ = try await disabled.call( + name: "sftp_copy", arguments: fixture.invocation().call.arguments, + runID: UUID(), clientName: "Tests", callID: "1", secretBroker: SecretBroker() + ) + } + } + + @Test + func mcpCopiesWithFullAccessAndRecordsMetadataOnly() async throws { + let fixture = CopyFixture() + await fixture.source.put("/source/a.bin", bytes: Data("PRIVATE-CONTENT".utf8)) + let audit = ToolAuditModel() + let server = fixture.server(audit: audit) + let result = try await server.call( + name: "sftp_copy", arguments: fixture.invocation().call.arguments, + runID: UUID(), clientName: "External Client", callID: "1", secretBroker: SecretBroker() + ) + #expect(!result.isError) + #expect(!result.content.contains("PRIVATE-CONTENT")) + try await copyWait { audit.entries().first?.finishedAt != nil } + let row = try #require(audit.entries().first) + #expect(row.clientName == "External Client") + #expect(row.bytesRead == 15) + #expect(row.bytesWritten == 15) + #expect(row.summary.contains(fixture.sourceID.uuidString)) + #expect(!row.summary.contains("PRIVATE-CONTENT")) + } + + @Test + func endingMcpSessionCancelsItsActiveCopy() async throws { + let fixture = CopyFixture() + await fixture.source.put("/source/a.bin", bytes: Data([1])) + await fixture.source.blockCopy() + let server = fixture.server() + let runID = UUID() + let broker = SecretBroker() + let task = Task { + try await server.call( + name: "sftp_copy", arguments: fixture.invocation().call.arguments, + runID: runID, clientName: "Tests", callID: "1", secretBroker: broker + ) + } + try await copyWait { await fixture.source.copyAttempts == 1 } + await server.endSession(runID: runID, secretBroker: broker) + #expect(try await task.value.isError) + #expect(await fixture.pool.openConnectionCount == 0) + } + + @Test + func manualDenialPreventsConnections() async throws { + let fixture = CopyFixture() + let runtime = AgentRuntime( + model: CopyCallingModel(arguments: fixture.invocation().call.arguments), + tools: [AnyAgentTool(fixture.tool)], + approval: AgentApprovalBroker(reviewer: nil, manualApproval: { _ in .deny("Denied") }) + ) + for await _ in await runtime.start( + AgentRunRequest(conversationID: UUID(), prompt: "copy", permissionMode: .askForApproval) + ) {} + #expect(await fixture.backend.opens == 0) + } + + @Test(arguments: ["en", "zh-Hans", "zh-Hant"]) + func approvalDetailsLocalizeBothEndpointsAndDestructiveWarning(_ language: String) throws { + let fixture = CopyFixture() + let items = SFTPCopyTool.presentArguments( + .init( + arguments: try #require(fixture.invocation(policy: "replace").call.arguments.objectValue), + locale: Locale(identifier: language) + )) + let labels = items.compactMap { item -> String? in + if case .field(let field) = item { return field.label } + return nil + } + #expect(labels.contains(language == "en" ? "Source host" : language == "zh-Hant" ? "來源主機" : "源主机")) + #expect(labels.contains(language == "en" ? "Destination host" : language == "zh-Hant" ? "目標主機" : "目标主机")) + #expect( + items.contains { + if case .message(_, .warning) = $0 { return true } + return false + }) + } +} + +private struct CopyCallingModel: AgentModelStreaming { + var arguments: AgentJSONValue + func stream(_ request: AgentModelRequest) -> AsyncThrowingStream { + AsyncThrowingStream { continuation in + if request.messages.contains(where: { $0.role == .tool }) { + continuation.yield(.finished(.completed)) + } else { + continuation.yield( + .toolCallSnapshot( + id: "1", providerItemID: nil, name: "sftp_copy", arguments: arguments.encodedString + )) + continuation.yield(.finished(.toolCalls)) + } + continuation.finish() + } + } +} + +@MainActor +private final class CopyFixture { + let sourceID = UUID() + let destinationID: UUID + let source = CopyFileSystem() + let destination = CopyFileSystem(target: true) + let backend: CopyBackend + let pool: AgentConnectionPool + let resolver: AgentHostResolver + let authentication: AgentSSHAuthenticationRegistry + var tool: SFTPCopyTool { .init(hosts: resolver, authentication: authentication, connections: pool) } + + init(sameHost: Bool = false) { + destinationID = sameHost ? sourceID : UUID() + let sourceID = sourceID + let destinationID = destinationID + let ids = Set([sourceID, destinationID]) + @Sendable nonisolated func configuration(_ id: UUID) -> SSHConnectionConfiguration { + SSHConnectionConfiguration( + hostID: id, host: id.uuidString, port: 22, username: "test", + credentials: [.password("test")], expectedHostKeys: [], proxy: nil, jumps: [] + ) + } + resolver = AgentHostResolver( + store: HostLibraryModel(), + temporaryTargets: { + ids.map { + .init( + id: $0, displayName: $0 == sourceID ? "Source" : "Destination", address: "test", + tags: [], connectionState: "connected", isTemporary: true) + } + }, + temporaryConfiguration: { configuration($0) } + ) + authentication = AgentSSHAuthenticationRegistry(resolver: resolver) + let backend = CopyBackend( + filesystems: sameHost ? [sourceID: source] : [sourceID: source, destinationID: destination]) + self.backend = backend + pool = AgentConnectionPool( + configuration: { configuration($0) }, routeFingerprint: { _ in "test" }, + openExec: { _ in throw AgentToolError.invalidHost }, + openSFTP: { try await backend.open($0.configuration.hostID) }, sweepInterval: nil + ) + } + + func invocation(path: String = "/source/a.bin", destination: String = "/target", policy: String? = nil) + -> AgentToolInvocation + { + var arguments: [String: AgentJSONValue] = [ + "source_host_id": .string(sourceID.uuidString), "source_path": .string(path), + "destination_host_id": .string(destinationID.uuidString), "destination_path": .string(destination), + ] + if let policy { arguments["conflict_policy"] = .string(policy) } + return .init(runID: UUID(), call: .init(id: "1", name: "sftp_copy", arguments: .object(arguments))) + } + + func copy(path: String = "/source/a.bin", destination: String = "/target", policy: String? = nil) async throws + -> AgentToolResult + { + let invocation = invocation(path: path, destination: destination, policy: policy) + let preflight = try await tool.preflight(invocation) + return try await tool.execute( + preflight.invocation, + context: .init( + runID: invocation.runID, userIntent: "copy", secretBroker: SecretBroker(), + userInteraction: UnavailableAgentUserInteraction(), preflightMetadata: preflight.executionMetadata, + authorize: { _ in .allow } + )) + } + + func object(_ result: AgentToolResult) throws -> [String: AgentJSONValue] { + let payload = try #require(result.untrustedPayload) + let value = try AgentJSONValue.decode(Data(payload.utf8)) + return try #require(value.objectValue) + } + + func server(disabled: Set = [], audit: ToolAuditModel = ToolAuditModel()) -> MCPServerTools { + MCPServerTools( + hosts: resolver, metrics: TestMetricsReader(), connections: pool, authentication: authentication, + background: AgentBackgroundCommands(sweepInterval: nil), externalAuthentication: ExternalAuthCoordinator(), + history: audit, disabledTools: disabled + ) + } +} + +private actor CopyFileSystem { + var nodes: [String: RemoteFile] = [:] + var data: [String: Data] = [:] + var copyAttempts = 0 + var blocked = false + var failing = false + var changing: String? + var stats: [String: Int] = [:] + + init(target: Bool = false) { + if target { + nodes["/target"] = Self.file("/target", kind: .directory) + } else { + nodes["/source"] = Self.file("/source", kind: .directory) + } + } + static func file(_ path: String, kind: RemoteFile.Kind, size: Int64 = 0) -> RemoteFile { + .init( + path: path, name: (path as NSString).lastPathComponent, kind: kind, size: size, + permissions: 0o600, ownerUID: 0, ownerGID: 0) + } + func node(_ path: String, kind: RemoteFile.Kind) { nodes[path] = Self.file(path, kind: kind) } + func directory(_ path: String) { node(path, kind: .directory) } + func put(_ path: String, bytes: Data) { + nodes[path] = Self.file(path, kind: .file, size: Int64(bytes.count)) + data[path] = bytes + } + func stat(_ path: String) -> RemoteFile? { + stats[path, default: 0] += 1 + if changing == path && stats[path] == 2 { nodes[path]?.size += 1 } + return nodes[path] + } + func children(_ path: String) -> [RemoteFile] { + let children = nodes.values.filter { RemotePath.parent(of: $0.path) == path }.sorted { $0.path < $1.path } + for child in children { stats[child.path, default: 0] += 1 } + return children + } + func remove(_ path: String) { + nodes[path] = nil + data[path] = nil + } + func blockCopy() { blocked = true } + func failCopy() { failing = true } + func changeOnSecondStat(_ path: String) { changing = path } + func startCopy() { copyAttempts += 1 } +} + +private actor CopyBackend { + var filesystems: [UUID: CopyFileSystem] + var opens = 0 + var disconnects = 0 + init(filesystems: [UUID: CopyFileSystem]) { self.filesystems = filesystems } + func open(_ id: UUID) throws -> CopyClient { + opens += 1 + return CopyClient(filesystem: try #require(filesystems[id]), backend: self) + } + func disconnected() { disconnects += 1 } +} + +private actor CopyClient: SFTPClientProtocol { + let filesystem: CopyFileSystem + let backend: CopyBackend + init(filesystem: CopyFileSystem, backend: CopyBackend) { + self.filesystem = filesystem + self.backend = backend + } + func connect(using configuration: SSHConnectionConfiguration, progress: SSHRouteProgressHandler?) async throws {} + func disconnect() async { await backend.disconnected() } + func homeDirectory() async throws -> String { "/" } + func listDirectory(at path: String) async throws -> [RemoteFile] { await filesystem.children(path) } + func attributes(at path: String) async throws -> RemoteFile? { await filesystem.stat(path) } + func createDirectory(at path: String) async throws { await filesystem.directory(path) } + func setPermissions(at path: String, to permissions: UInt32) async throws {} + func rename(from source: String, to destination: String, overwrite: Bool) async throws { + throw SFTPError.notConnected + } + func removeFile(at path: String) async throws { await filesystem.remove(path) } + func removeDirectory(at path: String) async throws { await filesystem.remove(path) } + func uploadFile( + from localURL: URL, to remotePath: String, overwrite: Bool, cancellation: TransferCancellationToken, + progress: @escaping @Sendable (SFTPProgress) -> Void + ) async throws { throw SFTPError.notConnected } + func downloadFile( + from remotePath: String, to localURL: URL, overwrite: Bool, cancellation: TransferCancellationToken, + progress: @escaping @Sendable (SFTPProgress) -> Void + ) async throws { throw SFTPError.notConnected } + func copyFile( + from sourcePath: String, to destinationPath: String, on destination: any SFTPClientProtocol, + byteCount: Int64, overwrite: Bool, cancellation: TransferCancellationToken, + progress: @escaping @Sendable (SFTPProgress) -> Void + ) async throws { + await filesystem.startCopy() + while await filesystem.blocked { + guard !cancellation.isCancelled else { throw SFTPError.cancelled } + try await Task.sleep(for: .milliseconds(10)) + } + if await filesystem.failing { throw SFTPError.connectionLost("Disconnected") } + let destination = try #require(destination as? CopyClient) + let bytes = try #require(await filesystem.data[sourcePath]) + progress(.init(completedBytes: byteCount, totalBytes: byteCount, bytesPerSecond: 1)) + await destination.filesystem.put(destinationPath, bytes: bytes) + } +} + +@MainActor +private func copyWait(_ condition: () async -> Bool) async throws { + let deadline = ContinuousClock.now.advanced(by: .seconds(5)) + while await !condition() { + guard ContinuousClock.now < deadline else { throw SFTPError.operation("Test timed out") } + try await Task.sleep(for: .milliseconds(10)) + } +} diff --git a/TFMTests/SFTPSocketRelayTests.swift b/TFMTests/SFTPSocketRelayTests.swift new file mode 100644 index 00000000..d902b468 --- /dev/null +++ b/TFMTests/SFTPSocketRelayTests.swift @@ -0,0 +1,60 @@ +import Darwin +import Foundation +import Testing + +@testable import TermindRemoteInfrastructure + +@Suite +struct SFTPSocketRelayTests { + @Test + func successfulSourceCompletionMarksEOFAfterBytes() throws { + let relay = try SFTPSocketRelay() + defer { relay.close() } + var sent: UInt8 = 255 + #expect(Darwin.send(relay.writer, &sent, 1, 0) == 1) + #expect(!relay.sourceSucceeded) + relay.finishWriting() + var received: UInt8 = 0 + #expect(Darwin.recv(relay.reader, &received, 1, 0) == 1) + #expect(received == sent) + #expect(Darwin.recv(relay.reader, &received, 1, 0) == 0) + #expect(relay.sourceSucceeded) + } + + @Test + func failedOrCancelledSourceEOFDoesNotAuthorizeCommit() throws { + let relay = try SFTPSocketRelay() + defer { relay.close() } + let cancellation = TransferCancellationToken() + let observer = cancellation.observeCancellation { relay.stop() } + defer { cancellation.removeObserver(observer) } + cancellation.cancel() + var byte: UInt8 = 0 + #expect(Darwin.recv(relay.reader, &byte, 1, 0) == 0) + #expect(!relay.sourceSucceeded) + } + + @Test + func cancellationAfterBytesStillDoesNotMarkSuccessfulSourceEOF() throws { + let relay = try SFTPSocketRelay() + defer { relay.close() } + var byte: UInt8 = 1 + #expect(Darwin.send(relay.writer, &byte, 1, 0) == 1) + #expect(Darwin.recv(relay.reader, &byte, 1, 0) == 1) + relay.stop() + #expect(!relay.sourceSucceeded) + #expect(Darwin.recv(relay.reader, &byte, 1, 0) == 0) + } + + @Test + func alreadyCancelledTokensImmediatelyStopNewRelays() throws { + let cancellation = TransferCancellationToken() + cancellation.cancel() + let relay = try SFTPSocketRelay() + defer { relay.close() } + let observer = cancellation.observeCancellation { relay.stop() } + defer { cancellation.removeObserver(observer) } + var byte: UInt8 = 0 + #expect(Darwin.recv(relay.reader, &byte, 1, 0) == 0) + } +}