From 2d3d68196de3241a199ee267438967e9252abbf3 Mon Sep 17 00:00:00 2001 From: onevcat Date: Fri, 29 May 2026 19:38:33 +0900 Subject: [PATCH] Propagate Task cancellation into ShellClient processes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ShellClient.runProcessStream used to spawn its worker on Task.detached and block on Process.waitUntilExit(), which means structured-concurrency timeouts that wrap a ShellClient call could not actually interrupt the running command — the catch path only ran after `gh` (or whatever shell program) terminated of its own accord. PullRequestRefreshCoordinator's soft timeout was the immediate motivation, but every shell-backed client inherits the same hole. Three changes wire cancellation through: 1. Replace `process.waitUntilExit()` with `waitForExit(of:)`, an async wait built on `process.terminationHandler` plus a CheckedContinuation. A `LockIsolated` guards against double-resume in the race where the process exits between `terminationHandler` assignment and the `isRunning` sanity check. 2. Wrap the wait in `withTaskCancellationHandler { … } onCancel: { process.terminate() }` so a cancelled worker Task immediately sends SIGTERM rather than waiting for natural exit. 3. Hold the running Process in a LockIsolated box and call `process.terminate()` plus `workerTask.cancel()` from the stream's `onTermination` handler, so when the stream consumer goes away (cancel, throw, explicit finish) the underlying process is torn down too. Switch the worker / pipe-reader Tasks from `.detached` to inheriting Tasks while we're here — they were already isolated by virtue of being created inside the stream closure, but the explicit form removes a small footgun. Tests cover the consumer-cancel and Task-cancel paths against a real `/bin/sleep 30`; the consumer used to block for the full 30 seconds, now it returns in well under a second. A regression test for a fast-exiting (`/bin/sh -c true`) command guards the terminationHandler/isRunning race. Closes #368. --- supacode/Clients/Shell/ShellClient.swift | 65 +++++++++++------ supacodeTests/ShellClientStreamingTests.swift | 69 +++++++++++++++++++ 2 files changed, 114 insertions(+), 20 deletions(-) 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( -- 2.51.2