1176 lines · 43.8 KB · Swift
RawDownload
#if canImport(Darwin)import Darwin#elseif canImport(Glibc)import Glibc#elseif canImport(Musl)import Musl#endifimport Foundation private enum TTYCommandRunnerActiveProcessRegistry {    private static let condition = NSCondition()    private nonisolated(unsafe) static var processes: [pid_t: ProcessInfo] = [:]    private nonisolated(unsafe) static var isShuttingDown = false    private nonisolated(unsafe) static var launchesInProgress = 0     private struct ProcessInfo {        let binary: String        var processGroup: pid_t?    }     @discardableResult    static func register(pid: pid_t, binary: String) -> Bool {        guard pid > 0 else { return false }        self.condition.lock()        defer { self.condition.unlock() }        guard !self.isShuttingDown else { return false }        self.processes[pid] = ProcessInfo(binary: binary, processGroup: nil)        return true    }     static func beginLaunch() -> Bool {        self.condition.lock()        defer { self.condition.unlock() }        guard !self.isShuttingDown else { return false }        self.launchesInProgress += 1        return true    }     static func endLaunch() {        self.condition.lock()        self.launchesInProgress = max(0, self.launchesInProgress - 1)        if self.launchesInProgress == 0 {            self.condition.broadcast()        }        self.condition.unlock()    }     static func updateProcessGroup(pid: pid_t, processGroup: pid_t?) {        guard pid > 0 else { return }        self.condition.lock()        guard var existing = self.processes[pid] else {            self.condition.unlock()            return        }        existing.processGroup = processGroup        self.processes[pid] = existing        self.condition.unlock()    }     static func unregister(pid: pid_t) {        guard pid > 0 else { return }        self.condition.lock()        self.processes.removeValue(forKey: pid)        self.condition.unlock()    }     static func drainForShutdown(        onFenceSet: (() -> Void)? = nil)        -> [(pid: pid_t, binary: String, processGroup: pid_t?)]    {        self.condition.lock()        self.isShuttingDown = true        onFenceSet?()        while self.launchesInProgress > 0 {            self.condition.wait()        }        let drained = self.processes.map {            (pid: $0.key, binary: $0.value.binary, processGroup: $0.value.processGroup)        }        self.processes.removeAll()        self.condition.unlock()        return drained    }     static func reset() {        self.condition.lock()        self.processes.removeAll()        self.isShuttingDown = false        self.launchesInProgress = 0        self.condition.broadcast()        self.condition.unlock()    }     static func count() -> Int {        self.condition.lock()        let count = self.processes.count        self.condition.unlock()        return count    }     static func testTrackProcess(pid: pid_t, binary: String, processGroup: pid_t?) {        guard pid > 0 else { return }        self.condition.lock()        self.processes[pid] = ProcessInfo(binary: binary, processGroup: processGroup)        self.condition.unlock()    }} enum TTYProcessTreeTerminator {    struct ProcessIdentity: Hashable {        let pid: pid_t        let startToken: UInt64    }     static func descendantPIDs(        of rootPID: pid_t,        childResolver: (pid_t) -> [pid_t] = Self.currentChildPIDs(of:)) -> [pid_t]    {        guard rootPID > 0 else { return [] }         var seen: Set<pid_t> = [rootPID]        var pending = childResolver(rootPID)        var descendants: [pid_t] = []         while let pid = pending.popLast() {            guard pid > 0, seen.insert(pid).inserted else { continue }            descendants.append(pid)            pending.append(contentsOf: childResolver(pid))        }         return descendants    }     static func currentChildPIDs(of parentPID: pid_t) -> [pid_t] {        guard parentPID > 0 else { return [] }         #if canImport(Darwin)        var pids = [pid_t](repeating: 0, count: 128)        let byteCount = Int32(pids.count * MemoryLayout<pid_t>.stride)        let childCount = proc_listchildpids(parentPID, &pids, byteCount)        guard childCount > 0 else { return [] }        return Array(pids.prefix(min(Int(childCount), pids.count))).filter { $0 > 0 }        #else        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?,        signal: Int32,        knownDescendants: [pid_t] = [],        childResolver: (pid_t) -> [pid_t] = Self.currentChildPIDs(of:),        signalSender: (pid_t, Int32) -> Void = { kill($0, $1) })    {        guard rootPID > 0 else { return }         var seen: Set<pid_t> = [rootPID]        let descendants = knownDescendants + self.descendantPIDs(of: rootPID, childResolver: childResolver)        for pid in descendants where pid > 0 && seen.insert(pid).inserted {            signalSender(pid, signal)        }        if let processGroup {            signalSender(-processGroup, signal)        }        signalSender(rootPID, signal)    }} /// Executes an interactive CLI inside a pseudo-terminal and returns all captured text./// Keeps it minimal so we can reuse for Codex and Claude without tmux.public struct TTYCommandRunner {    private static let log = CodexBarLog.logger(LogCategories.ttyRunner)     public struct Result: Sendable {        public let text: String    }     public struct Options: Sendable {        public var rows: UInt16 = 50        public var cols: UInt16 = 160        public var timeout: TimeInterval = 20.0        /// Stop early once output has been idle for this long (only for non-Codex flows).        /// Useful for interactive TUIs that render once and then wait for input indefinitely.        public var idleTimeout: TimeInterval?        public var workingDirectory: URL?        public var extraArgs: [String] = []        public var baseEnvironment: [String: String]?        public var initialDelay: TimeInterval = 0.4        public var sendEnterEvery: TimeInterval?        public var sendOnSubstrings: [String: String]        public var stopOnURL: Bool        public var stopOnSubstrings: [String]        public var settleAfterStop: TimeInterval        public var forceCodexStatusMode: Bool        public var useClaudeProbeWorkingDirectory: Bool         public init(            rows: UInt16 = 50,            cols: UInt16 = 160,            timeout: TimeInterval = 20.0,            idleTimeout: TimeInterval? = nil,            workingDirectory: URL? = nil,            extraArgs: [String] = [],            baseEnvironment: [String: String]? = nil,            initialDelay: TimeInterval = 0.4,            sendEnterEvery: TimeInterval? = nil,            sendOnSubstrings: [String: String] = [:],            stopOnURL: Bool = false,            stopOnSubstrings: [String] = [],            settleAfterStop: TimeInterval = 0.25,            forceCodexStatusMode: Bool = false,            useClaudeProbeWorkingDirectory: Bool = false)        {            self.rows = rows            self.cols = cols            self.timeout = timeout            self.idleTimeout = idleTimeout            self.workingDirectory = workingDirectory            self.extraArgs = extraArgs            self.baseEnvironment = baseEnvironment            self.initialDelay = initialDelay            self.sendEnterEvery = sendEnterEvery            self.sendOnSubstrings = sendOnSubstrings            self.stopOnURL = stopOnURL            self.stopOnSubstrings = stopOnSubstrings            self.settleAfterStop = settleAfterStop            self.forceCodexStatusMode = forceCodexStatusMode            self.useClaudeProbeWorkingDirectory = useClaudeProbeWorkingDirectory        }    }     public enum Error: Swift.Error, LocalizedError, Sendable {        case binaryNotFound(String)        case launchFailed(String)        case timedOut         public var errorDescription: String? {            switch self {            case let .binaryNotFound(bin):                "Missing CLI '\(bin)'. Install it (e.g. npm i -g @openai/codex) or add it to PATH."            case let .launchFailed(msg): "Failed to launch process: \(msg)"            case .timedOut: "PTY command timed out."            }        }    }     public init() {}     public static func terminateActiveProcessesForAppShutdown() {        let targets = TTYCommandRunnerActiveProcessRegistry.drainForShutdown()        guard !targets.isEmpty else { return }         let resolvedTargets = self.resolveShutdownTargets(            targets,            hostProcessGroup: getpgrp(),            groupResolver: { getpgid($0) })         for target in resolvedTargets where target.pid > 0 {            TTYProcessTreeTerminator.terminateProcessTree(                rootPID: target.pid,                processGroup: target.processGroup,                signal: SIGTERM)        }         for target in resolvedTargets where target.pid > 0 {            TTYProcessTreeTerminator.terminateProcessTree(                rootPID: target.pid,                processGroup: target.processGroup,                signal: SIGKILL)        }    }     private static func resolveShutdownTargets(        _ targets: [(pid: pid_t, binary: String, processGroup: pid_t?)],        hostProcessGroup: pid_t,        groupResolver: (pid_t) -> pid_t) -> [(pid: pid_t, binary: String, processGroup: pid_t?)]    {        var resolvedTargets: [(pid: pid_t, binary: String, processGroup: pid_t?)] = []        resolvedTargets.reserveCapacity(targets.count)         for target in targets {            var resolvedGroup = target.processGroup            if resolvedGroup == nil {                let pgid = groupResolver(target.pid)                if pgid > 0, pgid != hostProcessGroup {                    resolvedGroup = pgid                }            } else if resolvedGroup == hostProcessGroup {                resolvedGroup = nil            }             resolvedTargets.append((pid: target.pid, binary: target.binary, processGroup: resolvedGroup))        }        return resolvedTargets    }     struct RollingBuffer {        private let maxNeedle: Int        private var tail = Data()         init(maxNeedle: Int) {            self.maxNeedle = max(0, maxNeedle)        }         mutating func append(_ data: Data) -> Data {            guard !data.isEmpty else { return Data() }             var combined = Data()            combined.reserveCapacity(self.tail.count + data.count)            combined.append(self.tail)            combined.append(data)             if self.maxNeedle > 1 {                if combined.count >= self.maxNeedle - 1 {                    self.tail = combined.suffix(self.maxNeedle - 1)                } else {                    self.tail = combined                }            } else {                self.tail.removeAll(keepingCapacity: true)            }             return combined        }         mutating func reset() {            self.tail.removeAll(keepingCapacity: true)        }    }     enum DrainReadResult {        case data(Data)        case wouldBlock        case closed    }     static func lowercasedASCII(_ data: Data) -> Data {        guard !data.isEmpty else { return data }        var out = Data(count: data.count)        out.withUnsafeMutableBytes { dest in            data.withUnsafeBytes { source in                let src = source.bindMemory(to: UInt8.self)                let dst = dest.bindMemory(to: UInt8.self)                for idx in 0..<src.count {                    var byte = src[idx]                    if byte >= 65, byte <= 90 { byte += 32 }                    dst[idx] = byte                }            }        }        return out    }     static func drainRemainingOutput(        until drainDeadline: Date,        readChunk: () -> DrainReadResult,        processChunk: (Data) -> Void,        sleep: (UInt32) -> Void = { usleep($0) })    {        while Date() < drainDeadline {            switch readChunk() {            case let .data(newData):                processChunk(newData)            case .wouldBlock:                sleep(20000)            case .closed:                return            }        }    }     static func drainReadResult(for data: Data, terminalRead: Int, errno err: Int32) -> DrainReadResult {        if !data.isEmpty { return .data(data) }         if terminalRead == 0 {            return .closed        }         if terminalRead < 0 {            if err == EAGAIN || err == EWOULDBLOCK || err == EINTR {                return .wouldBlock            }            if err == EIO {                return .closed            }        }         return .closed    }     static func locateBundledHelper(_ name: String) -> String? {        let fm = FileManager.default         func isExecutable(_ path: String) -> Bool {            fm.isExecutableFile(atPath: path)        }         if let override = ProcessInfo.processInfo.environment["CODEXBAR_HELPER_\(name.uppercased())"],           isExecutable(override)        {            return override        }         func candidate(inAppBundleURL appURL: URL) -> String? {            let path = appURL                .appendingPathComponent("Contents", isDirectory: true)                .appendingPathComponent("Helpers", isDirectory: true)                .appendingPathComponent(name, isDirectory: false)                .path            return isExecutable(path) ? path : nil        }         let mainURL = Bundle.main.bundleURL        if mainURL.pathExtension == "app", let found = candidate(inAppBundleURL: mainURL) { return found }         if let argv0 = CommandLine.arguments.first {            var url = URL(fileURLWithPath: argv0)            if !argv0.hasPrefix("/") {                url = URL(fileURLWithPath: FileManager.default.currentDirectoryPath).appendingPathComponent(argv0)            }            var probe = url            for _ in 0..<6 {                let parent = probe.deletingLastPathComponent()                if parent.pathExtension == "app", let found = candidate(inAppBundleURL: parent) { return found }                if parent.path == probe.path { break }                probe = parent            }        }         return nil    }     // swiftlint:disable function_body_length    // swiftlint:disable:next cyclomatic_complexity    public func run(        binary: String,        send script: String,        options: Options = Options(),        onURLDetected: (@Sendable () -> Void)? = nil) throws -> Result    {        let resolved: String        if FileManager.default.isExecutableFile(atPath: binary) {            resolved = binary        } else if let hit = Self.which(binary) {            resolved = hit        } else {            Self.log.warning("PTY binary not found", metadata: ["binary": binary])            throw Error.binaryNotFound(binary)        }         let binaryName = URL(fileURLWithPath: resolved).lastPathComponent        Self.log.debug(            "PTY start",            metadata: [                "binary": binaryName,                "timeout": "\(options.timeout)",                "rows": "\(options.rows)",                "cols": "\(options.cols)",                "args": "\(options.extraArgs.count)",            ])

Lines 1–500 of 1176.