fix: reclaim generation slot on error and stream SSE headers before prefill

[CodeAndCanvas728]

Problem

Three related reliability failures around the --parallel semaphore and streaming TTFB:

1. Slot leak → permanent wedge. Any throw after semaphore.wait() (container.prepare, container.perform, …) skipped every manual signal() call site. With the default --parallel 1, one failed request wedged the server for all subsequent generations until restart.
2. Silent prefill → client ECONNRESET/retry loops. The streaming handler ran the whole prefill inside container.perform before returning response headers, so clients saw zero bytes for the entire prefill. Bun's fetch 10s idle timeout then killed the connection and retried — with an ever-grown payload (opencode retry-history growth; their side is tracked in opencode PR #26167, this is the server-side contributor).
3. Suspected /v1/models stall during generation — diagnosed with probes before any fix (see below).

Probe findings (before changing anything)

• Fresh-connection GET /v1/models and GET /health answer in ~0 ms even mid-prefill — no server-wide stall, no Hummingbird idleTimeout change needed.
• Same-connection pipelining is normal HTTP/1.1 serialization, not a bug.
• TTFB = end of prefill (60s+, 90s+ on large prompts) — confirms silent prefill as the root cause of the ECONNRESET/retry loop. A dead client's huge prompt still held the only slot through minutes of prefill.

Fix

• GenerationSlot — release-once guard around the semaphore: explicit release() on all success paths, deinit as fallback, whichever fires first signals exactly once (local bind so the Task never captures self). The four scattered await semaphore.signal() sites are gone.
• Deferred startGeneration — for streaming, container.perform (prefill) now runs inside the SSE consumer task after the handler returns, so headers + a : connected SSE comment hit the wire first; prefill heartbeat hook is installed up front; client disconnect cancels the consumer task via onTermination.
• Non-streaming behavior unchanged (awaits generation inline, slot released on all exits).

Behavior change

• Streaming request that fails during prefill now returns HTTP 200 + an SSE error event, then data: [DONE] — previously a 500. Headers are already on the wire by the time prefill runs, so the failure can only be reported in-stream. This matches OpenAI's streaming error behavior. The event payload is {"error":{"message":…,"type":"server_error","code":"internal_error"}}, JSON-encoded via JSONSerialization so any error text stays valid JSON.
• Non-streaming requests still return 500 on failure.

Measured

• Large-prompt streaming: first byte 0.126 s vs 54.76 s prefill (42,038-token prompt, Qwen2.5-0.5B), : connected sent, stream completed.
• Control-plane probes during in-flight generation: /v1/models + /health both answered within 3 s.

Tests

• GenerationSlotTests (new): deinit reclaims without explicit release; explicit release + deinit doesn't double-signal; repeated release() is idempotent — the regression is the leak, and no deterministic post-wait() error trigger exists for text-only test models, so the slot contract is unit-tested rather than e2e-forced.
• ServerSSETests.testErrorChunkEncodesMessageAsValidJSON (new): an error message containing ", \ and a newline round-trips through the SSE error event as valid JSON.
• test-server.sh Test 37: /v1/models + /health respond during in-flight generation on the default --parallel 1 server.
• test-server.sh Test 38: streaming TTFB < 2 s with : connected before a multi-second prefill.
• swift test: 178/178. ./tests/test-server.sh: 76/76. Release build: swift build -c release --product SwiftLM (full-product build still hits the known SwiftSoup/SwiftBuddy frontend crash — pre-existing, unrelated).

Known limitations (unchanged / out of scope)

• Prefill itself cannot be cancelled mid-eval() (mlx-swift-lm); disconnect during prefill still holds the slot until prefill completes — decode phase does cancel.
• Server-wide idle timeouts were deliberately not added: probes showed the control plane is healthy during prefill.

[CodeAndCanvas728]

Closing temporarily — holding for local testing before review.

[solderzzc]

@CodeAndCanvas728 Thanks for your contribution and verification, I've been also running test, since you are working on it, I'll let my agent wait for your commit.

[CodeAndCanvas728]

Manual soak re-run under healthy-machine conditions after diagnosing an unrelated OS-level low-swap SIGKILL (not this PR).

• GenerationSlotTests: 3/3
• Full unit suites: green (0 failures)
• e2e test-server.sh: 76/76 (incl. Tests 37 & 38)
• Tiel soak A–F: early TTFB 0.084s with : connected before prefill; slot reclaimed after success, mid-prefill disconnect, and error; sequential x3 OK; no process death / no exit 137

Reopening for review.

[solderzzc]

Thanks for the thorough soak run!

One blocker: Server.swift lines 208 and 231 use any Sendable?, which doesn't compile under Xcode 26.3 / Swift 6 (the version CI pins). It needs to be (any Sendable)?:

