Sources/OrgApp/BabelSessions.swift

  1#if os(macOS)
  2import Foundation
  3import OrgCore
  4
  5/// Interpreters kept running for `:session`, one per file folder, interpreter and session name.
  6/// Each run sends the code, then a line that prints a marker; the output is what came before it.
  7public final class BabelSessions: @unchecked Sendable {
  8    public static let shared = BabelSessions()
  9
 10    private let lock = NSLock()
 11    private var sessions: [String: Session] = [:]
 12
 13    final class Session: @unchecked Sendable {
 14        let process = Process()
 15        private let input: FileHandle
 16        private let lock = NSLock()
 17        private var stdout = Data()
 18        private var stderr = Data()
 19        private var waiting: (marker: String, continuation: CheckedContinuation<(String, String, Int32), Error>)?
 20
 21        init(executable: String, arguments: [String], directory: URL) throws {
 22            let inPipe = Pipe(), outPipe = Pipe(), errPipe = Pipe()
 23            process.executableURL = URL(fileURLWithPath: "/usr/bin/env")
 24            process.arguments = [executable] + arguments
 25            process.environment = BabelRunner.environment
 26            process.currentDirectoryURL = directory
 27            process.standardInput = inPipe
 28            process.standardOutput = outPipe
 29            process.standardError = errPipe
 30            input = inPipe.fileHandleForWriting
 31            outPipe.fileHandleForReading.readabilityHandler = { [weak self] handle in self?.received(handle.availableData, error: false) }
 32            errPipe.fileHandleForReading.readabilityHandler = { [weak self] handle in self?.received(handle.availableData, error: true) }
 33            process.terminationHandler = { [weak self] _ in self?.ended() }
 34            try process.run()
 35        }
 36
 37        /// Sends `code` and waits for `marker`, printed with the last status after it.
 38        func run(_ code: String, marker: String) async throws -> (stdout: String, stderr: String, status: Int32) {
 39            let result: (String, String, Int32) = try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<(String, String, Int32), Error>) in
 40                lock.withLock {
 41                    stdout = Data()
 42                    stderr = Data()
 43                    waiting = (marker, continuation)
 44                }
 45                do {
 46                    try input.write(contentsOf: Data(code.utf8))
 47                } catch {
 48                    finish(.failure(error))
 49                }
 50            }
 51            return (result.0, result.1, result.2)
 52        }
 53
 54        private func received(_ data: Data, error: Bool) {
 55            guard !data.isEmpty else { return }
 56            var result: (String, String, Int32)?
 57            lock.withLock {
 58                if error { stderr.append(data) } else { stdout.append(data) }
 59                guard let marker = waiting?.marker, !error else { return }
 60                let text = String(decoding: stdout, as: UTF8.self)
 61                guard let found = text.range(of: marker + " ", options: .backwards),
 62                      let lineEnd = text[found.upperBound...].firstIndex(of: "\n") else { return }
 63                let status = Int32(text[found.upperBound..<lineEnd].trimmingCharacters(in: .whitespaces)) ?? 0
 64                result = (String(text[..<found.lowerBound]), String(decoding: stderr, as: UTF8.self), status)
 65            }
 66            if let result { finish(.success(result)) }
 67        }
 68
 69        private func ended() {
 70            finish(.failure(BabelRunner.Failure(description: "The session ended.")))
 71        }
 72
 73        private func finish(_ result: Result<(String, String, Int32), Error>) {
 74            let continuation = lock.withLock { () -> CheckedContinuation<(String, String, Int32), Error>? in
 75                defer { waiting = nil }
 76                return waiting?.continuation
 77            }
 78            continuation?.resume(with: result)
 79        }
 80
 81        func stop() { if process.isRunning { process.terminate() } }
 82    }
 83
 84    /// Runs a session job: the session starts in `directory` the first time.
 85    func run(_ job: BabelJob, name: String, executable: String, python: Bool, directory: URL, resultFile: URL) async throws -> BabelOutput {
 86        let key = [directory.path, executable, name].joined(separator: "\u{1}")
 87        let session: Session = try lock.withLock {
 88            if let existing = sessions[key], existing.process.isRunning { return existing }
 89            let created = try Session(executable: executable, arguments: python ? ["-u", "-i", "-q"] : [], directory: directory)
 90            sessions[key] = created
 91            return created
 92        }
 93        let marker = "__orgstar_done_\(UUID().uuidString.replacingOccurrences(of: "-", with: ""))__"
 94        let program = job.input.replacingOccurrences(of: Babel.resultFileToken, with: resultFile.path)
 95        let code: String
 96        if python {
 97            let file = resultFile.deletingLastPathComponent().appendingPathComponent("program.py")
 98            try program.write(to: file, atomically: true, encoding: .utf8)
 99            let path = file.path.replacingOccurrences(of: "'", with: "\\'")
100            code = "try:\n    exec(compile(open('\(path)').read(), '<org-babel>', 'exec'), globals())\nfinally:\n    print('\\n\(marker) 0', flush=True)\n\n"
101        } else {
102            code = program + (program.hasSuffix("\n") ? "" : "\n") + "printf '\\n%s %d\\n' \(marker) $?\n"
103        }
104        let (stdout, stderr, status) = try await withTaskCancellationHandler {
105            try await session.run(code, marker: marker)
106        } onCancel: {
107            session.stop()
108        }
109        // The marker starts on a line of its own; the newline before it isn't output.
110        var output = stdout
111        if output.hasSuffix("\n") { output.removeLast() }
112        // Interactive Python prompts go to standard error.
113        let problems = stderr.split(separator: "\n").filter { $0.range(of: "^(>>> |\\.\\.\\. )*$", options: .regularExpression) == nil }
114        return BabelOutput(stdout: output, resultFile: try? String(contentsOf: resultFile, encoding: .utf8),
115                           stderr: problems.joined(separator: "\n"), status: status)
116    }
117
118    /// Stops every session, as when the app quits.
119    public func stopAll() {
120        lock.withLock {
121            sessions.values.forEach { $0.stop() }
122            sessions = [:]
123        }
124    }
125}
126#endif