Add async request queue to llm-server

swift/Sources/CoreAILMCommon/RequestQueue.swift

@@ -0,0 +1,163 @@
+// Copyright 2026 Apple Inc.
+//
+// Use of this source code is governed by a BSD-3-clause license that can
+// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
+
+import Foundation
+import Synchronization
+
+// MARK: - Request Queue
+
+/// Async semaphore gating concurrency at one in-flight request. A second request
+/// waits in a FIFO queue instead of getting an immediate 429; it is rejected only
+/// once `maxDepth` requests are already waiting.
+///
+/// `acquire()` returns a `QueuePermit` that frees the slot exactly once. Callers
+/// hold the permit for the lifetime of the request and let it release on scope
+/// exit (or hand it to a streaming body closure); the permit also releases on
+/// deallocation, so the slot is never leaked if a caller drops it.
+public final class RequestQueue: Sendable {
+    public let maxDepth: Int
+    private let state = Mutex<QueueState>(QueueState())
+
+    private struct QueueState {
+        var isActive: Bool = false
+        // Waiters carry a token so the cancellation handler can find and remove
+        // its own continuation (CheckedContinuation is not Equatable).
+        var waiters: [(id: UInt64, cont: CheckedContinuation<Void, any Error>)] = []
+        var nextID: UInt64 = 0
+    }
+
+    public init(maxDepth: Int) {
+        self.maxDepth = max(0, maxDepth)
+    }
+
+    public var depth: Int {
+        state.withLock { $0.waiters.count + ($0.isActive ? 1 : 0) }
+    }
+
+    public var queuedCount: Int {
+        state.withLock { $0.waiters.count }
+    }
+
+    public var isActive: Bool {
+        state.withLock { $0.isActive }
+    }
+
+    /// Acquire the exclusive slot, waiting in FIFO order if it is held. Throws
+    /// `ServerError.queueFull` when `maxDepth` waiters are already queued, or
+    /// `CancellationError` if the awaiting task is cancelled while queued.
+    public func acquire() async throws -> QueuePermit {
+        let needsWait: Bool = state.withLock { s in
+            if !s.isActive {
+                s.isActive = true
+                return false
+            }
+            return true
+        }
+        if !needsWait { return QueuePermit(self) }
+
+        // Reserve a token before suspending so the cancellation handler below can
+        // match this waiter even if it runs before the continuation is installed.
+        let id = state.withLock { s -> UInt64 in
+            defer { s.nextID &+= 1 }
+            return s.nextID
+        }
+
+        try await withTaskCancellationHandler {
+            try await withCheckedThrowingContinuation { (cont: CheckedContinuation<Void, any Error>) in
+                enum Action { case acquired, rejected, queued, cancelled }
+                let action: Action = state.withLock { s -> Action in
+                    if Task.isCancelled { return .cancelled }
+                    if !s.isActive {
+                        s.isActive = true
+                        return .acquired
+                    }
+                    if s.waiters.count >= maxDepth {
+                        return .rejected
+                    }
+                    s.waiters.append((id: id, cont: cont))
+                    return .queued
+                }
+                switch action {
+                case .acquired:
+                    cont.resume()
+                case .rejected:
+                    cont.resume(throwing: ServerError.queueFull(depth: maxDepth))
+                case .cancelled:
+                    cont.resume(throwing: CancellationError())
+                case .queued:
+                    break
+                }
+            }
+        } onCancel: {
+            // Remove our still-queued waiter and resume it throwing. The shared
+            // Mutex serializes this against release(), so a given waiter is popped
+            // by exactly one of the two and resumed exactly once.
+            let cont: CheckedContinuation<Void, any Error>? = state.withLock { s in
+                guard let idx = s.waiters.firstIndex(where: { $0.id == id }) else { return nil }
+                return s.waiters.remove(at: idx).cont
+            }
+            cont?.resume(throwing: CancellationError())
+        }
+        return QueuePermit(self)
+    }
+
+    /// Hand the slot to the next FIFO waiter, or clear the active flag. Internal;
+    /// callers release through `QueuePermit`.
+    fileprivate func release() {
+        let next: CheckedContinuation<Void, any Error>? = state.withLock { s in
+            if !s.waiters.isEmpty {
+                return s.waiters.removeFirst().cont
+            }
+            s.isActive = false
+            return nil
+        }
+        next?.resume()
+    }
+}
+
+// MARK: - Queue Permit
+
+/// A held `RequestQueue` slot. Releases exactly once — on the first `release()`
+/// call or, as a safety net, on deallocation (e.g. if a streaming response body
+/// is dropped by the server before its writer closure ever runs).
+public final class QueuePermit: Sendable {
+    private let queue: RequestQueue
+    private let released = Atomic<Bool>(false)
+
+    fileprivate init(_ queue: RequestQueue) {
+        self.queue = queue
+    }
+
+    public func release() {
+        // Flip false -> true atomically; exactly one caller wins the exchange and
+        // hands the slot back. `.relaxed` suffices: the exchange's atomicity (not
+        // its ordering) picks the single winner, and queue.release() takes the
+        // RequestQueue Mutex, which provides the actual cross-thread handoff.
+        let (won, _) = released.compareExchange(
+            expected: false, desired: true, ordering: .relaxed)
+        if won { queue.release() }
+    }
+
+    deinit { release() }
+}
+
+// MARK: - Server Errors
+
+public enum ServerError: Error, LocalizedError {
+    case badRequest(String)
+    case queueFull(depth: Int)
+
+    public var isBadRequest: Bool {
+        if case .badRequest = self { return true }
+        return false
+    }
+
+    public var errorDescription: String? {
+        switch self {
+        case .badRequest(let msg): return msg
+        case .queueFull(let depth): return "Queue full (depth: \(depth))"
+        }
+    }
+}

swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift

@@ -68,6 +68,8 @@ final class CoreAIPipelinedEngine: InferenceEngine, ConstrainedGenerationCapable
 
     var processedTokenCount: Int { engine.processedTokenCount }
 
+    var hasRecurrentState: Bool { engine.hasNonTruncatableStates }
+
     init(
         config: ModelConfig,
         preparedModel: PreparedModel,

swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift

@@ -35,6 +35,7 @@ public final class CoreAISequentialEngine: InferenceEngine, @unchecked Sendable
 
     public var supportsLogits: Bool { true }
     public var vocabSize: Int { config.vocabSize }
+    public var hasRecurrentState: Bool { hasNonTruncatableStates }
     public let config: ModelConfig
 
     // Core AI function handle

swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift

@@ -143,6 +143,15 @@ public protocol InferenceEngine: Sendable {
     /// Useful for debugging multi-turn efficiency and verifying prefix caching behavior.
     var lastPrefixHitCount: Int { get }
 
+    /// Whether the engine carries recurrent state (SSM / hybrid models).
+    ///
+    /// Recurrent/conv states summarize the entire prefix and cannot be truncated
+    /// by moving a KV cursor, so token-prefix reuse is impossible: on any rewind
+    /// the engine full-resets and replays the whole prompt. Callers (e.g. the
+    /// server's prefix-reuse accounting) should short-circuit when this is true.
+    /// Defaults to `false`.
+    var hasRecurrentState: Bool { get }
+
     // MARK: - Configuration
 
     associatedtype ConfigType: Codable, InferenceConfiguration
@@ -198,6 +207,10 @@ extension InferenceEngine {
     public var lastPrefixHitCount: Int { 0 }
 }
 
+extension InferenceEngine {
+    public var hasRecurrentState: Bool { false }
+}
+
 extension InferenceEngine {
     /// Default: engine is not busy.
     public var isBusy: Bool { false }

swift/Sources/Tools/llm-server/ChatHandler.swift

@@ -81,6 +81,9 @@ func startServer(state: ServerState, port: Int) async throws {
         configuration: .init(
             address: .hostname("127.0.0.1", port: port)
         ))
+
+    // Process lifecycle is left to the supervising process. TODO: any in-process
+    // idle-exit must first drain in-flight engine work (no drain hook today).
     try await app.run()
 }
 
@@ -100,8 +103,15 @@ private func handleAutoRoute(request: Request, state: ServerState) async throws
 private func handleChatCompletionsFromBody(body: ByteBuffer, state: ServerState, sessionID: String? = nil) async throws
     -> Response
 {
-    guard state.tryAcquire() else {
-        let err = ErrorResponse(error: .init(message: "Server is busy.", type: "server_error", code: "busy"))
+    // Non-streaming paths release the permit at each return; the streaming path
+    // hands it to the body closure. Release-on-deinit backstops a dropped response.
+    let permit: QueuePermit
+    do {
+        permit = try await state.queue.acquire()
+    } catch let error as ServerError {
+        let err = ErrorResponse(
+            error: .init(
+                message: error.localizedDescription, type: "server_error", code: "queue_full"))
         let data = try JSONEncoder().encode(err)
         return Response(
             status: .tooManyRequests, headers: [.contentType: "application/json"],
@@ -111,7 +121,7 @@ private func handleChatCompletionsFromBody(body: ByteBuffer, state: ServerState,
     do {
         chatRequest = try JSONDecoder().decode(ChatCompletionRequest.self, from: body)
     } catch {
-        state.release()
+        permit.release()
         let err = ErrorResponse(error: .init(message: "\(error)", type: "invalid_request_error", code: nil))
         let data = try JSONEncoder().encode(err)
         return Response(
@@ -121,23 +131,25 @@ private func handleChatCompletionsFromBody(body: ByteBuffer, state: ServerState,
     do {
         let shouldStream = chatRequest.stream ?? false
         if shouldStream {
-            return try await handleStreamingRequest(chatRequest: chatRequest, state: state, sessionID: sessionID)
+            // Ownership of the permit passes to the streaming body closure.
+            return try await handleStreamingRequest(
+                chatRequest: chatRequest, state: state, sessionID: sessionID, permit: permit)
         } else {
             let response = try await handleNonStreamingRequest(
                 chatRequest: chatRequest, state: state, sessionID: sessionID)
-            state.release()
+            permit.release()
             return response
         }
     } catch let error as ServerError {
-        state.release()
+        permit.release()
         let status: HTTPResponse.Status = error.isBadRequest ? .badRequest : .internalServerError
         let err = ErrorResponse(error: .init(message: "\(error)", type: "invalid_request_error", code: nil))
         let data = try JSONEncoder().encode(err)
         return Response(
             status: status, headers: [.contentType: "application/json"], body: .init(byteBuffer: ByteBuffer(data: data))
         )
     } catch {
-        state.release()
+        permit.release()
         let err = ErrorResponse(error: .init(message: "\(error)", type: "server_error", code: nil))
         let data = try JSONEncoder().encode(err)
         return Response(
@@ -327,7 +339,9 @@ private func handleNonStreamingRequest(chatRequest: ChatCompletionRequest, state
 
 // MARK: - Streaming (SSE)
 
-private func handleStreamingRequest(chatRequest: ChatCompletionRequest, state: ServerState, sessionID: String? = nil)
+private func handleStreamingRequest(
+    chatRequest: ChatCompletionRequest, state: ServerState, sessionID: String? = nil, permit: QueuePermit
+)
     async throws -> Response
 {
     let requestMaxTokens = chatRequest.maxCompletionTokens ?? chatRequest.maxTokens ?? state.config.defaultMaxTokens
@@ -364,7 +378,9 @@ private func handleStreamingRequest(chatRequest: ChatCompletionRequest, state: S
     }
 
     let responseBody = ResponseBody { writer in
-        defer { state.release() }
+        // The closure owns the permit for the stream's lifetime; release-on-deinit
+        // covers a response the framework drops before entering this closure.
+        defer { permit.release() }
         do {
             let encoder = JSONEncoder()
             let genStart = SuspendingClock().now

swift/Sources/Tools/llm-server/CompletionHandler.swift

@@ -30,14 +30,19 @@ func handleCompletionsFromBody(body: ByteBuffer, state: ServerState) async throw
             status: .notImplemented, headers: [.contentType: "application/json"],
             body: .init(byteBuffer: ByteBuffer(data: data)))
     }
-    guard state.tryAcquire() else {
-        let err = ErrorResponse(error: .init(message: "Server is busy.", type: "server_error", code: "busy"))
+    let permit: QueuePermit
+    do {
+        permit = try await state.queue.acquire()
+    } catch let error as ServerError {
+        let err = ErrorResponse(
+            error: .init(
+                message: error.localizedDescription, type: "server_error", code: "queue_full"))
         let data = try JSONEncoder().encode(err)
         return Response(
             status: .tooManyRequests, headers: [.contentType: "application/json"],
             body: .init(byteBuffer: ByteBuffer(data: data)))
     }
-    defer { state.release() }
+    defer { permit.release() }
     let req: CompletionRequest
     do {
         req = try JSONDecoder().decode(CompletionRequest.self, from: body)

swift/Sources/Tools/llm-server/LLMServerMain.swift

@@ -65,6 +65,9 @@ struct LLMServer: AsyncParsableCommand {
     )
     var chunkThreshold: Int?
 
+    @Option(name: .customLong("max-queue-depth"), help: "Max requests queued before returning 429 (default: 16)")
+    var maxQueueDepth: Int = 16
+
     @Flag(name: .customLong("no-thinking"), help: "Disable thinking/reasoning (appends /no_think or sets template)")
     var noThinking: Bool = false
 
@@ -77,6 +80,12 @@ struct LLMServer: AsyncParsableCommand {
     @Flag(help: "Enable verbose logging")
     var verbose: Bool = false
 
+    func validate() throws {
+        guard maxQueueDepth >= 0 else {
+            throw ValidationError("--max-queue-depth must be >= 0 (got \(maxQueueDepth))")
+        }
+    }
+
     func run() async throws {
         CLILogger.level = verbose ? 1 : 0
 
@@ -188,7 +197,8 @@ struct LLMServer: AsyncParsableCommand {
             supportsLogprobs: supportsLogprobs,
             maxContextLength: bundle.maxContextLength,
             vocabSize: bundle.vocabSize,
-            additionalEosTokenIds: additionalEosTokenIds
+            additionalEosTokenIds: additionalEosTokenIds,
+            maxQueueDepth: maxQueueDepth
         )
 
         let state = ServerState(
@@ -203,6 +213,7 @@ struct LLMServer: AsyncParsableCommand {
             print("  Logprobs: \(supportsLogprobs ? "supported" : "not supported (use --variant coreai-sequential)")")
             print("  Context: \(bundle.maxContextLength) tokens")
             print("  No-thinking: \(noThinking)")
+            print("  Max queue depth: \(maxQueueDepth)")
             let topKStr = topK.map { "\($0)" } ?? "nil"
             let topPStr = topP.map { "\($0)" } ?? "nil"
             print("  Sampling: temperature=\(temperature), topK=\(topKStr), topP=\(topPStr)")