diff --git a/Sources/CodexBar/CodexbarApp.swift b/Sources/CodexBar/CodexbarApp.swift index f870a2405..14035ecec 100644 --- a/Sources/CodexBar/CodexbarApp.swift +++ b/Sources/CodexBar/CodexbarApp.swift @@ -279,6 +279,15 @@ final class AppDelegate: NSObject, NSApplicationDelegate { } } + func applicationShouldTerminate(_ sender: NSApplication) -> NSApplication.TerminateReply { + ProcessRegistry.shared.terminateAll() + return .terminateNow + } + + func applicationWillTerminate(_ notification: Notification) { + ProcessRegistry.shared.terminateAll() + } + /// Use the classic (non-Liquid Glass) app icon on macOS versions before 26. private func configureAppIconForMacOSVersion() { if #unavailable(macOS 26) { diff --git a/Sources/CodexBarCore/Host/PTY/TTYCommandRunner.swift b/Sources/CodexBarCore/Host/PTY/TTYCommandRunner.swift index 5178696ca..fca51a624 100644 --- a/Sources/CodexBarCore/Host/PTY/TTYCommandRunner.swift +++ b/Sources/CodexBarCore/Host/PTY/TTYCommandRunner.swift @@ -265,6 +265,7 @@ public struct TTYCommandRunner { var cleanedUp = false var didLaunch = false var processGroup: pid_t? + var registryToken: UUID? /// Always tear down the PTY child (and its process group) even if we throw early /// while bootstrapping the CLI (e.g. when it prompts for login/telemetry). func cleanup() { @@ -301,6 +302,10 @@ public struct TTYCommandRunner { if didLaunch { proc.waitUntilExit() } + + if let registryToken { + ProcessRegistry.shared.unregister(registryToken) + } } // Ensure the PTY process is always torn down, even when we throw early (e.g. login prompt). @@ -324,6 +329,7 @@ public struct TTYCommandRunner { if setpgid(pid, pid) == 0 { processGroup = pid } + registryToken = ProcessRegistry.shared.register(process: proc, processGroup: processGroup) func send(_ text: String) throws { guard let data = text.data(using: .utf8) else { return } diff --git a/Sources/CodexBarCore/Host/Process/ProcessRegistry.swift b/Sources/CodexBarCore/Host/Process/ProcessRegistry.swift new file mode 100644 index 000000000..7d45af6c3 --- /dev/null +++ b/Sources/CodexBarCore/Host/Process/ProcessRegistry.swift @@ -0,0 +1,71 @@ +import Foundation +#if canImport(Darwin) +import Darwin +#else +import Glibc +#endif + +public final class ProcessRegistry: @unchecked Sendable { + public static let shared = ProcessRegistry() + + private struct Entry { + let process: Process + let processGroup: pid_t? + } + + private let lock = NSLock() + private var entries: [UUID: Entry] = [:] + + private init() {} + + @discardableResult + public func register(process: Process, processGroup: pid_t?) -> UUID { + let token = UUID() + self.lock.lock() + self.entries[token] = Entry(process: process, processGroup: processGroup) + self.lock.unlock() + return token + } + + public func unregister(_ token: UUID) { + self.lock.lock() + self.entries[token] = nil + self.lock.unlock() + } + + public func terminateAll() { + self.lock.lock() + let active = Array(self.entries.values) + self.entries.removeAll() + self.lock.unlock() + + for entry in active { + let pid = entry.process.processIdentifier + let shouldSignalGroup: Bool = { + guard let pgid = entry.processGroup, pgid > 0 else { return false } + guard entry.process.isRunning else { return false } + return getpgid(pid) == pgid + }() + + if entry.process.isRunning { + entry.process.terminate() + } + + if shouldSignalGroup, let pgid = entry.processGroup { + kill(-pgid, SIGTERM) + } + + let waitDeadline = Date().addingTimeInterval(1.0) + while entry.process.isRunning, Date() < waitDeadline { + usleep(100_000) + } + + if entry.process.isRunning { + if shouldSignalGroup, let pgid = entry.processGroup { + kill(-pgid, SIGKILL) + } + kill(pid, SIGKILL) + } + } + } +} diff --git a/Sources/CodexBarCore/Providers/CLIProbeSessionResetter.swift b/Sources/CodexBarCore/Providers/CLIProbeSessionResetter.swift index fd3300653..1eae2ed68 100644 --- a/Sources/CodexBarCore/Providers/CLIProbeSessionResetter.swift +++ b/Sources/CodexBarCore/Providers/CLIProbeSessionResetter.swift @@ -4,5 +4,6 @@ public enum CLIProbeSessionResetter { public static func resetAll() async { await ClaudeCLISession.shared.reset() await CodexCLISession.shared.reset() + ProcessRegistry.shared.terminateAll() } } diff --git a/Sources/CodexBarCore/Providers/Claude/ClaudeCLISession.swift b/Sources/CodexBarCore/Providers/Claude/ClaudeCLISession.swift index 9ae1bc052..6473b8638 100644 --- a/Sources/CodexBarCore/Providers/Claude/ClaudeCLISession.swift +++ b/Sources/CodexBarCore/Providers/Claude/ClaudeCLISession.swift @@ -32,6 +32,7 @@ actor ClaudeCLISession { private var processGroup: pid_t? private var binaryPath: String? private var startedAt: Date? + private var registryToken: UUID? private let promptSends: [String: String] = [ "Do you trust the files in this folder?": "y\r", @@ -302,6 +303,7 @@ actor ClaudeCLISession { if setpgid(pid, pid) == 0 { processGroup = pid } + let registryToken = ProcessRegistry.shared.register(process: proc, processGroup: processGroup) self.process = proc self.primaryFD = primaryFD @@ -310,6 +312,7 @@ actor ClaudeCLISession { self.processGroup = processGroup self.binaryPath = binary self.startedAt = Date() + self.registryToken = registryToken } private static func scrubbedClaudeEnvironment(from base: [String: String]) -> [String: String] { @@ -362,6 +365,10 @@ actor ClaudeCLISession { self.primaryFD = -1 self.processGroup = nil self.startedAt = nil + if let registryToken = self.registryToken { + ProcessRegistry.shared.unregister(registryToken) + } + self.registryToken = nil } private func readChunk() -> Data { diff --git a/Sources/CodexBarCore/Providers/Codex/CodexCLISession.swift b/Sources/CodexBarCore/Providers/Codex/CodexCLISession.swift index 8c7f3b245..9ee9c8844 100644 --- a/Sources/CodexBarCore/Providers/Codex/CodexCLISession.swift +++ b/Sources/CodexBarCore/Providers/Codex/CodexCLISession.swift @@ -31,6 +31,7 @@ actor CodexCLISession { private var startedAt: Date? private var ptyRows: UInt16 = 0 private var ptyCols: UInt16 = 0 + private var registryToken: UUID? private struct RollingBuffer { private let maxNeedle: Int @@ -289,6 +290,7 @@ actor CodexCLISession { if setpgid(pid, pid) == 0 { processGroup = pid } + let registryToken = ProcessRegistry.shared.register(process: proc, processGroup: processGroup) self.process = proc self.primaryFD = primaryFD @@ -299,6 +301,7 @@ actor CodexCLISession { self.startedAt = Date() self.ptyRows = rows self.ptyCols = cols + self.registryToken = registryToken } private func cleanup() { @@ -336,6 +339,10 @@ actor CodexCLISession { self.startedAt = nil self.ptyRows = 0 self.ptyCols = 0 + if let registryToken = self.registryToken { + ProcessRegistry.shared.unregister(registryToken) + } + self.registryToken = nil } private func readChunk() -> Data { diff --git a/Sources/CodexBarCore/UsageFetcher.swift b/Sources/CodexBarCore/UsageFetcher.swift index ca300ea9f..e7cb569a1 100644 --- a/Sources/CodexBarCore/UsageFetcher.swift +++ b/Sources/CodexBarCore/UsageFetcher.swift @@ -301,6 +301,8 @@ private final class CodexRPCClient: @unchecked Sendable { private let stdoutLineStream: AsyncStream private let stdoutLineContinuation: AsyncStream.Continuation private var nextID = 1 + private var registryToken: UUID? + private var processGroup: pid_t? private final class LineBuffer: @unchecked Sendable { private let lock = NSLock() @@ -367,6 +369,15 @@ private final class CodexRPCClient: @unchecked Sendable { throw RPCWireError.startFailed(error.localizedDescription) } + let pid = self.process.processIdentifier + if setpgid(pid, pid) == 0 { + self.processGroup = pid + } else { + self.processGroup = nil + } + + self.registryToken = ProcessRegistry.shared.register(process: self.process, processGroup: self.processGroup) + let stdoutHandle = self.stdoutPipe.fileHandleForReading let stdoutLineContinuation = self.stdoutLineContinuation let stdoutBuffer = LineBuffer() @@ -419,10 +430,39 @@ private final class CodexRPCClient: @unchecked Sendable { } func shutdown() { + let pid = self.process.processIdentifier + let shouldSignalGroup: Bool = { + guard let pgid = self.processGroup, pgid > 0 else { return false } + guard self.process.isRunning else { return false } + return getpgid(pid) == pgid + }() + if self.process.isRunning { Self.log.debug("Codex RPC stopping") self.process.terminate() } + + if shouldSignalGroup, let pgid = self.processGroup { + kill(-pgid, SIGTERM) + } + + let waitDeadline = Date().addingTimeInterval(1.0) + while self.process.isRunning, Date() < waitDeadline { + usleep(100_000) + } + + if self.process.isRunning { + if shouldSignalGroup, let pgid = self.processGroup { + kill(-pgid, SIGKILL) + } + kill(pid, SIGKILL) + } + + if !self.process.isRunning, let registryToken = self.registryToken { + ProcessRegistry.shared.unregister(registryToken) + self.registryToken = nil + } + self.processGroup = nil } // MARK: - JSON-RPC helpers