428 lines
15 KiB
Swift
428 lines
15 KiB
Swift
#if canImport(Darwin)
|
|
import Darwin
|
|
#elseif canImport(Glibc)
|
|
import Glibc
|
|
#elseif canImport(Musl)
|
|
import Musl
|
|
#endif
|
|
import Foundation
|
|
|
|
actor CodexCLISession {
|
|
static let shared = CodexCLISession()
|
|
|
|
enum SessionError: LocalizedError {
|
|
case launchFailed(String)
|
|
case timedOut
|
|
case processExited
|
|
|
|
var errorDescription: String? {
|
|
switch self {
|
|
case let .launchFailed(msg): "Failed to launch Codex CLI session: \(msg)"
|
|
case .timedOut: "Codex CLI session timed out."
|
|
case .processExited: "Codex CLI session exited."
|
|
}
|
|
}
|
|
}
|
|
|
|
private var process: Process?
|
|
private var primaryFD: Int32 = -1
|
|
private var primaryHandle: FileHandle?
|
|
private var secondaryHandle: FileHandle?
|
|
private var processGroup: pid_t?
|
|
private var binaryPath: String?
|
|
private var startedAt: Date?
|
|
private var ptyRows: UInt16 = 0
|
|
private var ptyCols: UInt16 = 0
|
|
private var sessionEnvironment: [String: String]?
|
|
private var sessionArguments: [String] = []
|
|
private var sessionWorkingDirectory: URL?
|
|
|
|
struct CaptureOptions {
|
|
let timeout: TimeInterval
|
|
let rows: UInt16
|
|
let cols: UInt16
|
|
let environment: [String: String]
|
|
let extraArgs: [String]
|
|
let workingDirectory: URL?
|
|
}
|
|
|
|
private 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)
|
|
}
|
|
}
|
|
|
|
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
|
|
}
|
|
|
|
// swiftlint:disable cyclomatic_complexity
|
|
func captureStatus(
|
|
binary: String,
|
|
options: CaptureOptions) async throws -> String
|
|
{
|
|
try self.ensureStarted(binary: binary, options: options)
|
|
if let startedAt {
|
|
let sinceStart = Date().timeIntervalSince(startedAt)
|
|
if sinceStart < 0.4 {
|
|
let delay = UInt64((0.4 - sinceStart) * 1_000_000_000)
|
|
try await Task.sleep(nanoseconds: delay)
|
|
}
|
|
}
|
|
self.drainOutput()
|
|
|
|
let script = "/status"
|
|
let cursorQuery = Data([0x1B, 0x5B, 0x36, 0x6E])
|
|
let statusMarkers = [
|
|
"Credits:",
|
|
"5h limit",
|
|
"5-hour limit",
|
|
"Weekly limit",
|
|
].map { Data($0.utf8) }
|
|
let updateNeedles = ["Update available!", "Run bun install -g @openai/codex", "0.60.1 ->"]
|
|
let updateNeedlesLower = updateNeedles.map { Data($0.lowercased().utf8) }
|
|
let statusNeedleLengths = statusMarkers.map(\.count)
|
|
let updateNeedleLengths = updateNeedlesLower.map(\.count)
|
|
let statusMaxNeedle = ([cursorQuery.count] + statusNeedleLengths).max() ?? cursorQuery.count
|
|
let updateMaxNeedle = updateNeedleLengths.max() ?? 0
|
|
var statusScanBuffer = RollingBuffer(maxNeedle: statusMaxNeedle)
|
|
var updateScanBuffer = RollingBuffer(maxNeedle: updateMaxNeedle)
|
|
|
|
var buffer = Data()
|
|
let deadline = Date().addingTimeInterval(options.timeout)
|
|
var nextCursorCheckAt = Date(timeIntervalSince1970: 0)
|
|
|
|
var skippedCodexUpdate = false
|
|
var sentScript = false
|
|
var updateSkipAttempts = 0
|
|
var lastEnter = Date(timeIntervalSince1970: 0)
|
|
var scriptSentAt: Date?
|
|
var resendStatusRetries = 0
|
|
var enterRetries = 0
|
|
var sawCodexStatus = false
|
|
var sawCodexUpdatePrompt = false
|
|
|
|
while Date() < deadline {
|
|
let newData = self.readChunk()
|
|
if !newData.isEmpty {
|
|
buffer.append(newData)
|
|
}
|
|
let scanData = statusScanBuffer.append(newData)
|
|
if Date() >= nextCursorCheckAt,
|
|
!scanData.isEmpty,
|
|
scanData.range(of: cursorQuery) != nil
|
|
{
|
|
try? self.send("\u{1b}[1;1R")
|
|
nextCursorCheckAt = Date().addingTimeInterval(1.0)
|
|
}
|
|
if !scanData.isEmpty, !sawCodexStatus {
|
|
if statusMarkers.contains(where: { scanData.range(of: $0) != nil }) {
|
|
sawCodexStatus = true
|
|
}
|
|
}
|
|
|
|
if !skippedCodexUpdate, !sawCodexUpdatePrompt, !newData.isEmpty {
|
|
let lowerData = Self.lowercasedASCII(newData)
|
|
let lowerScan = updateScanBuffer.append(lowerData)
|
|
if updateNeedlesLower.contains(where: { lowerScan.range(of: $0) != nil }) {
|
|
sawCodexUpdatePrompt = true
|
|
}
|
|
}
|
|
|
|
if !skippedCodexUpdate, sawCodexUpdatePrompt {
|
|
try? self.send("\u{1b}[B")
|
|
try await Task.sleep(nanoseconds: 120_000_000)
|
|
try? self.send("\r")
|
|
try await Task.sleep(nanoseconds: 150_000_000)
|
|
try? self.send("\r")
|
|
try? self.send(script)
|
|
try? self.send("\r")
|
|
updateSkipAttempts += 1
|
|
if updateSkipAttempts >= 1 {
|
|
skippedCodexUpdate = true
|
|
sentScript = false
|
|
scriptSentAt = nil
|
|
buffer.removeAll()
|
|
statusScanBuffer.reset()
|
|
updateScanBuffer.reset()
|
|
sawCodexStatus = false
|
|
}
|
|
try await Task.sleep(nanoseconds: 300_000_000)
|
|
}
|
|
|
|
if !sentScript, !sawCodexUpdatePrompt || skippedCodexUpdate {
|
|
try? self.send(script)
|
|
try? self.send("\r")
|
|
sentScript = true
|
|
scriptSentAt = Date()
|
|
lastEnter = Date()
|
|
try await Task.sleep(nanoseconds: 200_000_000)
|
|
continue
|
|
}
|
|
if sentScript, !sawCodexStatus {
|
|
if Date().timeIntervalSince(lastEnter) >= 1.2, enterRetries < 6 {
|
|
try? self.send("\r")
|
|
enterRetries += 1
|
|
lastEnter = Date()
|
|
try await Task.sleep(nanoseconds: 120_000_000)
|
|
continue
|
|
}
|
|
if let sentAt = scriptSentAt,
|
|
Date().timeIntervalSince(sentAt) >= 3.0,
|
|
resendStatusRetries < 2
|
|
{
|
|
try? self.send(script)
|
|
try? self.send("\r")
|
|
resendStatusRetries += 1
|
|
buffer.removeAll()
|
|
statusScanBuffer.reset()
|
|
updateScanBuffer.reset()
|
|
sawCodexStatus = false
|
|
scriptSentAt = Date()
|
|
lastEnter = Date()
|
|
try await Task.sleep(nanoseconds: 220_000_000)
|
|
continue
|
|
}
|
|
}
|
|
if sawCodexStatus { break }
|
|
if let proc = self.process, !proc.isRunning {
|
|
throw SessionError.processExited
|
|
}
|
|
try await Task.sleep(nanoseconds: 120_000_000)
|
|
}
|
|
|
|
if sawCodexStatus {
|
|
let settleDeadline = Date().addingTimeInterval(2.0)
|
|
while Date() < settleDeadline {
|
|
let newData = self.readChunk()
|
|
if !newData.isEmpty {
|
|
buffer.append(newData)
|
|
}
|
|
let scanData = statusScanBuffer.append(newData)
|
|
if Date() >= nextCursorCheckAt,
|
|
!scanData.isEmpty,
|
|
scanData.range(of: cursorQuery) != nil
|
|
{
|
|
try? self.send("\u{1b}[1;1R")
|
|
nextCursorCheckAt = Date().addingTimeInterval(1.0)
|
|
}
|
|
try await Task.sleep(nanoseconds: 100_000_000)
|
|
}
|
|
}
|
|
|
|
guard !buffer.isEmpty, let text = String(data: buffer, encoding: .utf8) else {
|
|
throw SessionError.timedOut
|
|
}
|
|
return text
|
|
}
|
|
|
|
// swiftlint:enable cyclomatic_complexity
|
|
|
|
func reset() {
|
|
self.cleanup()
|
|
}
|
|
|
|
private func ensureStarted(
|
|
binary: String,
|
|
options: CaptureOptions) throws
|
|
{
|
|
if let proc = self.process,
|
|
proc.isRunning,
|
|
self.binaryPath == binary,
|
|
self.ptyRows == options.rows,
|
|
self.ptyCols == options.cols,
|
|
self.sessionEnvironment == options.environment,
|
|
self.sessionArguments == options.extraArgs,
|
|
self.sessionWorkingDirectory == options.workingDirectory
|
|
{
|
|
return
|
|
}
|
|
self.cleanup()
|
|
|
|
var primaryFD: Int32 = -1
|
|
var secondaryFD: Int32 = -1
|
|
var win = winsize(ws_row: options.rows, ws_col: options.cols, ws_xpixel: 0, ws_ypixel: 0)
|
|
guard openpty(&primaryFD, &secondaryFD, nil, nil, &win) == 0 else {
|
|
throw SessionError.launchFailed("openpty failed")
|
|
}
|
|
_ = fcntl(primaryFD, F_SETFL, O_NONBLOCK)
|
|
|
|
let primaryHandle = FileHandle(fileDescriptor: primaryFD, closeOnDealloc: true)
|
|
let secondaryHandle = FileHandle(fileDescriptor: secondaryFD, closeOnDealloc: true)
|
|
|
|
let proc = Process()
|
|
let resolvedURL = URL(fileURLWithPath: binary)
|
|
proc.executableURL = resolvedURL
|
|
proc.arguments = options.extraArgs
|
|
proc.standardInput = secondaryHandle
|
|
proc.standardOutput = secondaryHandle
|
|
proc.standardError = secondaryHandle
|
|
proc.currentDirectoryURL = options.workingDirectory
|
|
|
|
let env = TTYCommandRunner.enrichedEnvironment(
|
|
baseEnv: options.environment,
|
|
home: options.environment["HOME"] ?? NSHomeDirectory())
|
|
proc.environment = env
|
|
|
|
guard TTYCommandRunner.beginActiveProcessLaunchForAppShutdown() else {
|
|
try? primaryHandle.close()
|
|
try? secondaryHandle.close()
|
|
throw SessionError.launchFailed("App shutdown in progress")
|
|
}
|
|
defer { TTYCommandRunner.endActiveProcessLaunchForAppShutdown() }
|
|
|
|
do {
|
|
try proc.run()
|
|
} catch {
|
|
try? primaryHandle.close()
|
|
try? secondaryHandle.close()
|
|
throw SessionError.launchFailed(error.localizedDescription)
|
|
}
|
|
|
|
let pid = proc.processIdentifier
|
|
guard TTYCommandRunner.registerActiveProcessForAppShutdown(
|
|
pid: pid,
|
|
binary: resolvedURL.lastPathComponent)
|
|
else {
|
|
proc.terminate()
|
|
kill(pid, SIGKILL)
|
|
try? primaryHandle.close()
|
|
try? secondaryHandle.close()
|
|
throw SessionError.launchFailed("App shutdown in progress")
|
|
}
|
|
|
|
var processGroup: pid_t?
|
|
if setpgid(pid, pid) == 0 {
|
|
processGroup = pid
|
|
TTYCommandRunner.updateActiveProcessGroupForAppShutdown(pid: pid, processGroup: processGroup)
|
|
}
|
|
|
|
self.process = proc
|
|
self.primaryFD = primaryFD
|
|
self.primaryHandle = primaryHandle
|
|
self.secondaryHandle = secondaryHandle
|
|
self.processGroup = processGroup
|
|
self.binaryPath = binary
|
|
self.startedAt = Date()
|
|
self.ptyRows = options.rows
|
|
self.ptyCols = options.cols
|
|
self.sessionEnvironment = options.environment
|
|
self.sessionArguments = options.extraArgs
|
|
self.sessionWorkingDirectory = options.workingDirectory
|
|
}
|
|
|
|
private func cleanup() {
|
|
if let proc = self.process, proc.isRunning, let handle = self.primaryHandle {
|
|
try? handle.write(contentsOf: Data("/exit\n".utf8))
|
|
}
|
|
try? self.primaryHandle?.close()
|
|
try? self.secondaryHandle?.close()
|
|
|
|
let descendants = self.process.map { TTYProcessTreeTerminator.descendantPIDs(of: $0.processIdentifier) } ?? []
|
|
if let proc = self.process, proc.isRunning {
|
|
proc.terminate()
|
|
}
|
|
if let proc = self.process {
|
|
TTYProcessTreeTerminator.terminateProcessTree(
|
|
rootPID: proc.processIdentifier,
|
|
processGroup: self.processGroup,
|
|
signal: SIGTERM,
|
|
knownDescendants: descendants)
|
|
}
|
|
let waitDeadline = Date().addingTimeInterval(1.0)
|
|
if let proc = self.process {
|
|
while proc.isRunning, Date() < waitDeadline {
|
|
usleep(100_000)
|
|
}
|
|
if proc.isRunning {
|
|
TTYProcessTreeTerminator.terminateProcessTree(
|
|
rootPID: proc.processIdentifier,
|
|
processGroup: self.processGroup,
|
|
signal: SIGKILL,
|
|
knownDescendants: descendants)
|
|
} else {
|
|
for pid in descendants where pid > 0 {
|
|
kill(pid, SIGKILL)
|
|
}
|
|
}
|
|
TTYCommandRunner.unregisterActiveProcessForAppShutdown(pid: proc.processIdentifier)
|
|
}
|
|
|
|
self.process = nil
|
|
self.primaryHandle = nil
|
|
self.secondaryHandle = nil
|
|
self.primaryFD = -1
|
|
self.processGroup = nil
|
|
self.binaryPath = nil
|
|
self.startedAt = nil
|
|
self.ptyRows = 0
|
|
self.ptyCols = 0
|
|
self.sessionEnvironment = nil
|
|
self.sessionArguments = []
|
|
self.sessionWorkingDirectory = nil
|
|
}
|
|
|
|
private func readChunk() -> Data {
|
|
guard self.primaryFD >= 0 else { return Data() }
|
|
var appended = Data()
|
|
while true {
|
|
var tmp = [UInt8](repeating: 0, count: 8192)
|
|
let n = read(self.primaryFD, &tmp, tmp.count)
|
|
if n > 0 {
|
|
appended.append(contentsOf: tmp.prefix(n))
|
|
continue
|
|
}
|
|
break
|
|
}
|
|
return appended
|
|
}
|
|
|
|
private func drainOutput() {
|
|
_ = self.readChunk()
|
|
}
|
|
|
|
private func send(_ text: String) throws {
|
|
guard let data = text.data(using: .utf8) else { return }
|
|
guard let handle = self.primaryHandle else { throw SessionError.processExited }
|
|
try handle.write(contentsOf: data)
|
|
}
|
|
}
|