446 lines
21 KiB
Swift
446 lines
21 KiB
Swift
// SPDX-FileCopyrightText: 2026 Kaede Akino
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
import AgentKit
|
|
import Foundation
|
|
import TermindApplication
|
|
import TermindRemoteInfrastructure
|
|
import Testing
|
|
|
|
private nonisolated func exited(_ code: Int32 = 0) -> SSHExecResult {
|
|
SSHExecResult(
|
|
stdout: "", stderr: "", exitCode: code, signal: nil,
|
|
stdoutTruncated: false, stderrTruncated: false)
|
|
}
|
|
|
|
@Suite
|
|
struct AgentBackgroundCommandTests {
|
|
@Test
|
|
func outputToolCapturesEachFullSnapshotWithItsOwnNameUnderTheConversation() async throws {
|
|
let directory = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString)
|
|
defer { try? FileManager.default.removeItem(at: directory) }
|
|
let workspace = AgentScratchWorkspace(base: directory)
|
|
let scope = UUID()
|
|
let accumulator = AgentOutputAccumulator(
|
|
workspace: workspace,
|
|
path: {
|
|
"\(scope.uuidString.lowercased())/outputs/ssh-exec-\(UUID().uuidString.lowercased()).txt"
|
|
})
|
|
let stream = AgentCommandOutputStream(accumulator: accumulator) { _, _ in }
|
|
let commands = AgentBackgroundCommands(sweepInterval: nil)
|
|
let source = "first\n" + String(repeating: "output\n", count: 2_100) + "last>"
|
|
let handle = try await commands.start(
|
|
scope: scope, hostID: UUID(), command: "build", keepsStdinOpen: false, stream: stream
|
|
) { session in
|
|
session.start { _ in 0 }
|
|
stream.handler(Data(source.utf8), false)
|
|
return exited()
|
|
}
|
|
try await commands.waitForCompletion(handle, scope: scope)
|
|
let run = await stream.snapshot()
|
|
#expect(run.capture?.path?.hasPrefix("\(scope.uuidString.lowercased())/outputs/ssh-exec-") == true)
|
|
let tool = SSHExecOutputTool(background: commands, scope: scope)
|
|
let call = AgentToolInvocation(
|
|
runID: UUID(),
|
|
call: .init(
|
|
id: "output", name: "ssh_exec_output",
|
|
arguments: .object([
|
|
"handle": .string(handle.uuidString)
|
|
])))
|
|
let context = AgentToolExecutionContext(
|
|
runID: call.runID, conversationID: scope, userIntent: "read output", secretBroker: SecretBroker(),
|
|
userInteraction: UnavailableAgentUserInteraction(), preflightMetadata: [:],
|
|
authorize: { _ in .allow }, workspace: workspace)
|
|
let first = try await tool.execute(call, context: context)
|
|
let second = try await tool.execute(call, context: context)
|
|
let path = try #require(first.outputCapture?.path)
|
|
#expect(path.hasPrefix("\(scope.uuidString.lowercased())/outputs/ssh-exec-output-"))
|
|
#expect(path.hasSuffix(".txt"))
|
|
#expect(second.outputCapture?.path != path)
|
|
#expect(try await workspace.data(at: path).data == Data(source.utf8))
|
|
#expect(first.metadata["output_total_lines"] == .number(2_102))
|
|
await commands.closeAll()
|
|
}
|
|
|
|
@Test
|
|
func outputIsImmediateCumulativeAndRepeatableAfterCompletion() async throws {
|
|
let commands = AgentBackgroundCommands(sweepInterval: nil)
|
|
let scope = UUID()
|
|
let release = AsyncRelease()
|
|
let stream = AgentCommandOutputStream(
|
|
accumulator: AgentOutputAccumulator(path: { "tests/outputs/ssh-exec-\(UUID().uuidString).txt" })
|
|
) { _, _ in }
|
|
let handle = try await commands.start(
|
|
scope: scope, hostID: UUID(), command: "build", keepsStdinOpen: false, stream: stream
|
|
) { session in
|
|
session.start { _ in 0 }
|
|
await release.wait()
|
|
return exited(7)
|
|
}
|
|
stream.handler(Data("step one\n".utf8), false)
|
|
let running = try await commands.output(handle, scope: scope)
|
|
#expect(running.objectValue?["status"] == .string("running"))
|
|
#expect(running.objectValue?["output"] == .string("step one\n"))
|
|
#expect(running.objectValue?["stdin_closed"] == .bool(true))
|
|
#expect(try await commands.output(handle, scope: scope) == running)
|
|
|
|
stream.handler(Data("step two\n".utf8), false)
|
|
stream.handler(Data("failure\n".utf8), true)
|
|
await release.release()
|
|
try await commands.waitForCompletion(handle, scope: scope)
|
|
let finished = try await commands.output(handle, scope: scope)
|
|
#expect(finished.objectValue?["status"] == .string("exited"))
|
|
#expect(finished.objectValue?["exit_code"] == .number(7))
|
|
#expect(finished.objectValue?["output"] == .string("step one\nstep two\nfailure\n"))
|
|
#expect(try await commands.output(handle, scope: scope) == finished)
|
|
#expect(await commands.count == 1)
|
|
await commands.closeAll()
|
|
}
|
|
|
|
@Test
|
|
func failedStartupLeavesNoHandle() async {
|
|
let commands = AgentBackgroundCommands(sweepInterval: nil)
|
|
await #expect(throws: CocoaError.self) {
|
|
try await commands.start(
|
|
scope: UUID(), hostID: UUID(), command: "build", keepsStdinOpen: true,
|
|
stream: AgentCommandOutputStream(
|
|
accumulator: AgentOutputAccumulator(path: { "tests/outputs/ssh-exec-\(UUID().uuidString).txt" })
|
|
) { _, _ in }
|
|
) { _ in throw CocoaError(.fileReadNoPermission) }
|
|
}
|
|
#expect(await commands.count == 0)
|
|
}
|
|
|
|
@Test
|
|
func cancellationDuringStartupDoesNotWaitForAuthenticationOrAllowLateStartup() async throws {
|
|
let commands = AgentBackgroundCommands(sweepInterval: nil)
|
|
let scope = UUID()
|
|
let connecting = AsyncRelease()
|
|
let release = AsyncRelease()
|
|
let finished = AsyncRelease()
|
|
let task = Task {
|
|
try await commands.start(
|
|
scope: scope, hostID: UUID(), command: "build", keepsStdinOpen: true,
|
|
stream: AgentCommandOutputStream(
|
|
accumulator: AgentOutputAccumulator(path: { "tests/outputs/ssh-exec-\(UUID().uuidString).txt" })
|
|
) { _, _ in }
|
|
) { session in
|
|
await connecting.release()
|
|
await release.wait()
|
|
session.start { _ in 0 }
|
|
#expect(!session.started)
|
|
await finished.release()
|
|
return exited()
|
|
}
|
|
}
|
|
await connecting.wait()
|
|
#expect(await commands.list(scope: scope).objectValue?["commands"] == .array([]))
|
|
task.cancel()
|
|
await #expect(throws: CancellationError.self) { try await task.value }
|
|
#expect(await commands.count == 0)
|
|
await release.release()
|
|
await finished.wait()
|
|
}
|
|
|
|
@Test
|
|
func runningCommandsNeverExpireAndFinishedReadsDoNotExtendRetention() async throws {
|
|
let commands = AgentBackgroundCommands(idleTimeout: .seconds(600), sweepInterval: nil)
|
|
let scope = UUID()
|
|
let release = AsyncRelease()
|
|
let handle = try await commands.start(
|
|
scope: scope, hostID: UUID(), command: "build", keepsStdinOpen: true,
|
|
stream: AgentCommandOutputStream(
|
|
accumulator: AgentOutputAccumulator(path: { "tests/outputs/ssh-exec-\(UUID().uuidString).txt" })
|
|
) { _, _ in }
|
|
) { session in
|
|
session.start { _ in 0 }
|
|
await release.wait()
|
|
return exited()
|
|
}
|
|
#expect(await commands.sweepIdle(now: .now.advanced(by: .seconds(86_400))) == 1)
|
|
await release.release()
|
|
try await commands.waitForCompletion(handle, scope: scope)
|
|
_ = try await commands.output(handle, scope: scope)
|
|
#expect(await commands.list(scope: scope).objectValue?["commands"] == .array([]))
|
|
#expect(
|
|
await commands.list(scope: scope, includeFinished: true).objectValue?["commands"]?.arrayValue?.count == 1)
|
|
#expect(await commands.sweepIdle(now: .now.advanced(by: .seconds(601))) == 0)
|
|
await #expect(throws: AgentBackgroundCommandError.self) {
|
|
try await commands.output(handle, scope: scope)
|
|
}
|
|
}
|
|
|
|
@Test
|
|
func scopeIsolationAndCancellationRetireOnlyTheLocalHandle() async throws {
|
|
let commands = AgentBackgroundCommands(sweepInterval: nil)
|
|
let scope = UUID()
|
|
let otherScope = UUID()
|
|
let host = UUID()
|
|
let release = AsyncRelease()
|
|
let handle = try await commands.start(
|
|
scope: scope, hostID: host, command: "build", keepsStdinOpen: true,
|
|
stream: AgentCommandOutputStream(
|
|
accumulator: AgentOutputAccumulator(path: { "tests/outputs/ssh-exec-\(UUID().uuidString).txt" })
|
|
) { _, _ in }
|
|
) { session in
|
|
session.start { _ in 0 }
|
|
await release.wait()
|
|
return exited()
|
|
}
|
|
#expect(await commands.list(scope: otherScope).objectValue?["commands"] == .array([]))
|
|
#expect(await commands.list(scope: scope, hostID: UUID()).objectValue?["commands"] == .array([]))
|
|
let rows = await commands.list(scope: scope, hostID: host).objectValue?["commands"]?.arrayValue
|
|
#expect(rows?.count == 1)
|
|
#expect(rows?.first?.objectValue?["output"] == nil)
|
|
await #expect(throws: AgentBackgroundCommandError.self) {
|
|
try await commands.cancel(handle, scope: otherScope)
|
|
}
|
|
let result = try await commands.cancel(handle, scope: scope)
|
|
#expect(result.objectValue?["status"] == .string("cancelled"))
|
|
#expect(result.objectValue?["was_running"] == .bool(true))
|
|
#expect(result.objectValue?["remote_termination_confirmed"] == .bool(false))
|
|
await #expect(throws: AgentBackgroundCommandError.self) {
|
|
try await commands.cancel(handle, scope: scope)
|
|
}
|
|
await release.release()
|
|
#expect(await commands.count == 0)
|
|
}
|
|
|
|
@Test
|
|
func stdinPreservesExactBytesAndEOFIsIdempotent() async throws {
|
|
let commands = AgentBackgroundCommands(sweepInterval: nil)
|
|
let scope = UUID()
|
|
let release = AsyncRelease()
|
|
let writes = InputLog()
|
|
let handle = try await commands.start(
|
|
scope: scope, hostID: UUID(), command: "cat", keepsStdinOpen: true,
|
|
stream: AgentCommandOutputStream(
|
|
accumulator: AgentOutputAccumulator(path: { "tests/outputs/ssh-exec-\(UUID().uuidString).txt" })
|
|
) { _, _ in }
|
|
) { session in
|
|
session.start { try await writes.write($0) }
|
|
await release.wait()
|
|
return exited()
|
|
}
|
|
let result = try await commands.write(handle, scope: scope, input: .text("abc"))
|
|
#expect(result.objectValue?["bytes_written"] == .number(3))
|
|
_ = try await commands.write(handle, scope: scope, input: .text(""))
|
|
_ = try await commands.write(handle, scope: scope, input: nil)
|
|
let repeated = try await commands.write(handle, scope: scope, input: nil)
|
|
#expect(repeated.objectValue?["stdin_closed"] == .bool(true))
|
|
#expect(await writes.values == ["abc", "", nil])
|
|
await #expect(throws: (any Error).self) {
|
|
try await commands.write(handle, scope: scope, input: .text("too late"))
|
|
}
|
|
await release.release()
|
|
try await commands.waitForCompletion(handle, scope: scope)
|
|
#expect(try await commands.cancel(handle, scope: scope).objectValue?["status"] == .string("released"))
|
|
}
|
|
|
|
@Test
|
|
func interjectionReturnsRunningSnapshotWithoutCancellingCommand() async throws {
|
|
let commands = AgentBackgroundCommands(sweepInterval: nil)
|
|
let scope = UUID()
|
|
let release = AsyncRelease()
|
|
let signal = AgentInterjectionSignal()
|
|
let handle = try await commands.start(
|
|
scope: scope, hostID: UUID(), command: "build", keepsStdinOpen: false,
|
|
stream: AgentCommandOutputStream(
|
|
accumulator: AgentOutputAccumulator(path: { "tests/outputs/ssh-exec-\(UUID().uuidString).txt" })
|
|
) { _, _ in }
|
|
) { session in
|
|
session.start { _ in 0 }
|
|
await release.wait()
|
|
#expect(!Task.isCancelled)
|
|
return exited()
|
|
}
|
|
signal.arm()
|
|
try await commands.waitForCompletion(handle, scope: scope, interjection: { await signal.wait() })
|
|
#expect(try await commands.output(handle, scope: scope).objectValue?["status"] == .string("running"))
|
|
await release.release()
|
|
try await commands.waitForCompletion(handle, scope: scope)
|
|
await commands.closeAll()
|
|
}
|
|
|
|
@Test
|
|
func outputAndListAreReadOnlyParallelToolsWithoutTimeouts() {
|
|
let commands = AgentBackgroundCommands(sweepInterval: nil)
|
|
let scope = UUID()
|
|
let tools = [
|
|
AnyAgentTool(SSHExecOutputTool(background: commands, scope: scope)),
|
|
AnyAgentTool(SSHExecListTool(background: commands, scope: scope)),
|
|
]
|
|
for tool in tools {
|
|
#expect(tool.descriptor.approvalPolicy == .approve)
|
|
#expect(tool.descriptor.concurrency == .parallel)
|
|
#expect(tool.descriptor.inputSchema.objectValue?["properties"]?.objectValue?["timeout_seconds"] == nil)
|
|
}
|
|
}
|
|
|
|
@Test(arguments: ["text", "empty", "secret", "eof"], [false, true])
|
|
func stdinToolAcceptsOmittedOrNullUnusedPayloads(mode: String, explicitNulls: Bool) async throws {
|
|
let secretText = "stdin-regression-secret"
|
|
let commands = AgentBackgroundCommands(sweepInterval: nil)
|
|
let scope = UUID()
|
|
let hostID = UUID()
|
|
let release = AsyncRelease()
|
|
let writes = InputLog()
|
|
let handle = try await commands.start(
|
|
scope: scope, hostID: hostID, command: "cat", keepsStdinOpen: true,
|
|
stream: AgentCommandOutputStream(
|
|
accumulator: AgentOutputAccumulator(path: { "tests/outputs/ssh-exec-\(UUID().uuidString).txt" })
|
|
) { _, _ in }
|
|
) { session in
|
|
session.start { try await writes.write($0) }
|
|
await release.wait()
|
|
return exited()
|
|
}
|
|
let broker = SecretBroker()
|
|
do {
|
|
let tool = SSHExecStdinTool(background: commands, scope: scope)
|
|
var arguments: [String: AgentJSONValue] = [
|
|
"handle": .string(handle.uuidString), "type": .string(mode == "empty" ? "text" : mode),
|
|
]
|
|
if explicitNulls {
|
|
arguments["stdin"] = .null
|
|
arguments["secret_handle"] = .null
|
|
}
|
|
switch mode {
|
|
case "text", "empty": arguments["stdin"] = .string(mode == "empty" ? "" : "input")
|
|
case "secret":
|
|
var bytes = Data(secretText.utf8)
|
|
let secret = try await broker.issue(AgentSecret(consumingUTF8: &bytes))
|
|
arguments["secret_handle"] = .string(secret.id.uuidString)
|
|
default: break
|
|
}
|
|
let invocation = AgentToolInvocation(
|
|
runID: UUID(), call: .init(id: "stdin", name: "ssh_exec_stdin", arguments: .object(arguments)),
|
|
targetHostID: UUID())
|
|
try AgentJSONSchemaValidator.validate(invocation.call.arguments, against: tool.descriptor.inputSchema)
|
|
let checked = try await tool.preflight(invocation)
|
|
#expect(checked.invocation.targetHostID == hostID)
|
|
#expect(checked.approvalPolicy == .ask)
|
|
let review = GuardianPolicy.reviewPrompt(
|
|
.init(
|
|
invocation: checked.invocation, descriptor: tool.descriptor, userIntent: "Provide command input",
|
|
localReasons: checked.reasons),
|
|
evidence: [])
|
|
#expect(review.contains(hostID.uuidString))
|
|
#expect(review.contains("Write input to: cat"))
|
|
#expect(review.contains("The plaintext is intentionally hidden from review") == (mode == "secret"))
|
|
#expect(!review.contains(secretText))
|
|
if mode == "secret" {
|
|
#expect(await broker.count == 1)
|
|
#expect(await writes.secretWrites == 0)
|
|
}
|
|
let result = try await tool.execute(
|
|
checked.invocation,
|
|
context: .init(
|
|
runID: invocation.runID, userIntent: "Provide command input", secretBroker: broker,
|
|
userInteraction: UnavailableAgentUserInteraction(), preflightMetadata: [:],
|
|
authorize: { _ in .allow }))
|
|
let payload = try AgentJSONValue.decode(Data(try #require(result.untrustedPayload).utf8))
|
|
#expect(!result.isError)
|
|
#expect(payload.objectValue?["status"] == .string(mode == "eof" ? "closed" : "written"))
|
|
#expect(payload.objectValue?["stdin_closed"] == .bool(mode == "eof"))
|
|
switch mode {
|
|
case "text": #expect(await writes.values == ["input"])
|
|
case "empty": #expect(await writes.values == [""])
|
|
case "secret":
|
|
#expect(await writes.secretWrites == 1)
|
|
#expect(await broker.count == 0)
|
|
#expect(payload.objectValue?["bytes_written"] == .number(Double(secretText.utf8.count)))
|
|
default: #expect(await writes.values == [nil])
|
|
}
|
|
} catch {
|
|
await release.release()
|
|
await commands.closeAll()
|
|
await broker.discardAll()
|
|
throw error
|
|
}
|
|
await release.release()
|
|
await commands.closeAll()
|
|
}
|
|
|
|
@Test(arguments: [false, true])
|
|
func listToolAcceptsOmittedOrNullFilters(explicitNulls: Bool) async throws {
|
|
let tool = SSHExecListTool(background: AgentBackgroundCommands(sweepInterval: nil), scope: UUID())
|
|
let invocation = AgentToolInvocation(
|
|
runID: UUID(),
|
|
call: .init(
|
|
id: "list", name: "ssh_exec_list",
|
|
arguments: .object(explicitNulls ? ["host_id": .null, "include_finished": .null] : [:])))
|
|
try AgentJSONSchemaValidator.validate(invocation.call.arguments, against: tool.descriptor.inputSchema)
|
|
let result = try await tool.execute(
|
|
invocation,
|
|
context: .init(
|
|
runID: invocation.runID, userIntent: "List commands", secretBroker: SecretBroker(),
|
|
userInteraction: UnavailableAgentUserInteraction(), preflightMetadata: [:], authorize: { _ in .allow }))
|
|
let payload = try AgentJSONValue.decode(Data(try #require(result.untrustedPayload).utf8))
|
|
#expect(payload == .object(["commands": .array([])]))
|
|
}
|
|
|
|
@Test(arguments: [
|
|
["type": AgentJSONValue.string("text"), "stdin": .null, "secret_handle": .null],
|
|
["type": .string("secret"), "stdin": .null, "secret_handle": .null],
|
|
["type": .string("secret"), "stdin": .null, "secret_handle": .string("invalid")],
|
|
["type": .string("text"), "stdin": .string("input"), "secret_handle": .string(UUID().uuidString)],
|
|
["type": .string("secret"), "stdin": .string(""), "secret_handle": .string(UUID().uuidString)],
|
|
["type": .string("eof"), "stdin": .string(""), "secret_handle": .null],
|
|
["type": .string("eof"), "stdin": .null, "secret_handle": .string(UUID().uuidString)],
|
|
["type": .string("eof"), "stdin": .bool(false), "secret_handle": .null],
|
|
])
|
|
func stdinToolStillRejectsMissingRequiredOrConflictingPayloads(arguments: [String: AgentJSONValue]) async {
|
|
let tool = SSHExecStdinTool(background: AgentBackgroundCommands(sweepInterval: nil), scope: UUID())
|
|
var arguments = arguments
|
|
arguments["handle"] = .string(UUID().uuidString)
|
|
let invocation = AgentToolInvocation(
|
|
runID: UUID(), call: .init(id: "stdin", name: "ssh_exec_stdin", arguments: .object(arguments)))
|
|
do {
|
|
_ = try await tool.preflight(invocation)
|
|
Issue.record("Invalid input must be rejected before command lookup")
|
|
} catch AgentToolError.invalidArguments {
|
|
} catch {
|
|
Issue.record("Unexpected error: \(error)")
|
|
}
|
|
}
|
|
}
|
|
|
|
private actor AsyncRelease {
|
|
private var isReleased = false
|
|
private var waiters: [CheckedContinuation<Void, Never>] = []
|
|
|
|
func wait() async {
|
|
if isReleased { return }
|
|
await withCheckedContinuation { waiters.append($0) }
|
|
}
|
|
|
|
func release() {
|
|
isReleased = true
|
|
let pending = waiters
|
|
waiters.removeAll()
|
|
for waiter in pending { waiter.resume() }
|
|
}
|
|
}
|
|
|
|
private actor InputLog {
|
|
var values: [String?] = []
|
|
var secretWrites = 0
|
|
func write(_ input: SSHCommandInput?) throws -> Int {
|
|
switch input {
|
|
case .text(let text):
|
|
values.append(text)
|
|
return text.utf8.count
|
|
case .plain(let data):
|
|
values.append(String(decoding: data, as: UTF8.self))
|
|
return data.count
|
|
case nil:
|
|
values.append(nil)
|
|
return 0
|
|
case .secret(let secret):
|
|
secretWrites += 1
|
|
return try secret.withUnsafeBytes { $0.count }
|
|
}
|
|
}
|
|
}
|