From c069c173691c4854ca9a281948af2729d1cb3ca8 Mon Sep 17 00:00:00 2001 From: Hamed Shams Date: Fri, 16 Oct 2026 08:58:57 +0000 Subject: [PATCH] fix: flush the JSON-mode buffer when a stream ends, keep nested objects whole handleChatStreaming buffers JSON-mode output so a hallucinated prefix object can be dropped, and only flushed the buffer from inside a chunk, once 32 chunks had arrived or a second "{" showed up. Nothing flushed it when the stream ended. The generator holds a "{"-led reply as a possible bare tool call and hands it over as one chunk, so the thresholds never tripped and a JSON-mode stream carried no content at all. The flush is now one function, called at the threshold as before and at the end of the stream, in the .info case and in the post-loop safety net. The prefix strip drops only a complete object that precedes a second one; the old rule cut at any second "{" and turned {"user": {"name": "Bob"}} into {"name": "Bob"}}. JSONModeStreamTests drives the handler with a synthetic stream, and contract Test 13 checks that a JSON-mode stream assembles to the non-streaming reply. Closes #233. --- Sources/SwiftLM/Server.swift | 87 ++++++++---- tests/SwiftLMTests/JSONModeStreamTests.swift | 139 +++++++++++++++++++ tests/test-contract.sh | 28 ++++ 3 files changed, 228 insertions(+), 26 deletions(-) create mode 100644 tests/SwiftLMTests/JSONModeStreamTests.swift diff --git a/Sources/SwiftLM/Server.swift b/Sources/SwiftLM/Server.swift index 3784bfe..5272aa8 100644 --- a/Sources/SwiftLM/Server.swift +++ b/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, (() 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 = "" diff --git a/tests/SwiftLMTests/JSONModeStreamTests.swift b/tests/SwiftLMTests/JSONModeStreamTests.swift new file mode 100644 index 0000000..592a21a --- /dev/null +++ b/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.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}"]) + } +} diff --git a/tests/test-contract.sh b/tests/test-contract.sh index 1f207a1..25fbb74 100755 --- a/tests/test-contract.sh +++ b/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 "═══════════════════════════════════════"