terminate helper reliably

Author: stephanos

Sources/ClaudeCost/AppDelegate.swift

@@ -7,6 +7,7 @@ final class AppDelegate: NSObject, NSApplicationDelegate, NSMenuDelegate {
 
   private var statusItem: NSStatusItem?
   private var timer: Timer?
+  private var refreshTask: Task<Void, Never>?
   private var state = AppState()
   private let loginItemManager = LoginItemManager()
   private var startAtLoginViewState = StartAtLoginViewState.make(status: .notRegistered)
@@ -35,6 +36,9 @@ final class AppDelegate: NSObject, NSApplicationDelegate, NSMenuDelegate {
   func applicationWillTerminate(_ notification: Notification) {
     timer?.invalidate()
     timer = nil
+    refreshTask?.cancel()
+    refreshTask = nil
+    UsageFetcher.cancelActiveHelper()
   }
 
   func menuNeedsUpdate(_ menu: NSMenu) {
@@ -71,13 +75,19 @@ final class AppDelegate: NSObject, NSApplicationDelegate, NSMenuDelegate {
     state = nextState
     renderTitle()
 
-    Task {
+    refreshTask?.cancel()
+    refreshTask = Task {
       do {
         let snapshot = try await UsageFetcher.fetchUsage()
         applyRefreshSuccess(snapshot)
       } catch {
+        guard !Task.isCancelled else {
+          return
+        }
         applyRefreshFailure(error)
       }
+
+      refreshTask = nil
     }
   }
 

Sources/ClaudeCost/UsageFetcher.swift

@@ -3,6 +3,7 @@ import Foundation
 private final class ProcessCompletionState: @unchecked Sendable {
   private let lock = NSLock()
   private var hasResumed = false
+  private var timedOut = false
 
   func tryMarkResumed() -> Bool {
     lock.lock()
@@ -15,6 +16,46 @@ private final class ProcessCompletionState: @unchecked Sendable {
     hasResumed = true
     return true
   }
+
+  func markTimedOut() {
+    lock.lock()
+    defer { lock.unlock() }
+    timedOut = true
+  }
+
+  func didTimeOut() -> Bool {
+    lock.lock()
+    defer { lock.unlock() }
+    return timedOut
+  }
+}
+
+private final class HelperProcessRegistry: @unchecked Sendable {
+  private let lock = NSLock()
+  private var process: Process?
+
+  func set(_ process: Process) {
+    lock.lock()
+    defer { lock.unlock() }
+    self.process = process
+  }
+
+  func clear(_ process: Process) {
+    lock.lock()
+    defer { lock.unlock() }
+
+    guard self.process === process else {
+      return
+    }
+
+    self.process = nil
+  }
+
+  func currentProcess() -> Process? {
+    lock.lock()
+    defer { lock.unlock() }
+    return process
+  }
 }
 
 struct UsagePayload: Codable {
@@ -63,19 +104,38 @@ enum UsageFetcherError: LocalizedError {
 
 enum UsageFetcher {
   private static let helperTimeout: TimeInterval = 30
+  private static let helperTerminationGracePeriod: TimeInterval = 2
+  private static let helperTerminationPollInterval: UInt32 = 50_000
+  private static let helperRegistry = HelperProcessRegistry()
 
   static func fetchUsage() async throws -> UsageSnapshot {
-    let monthStart = currentMonthStartString()
-    let helperURL = try usageHelperURL()
+    try await withTaskCancellationHandler {
+      let monthStart = currentMonthStartString()
+      let helperURL = try usageHelperURL()
 
-    let output = try await runHelper(at: helperURL, arguments: [monthStart])
-    return try UsagePayloadParser.decodeSnapshot(from: output)
+      let output = try await runHelper(at: helperURL, arguments: [monthStart])
+      return try UsagePayloadParser.decodeSnapshot(from: output)
+    } onCancel: {
+      cancelActiveHelper()
+    }
   }
 
   static func shouldTimeOut(processIsRunning: Bool) -> Bool {
     processIsRunning
   }
 
+  static func shouldEscalateTermination(processIsRunning: Bool, waitedEnough: Bool) -> Bool {
+    processIsRunning && waitedEnough
+  }
+
+  static func cancelActiveHelper() {
+    guard let process = helperRegistry.currentProcess() else {
+      return
+    }
+
+    terminateAndReap(process)
+  }
+
   private static func usageHelperURL() throws -> URL {
     if let override = ProcessInfo.processInfo.environment["CLAUDECOST_HELPER_PATH"],
       !override.isEmpty
@@ -121,12 +181,19 @@ enum UsageFetcher {
       process.standardError = stderr
 
       process.terminationHandler = { process in
+        helperRegistry.clear(process)
+
         let stdoutData = stdout.fileHandleForReading.readDataToEndOfFile()
         let stderrData = stderr.fileHandleForReading.readDataToEndOfFile()
         let stderrString =
           String(data: stderrData, encoding: .utf8)?
           .trimmingCharacters(in: .whitespacesAndNewlines) ?? ""
 
+        if completionState.didTimeOut() {
+          resumeOnce(with: .failure(UsageFetcherError.timedOut(helperURL.path)))
+          return
+        }
+
         guard process.terminationStatus == 0 else {
           resumeOnce(
             with: .failure(UsageFetcherError.commandFailed(helperURL.path, stderrString))
@@ -144,6 +211,7 @@ enum UsageFetcher {
 
       do {
         try process.run()
+        helperRegistry.set(process)
       } catch {
         resumeOnce(with: .failure(UsageFetcherError.failedToStart))
         return
@@ -154,12 +222,41 @@ enum UsageFetcher {
           return
         }
 
-        process.terminate()
+        completionState.markTimedOut()
+        terminateAndReap(process)
         resumeOnce(with: .failure(UsageFetcherError.timedOut(helperURL.path)))
       }
     }
   }
 
+  private static func terminateAndReap(_ process: Process) {
+    guard process.processIdentifier > 0 else {
+      return
+    }
+
+    if process.isRunning {
+      process.terminate()
+    }
+
+    waitForExit(of: process, timeout: helperTerminationGracePeriod)
+
+    if shouldEscalateTermination(
+      processIsRunning: process.isRunning,
+      waitedEnough: true
+    ) {
+      kill(process.processIdentifier, SIGKILL)
+      waitForExit(of: process, timeout: helperTerminationGracePeriod)
+    }
+  }
+
+  private static func waitForExit(of process: Process, timeout: TimeInterval) {
+    let deadline = Date().addingTimeInterval(timeout)
+
+    while process.isRunning && Date() < deadline {
+      usleep(helperTerminationPollInterval)
+    }
+  }
+
   private static func currentMonthStartString(now: Date = Date()) -> String {
     let formatter = DateFormatter()
     formatter.calendar = Calendar.current

Tests/UsageFetcherHarness.swift

@@ -9,4 +9,16 @@ func testUsageFetcherTimeoutDecision() throws {
     !UsageFetcher.shouldTimeOut(processIsRunning: false),
     "completed helper should not trigger timeout handling"
   )
+  try expect(
+    UsageFetcher.shouldEscalateTermination(processIsRunning: true, waitedEnough: true),
+    "running helper should escalate after the grace period"
+  )
+  try expect(
+    !UsageFetcher.shouldEscalateTermination(processIsRunning: true, waitedEnough: false),
+    "helper should not escalate before the grace period expires"
+  )
+  try expect(
+    !UsageFetcher.shouldEscalateTermination(processIsRunning: false, waitedEnough: true),
+    "stopped helper should not escalate"
+  )
 }