1270 lines · 47.6 KB · Swift
RawDownload
#if canImport(Darwin)import Darwin#elseif canImport(Glibc)import Glibc#elseif canImport(Musl)import Musl#endifimport Foundation // MARK: - Antigravity CLI Process Abstractions protocol AntigravityCLIProcessHandle: AnyObject, Sendable {    var pid: pid_t { get }    var isRunning: Bool { get }    var processGroup: pid_t? { get }     func assignProcessGroup() -> pid_t?    func sendExit() throws    func closePTY()    func terminateRoot()    func killRoot()    func descendantPIDs() -> [pid_t]    func terminateTree(signal: Int32, knownDescendants: [pid_t])    func killDescendants(_ descendants: [pid_t])    func drainOutput() -> Data} protocol AntigravityCLIProcessLaunching: Sendable {    func launch(binary: String) throws -> any AntigravityCLIProcessHandle} enum AntigravityCLIAuthenticationPrompt {    static let evidence = Data("Select login method:".utf8)    private static let promptPattern = #"select\s+login\s+method\s*:?"#     static func contains(_ output: Data) -> Bool {        // `agy` briefly prints "You are currently not signed in" before it        // auto-refreshes an existing login. The actual blocking state is the        // interactive login-method prompt. The exact prompt is a CLI-owned TUI        // string, so keep matching tolerant to casing and whitespace changes.        if output.range(of: self.evidence) != nil {            return true        }        let asciiBytes = output.map { $0 < 0x80 ? $0 : 0x20 }        let text = String(bytes: asciiBytes, encoding: .utf8) ?? ""        return text.range(            of: self.promptPattern,            options: [.regularExpression, .caseInsensitive]) != nil    }} struct AntigravityCLIProcessIdentity: Equatable {    let executablePath: String    let startEpoch: TimeInterval} protocol AntigravityCLIProcessIdentityProviding: Sendable {    func identity(for pid: pid_t) -> AntigravityCLIProcessIdentity?} struct AntigravityCLISessionRecord: Codable, Equatable {    let pid: pid_t    let requestedBinaryPath: String    let executablePath: String    let startEpoch: TimeInterval    let processGroup: pid_t?    let ownerPID: pid_t?    let ownerExecutablePath: String?    let ownerStartEpoch: TimeInterval?     init(        pid: pid_t,        requestedBinaryPath: String,        executablePath: String,        startEpoch: TimeInterval,        processGroup: pid_t?,        ownerPID: pid_t? = nil,        ownerExecutablePath: String? = nil,        ownerStartEpoch: TimeInterval? = nil)    {        self.pid = pid        self.requestedBinaryPath = requestedBinaryPath        self.executablePath = executablePath        self.startEpoch = startEpoch        self.processGroup = processGroup        self.ownerPID = ownerPID        self.ownerExecutablePath = ownerExecutablePath        self.ownerStartEpoch = ownerStartEpoch    }} protocol AntigravityCLISessionRecordStoring: Sendable {    func load() throws -> [AntigravityCLISessionRecord]    func save(_ record: AntigravityCLISessionRecord) throws    func remove(_ record: AntigravityCLISessionRecord) throws} protocol AntigravityCLISessionLaunchLocking: Sendable {    func withLock<T>(_ operation: () throws -> T) throws -> T} // MARK: - AntigravityCLISession /// Manages a bounded background ``agy`` process whose embedded localhost server/// provides the same ``GetUserStatus`` endpoint as the desktop Antigravity app's/// ``language_server``. The CLI is kept alive in a PTY so its daemon stays bound/// to a local port - this lets CodexBar read Claude + Gemini quotas even when/// the desktop Antigravity app is closed.////// The session intentionally does not scrape TUI output. It only launches and/// keeps the process reachable for HTTPS probing, drains discarded PTY output so/// the CLI cannot block on a full terminal buffer, and bounds the warm lifetime/// with an idle timer so CodexBar does not run an IDE backend forever.actor AntigravityCLISession {    static let shared = AntigravityCLISession()    private static let log = CodexBarLog.logger(LogCategories.antigravity)     enum ResetCause: Int {        case deferred        case oneShotCLI        case unhealthy        case authentication         var message: String {            switch self {            case .deferred: "deferred reset"            case .oneShotCLI: "one-shot CLI fetch"            case .unhealthy: "unhealthy CLI HTTPS session"            case .authentication: "authentication required"            }        }    }     private struct LaunchOutcome {        let pid: pid_t        let rejectedProcess: (any AntigravityCLIProcessHandle)?        let rejectionMessage: String?        let holdsLaunchReservation: Bool    }     struct Dependencies {        var launcher: any AntigravityCLIProcessLaunching        var identityProvider: any AntigravityCLIProcessIdentityProviding        var recordStore: any AntigravityCLISessionRecordStoring        var launchLock: any AntigravityCLISessionLaunchLocking        var beginAppShutdownTrackedLaunch: @Sendable () -> Bool        var endAppShutdownTrackedLaunch: @Sendable () -> Void        var registerForAppShutdown: @Sendable (pid_t, String) -> Bool        var updateAppShutdownProcessGroup: @Sendable (pid_t, pid_t?) -> Void        var unregisterForAppShutdown: @Sendable (pid_t) -> Void        var descendantPIDs: @Sendable (pid_t) -> [pid_t]        var terminateProcessTree: @Sendable (pid_t, pid_t?, Int32, [pid_t]) -> Void        var currentProcessID: @Sendable () -> pid_t        var now: @Sendable () -> Date        var sleep: @Sendable (UInt64) async throws -> Void        var idleWindow: TimeInterval        var failureRelaunchThreshold: Int        var terminationGracePeriod: TimeInterval         static func live() -> Self {            Self(                launcher: AntigravityPTYProcessLauncher(),                identityProvider: AntigravityProcessIdentityProvider(),                recordStore: AntigravityFileCLISessionRecordStore(),                launchLock: AntigravityFileCLISessionLaunchLock(),                beginAppShutdownTrackedLaunch: {                    TTYCommandRunner.beginActiveProcessLaunchForAppShutdown()                },                endAppShutdownTrackedLaunch: {                    TTYCommandRunner.endActiveProcessLaunchForAppShutdown()                },                registerForAppShutdown: { pid, binary in                    TTYCommandRunner.registerActiveProcessForAppShutdown(pid: pid, binary: binary)                },                updateAppShutdownProcessGroup: { pid, group in                    TTYCommandRunner.updateActiveProcessGroupForAppShutdown(pid: pid, processGroup: group)                },                unregisterForAppShutdown: { pid in                    TTYCommandRunner.unregisterActiveProcessForAppShutdown(pid: pid)                },                descendantPIDs: { pid in                    TTYProcessTreeTerminator.descendantPIDs(of: pid)                },                terminateProcessTree: { pid, group, signal, knownDescendants in                    TTYProcessTreeTerminator.terminateProcessTree(                        rootPID: pid,                        processGroup: group,                        signal: signal,                        knownDescendants: knownDescendants)                },                currentProcessID: getpid,                now: Date.init,                sleep: { nanoseconds in                    try await Task.sleep(nanoseconds: nanoseconds)                },                idleWindow: 180,                failureRelaunchThreshold: 2,                terminationGracePeriod: 1)        }    }     // MARK: State     private let dependencies: Dependencies    private var process: (any AntigravityCLIProcessHandle)?    private var binaryPath: String?    private var sessionIdleWindow: TimeInterval    private var activeProbeCount = 0    private var activeSessionProbeCount = 0    private var resetRequestedWhenIdle = false    private var hardResetRequestedWhenIdle = false    private var pendingResetCause: ResetCause?    private var idleTask: Task<Void, Never>?    private var sessionGeneration: UInt64 = 0    private var consecutiveProbeFailures = 0    private var persistedProcessIdentity: AntigravityCLIProcessIdentity?    private var recentOutput = Data()    private var authenticationPromptObserved = false    private var lastStopReason: String?    private var lifecycleOperationInProgress = false    private var lifecycleWaiters: [CheckedContinuation<Void, Never>] = []    private var exclusiveProbeWaiters: [CheckedContinuation<Void, Never>] = []     init(dependencies: Dependencies = .live()) {        self.dependencies = dependencies        self.sessionIdleWindow = dependencies.idleWindow    }     /// The pid of the running ``agy`` process, exposed so callers can discover    /// its listening ports via `lsof`.    var pid: pid_t? {        guard let proc = self.process, proc.isRunning else { return nil }        return proc.pid    }     /// Whether the managed process is alive and matches ``binaryPath``.    var isRunning: Bool {        guard let proc = self.process, proc.isRunning, self.binaryPath != nil else { return false }        return true    }     var failureCountForTesting: Int {        self.consecutiveProbeFailures    }     var idleWindowForTesting: TimeInterval {        self.sessionIdleWindow    }     var activeProbeCountForTesting: Int {        self.activeProbeCount    }     var lastStopReasonForTesting: String? {        self.lastStopReason    }     // MARK: Lifecycle     /// Mark a probe as active and ensure a warm ``agy`` is running on the given binary path.    ///    /// Callers must balance this with ``finishProbe(success:resetAfterFetch:)`` so    /// idle/reset cleanup cannot kill the process while its ports are being probed.    /// If previous probes repeatedly failed while the process stayed alive, this    /// force-relaunches instead of reusing a wedged HTTPS server forever.    func beginProbe(binary: String, idleWindow: TimeInterval? = nil) async throws -> pid_t {        self.activeProbeCount += 1        if let idleWindow, idleWindow > 0 {            self.sessionIdleWindow = max(self.dependencies.idleWindow, idleWindow)        }        self.cancelIdleTimer()        do {            return try await self.withLifecycleOperation {                let pid = try await self.ensureStartedLocked(binary: binary)                self.activeSessionProbeCount += 1                return pid            }        } catch {            self.activeProbeCount = max(0, self.activeProbeCount - 1)            self.notifyExclusiveProbeWaitersIfNeeded()            if self.activeProbeCount == 0, self.resetRequestedWhenIdle {                await self.withLifecycleOperation {                    guard self.activeProbeCount == 0 else {                        self.resetRequestedWhenIdle = true                        return                    }                    let forceTerminate = self.hardResetRequestedWhenIdle || self.consecutiveProbeFailures > 0                    self.resetRequestedWhenIdle = false                    self.hardResetRequestedWhenIdle = false                    await self.stopCurrentSessionLocked(                        reason: "deferred reset after failed begin",                        clearRecord: true,                        graceful: !forceTerminate)                }            }            throw error        }    }     /// Record probe completion and either keep the session warm for the bounded    /// idle window or tear it down immediately for one-shot CLI invocations.    func finishProbe(success: Bool, resetAfterFetch: Bool, forceTerminate: Bool = false) async {        if success {            self.consecutiveProbeFailures = 0        } else {            self.consecutiveProbeFailures += 1        }         self.activeProbeCount = max(0, self.activeProbeCount - 1)        self.activeSessionProbeCount = max(0, self.activeSessionProbeCount - 1)        self.notifyExclusiveProbeWaitersIfNeeded()        let shouldForceStopUnhealthy = !success &&            self.consecutiveProbeFailures >= max(1, self.dependencies.failureRelaunchThreshold)        let shouldReset = forceTerminate || resetAfterFetch || self.resetRequestedWhenIdle || shouldForceStopUnhealthy        let currentResetCause = Self.resetCause(            authenticationRequired: forceTerminate,            resetAfterFetch: resetAfterFetch,            shouldForceStopUnhealthy: shouldForceStopUnhealthy)         guard self.activeProbeCount == 0 else {            if shouldReset {                self.resetRequestedWhenIdle = true                self.pendingResetCause = Self.strongestResetCause(self.pendingResetCause, currentResetCause)            }            if shouldReset, !success || forceTerminate {                self.hardResetRequestedWhenIdle = true            }            return        }         if shouldReset {            let shouldForceTerminate = !success || forceTerminate || self.hardResetRequestedWhenIdle            let resetCause = Self.strongestResetCause(self.pendingResetCause, currentResetCause)            await self.withLifecycleOperation {                guard self.activeProbeCount == 0 else {                    self.resetRequestedWhenIdle = true                    self.pendingResetCause = Self.strongestResetCause(self.pendingResetCause, resetCause)                    if shouldForceTerminate {                        self.hardResetRequestedWhenIdle = true                    }                    return                }                self.resetRequestedWhenIdle = false                self.hardResetRequestedWhenIdle = false                self.pendingResetCause = nil                await self.stopCurrentSessionLocked(                    reason: resetCause.message,                    clearRecord: true,                    graceful: !shouldForceTerminate)            }        } else {            self.armIdleTimer()        }    }     static func resetCause(        authenticationRequired: Bool,        resetAfterFetch: Bool,        shouldForceStopUnhealthy: Bool) -> ResetCause    {        if authenticationRequired {            .authentication        } else if shouldForceStopUnhealthy {            .unhealthy        } else if resetAfterFetch {            .oneShotCLI        } else {            .deferred        }    }     static func strongestResetCause(_ existing: ResetCause?, _ incoming: ResetCause) -> ResetCause {        guard let existing else { return incoming }        return existing.rawValue >= incoming.rawValue ? existing : incoming    }     /// Ensure a warm ``agy`` is running on the given binary path.    ///    /// - If the process is already alive with the same binary, this returns immediately.    /// - If the process died, the binary changed, or repeated probes failed, it tears down the old one first.    /// - Returns the process identifier for port discovery.    func ensureStarted(binary: String) async throws -> pid_t {        try await self.withLifecycleOperation {            try await self.ensureStartedLocked(binary: binary)        }    }     /// Request teardown. If a probe is in flight, cleanup is deferred until the    /// matching ``finishProbe(success:resetAfterFetch:)`` call.    func reset() async {        self.cancelIdleTimer()        guard self.activeProbeCount == 0 else {            self.resetRequestedWhenIdle = true            return        }        await self.withLifecycleOperation {            guard self.activeProbeCount == 0 else {                self.resetRequestedWhenIdle = true                return            }            self.resetRequestedWhenIdle = false            let forceTerminate = self.hardResetRequestedWhenIdle || self.consecutiveProbeFailures > 0            self.hardResetRequestedWhenIdle = false            await self.stopCurrentSessionLocked(                reason: "manual reset",                clearRecord: true,                graceful: !forceTerminate)        }    }     /// Drain PTY output into one rolling buffer shared by concurrent probes.    func drainOutput() -> Data {        var searchableOutput = self.recentOutput        if let output = self.process?.drainOutput(), !output.isEmpty {            searchableOutput.append(output)            self.recentOutput = Data(searchableOutput.suffix(4096))        }         if !self.authenticationPromptObserved,           AntigravityCLIAuthenticationPrompt.contains(searchableOutput)        {            self.authenticationPromptObserved = true        }        if self.authenticationPromptObserved,           !AntigravityCLIAuthenticationPrompt.contains(searchableOutput)        {            searchableOutput.append(AntigravityCLIAuthenticationPrompt.evidence)        }        return searchableOutput    }     // MARK: Errors     enum SessionError: LocalizedError {        case launchFailed(String)         var errorDescription: String? {            switch self {            case let .launchFailed(msg): "Failed to launch Antigravity CLI session: \(msg)"            }        }    }     // MARK: Private     private func withLifecycleOperation<T>(_ operation: () async throws -> T) async throws -> T {        await self.acquireLifecycleOperation()        defer { self.releaseLifecycleOperation() }        return try await operation()    }     private func withLifecycleOperation(_ operation: () async -> Void) async {        await self.acquireLifecycleOperation()        defer { self.releaseLifecycleOperation() }        await operation()    }     private func acquireLifecycleOperation() async {        guard self.lifecycleOperationInProgress else {            self.lifecycleOperationInProgress = true            return        }        await withCheckedContinuation { continuation in            self.lifecycleWaiters.append(continuation)        }    }     private func releaseLifecycleOperation() {        guard !self.lifecycleWaiters.isEmpty else {            self.lifecycleOperationInProgress = false            return        }        let next = self.lifecycleWaiters.removeFirst()        next.resume()    }     private func waitForExclusiveProbeIfNeeded() async {        while self.activeSessionProbeCount > 0 {            await withCheckedContinuation { continuation in                self.exclusiveProbeWaiters.append(continuation)            }        }    }     private func notifyExclusiveProbeWaitersIfNeeded() {        guard self.activeSessionProbeCount == 0, !self.exclusiveProbeWaiters.isEmpty else { return }        let waiters = self.exclusiveProbeWaiters        self.exclusiveProbeWaiters.removeAll()        for waiter in waiters {            waiter.resume()        }    }     private func ensureStartedLocked(binary: String) async throws -> pid_t {        while true {            if let proc = self.process,               proc.isRunning,               self.binaryPath == binary,               !self.resetRequestedWhenIdle,               !self.hardResetRequestedWhenIdle,

Lines 1–500 of 1270.