mirror of
https://github.com/AgentSeal/codeburn.git
synced 2026-08-21 14:34:32 +00:00
529 lines
22 KiB
Swift
529 lines
22 KiB
Swift
import Darwin
|
|
import Foundation
|
|
|
|
/// A resident `codeburn serve --stdio` child, held so payload fetches skip the
|
|
/// per-spawn cost (node boot + a 100MB+ session-cache parse on large corpora,
|
|
/// seconds per fetch at the CLI level). Requests are JSON lines `{id, args}`;
|
|
/// replies are `{id, ok, output}`. Mirrors the desktop app's client contract:
|
|
///
|
|
/// - Only `status` payload queries route here; anything else spawns as before.
|
|
/// - The first real status request is also the warm-up. It may be written
|
|
/// before the child announces READY; the pipe buffers it until serve reads
|
|
/// stdin, avoiding a second one-shot process that parses the same cache.
|
|
/// - Transport/protocol failures fall back to the spawn path for that call;
|
|
/// resource-policy failures remain terminal. Three child deaths disable
|
|
/// serve for this app run.
|
|
/// - The child's stdin closing (app quit, even SIGKILL) ends the server loop
|
|
/// on the CLI side, so no orphan survives the menubar.
|
|
actor ServeConnection {
|
|
static let shared = ServeConnection()
|
|
|
|
typealias ProcessFactory = ([String], QualityOfService) -> Process
|
|
typealias TimeoutSleep = @Sendable (UInt64) async throws -> Void
|
|
|
|
private struct QueuedRequest {
|
|
let token: Int
|
|
let args: [String]
|
|
let continuation: CheckedContinuation<Data, Error>
|
|
}
|
|
|
|
private struct ActiveRequest {
|
|
let token: Int
|
|
let id: Int
|
|
let args: [String]
|
|
let child: Process
|
|
}
|
|
|
|
private var process: Process?
|
|
private var stdinHandle: FileHandle?
|
|
private var nextId = 1
|
|
private var nextRequestToken = 1
|
|
private var queuedRequests: [QueuedRequest] = []
|
|
private var activeRequest: ActiveRequest?
|
|
private var pending: [Int: CheckedContinuation<Data, Error>] = [:]
|
|
private var requestTimeouts: [Int: Task<Void, Never>] = [:]
|
|
private var timeoutOwners: [Int: Process] = [:]
|
|
private var responseBytes: [Int: Int] = [:]
|
|
private var deaths = 0
|
|
private var buffer = Data()
|
|
private var receivedTerminalResponse = false
|
|
private var outputTasks: [ObjectIdentifier: Task<Void, Never>] = [:]
|
|
private var terminationTasks: [ObjectIdentifier: Task<Void, Never>] = [:]
|
|
private let makeProcess: ProcessFactory
|
|
private let timeoutSleep: TimeoutSleep
|
|
private let terminationGraceSleep: TimeoutSleep
|
|
private let responseLimitBytes: Int
|
|
|
|
private static let maxDeaths = 3
|
|
static let maxResponseBytes = 16 * 1024 * 1024
|
|
private static let stdoutReadChunkBytes = 64 * 1024
|
|
private static let terminationGraceNanoseconds: UInt64 = 1_000_000_000
|
|
private static let coldRequestTimeoutNanoseconds: UInt64 = 10 * 60 * 1_000_000_000
|
|
private static let warmRequestTimeoutNanoseconds: UInt64 = 60 * 1_000_000_000
|
|
|
|
struct ServeUnavailable: Error {}
|
|
enum FailureReason: Sendable, Equatable {
|
|
case generic
|
|
case outputTooLarge
|
|
}
|
|
struct ServeRequestFailed: Error, Sendable {
|
|
let message: String
|
|
let reason: FailureReason
|
|
|
|
init(message: String, reason: FailureReason = .generic) {
|
|
self.message = message
|
|
self.reason = reason
|
|
}
|
|
}
|
|
|
|
init(
|
|
makeProcess: @escaping ProcessFactory = CodeburnCLI.makeProcess,
|
|
timeoutSleep: @escaping TimeoutSleep = { nanoseconds in
|
|
try await Task<Never, Never>.sleep(nanoseconds: nanoseconds)
|
|
},
|
|
terminationGraceSleep: @escaping TimeoutSleep = { nanoseconds in
|
|
try await Task<Never, Never>.sleep(nanoseconds: nanoseconds)
|
|
},
|
|
responseLimitBytes: Int = ServeConnection.maxResponseBytes
|
|
) {
|
|
self.makeProcess = makeProcess
|
|
self.timeoutSleep = timeoutSleep
|
|
self.terminationGraceSleep = terminationGraceSleep
|
|
precondition(responseLimitBytes > 0)
|
|
self.responseLimitBytes = responseLimitBytes
|
|
}
|
|
|
|
static func isEligible(_ subcommand: [String]) -> Bool {
|
|
subcommand.first == "status"
|
|
}
|
|
|
|
/// Kick the child off (idempotent). Called from app startup and again by
|
|
/// the first request in case the startup task has not run yet.
|
|
func ensureStarted() {
|
|
guard process == nil, deaths < Self.maxDeaths else { return }
|
|
// This single resident serves both background and user-visible status
|
|
// requests. Its cold hydration replaces the old interactive one-shot,
|
|
// so keep the child at the same user-initiated QoS as visible fetches.
|
|
let child = makeProcess(["serve", "--stdio"], .userInitiated)
|
|
let stdinPipe = Pipe()
|
|
let stdinWriter = stdinPipe.fileHandleForWriting
|
|
// Suppress SIGPIPE only for this connection's write end. A process-wide
|
|
// SIG_IGN leaks into unrelated libraries and children; F_SETNOSIGPIPE
|
|
// keeps a closed child stdin on the normal throwable EPIPE path.
|
|
guard Darwin.fcntl(stdinWriter.fileDescriptor, F_SETNOSIGPIPE, 1) == 0 else {
|
|
deaths = Self.maxDeaths
|
|
return
|
|
}
|
|
let stdoutPipe = Pipe()
|
|
let stdoutReader = stdoutPipe.fileHandleForReading
|
|
child.standardInput = stdinPipe
|
|
child.standardOutput = stdoutPipe
|
|
child.standardError = FileHandle.nullDevice
|
|
do {
|
|
try child.run()
|
|
} catch {
|
|
deaths = Self.maxDeaths // spawn path can't produce the binary either better than makeProcess did
|
|
return
|
|
}
|
|
process = child
|
|
stdinHandle = stdinWriter
|
|
let generation = ObjectIdentifier(child)
|
|
// One blocking reader owns this generation's stdout. It never reads a
|
|
// second bounded chunk until the actor has consumed the first, giving
|
|
// the 16 MiB protocol limit real backpressure instead of accumulating
|
|
// an unbounded callback/AsyncStream backlog. EOF is observed only after
|
|
// the pipe's final bytes, so child death cannot overtake a split reply.
|
|
outputTasks[generation] = Task.detached { [weak self] in
|
|
var bytes = [UInt8](repeating: 0, count: Self.stdoutReadChunkBytes)
|
|
while !Task.isCancelled {
|
|
let count = Darwin.read(stdoutReader.fileDescriptor, &bytes, bytes.count)
|
|
if count > 0 {
|
|
guard let self else { break }
|
|
await self.consume(Data(bytes[0..<count]), from: child)
|
|
} else if count == -1, errno == EINTR {
|
|
continue
|
|
} else {
|
|
break
|
|
}
|
|
}
|
|
await self?.outputStreamEnded(for: child)
|
|
await self?.outputStreamFinished(for: child)
|
|
}
|
|
}
|
|
|
|
/// Send the first real payload through the resident child. A request does
|
|
/// not need to wait for the READY frame: stdin is safe to write as soon as
|
|
/// Process.run() succeeds, and serve serializes it after initialization.
|
|
func request(args: [String]) async throws -> Data {
|
|
try Task.checkCancellation()
|
|
ensureStarted()
|
|
guard process != nil else { throw ServeUnavailable() }
|
|
let token = nextRequestToken
|
|
nextRequestToken += 1
|
|
let response = try await withTaskCancellationHandler {
|
|
try await withCheckedThrowingContinuation { continuation in
|
|
queuedRequests.append(QueuedRequest(
|
|
token: token,
|
|
args: args,
|
|
continuation: continuation
|
|
))
|
|
startNextRequestIfPossible()
|
|
}
|
|
} onCancel: {
|
|
Task { await self.cancelRequest(token: token) }
|
|
}
|
|
try Task.checkCancellation()
|
|
return response
|
|
}
|
|
|
|
func shutdown() {
|
|
deaths = Self.maxDeaths
|
|
process?.terminate()
|
|
for task in terminationTasks.values { task.cancel() }
|
|
terminationTasks.removeAll()
|
|
cancelAllTimeouts()
|
|
failAllRequests()
|
|
process = nil
|
|
stdinHandle = nil
|
|
buffer = Data()
|
|
receivedTerminalResponse = false
|
|
}
|
|
|
|
// MARK: - internals
|
|
|
|
private func startNextRequestIfPossible() {
|
|
guard activeRequest == nil, !queuedRequests.isEmpty else { return }
|
|
ensureStarted()
|
|
guard let stdinHandle, let child = process else {
|
|
failQueuedRequests(error: ServeUnavailable())
|
|
return
|
|
}
|
|
// A Process can report not-running just before its termination callback
|
|
// reaches the ordered event stream. Keep the request queued for that
|
|
// event instead of writing to a generation which is already exiting.
|
|
guard child.isRunning else { return }
|
|
|
|
let request = queuedRequests.removeFirst()
|
|
let id = nextId
|
|
nextId += 1
|
|
let line: Data
|
|
do {
|
|
line = try JSONSerialization.data(withJSONObject: ["id": id, "args": request.args])
|
|
} catch {
|
|
request.continuation.resume(throwing: error)
|
|
startNextRequestIfPossible()
|
|
return
|
|
}
|
|
|
|
// The previous response can resume its caller just before EOF reaches
|
|
// this actor. Avoid admitting a successor to an already-reaped child;
|
|
// the reader's ordered EOF path will start it on a replacement.
|
|
guard child.isRunning else {
|
|
queuedRequests.insert(request, at: 0)
|
|
outputStreamEnded(for: child)
|
|
return
|
|
}
|
|
|
|
// Select and arm the timeout only when this request becomes the sole
|
|
// protocol request in flight. A queued request must not spend its own
|
|
// budget while its predecessor is still hydrating or draining.
|
|
let timeoutNanoseconds = receivedTerminalResponse
|
|
? Self.warmRequestTimeoutNanoseconds
|
|
: Self.coldRequestTimeoutNanoseconds
|
|
activeRequest = ActiveRequest(
|
|
token: request.token,
|
|
id: id,
|
|
args: request.args,
|
|
child: child
|
|
)
|
|
pending[id] = request.continuation
|
|
responseBytes[id] = 0
|
|
do {
|
|
try stdinHandle.write(contentsOf: line + Data("\n".utf8))
|
|
armTimeout(id: id, child: child, nanoseconds: timeoutNanoseconds)
|
|
} catch {
|
|
// The previous terminal frame can resume its caller just before
|
|
// EOF detaches that generation. Preserve this never-admitted
|
|
// request and retry it on the replacement instead of surfacing a
|
|
// transient EPIPE to the UI.
|
|
pending.removeValue(forKey: id)
|
|
responseBytes.removeValue(forKey: id)
|
|
activeRequest = nil
|
|
queuedRequests.insert(request, at: 0)
|
|
outputStreamEnded(for: child)
|
|
}
|
|
}
|
|
|
|
private func cancelRequest(token: Int) {
|
|
if let index = queuedRequests.firstIndex(where: { $0.token == token }) {
|
|
let request = queuedRequests.remove(at: index)
|
|
request.continuation.resume(throwing: CancellationError())
|
|
return
|
|
}
|
|
guard let activeRequest, activeRequest.token == token,
|
|
let continuation = pending.removeValue(forKey: activeRequest.id) else { return }
|
|
continuation.resume(throwing: CancellationError())
|
|
// Caller cancellation abandons only this response. The serialized serve
|
|
// child may still be doing the expensive first hydration, and killing it
|
|
// here lets tab switches and UI watchdogs restart that work indefinitely.
|
|
// Its independent request timeout remains armed: a command that never
|
|
// returns is still reaped, so it cannot wedge every later serialized call.
|
|
}
|
|
|
|
private func armTimeout(id: Int, child: Process, nanoseconds: UInt64) {
|
|
let sleep = timeoutSleep
|
|
timeoutOwners[id] = child
|
|
requestTimeouts[id] = Task.detached { [weak self] in
|
|
do {
|
|
try await sleep(nanoseconds)
|
|
} catch {
|
|
return
|
|
}
|
|
await self?.requestTimedOut(id: id)
|
|
}
|
|
}
|
|
|
|
private func requestTimedOut(id: Int) {
|
|
guard let child = timeoutOwners.removeValue(forKey: id) else { return }
|
|
requestTimeouts.removeValue(forKey: id)
|
|
responseBytes.removeValue(forKey: id)
|
|
if let continuation = pending.removeValue(forKey: id) {
|
|
continuation.resume(throwing: ServeRequestFailed(message: "serve timeout"))
|
|
}
|
|
// The waiter may already have been abandoned by caller cancellation.
|
|
// Timeout ownership is deliberately independent of that continuation:
|
|
// kill only the exact generation that received the timed-out request.
|
|
guard process === child else {
|
|
if activeRequest?.id == id { activeRequest = nil }
|
|
startNextRequestIfPossible()
|
|
return
|
|
}
|
|
// Retire the timed-out generation synchronously. Its stdout may never
|
|
// reach EOF (for example, a stuck child can ignore SIGTERM or a
|
|
// descendant can retain the pipe), so waiting for the reader would also
|
|
// spend every queued caller's timeout before it can even be admitted.
|
|
process = nil
|
|
stdinHandle = nil
|
|
buffer = Data()
|
|
receivedTerminalResponse = false
|
|
deaths += 1
|
|
if activeRequest?.id == id { activeRequest = nil }
|
|
cancelTimeouts(ownedBy: child)
|
|
terminateTimedOutChild(child)
|
|
// The waiter was removed above and cannot be requeued by stale EOF.
|
|
// A queued read starts on a replacement immediately, subject to the
|
|
// ordinary three-death budget.
|
|
startNextRequestIfPossible()
|
|
}
|
|
|
|
private func cancelTimeout(id: Int) {
|
|
timeoutOwners.removeValue(forKey: id)
|
|
requestTimeouts.removeValue(forKey: id)?.cancel()
|
|
responseBytes.removeValue(forKey: id)
|
|
}
|
|
|
|
private func cancelTimeouts(ownedBy child: Process) {
|
|
let ids = timeoutOwners.compactMap { id, owner in owner === child ? id : nil }
|
|
for id in ids { cancelTimeout(id: id) }
|
|
}
|
|
|
|
private func cancelAllTimeouts() {
|
|
for task in requestTimeouts.values { task.cancel() }
|
|
requestTimeouts.removeAll()
|
|
timeoutOwners.removeAll()
|
|
responseBytes.removeAll()
|
|
}
|
|
|
|
private func outputStreamFinished(for child: Process) {
|
|
outputTasks.removeValue(forKey: ObjectIdentifier(child))
|
|
if !child.isRunning {
|
|
terminationTasks.removeValue(forKey: ObjectIdentifier(child))?.cancel()
|
|
}
|
|
}
|
|
|
|
private func terminateTimedOutChild(_ child: Process) {
|
|
guard child.isRunning else { return }
|
|
child.terminate()
|
|
let generation = ObjectIdentifier(child)
|
|
let sleep = terminationGraceSleep
|
|
terminationTasks[generation] = Task.detached { [weak self] in
|
|
do {
|
|
try await sleep(Self.terminationGraceNanoseconds)
|
|
} catch {
|
|
// Cancellation means the owner stopped waiting: either shutdown
|
|
// (which must not orphan a SIGTERM-ignoring generation) or the
|
|
// child already died and the stream finished. Escalate either
|
|
// way; the isRunning guard makes the dead-child case a no-op.
|
|
await self?.forceKillAfterGrace(child)
|
|
return
|
|
}
|
|
await self?.forceKillAfterGrace(child)
|
|
}
|
|
}
|
|
|
|
private func forceKillAfterGrace(_ child: Process) {
|
|
terminationTasks.removeValue(forKey: ObjectIdentifier(child))
|
|
guard child.isRunning else { return }
|
|
_ = Darwin.kill(child.processIdentifier, SIGKILL)
|
|
}
|
|
|
|
private func outputStreamEnded(for child: Process) {
|
|
guard process === child else { return }
|
|
// EOF/read failure is a transport death even if the process has not
|
|
// reaped yet. Terminate that exact generation so a child which closed
|
|
// stdout cannot survive after the actor starts its replacement.
|
|
if child.isRunning { child.terminate() }
|
|
childDied(child)
|
|
}
|
|
|
|
// Internal so the generation guard can be exercised deterministically by
|
|
// tests without relying on Foundation callback scheduling at process exit.
|
|
func consume(_ data: Data, from child: Process) {
|
|
// A readability callback can already have queued its actor Task when the
|
|
// old process exits. If a replacement starts first, those late bytes must
|
|
// not repopulate the shared line buffer or mark the new child as warm.
|
|
guard process === child else { return }
|
|
var remaining = data[data.startIndex..<data.endIndex]
|
|
while !remaining.isEmpty {
|
|
if let newline = remaining.firstIndex(of: UInt8(ascii: "\n")) {
|
|
let fragment = remaining[remaining.startIndex..<newline]
|
|
guard fragment.count <= responseLimitBytes - buffer.count else {
|
|
outputOverflowed(child)
|
|
return
|
|
}
|
|
buffer.append(contentsOf: fragment)
|
|
let lineData = buffer
|
|
buffer = Data()
|
|
consumeLine(lineData, from: child)
|
|
guard process === child else { return }
|
|
remaining = remaining[remaining.index(after: newline)..<remaining.endIndex]
|
|
} else {
|
|
guard remaining.count <= responseLimitBytes - buffer.count else {
|
|
outputOverflowed(child)
|
|
return
|
|
}
|
|
buffer.append(contentsOf: remaining)
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
private func consumeLine(_ lineData: Data, from child: Process) {
|
|
guard !lineData.isEmpty,
|
|
let object = try? JSONSerialization.jsonObject(with: lineData) as? [String: Any] else { return }
|
|
if object["ready"] as? Bool == true {
|
|
return
|
|
}
|
|
guard let id = object["id"] as? Int,
|
|
responseBytes[id] != nil else { return }
|
|
// Desktop asks serve to stream cold-scan stderr as progress frames.
|
|
// Menubar has no progress UI, but must leave the request pending
|
|
// until the terminal response arrives if such a frame is emitted.
|
|
if let progress = object["progress"] as? String {
|
|
guard accountResponseBytes(Data(progress.utf8).count, id: id, child: child) else { return }
|
|
return
|
|
}
|
|
let succeeded = object["ok"] as? Bool == true
|
|
let payload = succeeded
|
|
? (object["output"] as? String)
|
|
: (object["error"] as? String)
|
|
if let payload {
|
|
guard accountResponseBytes(Data(payload.utf8).count, id: id, child: child) else { return }
|
|
}
|
|
// A refused/failed command can finish before any cache hydration.
|
|
// Only a successful terminal proves the resident is warm. Keep
|
|
// this before the waiter lookup so a successful orphan response
|
|
// still records the child as warm without resuming anything.
|
|
if succeeded { receivedTerminalResponse = true }
|
|
cancelTimeout(id: id)
|
|
let continuation = pending.removeValue(forKey: id)
|
|
if activeRequest?.id == id { activeRequest = nil }
|
|
if let continuation {
|
|
if succeeded, let output = object["output"] as? String {
|
|
continuation.resume(returning: Data(output.utf8))
|
|
} else {
|
|
let message = object["error"] as? String ?? "serve request failed"
|
|
continuation.resume(throwing: ServeRequestFailed(message: message))
|
|
}
|
|
}
|
|
// This also advances after an orphan terminal response whose caller
|
|
// was cancelled: cancellation removes only the waiter, not the
|
|
// active protocol lifecycle.
|
|
startNextRequestIfPossible()
|
|
}
|
|
|
|
private func accountResponseBytes(_ count: Int, id: Int, child: Process) -> Bool {
|
|
guard let current = responseBytes[id],
|
|
count <= responseLimitBytes - current else {
|
|
outputOverflowed(child)
|
|
return false
|
|
}
|
|
responseBytes[id] = current + count
|
|
return true
|
|
}
|
|
|
|
private func outputOverflowed(_ child: Process) {
|
|
guard process === child else { return }
|
|
// Detach this exact generation before terminating it. Its eventual exit
|
|
// and any already-scheduled stdout callbacks are then stale and cannot
|
|
// consume a second death or corrupt a replacement generation.
|
|
process = nil
|
|
stdinHandle = nil
|
|
buffer = Data()
|
|
receivedTerminalResponse = false
|
|
deaths += 1
|
|
cancelTimeouts(ownedBy: child)
|
|
failAllRequests(error: ServeRequestFailed(
|
|
message: "serve output exceeded \(responseLimitBytes) bytes",
|
|
reason: .outputTooLarge
|
|
))
|
|
if child.isRunning { child.terminate() }
|
|
}
|
|
|
|
private func childDied(_ child: Process) {
|
|
guard process === child else { return }
|
|
process = nil
|
|
stdinHandle = nil
|
|
buffer.removeAll()
|
|
receivedTerminalResponse = false
|
|
deaths += 1
|
|
cancelTimeouts(ownedBy: child)
|
|
if let activeRequest, activeRequest.child === child {
|
|
if let continuation = pending.removeValue(forKey: activeRequest.id) {
|
|
// Only read-only status requests enter this connection. If a
|
|
// generation exits after admission but before its terminal
|
|
// reply, retain the waiter and retry on the replacement rather
|
|
// than racing it into a one-shot fallback. A timed-out or
|
|
// cancelled waiter is already absent and is never retried.
|
|
queuedRequests.insert(QueuedRequest(
|
|
token: activeRequest.token,
|
|
args: activeRequest.args,
|
|
continuation: continuation
|
|
), at: 0)
|
|
}
|
|
self.activeRequest = nil
|
|
}
|
|
// Requests which were never written survive an ordinary child crash.
|
|
// They begin on a replacement only after this ordered death event.
|
|
startNextRequestIfPossible()
|
|
}
|
|
|
|
private func failAllRequests(
|
|
error: Error = ServeRequestFailed(message: "serve exited")
|
|
) {
|
|
for (_, continuation) in pending {
|
|
continuation.resume(throwing: error)
|
|
}
|
|
pending.removeAll()
|
|
activeRequest = nil
|
|
failQueuedRequests(error: error)
|
|
}
|
|
|
|
private func failQueuedRequests(error: Error) {
|
|
let requests = queuedRequests
|
|
queuedRequests.removeAll()
|
|
for request in requests {
|
|
request.continuation.resume(throwing: error)
|
|
}
|
|
}
|
|
}
|