diff --git a/supacode/Clients/Shell/ShellClient.swift b/supacode/Clients/Shell/ShellClient.swift index 94f0aeed..aca3adb3 100644 --- a/supacode/Clients/Shell/ShellClient.swift +++ b/supacode/Clients/Shell/ShellClient.swift @@ -168,7 +168,12 @@ nonisolated private func runProcessStream( currentDirectoryURL: URL? ) -> AsyncThrowingStream { AsyncThrowingStream { continuation in - Task.detached { + // Stored so onTermination can SIGTERM the process when the consumer of the + // stream goes away (Task cancellation, for-await throw, explicit finish). + // Without this hook structured-concurrency timeouts that wrap a ShellClient + // call would have to wait for the process to exit on its own. + let processBox = LockIsolated(nil) + let workerTask = Task { let outputAccumulator = ShellOutputAccumulator() let process = Process() process.executableURL = executableURL @@ -184,33 +189,24 @@ nonisolated private func runProcessStream( let command = ([executableURL.path(percentEncoded: false)] + arguments).joined(separator: " ") do { try process.run() - let stdoutTask = Task.detached { + processBox.withValue { $0 = process } + let stdoutTask = Task { for await line in lineStream(from: outputHandle) { await outputAccumulator.append(line, source: .stdout) - continuation.yield( - .line( - ShellStreamLine( - source: .stdout, - text: line - ) - ) - ) + continuation.yield(.line(ShellStreamLine(source: .stdout, text: line))) } } - let stderrTask = Task.detached { + let stderrTask = Task { for await line in lineStream(from: errorHandle) { await outputAccumulator.append(line, source: .stderr) - continuation.yield( - .line( - ShellStreamLine( - source: .stderr, - text: line - ) - ) - ) + continuation.yield(.line(ShellStreamLine(source: .stderr, text: line))) } } - process.waitUntilExit() + await withTaskCancellationHandler { + await waitForExit(of: process) + } onCancel: { + process.terminate() + } await stdoutTask.value await stderrTask.value let output = await outputAccumulator.output(exitCode: process.terminationStatus) @@ -231,6 +227,35 @@ nonisolated private func runProcessStream( continuation.finish(throwing: error) } } + continuation.onTermination = { _ in + processBox.withValue { $0?.terminate() } + workerTask.cancel() + } + } +} + +/// Waits asynchronously for `process` to exit using `terminationHandler` +/// instead of `process.waitUntilExit()` so the caller's Task cancellation +/// can be honoured. The handler is paired with a synchronous `isRunning` +/// check to cover the race where the process exits before the handler is +/// installed. +nonisolated private func waitForExit(of process: Process) async { + await withCheckedContinuation { (continuation: CheckedContinuation) in + let resumed = LockIsolated(false) + let resumeOnce: @Sendable () -> Void = { + let shouldResume = resumed.withValue { (value: inout Bool) -> Bool in + guard !value else { return false } + value = true + return true + } + if shouldResume { + continuation.resume() + } + } + process.terminationHandler = { _ in resumeOnce() } + if !process.isRunning { + resumeOnce() + } } } diff --git a/supacodeTests/ShellClientStreamingTests.swift b/supacodeTests/ShellClientStreamingTests.swift index 0217b5a9..f0d95b35 100644 --- a/supacodeTests/ShellClientStreamingTests.swift +++ b/supacodeTests/ShellClientStreamingTests.swift @@ -127,6 +127,75 @@ struct ShellClientStreamingTests { #expect(streamedLines.contains(where: { $0.source == .stderr && $0.text == "err" })) } + @Test func cancellingRunStreamConsumerTerminatesProcessQuickly() async throws { + let shell = ShellClient.liveValue + let commandURL = URL(fileURLWithPath: "/bin/sleep") + let stream = shell.runStream(commandURL, ["30"], nil) + + let consumer = Task { + var lines: [ShellStreamLine] = [] + do { + for try await event in stream { + if case .line(let line) = event { + lines.append(line) + } + } + } catch { + // CancellationError or ShellClientError after SIGTERM both indicate + // the cancellation pathway is connected. + } + return lines + } + + try await Task.sleep(for: .milliseconds(120)) + + let start = ContinuousClock.now + consumer.cancel() + _ = await consumer.value + let elapsed = ContinuousClock.now - start + + #expect( + elapsed < .seconds(2), + "consumer cancel should propagate to the shell process; took \(elapsed)" + ) + } + + @Test func runReturnsQuicklyWhenCallingTaskIsCancelled() async { + let shell = ShellClient.liveValue + let commandURL = URL(fileURLWithPath: "/bin/sleep") + + let runTask = Task { + try await shell.run(commandURL, ["30"], nil) + } + + try? await Task.sleep(for: .milliseconds(120)) + + let start = ContinuousClock.now + runTask.cancel() + _ = await runTask.result + let elapsed = ContinuousClock.now - start + + #expect( + elapsed < .seconds(2), + "run() should propagate cancellation to the process; took \(elapsed)" + ) + } + + @Test func runStreamSucceedsForShortLivedProcessAfterCancellationFixes() async throws { + // Regression guard: terminationHandler / isRunning race in waitForExit + // must not deadlock or double-resume on fast-exiting processes. + let shell = ShellClient.liveValue + let commandURL = URL(fileURLWithPath: "/bin/sh") + let stream = shell.runStream(commandURL, ["-c", "true"], nil) + var finished: ShellOutput? + for try await event in stream { + if case .finished(let output) = event { + finished = output + } + } + #expect(finished?.exitCode == 0) + } + @Test func runLoginStreamForwardsParameters() async throws { let recorder = LoginStreamCallRecorder() let shell = ShellClient(