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 "═══════════════════════════════════════"