fix: send SSE comment keepalives so slow streams are not idle-killed

Sources/SwiftLM/Server.swift

@@ -2333,6 +2333,27 @@ let forwardPrefillProgress: @Sendable (Int, Int) -> Void = { processed, total in
     activePrefillProgressHook?(processed, total)
 }
 
+/// Emits an SSE comment every `interval` until the stream terminates or the task is
+/// cancelled.
+///
+/// Without it a stream can sit silent for minutes on slow hardware: through a long
+/// prefill (the progress heartbeat is opt-in) and while a tool call is buffered, since
+/// no delta is sent until it parses. Node/Bun clients abort a body idle for ~300 s
+/// (`terminated`), and the agent's retry re-pays the whole prefill. SSE parsers ignore
+/// comment lines, so this is safe for every client. Each yield is a complete event, so
+/// a comment can never land inside another event.
+func startSSEKeepalive(
+    _ cont: AsyncStream<String>.Continuation, interval: Duration = .seconds(5)
+) -> Task<Void, Never> {
+    Task {
+        while !Task.isCancelled {
+            try? await Task.sleep(for: interval)
+            guard !Task.isCancelled else { return }
+            if case .terminated = cont.yield(": keepalive\r\n\r\n") { return }
+        }
+    }
+}
+
 func handleChatStreaming(
     startGeneration: @escaping () async throws -> (AsyncStream<Generation>, (() async -> Void)?),
     modelId: String,
@@ -2388,6 +2409,7 @@ func handleChatStreaming(
     // liveness so a long prefill cannot trip the client's idle timeout
     // (e.g. Bun fetch's 10s) into an ECONNRESET-and-retry loop.
     cont.yield(": connected\r\n\r\n")
+    let keepaliveTask = startSSEKeepalive(cont)
 
     let consumerTask: Task<Void, Never> = Task {
         var hasToolCalls = false
@@ -2410,6 +2432,7 @@ func handleChatStreaming(
         // generation slot is returned on ALL exit paths (normal completion,
         // startGeneration failure, client disconnect, or task cancellation).
         defer {
+            keepaliveTask.cancel()
             heartbeatTask?.cancel()
             heartbeatTask = nil
             activePrefillProgressHook = nil
@@ -2969,6 +2992,7 @@ func handleTextStreaming(
     }
     // First byte before any model work — same liveness guarantee as the chat path.
     cont.yield(": connected\r\n\r\n")
+    let keepaliveTask = startSSEKeepalive(cont)
     let consumerTask: Task<Void, Never> = Task {
         var completionTokenCount = 0
         var fullText = ""
@@ -2981,6 +3005,7 @@ func handleTextStreaming(
         // Unconditional cleanup: cancels the heartbeat and returns the generation
         // slot on ALL exit paths (completion, startGeneration failure, disconnect).
         defer {
+            keepaliveTask.cancel()
             heartbeatTask?.cancel()
             heartbeatTask = nil
             activePrefillProgressHook = nil

tests/SwiftLMTests/ServerSSETests.swift

@@ -173,4 +173,28 @@ final class ServerSSETests: XCTestCase {
         XCTAssertEqual(err["type"] as? String, "server_error")
         XCTAssertEqual(err["code"] as? String, "internal_error")
     }
+
+    // MARK: - Keepalive
+
+    /// A silent stream (long prefill, buffered tool call) must still carry bytes, as
+    /// SSE comments that parsers ignore, or Node/Bun clients abort it after ~300 s.
+    func testKeepaliveEmitsSSECommentsWhileStreamIsSilent() async {
+        let (stream, cont) = AsyncStream<String>.makeStream()
+        let task = startSSEKeepalive(cont, interval: .milliseconds(20))
+        var received: [String] = []
+        for await event in stream {
+            received.append(event)
+            if received.count == 2 { break }
+        }
+        task.cancel()
+        XCTAssertEqual(received, [": keepalive\r\n\r\n", ": keepalive\r\n\r\n"])
+    }
+
+    func testKeepaliveStopsWhenStreamFinishes() async {
+        let (_, cont) = AsyncStream<String>.makeStream()
+        let task = startSSEKeepalive(cont, interval: .milliseconds(10))
+        cont.finish()
+        // Returns on its own (yield reports .terminated), without being cancelled.
+        await task.value
+    }
 }

tests/test-server.sh

@@ -1161,13 +1161,14 @@ rm -f /tmp/mlx_inflight_models.json /tmp/mlx_inflight_health.json /tmp/mlx_infli
 # whole prefill inside container.perform before returning headers, so clients
 # saw no bytes for the entire prefill (Bun fetch 10s idle → ECONNRESET/retry).
 # A large prompt makes prefill multi-second; TTFB must stay well under that.
-# ~2k tokens. Gemma-4 prefill allocates O(n²) attention in one Metal buffer
+# Under 2k tokens: 270 numbers measure 1,916 with the gemma-4-e2b tokenizer and
+# chat wrapper (300 were 2,126). Gemma-4 prefill allocates O(n²) attention in one Metal buffer
 # (~149 bytes × n² for gemma-4-e2b) and the CI runner caps a single buffer at
 # 3.5 GB, i.e. ~4.9k tokens max. 7k tokens crashed the server with a 7.4 GB
 # malloc; 2k tokens needs ~0.6 GB and still gives a measurable prefill.
 log "Test 38: streaming TTFB — ': connected' arrives before model prefill"
 
-python3 -c "print(' '.join(f'{i:06d}' for i in range(300)))" > /tmp/mlx_ttfb_prompt.txt
+python3 -c "print(' '.join(f'{i:06d}' for i in range(270)))" > /tmp/mlx_ttfb_prompt.txt
 jq -nc --arg m "$MODEL" --rawfile p /tmp/mlx_ttfb_prompt.txt \
     '{model:$m, stream:true, max_tokens:5, messages:[{role:"user",content:("Summarize this list in one word:\n"+$p)}]}' \
     > /tmp/mlx_ttfb_body.json