Fix Kiro usage command pipe hangs (#1535)

* Fix Kiro usage command pipe hangs

* Fix Kiro probe lint failures

* fix: harden Kiro command cleanup

* fix: terminate escaped subprocess descendants

---------

Co-authored-by: Kiran Magic <kiran@Alices-Laptop.local>
Co-authored-by: kiranmagic7 <209323973+kiranmagic7@users.noreply.github.com>
Co-authored-by: Peter Steinberger <steipete@gmail.com>

Author: 4 people

CHANGELOG.md

@@ -24,6 +24,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!
 

Sources/CodexBarCore/Host/PTY/TTYCommandRunner.swift

@@ -105,6 +105,11 @@ private enum TTYCommandRunnerActiveProcessRegistry {
 }
 
 enum TTYProcessTreeTerminator {
+    struct ProcessIdentity: Equatable {
+        let pid: pid_t
+        let startToken: UInt64
+    }
+
     static func descendantPIDs(
         of rootPID: pid_t,
         childResolver: (pid_t) -> [pid_t] = Self.currentChildPIDs(of:)) -> [pid_t]
@@ -134,10 +139,49 @@ enum TTYProcessTreeTerminator {
         guard childCount > 0 else { return [] }
         return Array(pids.prefix(min(Int(childCount), pids.count))).filter { $0 > 0 }
         #else
-        return []
+        let taskPath = "/proc/\(parentPID)/task"
+        guard let taskIDs = try? FileManager.default.contentsOfDirectory(atPath: taskPath) else { return [] }
+
+        var children: Set<pid_t> = []
+        for taskID in taskIDs {
+            let childrenPath = "\(taskPath)/\(taskID)/children"
+            guard let text = try? String(contentsOfFile: childrenPath, encoding: .utf8) else { continue }
+            children.formUnion(text.split(whereSeparator: \.isWhitespace).compactMap { pid_t($0) })
+        }
+        return children.sorted()
+        #endif
+    }
+
+    static func processIdentity(for pid: pid_t) -> ProcessIdentity? {
+        guard pid > 0 else { return nil }
+
+        #if canImport(Darwin)
+        var info = proc_bsdinfo()
+        let size = proc_pidinfo(
+            pid,
+            PROC_PIDTBSDINFO,
+            0,
+            &info,
+            Int32(MemoryLayout<proc_bsdinfo>.stride))
+        guard size == Int32(MemoryLayout<proc_bsdinfo>.stride) else { return nil }
+        let startToken = UInt64(info.pbi_start_tvsec) * 1_000_000 + UInt64(info.pbi_start_tvusec)
+        return ProcessIdentity(pid: pid, startToken: startToken)
+        #else
+        guard let stat = try? String(contentsOfFile: "/proc/\(pid)/stat", encoding: .utf8),
+              let commandEnd = stat.lastIndex(of: ")")
+        else {
+            return nil
+        }
+        let fields = stat[stat.index(after: commandEnd)...].split(whereSeparator: \.isWhitespace)
+        guard fields.count > 19, let startToken = UInt64(fields[19]) else { return nil }
+        return ProcessIdentity(pid: pid, startToken: startToken)
         #endif
     }
 
+    static func isCurrent(_ identity: ProcessIdentity) -> Bool {
+        self.processIdentity(for: identity.pid) == identity
+    }
+
     static func terminateProcessTree(
         rootPID: pid_t,
         processGroup: pid_t?,

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()
     }

Sources/CodexBarCore/Host/Process/SubprocessRunner.swift

@@ -110,21 +110,36 @@ 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 {
-            kill(-pgid, SIGTERM)
-        }
+        let descendants = TTYProcessTreeTerminator.descendantPIDs(of: process.processIdentifier)
+        let descendantIdentities = descendants.compactMap(TTYProcessTreeTerminator.processIdentity(for:))
+        TTYProcessTreeTerminator.terminateProcessTree(
+            rootPID: process.processIdentifier,
+            processGroup: processGroup,
+            signal: SIGTERM,
+            knownDescendants: descendants)
         let killDeadline = Date().addingTimeInterval(0.4)
         while process.isRunning, Date() < killDeadline {
             usleep(50000)
         }
         if process.isRunning {
-            if let pgid = processGroup {
-                kill(-pgid, SIGKILL)
+            let currentDescendants = descendantIdentities
+                .filter(TTYProcessTreeTerminator.isCurrent(_:))
+                .map(\.pid)
+            TTYProcessTreeTerminator.terminateProcessTree(
+                rootPID: process.processIdentifier,
+                processGroup: processGroup,
+                signal: SIGKILL,
+                knownDescendants: currentDescendants)
+            let reapDeadline = Date().addingTimeInterval(0.4)
+            while process.isRunning, Date() < reapDeadline {
+                usleep(50000)
+            }
+        } else {
+            for identity in descendantIdentities where TTYProcessTreeTerminator.isCurrent(identity) {
+                kill(identity.pid, SIGKILL)
             }
-            kill(process.processIdentifier, SIGKILL)
         }
         return true
     }

Sources/CodexBarCore/Providers/Kiro/KiroStatusProbe.swift

@@ -1,4 +1,9 @@
 import Foundation
+#if canImport(Darwin)
+import Darwin
+#else
+import Glibc
+#endif
 
 public struct KiroUsageSnapshot: Sendable {
     public let planName: String
@@ -219,7 +224,15 @@ public enum KiroStatusProbeError: LocalizedError, Sendable {
 }
 
 public struct KiroStatusProbe: Sendable {
-    public init() {}
+    private let cliBinaryResolver: @Sendable () -> String?
+
+    public init() {
+        self.cliBinaryResolver = { TTYCommandRunner.which("kiro-cli") }
+    }
+
+    init(cliBinaryResolver: @escaping @Sendable () -> String?) {
+        self.cliBinaryResolver = cliBinaryResolver
+    }
 
     private static let logger = CodexBarLog.logger(LogCategories.kiro)
 
@@ -256,7 +269,7 @@ public struct KiroStatusProbe: Sendable {
             contextUsage: contextUsage)
     }
 
-    private struct KiroCLIResult {
+    struct KiroCLIResult {
         let stdout: String
         let stderr: String
         let terminationStatus: Int32
@@ -383,12 +396,12 @@ 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
     {
-        guard let binary = TTYCommandRunner.which("kiro-cli") else {
+        guard let binary = self.cliBinaryResolver() else {
             throw KiroStatusProbeError.cliNotFound
         }
 
@@ -406,124 +419,98 @@ 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
-            }
-
-            func appendStdout(_ data: Data) {
-                self.lock.lock()
-                defer { lock.unlock() }
-                self._stdoutData.append(data)
-                self._lastActivityAt = Date()
-                self._hasReceivedOutput = true
+                self.lock.withLock { self._hasReceivedOutput }
             }
 
-            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()
+        let stdoutCapture = ProcessPipeCapture(pipe: stdoutPipe, onData: { state.markActivity() })
+        let stderrCapture = ProcessPipeCapture(pipe: stderrPipe, onData: { state.markActivity() })
 
-        // Set up readability handlers to track activity
-        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)
-            }
+        do {
+            try process.run()
+        } catch {
+            stdoutCapture.stop()
+            stderrCapture.stop()
+            throw error
         }
+        stdoutCapture.start()
+        stderrCapture.start()
+        let pid = process.processIdentifier
+        let processGroup: pid_t? = setpgid(pid, pid) == 0 ? pid : nil
 
-        return try await withCheckedThrowingContinuation { continuation in
-            DispatchQueue.global().async {
-                do {
-                    try process.run()
-                } catch {
-                    stdoutPipe.fileHandleForReading.readabilityHandler = nil
-                    stderrPipe.fileHandleForReading.readabilityHandler = nil
-                    continuation.resume(throwing: error)
-                    return
-                }
+        let deadline = Date().addingTimeInterval(timeout)
+        var didHitDeadline = false
+        var didTerminateForIdle = false
 
-                let deadline = Date().addingTimeInterval(timeout)
-                var didHitDeadline = false
-                var didTerminateForIdle = false
-
-                while process.isRunning {
-                    if Date() >= deadline {
-                        didHitDeadline = true
-                        break
-                    }
-                    // Idle timeout: if we got output but then it went silent
-                    if state.hasReceivedOutput,
-                       Date().timeIntervalSince(state.lastActivityAt) >= idleTimeout
-                    {
-                        // Process went idle after producing output - likely done or stuck
-                        didTerminateForIdle = true
-                        break
-                    }
-                    Thread.sleep(forTimeInterval: 0.1)
+        do {
+            while process.isRunning {
+                try Task.checkCancellation()
+                if Date() >= deadline {
+                    didHitDeadline = true
+                    break
                 }
-
-                // Clean up handlers
-                stdoutPipe.fileHandleForReading.readabilityHandler = nil
-                stderrPipe.fileHandleForReading.readabilityHandler = nil
-
-                if process.isRunning {
-                    process.terminate()
-                    process.waitUntilExit()
-                    if didHitDeadline || !state.hasReceivedOutput {
-                        continuation.resume(throwing: KiroStatusProbeError.timeout)
-                        return
-                    }
+                if state.hasReceivedOutput,
+                   Date().timeIntervalSince(state.lastActivityAt) >= idleTimeout
+                {
+                    didTerminateForIdle = true
+                    break
                 }
+                try await Task.sleep(for: .milliseconds(100))
+            }
+        } catch {
+            await Self.terminateProcess(process, processGroup: processGroup)
+            stdoutCapture.stop()
+            stderrCapture.stop()
+            throw error
+        }
+
+        if process.isRunning {
+            await Self.terminateProcess(process, processGroup: processGroup)
+            guard !process.isRunning else {
+                stdoutCapture.stop()
+                stderrCapture.stop()
+                throw KiroStatusProbeError.timeout
+            }
+            if didHitDeadline || !state.hasReceivedOutput {
+                stdoutCapture.stop()
+                stderrCapture.stop()
+                throw KiroStatusProbeError.timeout
+            }
+        }
+
+        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: String(data: output.stdout, encoding: .utf8) ?? "",
+            stderr: String(data: output.stderr, encoding: .utf8) ?? "",
+            terminationStatus: process.terminationStatus,
+            terminatedForIdle: didTerminateForIdle)
+    }
 
-                // Read any remaining data
-                let remainingStdout = stdoutPipe.fileHandleForReading.readDataToEndOfFile()
-                let remainingStderr = stderrPipe.fileHandleForReading.readDataToEndOfFile()
-
-                var output = state.getOutput()
-                output.stdout.append(remainingStdout)
-                output.stderr.append(remainingStderr)
-
-                let stdoutOutput = String(data: output.stdout, encoding: .utf8) ?? ""
-                let stderrOutput = String(data: output.stderr, encoding: .utf8) ?? ""
-                continuation.resume(returning: KiroCLIResult(
-                    stdout: stdoutOutput,
-                    stderr: stderrOutput,
-                    terminationStatus: process.terminationStatus,
-                    terminatedForIdle: didTerminateForIdle))
+    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()
             }
         }
     }