diff --git a/Sources/SwiftLM/Server.swift b/Sources/SwiftLM/Server.swift index c69dcf7..55c4f45 100644 --- a/Sources/SwiftLM/Server.swift +++ b/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.Continuation, interval: Duration = .seconds(5) +) -> Task { + 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, (() 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 = 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 = 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 diff --git a/tests/SwiftLMTests/ServerSSETests.swift b/tests/SwiftLMTests/ServerSSETests.swift index 48de24f..c96b7d4 100644 --- a/tests/SwiftLMTests/ServerSSETests.swift +++ b/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.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.makeStream() + let task = startSSEKeepalive(cont, interval: .milliseconds(10)) + cont.finish() + // Returns on its own (yield reports .terminated), without being cancelled. + await task.value + } } diff --git a/tests/test-server.sh b/tests/test-server.sh index 110bfe9..95b0939 100755 --- a/tests/test-server.sh +++ b/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