func sanitizeForJinja(_ value: any Sendable) -> (any Sendable)? {
...
func mapValuesDeep(_ transform: (any Sendable) -> (any Sendable)?) -> [String: any Sendable] {

With that change applied locally, it builds on 26.3 and all 177 tests pass, including the new GenerationSlotTests.

The slot-reclaim logic looks good to me. It releases exactly once, deinit catches anything that slips through, and on a mid-prefill disconnect the slot stays held until prefill actually ends, so two generations never run at once.

Two non-blocking notes:

1. When a streaming request fails during prefill, it now returns HTTP 200 plus an SSE error event instead of a 500. That matches OpenAI's behavior, but please call it out in the PR description.
2. The error event's message only has " replaced with '. A backslash or newline in the error text would make the event invalid JSON. Consider encoding it with JSONSerialization.

Since #170 already contains everything in #169, I'd suggest merging #170 once the fix is in, and closing #169.

[CodeAndCanvas728]

Thanks for the review — all three addressed in ac26d57:

1. Blocker: sanitizeForJinja / mapValuesDeep now use (any Sendable)?. Those were the only two unparenthesized optional existentials in Sources/ and tests/. (My local toolchain is Xcode 27 / Swift 6.4, which accepts the bare form — that's how it slipped through.)
2. HTTP 200 + SSE error on prefill failure: called out under a new "Behavior change" section in the description.
3. Error event encoding: new sseErrorChunk(_:) builds the payload with JSONSerialization; both streaming catch sites (chat + text completions) use it. Added ServerSSETests.testErrorChunkEncodesMessageAsValidJSON covering ", \ and newline.

swift test: 178/178. test-server.sh: 76/76.

Side note, not touched here since it predates this PR: the non-streaming 500 handlers for /v1/chat/completions and /v1/completions build their error body the same way (" → ' only), so they have the same invalid-JSON risk. Happy to send a small follow-up if you want.

Agreed on merging #170 and closing #169 — closing #169 now.

[solderzzc]

Thanks @CodeAndCanvas728 , let me turn on CI to get it passed.

[solderzzc]

Thanks for the review — all three addressed in ac26d57:

1. Blocker: sanitizeForJinja / mapValuesDeep now use (any Sendable)?. Those were the only two unparenthesized optional existentials in Sources/ and tests/. (My local toolchain is Xcode 27 / Swift 6.4, which accepts the bare form — that's how it slipped through.)
2. HTTP 200 + SSE error on prefill failure: called out under a new "Behavior change" section in the description.
3. Error event encoding: new sseErrorChunk(_:) builds the payload with JSONSerialization; both streaming catch sites (chat + text completions) use it. Added ServerSSETests.testErrorChunkEncodesMessageAsValidJSON covering ", \ and newline.

swift test: 178/178. test-server.sh: 76/76.

Side note, not touched here since it predates this PR: the non-streaming 500 handlers for /v1/chat/completions and /v1/completions build their error body the same way (" → ' only), so they have the same invalid-JSON risk. Happy to send a small follow-up if you want.

Agreed on merging #170 and closing #169 — closing #169 now.

Yes please on the follow-up for the non-streaming 500 handlers. Thanks a lot for your contribution.

[CodeAndCanvas728]

Thanks! The red job (integration_matrix (server, -gemma-e2b …), 75/76) is Test 38, which this PR added. The PR's fix itself passed there: TTFB was 0.21s and : connected arrived first. What failed is that the server crashed during prefill, so the stream never reached [DONE]:

srv  slot_launch: id 0 | prompt=42017t | thinking=false | prefilling...
[metal::malloc] Attempting to allocate 28383323840 bytes which is greater than the maximum allowed buffer size of 3758096384 bytes.

The test prompt (6000 numbers) comes to about 42k tokens on Gemma-4. Gemma-4 prefill allocates full attention in a single buffer (28383323840 = 42017 × 42220 × 16 bytes, so it grows with the square of the prompt length). That goes well over the runner's 3.5 GB buffer limit. Qwen 0.5B passed the same test.

Fixed in fe9da73 by cutting the prompt to 1000 numbers (about 7k tokens, roughly 0.8 GB). The : connected-first check still catches the original regression. Locally, test-server.sh passes 76/76 and Test 38 reaches first byte in 0.04s. The new run is waiting for workflow approval ("action required").

Separately, a 42k-token prompt that fits in context still crashes the server on Gemma-4. That's a real bug but outside this PR, so I'll handle it separately.

The non-streaming 500 follow-up is in #171. It's stacked on this PR, so only its last commit is new.

[solderzzc]

Thanks! The red job (integration_matrix (server, -gemma-e2b …), 75/76) is Test 38, which this PR added. The PR's fix itself passed there: TTFB was 0.21s and : connected arrived first. What failed is that the server crashed during prefill, so the stream never reached [DONE]:

srv  slot_launch: id 0 | prompt=42017t | thinking=false | prefilling...
[metal::malloc] Attempting to allocate 28383323840 bytes which is greater than the maximum allowed buffer size of 3758096384 bytes.

The test prompt (6000 numbers) comes to about 42k tokens on Gemma-4. Gemma-4 prefill allocates full attention in a single buffer (28383323840 = 42017 × 42220 × 16 bytes, so it grows with the square of the prompt length). That goes well over the runner's 3.5 GB buffer limit. Qwen 0.5B passed the same test.

Fixed in fe9da73 by cutting the prompt to 1000 numbers (about 7k tokens, roughly 0.8 GB). The : connected-first check still catches the original regression. Locally, test-server.sh passes 76/76 and Test 38 reaches first byte in 0.04s. The new run is waiting for workflow approval ("action required").

Separately, a 42k-token prompt that fits in context still crashes the server on Gemma-4. That's a real bug but outside this PR, so I'll handle it separately.

The non-streaming 500 follow-up is in #171. It's stacked on this PR, so only its last commit is new.

Thanks, let me start the CI.

[solderzzc]

I pushed c51752e directly to this branch to shrink Test 38's prompt to ~2k tokens.

Why: fe9da73's ~7k-token prompt still exceeded the CI runner's limit. gemma-4-e2b's prefill tried to allocate 7.4 GB in a single Metal buffer, against the runner's 3.5 GB cap:

[metal::malloc] Attempting to allocate 7357857792 bytes which is greater than the maximum allowed buffer size of 3758096384 bytes

Prefill costs roughly 149 bytes × n² for this model, so the real ceiling is ~4.9k tokens, not ~15k as the comment said. It now uses range(300) (~2k tokens, ~0.6 GB), and the comment is corrected. Test-only change; no server code touched.

#171 is stacked on fe9da73, so it will need a rebase. This change doesn't touch anything in #171, so there should be no conflicts.