fix: JSON mode with stream: true sends its reply; nested objects stay whole

Sources/SwiftLM/Server.swift

@@ -3118,6 +3118,33 @@ func startSSEKeepalive(
     }
 }
 
+/// What a JSON-mode stream sends for its buffered opening: the reply itself, unless a
+/// complete object precedes a second one; then the first is a hallucinated prefix and
+/// only the text from the second `{` on goes out. The earlier rule cut at the second `{`
+/// wherever it was, which turned `{"user": {"name": "Bob"}}` into `{"name": "Bob"}}`
+/// (#233). Braces are counted without regard to strings: this guards against a stray
+/// leading object, it is not a JSON parser.
+func jsonModeFlushText(_ buffered: String) -> String {
+    let trimmed = buffered.trimmingCharacters(in: .whitespacesAndNewlines)
+    guard trimmed.hasPrefix("{") else { return trimmed }
+    var depth = 0
+    for i in trimmed.indices {
+        switch trimmed[i] {
+        case "{": depth += 1
+        case "}": depth -= 1
+        default: continue
+        }
+        if depth == 0 {
+            let rest = trimmed[trimmed.index(after: i)...]
+            if let next = rest.firstIndex(of: "{") {
+                return String(trimmed[next...])
+            }
+            return trimmed
+        }
+    }
+    return trimmed
+}
+
 func handleChatStreaming(
     startGeneration: @escaping () async throws -> (AsyncStream<Generation>, (() async -> Void)?),
     modelId: String,
@@ -3226,6 +3253,31 @@ func handleChatStreaming(
         // ── JSON mode streaming: buffer early tokens to strip hallucinated prefixes ──
         var jsonBuffering = jsonMode
         var jsonBuffer = ""
+        // Hands the buffered text to the client, once: at the flush threshold, or at the
+        // end of a stream that never reached it. The thresholds count chunks, and a reply
+        // that starts with `{` is held by the generator as a possible bare tool call until
+        // generation ends, so it usually arrives as one chunk; without the end-of-stream
+        // flush a JSON-mode stream carried no content at all (#233).
+        func flushJSONBuffer() {
+            guard jsonBuffering else { return }
+            jsonBuffering = false
+            let cleaned = jsonModeFlushText(jsonBuffer)
+            jsonBuffer = ""
+            let (rText, cText) = enableThinking ? tracker.process(cleaned) : ("", cleaned)
+            if !rText.isEmpty || !cText.isEmpty {
+                cont.yield(sseChunk(
+                    modelId: modelId,
+                    reasoningContent: rText.isEmpty ? nil : rText,
+                    content: cText.isEmpty ? nil : cText,
+                    finishReason: nil
+                ))
+            }
+            // Everything accumulated so far has now gone to the client, so
+            // the stop path must not re-emit it. Counting raw characters
+            // consumed (not the cleaned derivative that went on the wire)
+            // keeps this consistent with how the normal path counts.
+            emittedTextCount = fullText.count
+        }
 
         for await generation in stream {
             if stopped { break }
@@ -3265,32 +3317,8 @@ func handleChatStreaming(
                     let enoughTokens = completionTokenCount >= 3
                     let hitMax = completionTokenCount >= 32
                     let hasDoubleBrace = enoughTokens && trimmed.hasPrefix("{") && trimmed.dropFirst().contains("{")
-
                     if hitMax || hasDoubleBrace {
-                        var cleaned = trimmed
-                        if hasDoubleBrace {
-                            if let firstBrace = cleaned.firstIndex(of: "{") {
-                                let afterFirst = cleaned.index(after: firstBrace)
-                                if let secondBrace = cleaned[afterFirst...].firstIndex(of: "{") {
-                                    cleaned = String(cleaned[secondBrace...])
-                                }
-                            }
-                        }
-                        let (rText, cText) = enableThinking ? tracker.process(cleaned) : ("", cleaned)
-                        if !rText.isEmpty || !cText.isEmpty {
-                            cont.yield(sseChunk(
-                                modelId: modelId,
-                                reasoningContent: rText.isEmpty ? nil : rText,
-                                content: cText.isEmpty ? nil : cText,
-                                finishReason: nil
-                            ))
-                        }
-                        // Everything accumulated so far has now gone to the client, so
-                        // the stop path must not re-emit it. Counting raw characters
-                        // consumed (not the cleaned derivative that went on the wire)
-                        // keeps this consistent with how the normal path counts.
-                        emittedTextCount = fullText.count
-                        jsonBuffering = false
+                        flushJSONBuffer()
                     }
                     continue  // skip normal emit while buffering or just flushed
                 }
@@ -3400,6 +3428,9 @@ func handleChatStreaming(
                     case .cancelled, .stop:
                         reason = hasToolCalls ? "tool_calls" : "stop"
                     }
+                    // A JSON-mode reply shorter than the flush threshold is still in the
+                    // buffer; it goes out before the finish chunk.
+                    flushJSONBuffer()
                     // The withheld tail never became a stop sequence, so it is ordinary
                     // output and must not be dropped.
                     if !heldStopTail.isEmpty {
@@ -3460,7 +3491,11 @@ func handleChatStreaming(
         }
         // Safety net: the tail is normally released in `case .info`, but a stream that
         // ends without one — a generation error, or cancellation mid-decode — would
-        // otherwise discard up to maxStopLength-1 characters of real output.
+        // otherwise discard up to maxStopLength-1 characters of real output, and in
+        // JSON mode the whole buffered reply.
+        if !stopped {
+            flushJSONBuffer()
+        }
         if !stopped && !heldStopTail.isEmpty {
             let (r, c) = enableThinking ? tracker.process(heldStopTail) : ("", heldStopTail)
             heldStopTail = ""

tests/SwiftLMTests/JSONModeStreamTests.swift

@@ -0,0 +1,139 @@
+import XCTest
+import Foundation
+import Hummingbird
+import HTTPTypes
+import MLXLMCommon
+import NIOCore
+@testable import SwiftLM
+
+// MARK: - JSON mode on the streaming chat path (#233)
+//
+// JSON-mode output is buffered so a hallucinated prefix object can be dropped before
+// anything reaches the client. The buffer used to be flushed only from inside a chunk,
+// once 32 chunks had arrived or a second `{` showed up, and never when the stream ended.
+// A reply that starts with `{` is held by the generator as a possible bare tool call and
+// arrives as one chunk, so in practice a JSON-mode stream sent nothing. These drive the
+// handler with a synthetic generation stream, no model needed.
+final class JSONModeStreamTests: XCTestCase {
+
+    // MARK: The prefix rule
+
+    func testFlatObjectIsUntouched() {
+        XCTAssertEqual(jsonModeFlushText("{\"ok\": true}"), "{\"ok\": true}")
+    }
+
+    /// The old rule cut at the second `{` wherever it was, so a nested object lost its
+    /// outer key and kept an unbalanced brace.
+    func testNestedObjectIsNotCutAtItsInnerBrace() {
+        let nested = "{\"user\": {\"name\": \"Bob\"}}"
+        XCTAssertEqual(jsonModeFlushText(nested), nested)
+        let list = "{\"items\": [{\"id\": 1}, {\"id\": 2}]}"
+        XCTAssertEqual(jsonModeFlushText(list), list)
+    }
+
+    /// The reason the buffer exists: a complete object before the reply is dropped.
+    func testClosedPrefixObjectIsDropped() {
+        XCTAssertEqual(jsonModeFlushText("{\"thought\": \"ok\"}\n{\"answer\": 1}"), "{\"answer\": 1}")
+        XCTAssertEqual(jsonModeFlushText("{}{\"answer\": 1}"), "{\"answer\": 1}")
+    }
+
+    func testOpenObjectAndNonObjectTextPassThroughTrimmed() {
+        XCTAssertEqual(jsonModeFlushText("  {\"a\": {\"b\": "), "{\"a\": {\"b\":")
+        XCTAssertEqual(jsonModeFlushText("```json\n{\"a\": 1}\n"), "```json\n{\"a\": 1}")
+    }
+
+    // MARK: The handler
+
+    private final class BodySink: @unchecked Sendable {
+        var data = Data()
+    }
+
+    private struct CollectingWriter: ResponseBodyWriter {
+        let sink: BodySink
+        mutating func write(_ buffer: ByteBuffer) async throws {
+            sink.data.append(contentsOf: buffer.readableBytesView)
+        }
+        consuming func finish(_ trailingHeaders: HTTPFields?) async throws {}
+    }
+
+    private func info(stopReason: GenerateStopReason, tokens: Int) -> Generation {
+        .info(GenerateCompletionInfo(
+            promptTokenCount: 3, generationTokenCount: tokens,
+            promptTime: 0, generationTime: 0, stopReason: stopReason))
+    }
+
+    private func chunks(_ pieces: [String]) -> [Generation] {
+        pieces.enumerated().map { .chunk($1, tokenId: $0 + 1) }
+    }
+
+    /// Streams `events` through the chat handler and returns the content deltas it put
+    /// on the wire, in order.
+    private func contentDeltas(for events: [Generation], jsonMode: Bool = true) async throws -> [String] {
+        let (stream, continuation) = AsyncStream<Generation>.makeStream()
+        for event in events { continuation.yield(event) }
+        continuation.finish()
+
+        let semaphore = AsyncSemaphore(limit: 1)
+        await semaphore.wait()
+        let slot = GenerationSlot(semaphore: semaphore)
+        let response = handleChatStreaming(
+            startGeneration: { (stream, nil) }, modelId: "test-model", stopSequences: [],
+            includeUsage: false, promptTokenCount: 3, jsonMode: jsonMode, tokenEcho: false,
+            slot: slot, stats: ServerStats(), genStart: Date(), prefillStart: Date(),
+            emitPrefillProgress: false)
+
+        let sink = BodySink()
+        try await response.body.write(CollectingWriter(sink: sink))
+        let body = String(decoding: sink.data, as: UTF8.self)
+        var deltas: [String] = []
+        for raw in body.components(separatedBy: "\n") {
+            let line = raw.trimmingCharacters(in: .whitespacesAndNewlines)
+            guard line.hasPrefix("data: {") else { continue }
+            let event = try XCTUnwrap(
+                JSONSerialization.jsonObject(with: Data(line.dropFirst(6).utf8)) as? [String: Any])
+            let choices = try XCTUnwrap(event["choices"] as? [[String: Any]])
+            if let delta = choices.first?["delta"] as? [String: Any],
+               let text = delta["content"] as? String {
+                deltas.append(text)
+            }
+        }
+        return deltas
+    }
+
+    /// The bug itself: a reply shorter than the flush threshold, delivered as one chunk,
+    /// never reached the client.
+    func testReplyShorterThanTheFlushThresholdIsSentWhenTheStreamEnds() async throws {
+        let deltas = try await contentDeltas(for: [
+            .chunk("{\"ok\": true}", tokenId: 1), info(stopReason: .stop, tokens: 5),
+        ])
+        XCTAssertEqual(deltas.joined(), "{\"ok\": true}")
+    }
+
+    /// A stream that ends without `.info` (an error, a cancelled decode) keeps what it had.
+    func testBufferIsSentWhenTheStreamEndsWithoutInfo() async throws {
+        let deltas = try await contentDeltas(for: [.chunk("{\"ok\": true}", tokenId: 1)])
+        XCTAssertEqual(deltas.joined(), "{\"ok\": true}")
+    }
+
+    /// With per-token delivery the second `{` of a nested object used to trigger the
+    /// prefix cut, so the client saw `{"name": "Bob"}}`.
+    func testNestedObjectArrivingPerTokenIsForwardedWhole() async throws {
+        let pieces = ["{\"", "user", "\": {\"", "name", "\": \"", "Bob", "\"}}"]
+        let deltas = try await contentDeltas(for: chunks(pieces) + [info(stopReason: .stop, tokens: pieces.count)])
+        XCTAssertEqual(deltas.joined(), "{\"user\": {\"name\": \"Bob\"}}")
+    }
+
+    /// What the buffer is for, unchanged: a closed object ahead of the reply is dropped.
+    func testHallucinatedPrefixObjectIsStillDropped() async throws {
+        let pieces = ["{\"", "note", "\": 1}", "{\"", "answer", "\": 2}"]
+        let deltas = try await contentDeltas(for: chunks(pieces) + [info(stopReason: .stop, tokens: pieces.count)])
+        XCTAssertEqual(deltas.joined(), "{\"answer\": 2}")
+    }
+
+    func testOutsideJSONModeTheChunkIsForwardedAsIs() async throws {
+        let deltas = try await contentDeltas(
+            for: [.chunk("{\"ok\": true}", tokenId: 1), info(stopReason: .stop, tokens: 5)],
+            jsonMode: false)
+        XCTAssertEqual(deltas, ["{\"ok\": true}"])
+    }
+}

tests/test-contract.sh

@@ -379,6 +379,34 @@ else
     fi
 fi
 
+# ── 13. JSON mode streams its reply ─────────────────────────────────────────
+# A JSON-mode reply is buffered so a hallucinated prefix can be dropped, and the buffer
+# was never flushed when the stream ended. The generator holds a `{`-led reply as a
+# possible tool call and delivers it as one chunk, so a JSON-mode stream sent nothing
+# while the non-streaming twin was fine (#233).
+log "Test 13: json_object mode streams the content it returns non-streaming"
+JSON_REQ='"model":"x","messages":[{"role":"user","content":"Return exactly this JSON and nothing else: {\"ok\": true}"}],"max_tokens":60,"temperature":0,"response_format":{"type":"json_object"}'
+NS_JSON=$(chat "{$JSON_REQ}" | python3 -c 'import json,sys; print(json.load(sys.stdin)["choices"][0]["message"].get("content") or "")')
+ST_JSON=$(curl -sf -N "$URL/v1/chat/completions" -H 'Content-Type: application/json' \
+    -d "{$JSON_REQ,\"stream\":true}" | python3 -c '
+import json,sys
+out=[]
+for line in sys.stdin:
+    line=line.strip()
+    if not line.startswith("data: ") or line=="data: [DONE]": continue
+    try: d=json.loads(line[6:])
+    except Exception: continue
+    for c in d.get("choices",[]):
+        out.append(c.get("delta",{}).get("content") or "")
+print("".join(out))')
+if [ -z "$ST_JSON" ]; then
+    fail "json_object stream carried no content (non-streaming returned: $(echo "$NS_JSON" | head -c 60))"
+elif [ "$ST_JSON" = "$NS_JSON" ]; then
+    pass "json_object stream matches the non-streaming content: $(echo "$ST_JSON" | head -c 60)"
+else
+    fail "json_object stream differs from non-streaming: '$(echo "$ST_JSON" | head -c 60)' vs '$(echo "$NS_JSON" | head -c 60)'"
+fi
+
 log "═══════════════════════════════════════"
 log "Results: $PASS passed, $FAIL failed, $SKIP skipped, $TOTAL total"
 log "═══════════════════════════════════════"