From 2b552309404eac709b9b358e3feae3ba07acecd5 Mon Sep 17 00:00:00 2001 From: CodeAndCanvas728 Date: Thu, 24 Sep 2026 12:54:42 +0200 Subject: [PATCH 1/2] fix: send SSE comment keepalives so slow streams are not idle-killed On slow hardware a streaming response can carry no bytes for minutes: 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. Observed with little-coder/pi on a 16 GB M2: a 4k-token write tool call went silent after 16 tokens and died at +300 s, then every retry timed out. Emit ": keepalive" every 5 s on chat and text completion streams until the stream terminates. SSE parsers ignore comment lines, and every yield is a complete event, so a comment never lands inside another event. Co-Authored-By: Claude Opus 5.5 (1M context) --- Sources/SwiftLM/Server.swift | 25 +++++++++++++++++++++++++ tests/SwiftLMTests/ServerSSETests.swift | 24 ++++++++++++++++++++++++ 2 files changed, 49 insertions(+) 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 + } } From d0c6b135eb064150b775a91b87b516b8f6d764e6 Mon Sep 17 00:00:00 2001 From: CodeAndCanvas728 Date: Fri, 25 Sep 2026 13:27:39 +0200 Subject: [PATCH 2/2] test: keep Test 38 prompt under 2k tokens for gemma-4-e2b 300 six-digit numbers measure 2,126 tokens with the gemma-4-e2b tokenizer and chat wrapper; 270 measure 1,916, keeping CI below the 2k-token budget. Co-Authored-By: Claude Opus 5.5 (1M context) --- tests/test-server.sh | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) 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