Files
nucleic/Sources/NucleicCore/AFMRequestQueue.swift
T

494 lines
26 KiB
Swift

import Foundation
/// The session a piece of soft-AI work serves, so the queue inspector — and the sidebar's
/// omnisearch — can tie a queued request back to its chat. Threaded as a task-local
/// (``AFMWorkContext``) rather than through every provider method's signature.
public struct AFMSessionRef: Sendable, Equatable {
public let id: SessionID
public let title: String
public init(id: SessionID, title: String) {
self.id = id
self.title = title
}
}
/// Carries the originating session into the shared ``AFMRequestQueue`` without changing any
/// provider signature. `AppStore` binds it around a session-scoped generation
/// (``AFMWorkContext/session/withValue(_:_:)``); the queue reads it when it enqueues a request
/// and stamps the session onto the entry, so a request the search bar filters can be matched to
/// its chat. Task-locals propagate through the `await` into the queue's actor call, so the value
/// bound at the call site is visible where the entry is created.
public enum AFMWorkContext {
@TaskLocal public static var session: AFMSessionRef?
}
/// Serializes on-device Apple Foundation Model (AFM) requests through a small
/// concurrency limit, so a burst of soft-AI work can't saturate the Neural Engine and
/// stall the rest of the app. The on-device model is a single shared device resource —
/// firing several generations at once only adds contention and latency, never speed —
/// so every provider routes through the one ``shared`` queue.
///
/// Waiting requests are released in **priority** order: a generation tied to something
/// the user is watching right now (naming the chat they just started) always jumps
/// ahead of queued background work (triaging the backlog). Ties break FIFO, so same-
/// priority work still runs in arrival order. A higher priority only preempts *waiting*
/// work — it never interrupts a request already in flight, since AFM calls aren't
/// cancellable mid-generation.
public actor AFMRequestQueue {
/// Which user experience a request serves, highest priority first. The provider maps
/// each soft-AI feature onto one of these so the queue can order a backlog of work.
public enum Priority: Int, Comparable, Sendable {
/// Background backlog work the user isn't actively watching: backlog triage,
/// to-do gisting, and the non-Bash collapsed tool-call summaries (Grep/Glob folds).
case background = 0
/// Writing a completed session's glanceable summary card as it wraps up.
case completion = 1
/// Collapsed Bash tool-call summaries (the map/reduce gist of a run's commands):
/// they fold into a line the user reads in the open transcript, so they outrank the
/// session summary card — but still yield to turn classification and a live chat name.
case bashSummary = 2
/// Classifying whether a finished turn is done or awaits the user. It gates the
/// sidebar's Done/Awaiting state and, for an autoship session, whether the finished
/// turn ships — so it must not sit behind a burst of Bash-summary folds while a
/// completed turn waits to be shipped. Outranks everything but a live chat name.
case turnClassification = 3
/// Directly tied to a fresh user action and visible immediately: naming a chat
/// from its first message.
case interactive = 4
public static func < (lhs: Priority, rhs: Priority) -> Bool {
lhs.rawValue < rhs.rawValue
}
/// A short human label for the queue inspector ("Running", "Completing"…).
public var label: String {
switch self {
case .background: "Background"
case .completion: "Completion"
case .bashSummary: "Nash"
case .turnClassification: "Turn status"
case .interactive: "Interactive"
}
}
}
/// The on-device queue every provider shares for local model work. The Neural Engine is a
/// single device resource, so concurrency is capped at one — a burst of soft-AI work lines
/// up here instead of contending for the chip.
public static let shared = AFMRequestQueue(maxConcurrent: 1, isCloud: false)
/// The Private Cloud Compute queue. PCC runs on Apple's servers, not this device, so it
/// isn't resource-constrained the way the Neural Engine is — many requests can be in flight
/// at once. Work still routes through a queue so the inspector can show what the cloud is
/// generating and what it just finished; with an effectively unbounded concurrency cap it
/// fires requests concurrently and never makes anything wait — except while frozen for
/// inspection, where new requests line up as waiters and resume drains them all at once.
public static let sharedCloud = AFMRequestQueue(maxConcurrent: .max, isCloud: true)
private let maxConcurrent: Int
/// Whether this queue's work runs in Private Cloud Compute (Apple's servers) rather than on
/// this device's Neural Engine. Stamped into every ``Request`` so the merged inspector can
/// flag PCC rows with a cloud glyph; it doesn't affect ordering or concurrency.
private let isCloud: Bool
/// Monotonic ticket, stamped on each request to break same-priority ties FIFO and to
/// identify a request when it's released or cancelled.
private var nextSeq: UInt64 = 0
/// One pending or in-flight request: its priority, a human-readable description of
/// the work (for the queue inspector), and its FIFO ticket.
private struct Entry {
let seq: UInt64
let priority: Priority
let label: String
/// A richer, human-readable view of the actual work — typically a preview of the
/// prompt the model is processing — surfaced when a queue row is expanded.
let detail: String
/// Tags a class of redundant background work where only the newest still-waiting
/// instance is worth running: enqueuing one supersedes any earlier waiter sharing
/// the key (re-triaging the whole backlog covers what a queued triage of an older
/// snapshot would have). nil for one-off work that never coalesces.
let coalesceKey: String?
/// The session this work serves, captured from ``AFMWorkContext`` at enqueue, so the
/// inspector and the sidebar's omnisearch can tie the request to its chat. nil for
/// work not scoped to a single session (project-level triage, mesh jobs).
let session: AFMSessionRef?
/// When the request entered the queue.
let enqueuedAt: Date
/// When it was handed a slot and began generating; nil while still waiting. Drives
/// the "running for…" / "waiting…" timers in the inspector.
var startedAt: Date?
/// When the request finished and gave its slot back; nil while running or waiting.
/// Set only on entries moved into ``historyEntries``.
var finishedAt: Date?
/// The text the request produced, captured on release for the inspector's expanded
/// detail — what the model actually said (the Bash summary line). nil while running or
/// waiting, and for work that returns no text (structured triage).
var result: String?
}
/// Requests currently holding a slot — executing, or just handed one and about to
/// start. A queued (still-waiting) request is NOT here; it's in `waiters`.
private var runningEntries: [Entry] = []
private struct Waiter {
let entry: Entry
/// Resumes the suspended `acquire`. `true` hands over a slot (the waiter may
/// run); `false` means it was cancelled before getting one (no slot to release).
let resume: @Sendable (Bool) -> Void
}
private var waiters: [Waiter] = []
/// A bounded, most-recent-last log of requests that have finished, so the inspector can
/// show what the model just worked on (the queue moves fast and rows vanish the instant
/// they complete). Capped so a long-running app can't grow this without bound.
private var historyEntries: [Entry] = []
private let maxHistory = 50
/// When `true`, the queue stops handing slots to waiting requests: nothing new starts,
/// so the waiting list holds still while the user inspects it in the sidebar. A request
/// already in flight isn't interrupted (AFM can't be cancelled mid-generation) — it
/// finishes and moves to history, then its slot sits idle until the freeze lifts. New
/// requests still enqueue while frozen; they just wait their turn like everything else.
private var frozen = false
/// `maxConcurrent` defaults to 1: the on-device model effectively serializes on the
/// Neural Engine anyway, so one in-flight request at a time is the honest cap.
public init(maxConcurrent: Int = 1, isCloud: Bool = false) {
self.maxConcurrent = max(1, maxConcurrent)
self.isCloud = isCloud
}
/// Runs `operation` once a slot is free, with waiting requests ordered by `priority`.
/// The operation itself runs outside the queue's isolation, so a long generation
/// never blocks the bookkeeping here — only slot accounting is serialized. `label`
/// describes the work for the queue inspector (e.g. "Chat name", "Backlog triage").
///
/// `resultText`, when given, derives the request's output from its result once it returns
/// — for text work (a Bash summary), the model's own reply — and stores it on the request's
/// history entry, so the inspector's expanded detail can show what the model produced. A
/// nil/blank return leaves the entry with no recorded output.
public func run<T: Sendable>(
priority: Priority, label: String = "", detail: String = "",
resultText: (@Sendable (T) -> String?)? = nil,
_ operation: @Sendable () async -> T
) async -> T {
let seq = await acquire(priority: priority, label: label, detail: detail)
let result = await operation()
// Only release a slot we were actually handed — a cancelled waiter never held one.
if let seq { release(seq, result: resultText?(result)) }
return result
}
/// Like ``run``, but coalesces redundant background work tagged with `coalesceKey` — a
/// class where only the newest still-waiting instance is worth running because it
/// supersedes the rest (re-triaging the whole current backlog covers everything an
/// earlier, still-queued triage of an older snapshot would have). Enqueuing here cancels
/// every *waiting* request sharing the key — each returns `nil` without running its
/// operation — so a burst of identical background work collapses to a single pass. A
/// request already in flight is never touched (AFM isn't cancellable mid-generation, and
/// it's about to finish anyway). Returns the operation's result, or `nil` if this request
/// was itself superseded by a newer one before it was handed a slot.
public func runCoalescing<T: Sendable>(
priority: Priority, coalesceKey: String, label: String = "", detail: String = "",
resultText: (@Sendable (T) -> String?)? = nil,
_ operation: @Sendable () async -> T
) async -> T? {
guard let seq = await acquire(priority: priority, label: label, detail: detail,
coalesceKey: coalesceKey) else { return nil }
let result = await operation()
release(seq, result: resultText?(result))
return result
}
/// Suspends until a slot is free (immediately if one is). Returns the request's ticket
/// when a slot was granted, `nil` if the task was cancelled while waiting (so the
/// caller must not release). Cancellation-safe: a waiting request cancelled mid-queue
/// is pulled out and resumed rather than left to leak its continuation.
private func acquire(priority: Priority, label: String, detail: String,
coalesceKey: String? = nil) async -> UInt64? {
// The newest keyed request wins: drop any earlier still-waiting work of the same
// class, since this one supersedes it. (A request already holding a slot keeps it —
// AFM can't be cancelled mid-generation.) A free slot means there are no waiters to
// supersede, so this only ever fires on the waiting path.
if let coalesceKey { supersedeWaiters(coalesceKey: coalesceKey) }
let seq = nextSeq
nextSeq += 1
let now = Date()
var entry = Entry(seq: seq, priority: priority, label: label, detail: detail,
coalesceKey: coalesceKey, session: AFMWorkContext.session,
enqueuedAt: now, startedAt: nil)
// A frozen queue grants no slots, even an idle one — the request queues so the
// user sees it line up without anything starting to run.
if !frozen && runningEntries.count < maxConcurrent {
entry.startedAt = now
runningEntries.append(entry)
return seq
}
return await withTaskCancellationHandler {
await withCheckedContinuation { (continuation: CheckedContinuation<UInt64?, Never>) in
// Already cancelled before we suspended — don't enqueue at all. (Actor
// isolation guarantees this body runs to completion before the onCancel
// task below can touch `waiters`, so there's no append-after-cancel gap.)
if Task.isCancelled {
continuation.resume(returning: nil)
} else {
waiters.append(Waiter(entry: entry,
resume: { continuation.resume(returning: $0 ? seq : nil) }))
}
}
} onCancel: {
Task { await self.cancelWaiter(seq) }
}
}
/// Pulls a still-waiting request out of the queue on cancellation and resumes it with
/// "no slot". A no-op if it has already been granted a slot (it'll run and release
/// normally; the operation can observe cancellation itself).
private func cancelWaiter(_ seq: UInt64) {
guard let index = waiters.firstIndex(where: { $0.entry.seq == seq }) else { return }
let waiter = waiters.remove(at: index)
waiter.resume(false)
}
/// Cancels every still-waiting request tagged with `coalesceKey`, resuming each with
/// "no slot" so its ``runCoalescing`` returns nil and skips its operation. Called when a
/// newer request of the same class arrives and renders the earlier waiters redundant.
/// Only waiters are affected — a running request is mid-generation and can't be undone.
private func supersedeWaiters(coalesceKey: String) {
var kept: [Waiter] = []
var superseded: [Waiter] = []
for waiter in waiters {
if waiter.entry.coalesceKey == coalesceKey { superseded.append(waiter) }
else { kept.append(waiter) }
}
guard !superseded.isEmpty else { return }
waiters = kept
for waiter in superseded { waiter.resume(false) }
}
/// Frees the slot held by `seq`, then hands it to the highest-priority waiter (or lets
/// the slot go idle if none wait). Removing then re-adding keeps the running count at
/// `maxConcurrent` when a waiter is promoted, and drops it by one when none is.
///
/// `result`, when non-blank, is recorded on the finished entry before it lands in history
/// — so a text request (a Bash summary) carries the model's own output into the inspector's
/// expanded detail. Blank/whitespace records nothing.
private func release(_ seq: UInt64, result: String? = nil) {
if let index = runningEntries.firstIndex(where: { $0.seq == seq }) {
var finished = runningEntries.remove(at: index)
finished.finishedAt = Date()
if let result,
!result.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty {
finished.result = result
}
historyEntries.append(finished)
if historyEntries.count > maxHistory {
historyEntries.removeFirst(historyEntries.count - maxHistory)
}
}
// While frozen, leave the freed slot idle so the waiting queue holds still for
// inspection; lifting the freeze fills it again (see ``setFrozen``).
if !frozen { promoteNextWaiter() }
}
/// Promotes the highest-priority waiter into a free slot, stamping the moment it starts
/// running and resuming it. No-op when every slot is taken or no one is waiting; returns
/// whether a waiter was promoted so ``setFrozen`` can fill all idle slots on resume.
@discardableResult
private func promoteNextWaiter() -> Bool {
guard runningEntries.count < maxConcurrent, let index = nextWaiterIndex() else { return false }
let waiter = waiters.remove(at: index)
var promoted = waiter.entry // promote the waiter into the freed slot…
promoted.startedAt = Date() // …and stamp the moment it starts running
runningEntries.append(promoted)
waiter.resume(true)
return true
}
/// Freezes or resumes slot hand-off. While frozen, no waiting request is started, so the
/// queue the user is inspecting stops moving. Lifting the freeze immediately fills every
/// idle slot from the waiting queue, so processing picks up right where it left off.
public func setFrozen(_ value: Bool) {
guard frozen != value else { return }
frozen = value
if !frozen { while promoteNextWaiter() {} }
}
/// Whether slot hand-off is currently frozen.
public var isFrozen: Bool { frozen }
/// The index of the next waiter to run: highest priority, ties broken by lowest seq
/// (earliest arrival).
private func nextWaiterIndex() -> Int? {
guard !waiters.isEmpty else { return nil }
var best = waiters.startIndex
for index in waiters.indices where index != best {
let candidate = waiters[index].entry
let leader = waiters[best].entry
if candidate.priority > leader.priority
|| (candidate.priority == leader.priority && candidate.seq < leader.seq) {
best = index
}
}
return best
}
// MARK: - Inspection
/// A point-in-time view of the queue for the inspector UI: what's running now and
/// what's waiting (already ordered by release order — highest priority first, ties
/// FIFO). One queue is on-device, one is Private Cloud Compute; the inspector merges
/// both and marks each request by the backend it ran on (see ``Request/isCloud``).
public struct Snapshot: Sendable {
public var running: [Request]
public var waiting: [Request]
/// Recently finished requests, most-recent first. Surfaced as a grayed-out history
/// so the user can see what the model just did; doesn't count toward ``isEmpty``,
/// which tracks only live (running or waiting) work.
public var history: [Request]
/// Whether slot hand-off is frozen right now (see ``setFrozen``), so the inspector
/// can reflect the paused state and offer to resume.
public var frozen: Bool
public var isEmpty: Bool { running.isEmpty && waiting.isEmpty }
public init(running: [Request], waiting: [Request], history: [Request] = [],
frozen: Bool = false) {
self.running = running
self.waiting = waiting
self.history = history
self.frozen = frozen
}
}
/// One request in a ``Snapshot``: its priority, the work's short label, a richer
/// detail string (the prompt preview, for the expanded inspector), and the timestamps
/// that drive its live "running for…" / "waiting…" timers.
public struct Request: Sendable, Identifiable {
public let id: UInt64
public let priority: Priority
public let label: String
public let detail: String
public let enqueuedAt: Date
public let startedAt: Date?
/// When the request finished; non-nil only for history rows.
public let finishedAt: Date?
/// The text the request produced — the model's reply — shown in the expanded detail
/// when a finished row is clicked. nil while running/waiting and for non-text work.
public let result: String?
/// Whether this request ran in Private Cloud Compute rather than on-device — drives
/// the cloud glyph the merged inspector puts on PCC rows.
public let isCloud: Bool
/// The session this work serves (see ``Entry/session``), so the sidebar's omnisearch
/// can filter the queue to requests belonging to matching chats. nil for work not
/// scoped to one session.
public let session: AFMSessionRef?
public init(id: UInt64, priority: Priority, label: String, detail: String,
enqueuedAt: Date, startedAt: Date?, finishedAt: Date? = nil,
result: String? = nil, isCloud: Bool = false,
session: AFMSessionRef? = nil) {
self.id = id
self.priority = priority
self.label = label
self.detail = detail
self.enqueuedAt = enqueuedAt
self.startedAt = startedAt
self.finishedAt = finishedAt
self.result = result
self.isCloud = isCloud
self.session = session
}
/// A key unique across both queues. The per-queue `id` alone collides between the
/// on-device and cloud queues (they share no counter), so the merged inspector keys
/// row identity and expansion on this instead.
public var traceID: String { "\(isCloud ? "pcc" : "afm"):\(id)" }
/// Seconds since the request entered the queue.
public func queuedFor(now: Date) -> TimeInterval { max(0, now.timeIntervalSince(enqueuedAt)) }
/// Seconds since it started generating, or nil if it hasn't started yet.
public func runningFor(now: Date) -> TimeInterval? {
startedAt.map { max(0, now.timeIntervalSince($0)) }
}
/// How long the request ran end to end, or nil if it never started or hasn't
/// finished. Drives the "ran for…" line on history rows.
public var ranFor: TimeInterval? {
guard let startedAt, let finishedAt else { return nil }
return max(0, finishedAt.timeIntervalSince(startedAt))
}
}
/// A snapshot of the queue right now — running requests in start order, waiting ones
/// in the order they'll be released, and recently finished ones most-recent first.
public func snapshot() -> Snapshot {
func request(_ e: Entry) -> Request {
Request(id: e.seq, priority: e.priority, label: e.label, detail: e.detail,
enqueuedAt: e.enqueuedAt, startedAt: e.startedAt, finishedAt: e.finishedAt,
result: e.result, isCloud: isCloud, session: e.session)
}
let running = runningEntries.map(request)
let ordered = waiters.map(\.entry).sorted {
$0.priority != $1.priority ? $0.priority > $1.priority : $0.seq < $1.seq
}
let waiting = ordered.map(request)
let history = historyEntries.reversed().map(request) // most-recent first
return Snapshot(running: running, waiting: waiting, history: history, frozen: frozen)
}
}
/// Runs `operation`, returning its result, or `nil` if it hasn't produced one within
/// `seconds`. On timeout the operation is **abandoned** — left running detached rather
/// than awaited — so a call that never returns can't hold its caller past the deadline.
///
/// This exists because an on-device Foundation Models generation (``LanguageModelSession``
/// `respond`) is not cancellable mid-flight and, on rare occasions, never returns. Since
/// every soft-AI request serializes through the one-slot ``AFMRequestQueue/shared``, a
/// single wedged generation would otherwise hold that slot forever and silently stall every
/// later request behind it — including the turn classification autoship waits on before it
/// can ship, so autoship would quietly stop firing for the rest of the run with no error.
/// Bounding each generation lets the stuck call be abandoned, the slot released, and the
/// caller fall back to its deterministic heuristic: model providers already treat `nil` as
/// "no result", so a timeout degrades to the heuristic instead of wedging the queue.
///
/// The abandoned operation keeps running to completion (it can't be cancelled); only its
/// result is discarded. A `seconds` of zero or less is treated as "no timeout" so a caller
/// can opt out.
public func withTimeout<T: Sendable>(
_ seconds: TimeInterval, _ operation: @escaping @Sendable () async -> T?
) async -> T? {
guard seconds > 0 else { return await operation() }
// Serializes the single continuation resume: whichever child finishes first claims it;
// the loser's later attempt is dropped and its value discarded.
let gate = TimeoutResumeGate()
return await withCheckedContinuation { (continuation: CheckedContinuation<T?, Never>) in
let sleeper = Task {
try? await Task.sleep(nanoseconds: UInt64(seconds * 1_000_000_000))
if gate.claim() { continuation.resume(returning: nil) }
}
Task {
let value = await operation()
if gate.claim() {
sleeper.cancel() // operation won — stop the timer promptly
continuation.resume(returning: value)
}
}
}
}
/// One-shot guard for ``withTimeout``'s continuation: the first of the two racing child
/// tasks to `claim()` wins and resumes the continuation; the loser's later `claim()` returns
/// false and is dropped. A tiny lock-guarded flag rather than an actor, so the racing tasks
/// can check it synchronously without another suspension. (A type can't be nested in a
/// generic function, so it lives at file scope.)
private final class TimeoutResumeGate: @unchecked Sendable {
private let lock = NSLock()
private var done = false
func claim() -> Bool {
lock.lock(); defer { lock.unlock() }
if done { return false }
done = true
return true
}
}