fix: harden Kiro command cleanup

steipete · steipete · commit ad1453433cb2 · 2026-06-15T01:02:15.000-04:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -21,6 +21,7 @@
 - Provider probes: stop waiting indefinitely for inherited output pipes after subprocesses or CLI version checks exit (fixes #1531).
 - Menu bar: update visible usage values in place when a manual refresh completes instead of leaving the open provider card stale until the menu is reopened (fixes #1516).
 - Gemini: recognize the current `gemini-api-key` CLI auth setting so API-key sessions show the supported OAuth guidance instead of a misleading not-logged-in error (fixes #1511).
+- Kiro: keep usage refreshes bounded when CLI helpers retain output pipes, ignore termination, or are cancelled (fixes #1533). Thanks @kiranmagic7!
 - Xiaomi MiMo: cancel optional token-plan requests when the required balance request fails instead of delaying the error for up to 30 seconds.
 - Settings: make the cost history window directly editable by keyboard while preserving the existing stepper and 1–365 day bounds (fixes #1499). Thanks @kiranmagic7!
 
diff --git a/Sources/CodexBarCore/Host/Process/ProcessPipeCapture.swift b/Sources/CodexBarCore/Host/Process/ProcessPipeCapture.swift
@@ -2,15 +2,17 @@ import Foundation
 
 package final class ProcessPipeCapture: @unchecked Sendable {
     private let handle: FileHandle
+    private let onData: (@Sendable () -> Void)?
     private let condition = NSCondition()
     private var data = Data()
     private var activeCallbacks = 0
     private var isFinished = false
     private var isStopping = false
     private var continuation: CheckedContinuation<Void, Never>?
 
-    package init(pipe: Pipe) {
+    package init(pipe: Pipe, onData: (@Sendable () -> Void)? = nil) {
         self.handle = pipe.fileHandleForReading
+        self.onData = onData
     }
 
     package func start() {
@@ -70,6 +72,8 @@ package final class ProcessPipeCapture: @unchecked Sendable {
 
         if chunk.isEmpty {
             handle.readabilityHandler = nil
+        } else {
+            self.onData?()
         }
         continuation?.resume()
     }
diff --git a/Sources/CodexBarCore/Host/Process/SubprocessRunner.swift b/Sources/CodexBarCore/Host/Process/SubprocessRunner.swift
@@ -110,7 +110,7 @@ public enum SubprocessRunner {
     /// Terminates a process and its process group, escalating from SIGTERM to SIGKILL.
     /// Returns `true` if the process was actually killed, `false` if it had already exited.
     @discardableResult
-    private static func terminateProcess(_ process: Process, processGroup: pid_t?) -> Bool {
+    package static func terminateProcess(_ process: Process, processGroup: pid_t?) -> Bool {
         guard process.isRunning else { return false }
         process.terminate()
         if let pgid = processGroup {
@@ -125,6 +125,10 @@ public enum SubprocessRunner {
                 kill(-pgid, SIGKILL)
             }
             kill(process.processIdentifier, SIGKILL)
+            let reapDeadline = Date().addingTimeInterval(0.4)
+            while process.isRunning, Date() < reapDeadline {
+                usleep(50000)
+            }
         }
         return true
     }
diff --git a/Sources/CodexBarCore/Providers/Kiro/KiroStatusProbe.swift b/Sources/CodexBarCore/Providers/Kiro/KiroStatusProbe.swift
@@ -269,7 +269,7 @@ public struct KiroStatusProbe: Sendable {
             contextUsage: contextUsage)
     }
 
-    private struct KiroCLIResult {
+    struct KiroCLIResult {
         let stdout: String
         let stderr: String
         let terminationStatus: Int32
@@ -396,7 +396,7 @@ public struct KiroStatusProbe: Sendable {
         return self.parseContextUsage(output: output)
     }
 
-    private func runCommand(
+    func runCommand(
         arguments: [String],
         timeout: TimeInterval,
         idleTimeout: TimeInterval = 5.0) async throws -> KiroCLIResult
@@ -419,123 +419,100 @@ public struct KiroStatusProbe: Sendable {
         env["TERM"] = "xterm-256color"
         process.environment = env
 
-        // Thread-safe state for activity tracking
         final class ActivityState: @unchecked Sendable {
             private let lock = NSLock()
             private var _lastActivityAt = Date()
             private var _hasReceivedOutput = false
-            private var _stdoutData = Data()
-            private var _stderrData = Data()
 
             var lastActivityAt: Date {
-                self.lock.lock()
-                defer { lock.unlock() }
-                return self._lastActivityAt
+                self.lock.withLock { self._lastActivityAt }
             }
 
             var hasReceivedOutput: Bool {
-                self.lock.lock()
-                defer { lock.unlock() }
-                return self._hasReceivedOutput
+                self.lock.withLock { self._hasReceivedOutput }
             }
 
-            func appendStdout(_ data: Data) {
-                self.lock.lock()
-                defer { lock.unlock() }
-                self._stdoutData.append(data)
-                self._lastActivityAt = Date()
-                self._hasReceivedOutput = true
-            }
-
-            func appendStderr(_ data: Data) {
-                self.lock.lock()
-                defer { lock.unlock() }
-                self._stderrData.append(data)
-                self._lastActivityAt = Date()
-                self._hasReceivedOutput = true
-            }
-
-            func getOutput() -> (stdout: Data, stderr: Data) {
-                self.lock.lock()
-                defer { lock.unlock() }
-                return (self._stdoutData, self._stderrData)
+            func markActivity() {
+                self.lock.withLock {
+                    self._lastActivityAt = Date()
+                    self._hasReceivedOutput = true
+                }
             }
         }
 
         let state = ActivityState()
-        stdoutPipe.fileHandleForReading.readabilityHandler = { handle in
-            let data = handle.availableData
-            if !data.isEmpty {
-                state.appendStdout(data)
-            }
-        }
-        stderrPipe.fileHandleForReading.readabilityHandler = { handle in
-            let data = handle.availableData
-            if !data.isEmpty {
-                state.appendStderr(data)
-            }
-        }
+        let stdoutCapture = ProcessPipeCapture(pipe: stdoutPipe, onData: { state.markActivity() })
+        let stderrCapture = ProcessPipeCapture(pipe: stderrPipe, onData: { state.markActivity() })
 
         do {
             try process.run()
         } catch {
-            stdoutPipe.fileHandleForReading.readabilityHandler = nil
-            stderrPipe.fileHandleForReading.readabilityHandler = nil
+            stdoutCapture.stop()
+            stderrCapture.stop()
             throw error
         }
+        stdoutCapture.start()
+        stderrCapture.start()
+        let pid = process.processIdentifier
+        let processGroup: pid_t? = setpgid(pid, pid) == 0 ? pid : nil
 
         let deadline = Date().addingTimeInterval(timeout)
         var didHitDeadline = false
         var didTerminateForIdle = false
 
-        while process.isRunning {
-            if Date() >= deadline {
-                didHitDeadline = true
-                break
-            }
-            if state.hasReceivedOutput,
-               Date().timeIntervalSince(state.lastActivityAt) >= idleTimeout
-            {
-                didTerminateForIdle = true
-                break
+        do {
+            while process.isRunning {
+                try Task.checkCancellation()
+                if Date() >= deadline {
+                    didHitDeadline = true
+                    break
+                }
+                if state.hasReceivedOutput,
+                   Date().timeIntervalSince(state.lastActivityAt) >= idleTimeout
+                {
+                    didTerminateForIdle = true
+                    break
+                }
+                try await Task.sleep(for: .milliseconds(100))
             }
-            try await Task.sleep(for: .milliseconds(100))
+        } catch {
+            await Self.terminateProcess(process, processGroup: processGroup)
+            stdoutCapture.stop()
+            stderrCapture.stop()
+            throw error
         }
 
         if process.isRunning {
-            Self.terminateProcess(process)
+            await Self.terminateProcess(process, processGroup: processGroup)
+            guard !process.isRunning else {
+                stdoutCapture.stop()
+                stderrCapture.stop()
+                throw KiroStatusProbeError.timeout
+            }
             if didHitDeadline || !state.hasReceivedOutput {
-                stdoutPipe.fileHandleForReading.readabilityHandler = nil
-                stderrPipe.fileHandleForReading.readabilityHandler = nil
+                stdoutCapture.stop()
+                stderrCapture.stop()
                 throw KiroStatusProbeError.timeout
             }
         }
 
-        try await Task.sleep(for: .milliseconds(100))
-        stdoutPipe.fileHandleForReading.readabilityHandler = nil
-        stderrPipe.fileHandleForReading.readabilityHandler = nil
-
-        let output = state.getOutput()
-        let stdoutOutput = String(data: output.stdout, encoding: .utf8) ?? ""
-        let stderrOutput = String(data: output.stderr, encoding: .utf8) ?? ""
+        async let stdoutData = stdoutCapture.finish(timeout: .seconds(1))
+        async let stderrData = stderrCapture.finish(timeout: .seconds(1))
+        let output = await (stdout: stdoutData, stderr: stderrData)
         return KiroCLIResult(
-            stdout: stdoutOutput,
-            stderr: stderrOutput,
+            stdout: String(data: output.stdout, encoding: .utf8) ?? "",
+            stderr: String(data: output.stderr, encoding: .utf8) ?? "",
             terminationStatus: process.terminationStatus,
             terminatedForIdle: didTerminateForIdle)
     }
 
-    private static func terminateProcess(_ process: Process) {
-        guard process.isRunning else { return }
-        process.terminate()
-        let deadline = Date().addingTimeInterval(0.4)
-        while process.isRunning, Date() < deadline {
-            usleep(50000)
-        }
-        if process.isRunning {
-            kill(process.processIdentifier, SIGKILL)
+    private static func terminateProcess(_ process: Process, processGroup: pid_t?) async {
+        await withCheckedContinuation { continuation in
+            DispatchQueue.global(qos: .userInitiated).async {
+                SubprocessRunner.terminateProcess(process, processGroup: processGroup)
+                continuation.resume()
+            }
         }
-        process.waitUntilExit()
     }
 
     func parse(
diff --git a/Tests/CodexBarTests/KiroStatusProbeTests.swift b/Tests/CodexBarTests/KiroStatusProbeTests.swift
@@ -75,6 +75,86 @@ struct KiroStatusProbeTests {
         #expect(elapsed < 3, "Kiro usage capture should not wait for inherited pipe EOF, took \(elapsed)s")
     }
 
+    @Test
+    func `run command hard stops a process that ignores SIGTERM`() async throws {
+        let cliURL = try self.makeCLI(
+            """
+            #!/bin/sh
+            trap '' TERM
+            printf 'partial output\\n'
+            while true; do sleep 1; done
+            """)
+        defer { try? FileManager.default.removeItem(at: cliURL.deletingLastPathComponent()) }
+
+        let probe = KiroStatusProbe(cliBinaryResolver: { cliURL.path })
+        let start = Date()
+        let result = try await probe.runCommand(arguments: [], timeout: 2, idleTimeout: 0.1)
+        let elapsed = Date().timeIntervalSince(start)
+
+        #expect(result.terminatedForIdle)
+        #expect(result.stdout.contains("partial output"))
+        #expect(result.terminationStatus != 0)
+        #expect(elapsed < 2, "Ignored SIGTERM should escalate to SIGKILL, took \(elapsed)s")
+    }
+
+    @Test
+    func `run command preserves completed no-output failure status`() async throws {
+        let cliURL = try self.makeCLI(
+            """
+            #!/bin/sh
+            exit 23
+            """)
+        defer { try? FileManager.default.removeItem(at: cliURL.deletingLastPathComponent()) }
+
+        let probe = KiroStatusProbe(cliBinaryResolver: { cliURL.path })
+        let result = try await probe.runCommand(arguments: [], timeout: 2)
+
+        #expect(result.stdout.isEmpty)
+        #expect(result.stderr.isEmpty)
+        #expect(result.terminationStatus == 23)
+        #expect(!result.terminatedForIdle)
+    }
+
+    @Test
+    func `run command cancellation terminates the process`() async throws {
+        let pidFile = FileManager.default.temporaryDirectory
+            .appendingPathComponent("codexbar-kiro-cancel-\(UUID().uuidString).pid")
+        let cliURL = try self.makeCLI(
+            """
+            #!/bin/sh
+            printf '%s\\n' "$$" > "$1"
+            trap '' TERM
+            while true; do sleep 1; done
+            """)
+        defer {
+            try? FileManager.default.removeItem(at: cliURL.deletingLastPathComponent())
+            try? FileManager.default.removeItem(at: pidFile)
+        }
+
+        let probe = KiroStatusProbe(cliBinaryResolver: { cliURL.path })
+        let task = Task {
+            try await probe.runCommand(arguments: [pidFile.path], timeout: 20)
+        }
+        defer { task.cancel() }
+
+        var capturedProcessID: pid_t?
+        for _ in 0..<100 {
+            if let text = try? String(contentsOf: pidFile, encoding: .utf8) {
+                capturedProcessID = pid_t(text.trimmingCharacters(in: .whitespacesAndNewlines))
+                break
+            }
+            try await Task.sleep(for: .milliseconds(20))
+        }
+        let processID = try #require(capturedProcessID)
+        defer { _ = kill(processID, SIGKILL) }
+
+        task.cancel()
+        await #expect(throws: CancellationError.self) {
+            try await task.value
+        }
+        #expect(kill(processID, 0) == -1)
+    }
+
     // MARK: - Happy Path Parsing
 
     @Test
@@ -99,6 +179,16 @@ struct KiroStatusProbeTests {
         #expect(snapshot.resetsAt != nil)
     }
 
+    private func makeCLI(_ script: String) throws -> URL {
+        let root = FileManager.default.temporaryDirectory
+            .appendingPathComponent("codexbar-kiro-cli-\(UUID().uuidString)", isDirectory: true)
+        let cliURL = root.appendingPathComponent("kiro-cli")
+        try FileManager.default.createDirectory(at: root, withIntermediateDirectories: true)
+        try script.write(to: cliURL, atomically: true, encoding: .utf8)
+        try FileManager.default.setAttributes([.posixPermissions: 0o755], ofItemAtPath: cliURL.path)
+        return cliURL
+    }
+
     @Test
     func `parses output with bonus credits`() throws {
         let output = """

Original file line number	Diff line number	Diff line change
`@@ -110,7 +110,7 @@ public enum SubprocessRunner {`
`110`	`110`	`/// Terminates a process and its process group, escalating from SIGTERM to SIGKILL.`
`111`	`111`	/// Returns `true` if the process was actually killed, `false` if it had already exited.
`112`	`112`	`@discardableResult`
`113`		`- private static func terminateProcess(_ process: Process, processGroup: pid_t?) -> Bool {`
	`113`	`+ package static func terminateProcess(_ process: Process, processGroup: pid_t?) -> Bool {`
`114`	`114`	`guard process.isRunning else { return false }`
`115`	`115`	`process.terminate()`
`116`	`116`	`if let pgid = processGroup {`
`@@ -125,6 +125,10 @@ public enum SubprocessRunner {`
`125`	`125`	`kill(-pgid, SIGKILL)`
`126`	`126`	`}`
`127`	`127`	`kill(process.processIdentifier, SIGKILL)`
	`128`	`+ let reapDeadline = Date().addingTimeInterval(0.4)`
	`129`	`+ while process.isRunning, Date() < reapDeadline {`
	`130`	`+ usleep(50000)`
	`131`	`+ }`
`128`	`132`	`}`
`129`	`133`	`return true`
`130`	`134`	`}`