fix: stop a non-streaming generation when the client disconnects

Sources/SwiftLM/Server.swift

@@ -1498,7 +1498,8 @@ struct MLXServer: AsyncParsableCommand {
         }
 
         // ── Build Hummingbird router ──
-        let router = Router()
+        // The non-streaming handlers need the connection to notice a client leaving.
+        let router = Router(context: ConnectionRequestContext.self)
 
         // ── CORS middleware ──
         if let origin = corsOrigin {
@@ -1587,15 +1588,16 @@ struct MLXServer: AsyncParsableCommand {
             }
         }
         memoryPressure.resume()
-        router.post("/v1/chat/completions") { request, _ -> Response in
+        router.post("/v1/chat/completions") { request, context -> Response in
             do {
                 let bodyData = try await collectBody(request)
                 return try await handleChatCompletion(
                     request: request, bodyData: bodyData, config: config, container: container, semaphore: semaphore, stats: stats, promptCache: promptCache,
                     draftModelRef: draftModelRef, numDraftTokens: numDraftTokensConfig,
                     dflashModel: dflashModel, dflashBlockSize: dflashBlockSizeConfig,
                     dflashTargetModel: dflashTargetModel,
-                    mtpAssistant: loadedMTPAssistant
+                    mtpAssistant: loadedMTPAssistant,
+                    channel: context.channel
                 )
             } catch {
                 return Response(
@@ -1607,11 +1609,12 @@ struct MLXServer: AsyncParsableCommand {
         }
 
         // Text completions — handler extracted to avoid type-checker timeout
-        router.post("/v1/completions") { request, _ -> Response in
+        router.post("/v1/completions") { request, context -> Response in
             do {
                 let bodyData = try await collectBody(request)
                 return try await handleTextCompletion(
-                    request: request, bodyData: bodyData, config: config, container: container, semaphore: semaphore, stats: stats
+                    request: request, bodyData: bodyData, config: config, container: container, semaphore: semaphore, stats: stats,
+                    channel: context.channel
                 )
             } catch {
                 return Response(
@@ -2475,7 +2478,8 @@ func handleChatCompletion(
     dflashModel: DFlashDraftModel? = nil,
     dflashBlockSize: Int? = nil,
     dflashTargetModel: (any DFlashTargetModel)? = nil,
-    mtpAssistant: (any DualModelMTP)? = nil
+    mtpAssistant: (any DualModelMTP)? = nil,
+    channel: (any Channel)? = nil
 ) async throws -> Response {
     let chatReq: ChatCompletionRequest
     do {
@@ -2941,14 +2945,21 @@ func handleChatCompletion(
             emitPrefillProgress: emitPrefillProgress
         )
     } else {
-        let (stream, onPrefillDone) = try await startGeneration()
-        return try await handleChatNonStreaming(
-            stream: stream, modelId: modelId, stopSequences: stopSequences,
-            promptTokenCount: promptTokenCount, enableThinking: enableThinking,
-            thinkingPreOpened: thinkingPreOpened, jsonMode: jsonMode, tokenEcho: config.tokenEcho,
-            slot: slot, stats: stats, genStart: genStart, prefillStart: prefillStart,
-            onPrefillDone: onPrefillDone
-        )
+        // A client may give up during a long answer, and nothing on this path would
+        // notice: the SSE path learns of a disconnect from a failed write, but a
+        // non-streaming response writes nothing until the end. The generation then ran
+        // to max_tokens for nobody and held the slot the whole time (and with the
+        // default --parallel 1, every later request behind it).
+        return try await cancellingOnClientDisconnect(channel) { [enableThinking, thinkingPreOpened] in
+            let (stream, onPrefillDone) = try await startGeneration()
+            return try await handleChatNonStreaming(
+                stream: stream, modelId: modelId, stopSequences: stopSequences,
+                promptTokenCount: promptTokenCount, enableThinking: enableThinking,
+                thinkingPreOpened: thinkingPreOpened, jsonMode: jsonMode, tokenEcho: config.tokenEcho,
+                slot: slot, stats: stats, genStart: genStart, prefillStart: prefillStart,
+                onPrefillDone: onPrefillDone
+            )
+        }
     }
 }
 
@@ -3416,7 +3427,9 @@ func handleChatStreaming(
                 // it arrives. The prompt is saved here, not on the first chunk: that raced the
                 // decode step the producer had already started on the same KVCacheSimple. An
                 // early exit (stop sequence, disconnect) never reaches `.info` and skips the save.
-                if !firstToken, let onPrefillDone { await onPrefillDone() }
+                // The save does not wait for a chunk: a turn that is only a tool call never
+                // yields one, and its prompt is exactly what the next turn of the loop needs.
+                if let onPrefillDone { await onPrefillDone() }
                 // The chunk count undercounts: tokens the decoder buffers (tool-call
                 // bodies, partial text) emit no chunk. `info` carries the real count.
                 completionTokenCount = info.generationTokenCount
@@ -3592,7 +3605,7 @@ func handleChatNonStreaming(
         case .info(let info):
             // The producer's last yield; see the matching comment in `handleChatStreaming`.
             // The stop-sequence `break` above exits before `.info`, so it skips the save.
-            if !firstToken, let onPrefillDone { await onPrefillDone() }
+            if let onPrefillDone { await onPrefillDone() }
             generationStopReason = info.stopReason
             completionTokenCount = info.generationTokenCount  // real count; see handleChatStreaming
         }
@@ -3725,7 +3738,8 @@ func handleTextCompletion(
     config: ServerConfig,
     container: ModelContainer,
     semaphore: AsyncSemaphore,
-    stats: ServerStats
+    stats: ServerStats,
+    channel: (any Channel)? = nil
 ) async throws -> Response {
     let compReq: TextCompletionRequest
     do {
@@ -3790,11 +3804,14 @@ func handleTextCompletion(
             genStart: genStart, emitPrefillProgress: emitPrefillProgress
         )
     } else {
-        let stream = try await container.generate(input: lmInput, parameters: params)
-        return try await handleTextNonStreaming(
-            stream: stream, modelId: modelId, stopSequences: stopSequences,
-            promptTokenCount: promptTokenCount, slot: slot, stats: stats, genStart: genStart
-        )
+        // Same as the chat path: stop generating if the client goes away.
+        return try await cancellingOnClientDisconnect(channel) {
+            let stream = try await container.generate(input: lmInput, parameters: params)
+            return try await handleTextNonStreaming(
+                stream: stream, modelId: modelId, stopSequences: stopSequences,
+                promptTokenCount: promptTokenCount, slot: slot, stats: stats, genStart: genStart
+            )
+        }
     }
 }
 
@@ -4105,6 +4122,106 @@ final class GenerationSlot: @unchecked Sendable {
     }
 }
 
+// ── Client disconnect (non-streaming) ─────────────────────────────────────────
+
+/// Request context that keeps the connection's channel. `BasicRequestContext` drops it,
+/// and the non-streaming handlers need it to find out when the client goes away.
+struct ConnectionRequestContext: RequestContext {
+    var coreContext: CoreRequestContextStorage
+    let channel: any Channel
+
+    init(source: ApplicationRequestContextSource) {
+        coreContext = .init(source: source)
+        channel = source.channel
+    }
+}
+
+/// Reports, once, that the client of a connection has stopped listening.
+///
+/// Hummingbird enables remote half-closure on its connections, so a client that closes
+/// its socket mid-request shows up as `ChannelEvent.inputClosed`, not as a closed
+/// channel; a reset shows up as `channelInactive`. The streaming handlers need neither,
+/// because their next write fails and tears down the body stream, but a non-streaming
+/// handler writes nothing until the end, so this is its only signal.
+///
+/// Installed at the head of the pipeline for the span of one request and removed after
+/// it, so a keep-alive connection closing between requests is not taken for a client
+/// abandoning one. Both callbacks run on the channel's event loop, and `fired` is only
+/// touched there.
+final class ClientDisconnectWatch: ChannelInboundHandler, RemovableChannelHandler, @unchecked Sendable {
+    typealias InboundIn = ByteBuffer
+
+    private let onDisconnect: @Sendable () -> Void
+    private var fired = false
+
+    init(onDisconnect: @escaping @Sendable () -> Void) {
+        self.onDisconnect = onDisconnect
+    }
+
+    func userInboundEventTriggered(context: ChannelHandlerContext, event: Any) {
+        if (event as? ChannelEvent) == .inputClosed { fire() }
+        context.fireUserInboundEventTriggered(event)
+    }
+
+    func channelInactive(context: ChannelHandlerContext) {
+        fire()
+        context.fireChannelInactive()
+    }
+
+    private func fire() {
+        guard !fired else { return }
+        fired = true
+        onDisconnect()
+    }
+}
+
+/// A cancel request that may arrive before the task it is meant for exists.
+final class DeferredCancellation: @unchecked Sendable {
+    private let lock = NSLock()
+    private var cancelled = false
+    private var cancelTask: (@Sendable () -> Void)?
+
+    func cancel() {
+        lock.lock()
+        cancelled = true
+        let cancelTask = self.cancelTask
+        lock.unlock()
+        cancelTask?()
+    }
+
+    func attach<T: Sendable, E: Error>(_ task: Task<T, E>) {
+        lock.lock()
+        cancelTask = { task.cancel() }
+        let cancelled = self.cancelled
+        lock.unlock()
+        if cancelled { task.cancel() }
+    }
+}
+
+/// Runs `body`, a non-streaming generation, and cancels it if the client disconnects
+/// meanwhile. Cancellation reaches the producer through the generation stream, the
+/// same way it does when an SSE consumer is cancelled: the model stops after the token
+/// in flight, the handler falls out of its loop and releases the slot, and the response
+/// it builds goes to a closed socket. With no channel (tests), `body` just runs.
+func cancellingOnClientDisconnect<T: Sendable>(
+    _ channel: (any Channel)?, _ body: @escaping @Sendable () async throws -> T
+) async throws -> T {
+    guard let channel else { return try await body() }
+    let cancellation = DeferredCancellation()
+    let watch = ClientDisconnectWatch { cancellation.cancel() }
+    // Installed before the work starts, so a disconnect in between is not missed. If
+    // the channel is already closed the install fails, and the client has left.
+    do {
+        try await channel.pipeline.addHandler(watch, position: .first).get()
+    } catch {
+        cancellation.cancel()
+    }
+    defer { channel.pipeline.removeHandler(watch, promise: nil) }
+    let work = Task { try await body() }
+    cancellation.attach(work)
+    return try await work.value
+}
+
 // ── CORS Middleware ───────────────────────────────────────────────────────────
 
 struct CORSMiddleware<Context: RequestContext>: RouterMiddleware {