fix: harden Kiro CLI idle handling (#145) (thanks @chadneal)

Author: steipete

CHANGELOG.md

@@ -9,6 +9,7 @@
 - Vertex AI: harden quota usage parsing for edge-case responses.
 - Kiro: add CLI-based usage provider via kiro-cli. Thanks @neror!
 - Kiro: clean up provider wiring and show plan name in the menu.
+- Kiro: harden CLI idle handling to avoid partial usage snapshots (#145). Thanks @chadneal!
 - Augment: add provider with browser-cookie usage tracking.
 - Cursor: support legacy request-based plans and show individual on-demand usage (#125) — thanks @vltansky
 - Cursor: avoid Intel crash when opening login and harden WebKit teardown. Thanks @meghanto!

Sources/CodexBarCore/Providers/Kiro/KiroStatusProbe.swift

@@ -114,6 +114,7 @@ public struct KiroStatusProbe: Sendable {
         let stdout: String
         let stderr: String
         let terminationStatus: Int32
+        let terminatedForIdle: Bool
     }
 
     private func ensureLoggedIn() async throws {
@@ -147,7 +148,10 @@ public struct KiroStatusProbe: Sendable {
     }
 
     private func runUsageCommand() async throws -> String {
-        let result = try await self.runCommand(arguments: ["chat", "--no-interactive", "/usage"], timeout: 20.0)
+        let result = try await self.runCommand(
+            arguments: ["chat", "--no-interactive", "/usage"],
+            timeout: 20.0,
+            idleTimeout: 10.0)
         let trimmedStdout = result.stdout.trimmingCharacters(in: .whitespacesAndNewlines)
         let trimmedStderr = result.stderr.trimmingCharacters(in: .whitespacesAndNewlines)
         let combinedOutput = trimmedStderr.isEmpty ? trimmedStdout : trimmedStderr
@@ -162,6 +166,10 @@ public struct KiroStatusProbe: Sendable {
             throw KiroStatusProbeError.notLoggedIn
         }
 
+        if result.terminatedForIdle, !Self.isUsageOutputComplete(combinedOutput) {
+            throw KiroStatusProbeError.timeout
+        }
+
         if !trimmedStdout.isEmpty {
             return result.stdout
         }
@@ -183,8 +191,8 @@ public struct KiroStatusProbe: Sendable {
     private func runCommand(
         arguments: [String],
         timeout: TimeInterval,
-        idleTimeout: TimeInterval = 5.0
-    ) async throws -> KiroCLIResult {
+        idleTimeout: TimeInterval = 5.0) async throws -> KiroCLIResult
+    {
         guard let binary = TTYCommandRunner.which("kiro-cli") else {
             throw KiroStatusProbeError.cliNotFound
         }
@@ -212,37 +220,37 @@ public struct KiroStatusProbe: Sendable {
             private var _stderrData = Data()
 
             var lastActivityAt: Date {
-                lock.lock()
+                self.lock.lock()
                 defer { lock.unlock() }
-                return _lastActivityAt
+                return self._lastActivityAt
             }
 
             var hasReceivedOutput: Bool {
-                lock.lock()
+                self.lock.lock()
                 defer { lock.unlock() }
-                return _hasReceivedOutput
+                return self._hasReceivedOutput
             }
 
             func appendStdout(_ data: Data) {
-                lock.lock()
+                self.lock.lock()
                 defer { lock.unlock() }
-                _stdoutData.append(data)
-                _lastActivityAt = Date()
-                _hasReceivedOutput = true
+                self._stdoutData.append(data)
+                self._lastActivityAt = Date()
+                self._hasReceivedOutput = true
             }
 
             func appendStderr(_ data: Data) {
-                lock.lock()
+                self.lock.lock()
                 defer { lock.unlock() }
-                _stderrData.append(data)
-                _lastActivityAt = Date()
-                _hasReceivedOutput = true
+                self._stderrData.append(data)
+                self._lastActivityAt = Date()
+                self._hasReceivedOutput = true
             }
 
             func getOutput() -> (stdout: Data, stderr: Data) {
-                lock.lock()
+                self.lock.lock()
                 defer { lock.unlock() }
-                return (_stdoutData, _stderrData)
+                return (self._stdoutData, self._stderrData)
             }
         }
 
@@ -274,13 +282,20 @@ public struct KiroStatusProbe: Sendable {
                 }
 
                 let deadline = Date().addingTimeInterval(timeout)
+                var didHitDeadline = false
+                var didTerminateForIdle = false
 
-                while process.isRunning, Date() < deadline {
+                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)
@@ -293,8 +308,7 @@ public struct KiroStatusProbe: Sendable {
                 if process.isRunning {
                     process.terminate()
                     process.waitUntilExit()
-                    // Only throw timeout if we never got output
-                    if !state.hasReceivedOutput {
+                    if didHitDeadline || !state.hasReceivedOutput {
                         continuation.resume(throwing: KiroStatusProbeError.timeout)
                         return
                     }
@@ -313,7 +327,8 @@ public struct KiroStatusProbe: Sendable {
                 continuation.resume(returning: KiroCLIResult(
                     stdout: stdoutOutput,
                     stderr: stderrOutput,
-                    terminationStatus: process.terminationStatus))
+                    terminationStatus: process.terminationStatus,
+                    terminatedForIdle: didTerminateForIdle))
             }
         }
     }
@@ -342,7 +357,6 @@ public struct KiroStatusProbe: Sendable {
         }
 
         // Track which key patterns matched to detect format changes
-        var matchedPlanName = false
         var matchedPercent = false
         var matchedCredits = false
 
@@ -351,7 +365,6 @@ public struct KiroStatusProbe: Sendable {
         if let planMatch = stripped.range(of: #"\|\s*(KIRO\s+\w+)"#, options: .regularExpression) {
             let raw = String(stripped[planMatch]).replacingOccurrences(of: "|", with: "")
             planName = raw.trimmingCharacters(in: .whitespaces)
-            matchedPlanName = true
         }
 
         // Parse reset date from "resets on 01/01"
@@ -411,7 +424,7 @@ public struct KiroStatusProbe: Sendable {
         }
 
         // Require at least one key pattern to match to avoid silent failures
-        if !matchedPlanName, !matchedPercent, !matchedCredits {
+        if !matchedPercent, !matchedCredits {
             throw KiroStatusProbeError.parseError(
                 "No recognizable usage patterns found. Kiro CLI output format may have changed.")
         }
@@ -463,4 +476,11 @@ public struct KiroStatusProbe: Sendable {
         components.year = currentYear + 1
         return calendar.date(from: components)
     }
+
+    private static func isUsageOutputComplete(_ output: String) -> Bool {
+        let stripped = self.stripANSI(output).lowercased()
+        return stripped.contains("covered in plan")
+            || stripped.contains("resets on")
+            || stripped.contains("bonus credits")
+    }
 }

Tests/CodexBarTests/KiroStatusProbeTests.swift

@@ -110,6 +110,18 @@ struct KiroStatusProbeTests {
         #expect(snapshot.bonusExpiryDays == 1)
     }
 
+    @Test
+    func rejectsOutputMissingUsageMarkers() throws {
+        let output = """
+        | KIRO FREE                                          |
+        """
+
+        let probe = KiroStatusProbe()
+        #expect(throws: KiroStatusProbeError.self) {
+            try probe.parse(output: output)
+        }
+    }
+
     // MARK: - Snapshot Conversion
 
     @Test

docs/kiro.md

@@ -14,7 +14,7 @@ Kiro uses the AWS `kiro-cli` tool to fetch usage data. No browser cookies or OAu
 
 1) **CLI command** (primary and only strategy)
    - Command: `kiro-cli chat --no-interactive "/usage"`
-   - Timeout: 10 seconds.
+   - Timeout: 20 seconds (idle cutoff after ~10s of no output once the CLI starts responding).
    - Requires `kiro-cli` installed and logged in via AWS Builder ID.
    - Output is ANSI-decorated; CodexBar strips escape sequences before parsing.