CodexBar / Sources / CodexBarCore / TTYCommandRunner.swift

import Darwin
import Foundation

/// 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 {
    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
        public var extraArgs: [String] = []

        public init(rows: UInt16 = 50, cols: UInt16 = 160, timeout: TimeInterval = 20.0, extraArgs: [String] = []) {
            self.rows = rows
            self.cols = cols
            self.timeout = timeout
            self.extraArgs = extraArgs
        }
    }

    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): "Binary not found on PATH: \(bin)"
            case let .launchFailed(msg): "Failed to launch process: \(msg)"
            case .timedOut: "PTY command timed out."
            }
        }
    }

    public init() {}

    // swiftlint:disable function_body_length
    public func run(binary: String, send script: String, options: Options = Options()) throws -> Result {
        guard let resolved = Self.which(binary) else { throw Error.binaryNotFound(binary) }

        var primaryFD: Int32 = -1
        var secondaryFD: Int32 = -1
        var term = termios()
        var win = winsize(ws_row: options.rows, ws_col: options.cols, ws_xpixel: 0, ws_ypixel: 0)
        guard openpty(&primaryFD, &secondaryFD, nil, &term, &win) == 0 else {
            throw Error.launchFailed("openpty failed")
        }
        // Make primary side non-blocking so read loops don't hang when no data is available.
        _ = fcntl(primaryFD, F_SETFL, O_NONBLOCK)

        let primaryHandle = FileHandle(fileDescriptor: primaryFD, closeOnDealloc: true)
        let secondaryHandle = FileHandle(fileDescriptor: secondaryFD, closeOnDealloc: true)

        let proc = Process()
        proc.executableURL = URL(fileURLWithPath: resolved)
        proc.arguments = options.extraArgs
        proc.standardInput = secondaryHandle
        proc.standardOutput = secondaryHandle
        proc.standardError = secondaryHandle
        // Mirror RPC PATH seeding so CLIs installed via npm/nvm/fnm/bun still launch in hardened builds,
        // but keep the caller’s environment (HOME, LANG, BUN_INSTALL, etc.) so the CLIs can find their
        // auth/config files.
        proc.environment = Self.enrichedEnvironment()

        var cleanedUp = false
        var didLaunch = false
        var processGroup: pid_t?
        // 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() {
            guard !cleanedUp else { return }
            cleanedUp = true

            if didLaunch, proc.isRunning {
                let exitData = Data("/exit\n".utf8)
                try? primaryHandle.write(contentsOf: exitData)
            }

            try? primaryHandle.close()
            try? secondaryHandle.close()

            guard didLaunch else { return }

            if proc.isRunning {
                proc.terminate()
            }
            if let pgid = processGroup {
                kill(-pgid, SIGTERM)
            }
            let waitDeadline = Date().addingTimeInterval(2.0)
            while proc.isRunning, Date() < waitDeadline {
                usleep(100_000)
            }
            if proc.isRunning {
                if let pgid = processGroup {
                    kill(-pgid, SIGKILL)
                }
                kill(proc.processIdentifier, SIGKILL)
            }
            if didLaunch {
                proc.waitUntilExit()
            }
        }

        // Ensure the PTY process is always torn down, even when we throw early (e.g. login prompt).
        defer { cleanup() }

        try proc.run()
        didLaunch = true

        // Isolate the child into its own process group so descendant helpers can be
        // terminated together. If this fails (e.g. process already exec'ed), we
        // continue and fall back to single-PID termination.
        let pid = proc.processIdentifier
        if setpgid(pid, pid) == 0 {
            processGroup = pid
        }

        func send(_ text: String) throws {
            guard let data = text.data(using: .utf8) else { return }
            try primaryHandle.write(contentsOf: data)
        }

        let deadline = Date().addingTimeInterval(options.timeout)
        var buffer = Data()
        func readChunk() {
            var tmp = [UInt8](repeating: 0, count: 8192)
            let n = Darwin.read(primaryFD, &tmp, tmp.count)
            if n > 0 { buffer.append(contentsOf: tmp.prefix(n)) }
        }

        func containsCodexStatus() -> Bool {
            let markers = [
                "Credits:",
                "5h limit",
                "5-hour limit",
                "Weekly limit",
            ].map { Data($0.utf8) }
            return markers.contains { buffer.contains($0) }
        }

        func respondIfCursorQuerySeen() {
            let query = Data([0x1B, 0x5B, 0x36, 0x6E]) // ESC [ 6 n
            if buffer.contains(query) {
                // Pretend cursor is at 1;1, which is enough to satisfy Codex CLI's probe.
                try? send("\u{1b}[1;1R")
            }
        }

        func containsCodexUpdatePrompt() -> Bool {
            let needles = [
                "Update available!",
                "Run bun install -g @openai/codex",
                "0.60.1 ->",
            ]
            let lower = String(data: buffer, encoding: .utf8)?.lowercased() ?? ""
            return needles.contains { lower.contains($0.lowercased()) }
        }

        // Generic behavior (Codex /status and other commands).
        usleep(400_000) // small boot grace
        let delayInitialSend = script.trimmingCharacters(in: .whitespacesAndNewlines) == "/status"
        if !delayInitialSend {
            try send(script)
            try send("\r")
            usleep(150_000)
            try send("\r")
            try send("\u{1b}")
        }

        var skippedCodexUpdate = false
        var sentScript = !delayInitialSend
        var updateSkipAttempts = 0
        var lastEnter = Date(timeIntervalSince1970: 0)
        var scriptSentAt: Date? = sentScript ? Date() : nil
        var resendStatusRetries = 0
        var enterRetries = 0
        var sawCodexStatus = false

        while Date() < deadline {
            readChunk()
            respondIfCursorQuerySeen()
            if !skippedCodexUpdate, containsCodexUpdatePrompt() {
                // Prompt shows options: 1) Update now, 2) Skip, 3) Skip until next version.
                // Users report one Down + Enter is enough; follow with an extra Enter for safety, then re-run
                // /status.
                try? send("\u{1b}[B") // highlight option 2 (Skip)
                usleep(120_000)
                try? send("\r")
                usleep(150_000)
                try? send("\r") // if still focused on prompt, confirm again
                try? send("/status")
                try? send("\r")
                updateSkipAttempts += 1
                if updateSkipAttempts >= 1 {
                    skippedCodexUpdate = true
                    sentScript = false // re-send /status after dismissing
                    scriptSentAt = nil
                    buffer.removeAll()
                }
                usleep(300_000)
            }
            if !sentScript, !containsCodexUpdatePrompt() || skippedCodexUpdate {
                try? send(script)
                try? send("\r")
                sentScript = true
                scriptSentAt = Date()
                lastEnter = Date()
                usleep(200_000)
                continue
            }
            if sentScript, !containsCodexStatus() {
                if Date().timeIntervalSince(lastEnter) >= 1.2, enterRetries < 6 {
                    try? send("\r")
                    enterRetries += 1
                    lastEnter = Date()
                    usleep(120_000)
                    continue
                }
                if let sentAt = scriptSentAt,
                   Date().timeIntervalSince(sentAt) >= 3.0,
                   resendStatusRetries < 2
                {
                    try? send("/status")
                    try? send("\r")
                    resendStatusRetries += 1
                    buffer.removeAll()
                    scriptSentAt = Date()
                    lastEnter = Date()
                    usleep(220_000)
                    continue
                }
            }
            if containsCodexStatus() {
                sawCodexStatus = true
                break
            }
            usleep(120_000)
        }

        if sawCodexStatus {
            let settleDeadline = Date().addingTimeInterval(2.0)
            while Date() < settleDeadline {
                readChunk()
                respondIfCursorQuerySeen()
                usleep(100_000)
            }
        }

        guard let text = String(data: buffer, encoding: .utf8), !text.isEmpty else {
            throw Error.timedOut
        }

        return Result(text: text)
    }

    // swiftlint:enable function_body_length

    public static func which(_ tool: String) -> String? {
        if tool == "codex", let located = BinaryLocator.resolveCodexBinary() { return located }
        if tool == "claude", let located = BinaryLocator.resolveClaudeBinary() { return located }
        // First try system PATH
        if let path = runWhich(tool) { return path }
        // Fallback to common locations (Homebrew, local bins)
        let home = NSHomeDirectory()
        let candidates = [
            "/opt/homebrew/bin/\(tool)",
            "/usr/local/bin/\(tool)",
            "\(home)/.local/bin/\(tool)",
            "\(home)/bin/\(tool)",
        ]
        for c in candidates where FileManager.default.isExecutableFile(atPath: c) {
            return c
        }
        return nil
    }

    private static func runWhich(_ tool: String) -> String? {
        let proc = Process()
        proc.executableURL = URL(fileURLWithPath: "/usr/bin/which")
        proc.arguments = [tool]
        let pipe = Pipe()
        proc.standardOutput = pipe
        try? proc.run()
        proc.waitUntilExit()
        guard proc.terminationStatus == 0 else { return nil }
        let data = pipe.fileHandleForReading.readDataToEndOfFile()
        guard let path = String(data: data, encoding: .utf8)?
            .trimmingCharacters(in: .whitespacesAndNewlines),
            !path.isEmpty else { return nil }
        return path
    }

    /// Expands PATH with the same defaults we use for Codex RPC, so TTY probes can find CLIs installed via Homebrew,
    /// bun, nvm, fnm, or npm.
    public static func enrichedPath() -> String {
        PathBuilder.effectivePATH(
            purposes: [.tty, .nodeTooling],
            env: ProcessInfo.processInfo.environment)
    }

    static func enrichedEnvironment(
        baseEnv: [String: String] = ProcessInfo.processInfo.environment,
        home: String = NSHomeDirectory()) -> [String: String]
    {
        var env = baseEnv
        env["PATH"] = PathBuilder.effectivePATH(
            purposes: [.tty, .nodeTooling],
            env: baseEnv,
            home: home)
        if env["HOME"]?.isEmpty ?? true {
            env["HOME"] = home
        }
        if env["TERM"]?.isEmpty ?? true {
            env["TERM"] = "xterm-256color"
        }
        return env
    }
}