From 3803b75306270673da44ece30addc33ec83c3f7e Mon Sep 17 00:00:00 2001 From: Simba Zhang Date: Sun, 11 Oct 2026 12:18:55 -0700 Subject: [PATCH 1/2] fix: stop non-streaming generation at the first stop sequence The streaming handlers end generation when a stop sequence matches. The two non-streaming handlers (chat and text) kept generating to EOS or max_tokens and trimmed the text afterwards, so the same request spent more compute and reported a larger completion_tokens when it was not streamed. Contract tests: finish_reason is "length" on the chat and streaming-text paths (Test 11), and completion_tokens after a stop sequence agree between streaming and non-streaming for chat (Test 12) and drop for non-streaming completions (Test 12b). Tests 12 and 12b fail on main (38 vs 14 and 60 vs 60 tokens). Co-Authored-By: Claude Sonnet 5.5 --- Sources/SwiftLM/Server.swift | 21 ++++++++- tests/test-contract.sh | 84 ++++++++++++++++++++++++++++++++++++ 2 files changed, 103 insertions(+), 2 deletions(-) diff --git a/Sources/SwiftLM/Server.swift b/Sources/SwiftLM/Server.swift index 927f75e..77fbd8e 100644 --- a/Sources/SwiftLM/Server.swift +++ b/Sources/SwiftLM/Server.swift @@ -3444,11 +3444,21 @@ func handleChatNonStreaming( var tcIndex = 0 var generationStopReason: GenerateStopReason = .stop var firstToken = true - for await generation in stream { + let stopWindow = stopScanWindow(stopSequences) + generationLoop: for await generation in stream { switch generation { case .chunk(let text, _): fullText += text completionTokenCount += 1 + // End generation at the first stop sequence, as the streaming path does. The + // text is trimmed below either way; continuing only spends compute and bills + // the client for tokens that are discarded. `.info` never arrives on this + // exit, so the count stays the chunk count. + if !stopSequences.isEmpty, + checkStopSequences(fullText, stopSequences: stopSequences, lookback: stopWindow + text.count) != nil + { + break generationLoop + } // GPU yield: prevent Metal from starving macOS WindowServer if completionTokenCount % 8 == 0 { try? await Task.sleep(for: .microseconds(50)) @@ -3856,11 +3866,18 @@ func handleTextNonStreaming( ) async throws -> Response { var fullText = "" var completionTokenCount = 0 - for await generation in stream { + let stopWindow = stopScanWindow(stopSequences) + generationLoop: for await generation in stream { switch generation { case .chunk(let text, _): fullText += text completionTokenCount += 1 + // See handleChatNonStreaming: stop generating at the first stop sequence. + if !stopSequences.isEmpty, + checkStopSequences(fullText, stopSequences: stopSequences, lookback: stopWindow + text.count) != nil + { + break generationLoop + } // GPU yield: prevent Metal from starving macOS WindowServer if completionTokenCount % 8 == 0 { try? await Task.sleep(for: .microseconds(50)) diff --git a/tests/test-contract.sh b/tests/test-contract.sh index 35dd6d9..ce614f4 100755 --- a/tests/test-contract.sh +++ b/tests/test-contract.sh @@ -291,6 +291,90 @@ else fail "got HTTP $CODE with error '$ERR_TYPE' (expected 400 invalid_request_error stop)" fi +# ── 11. finish_reason says when the model was cut off ───────────────────────── +# A client cannot tell a truncated completion from a finished one except through +# `finish_reason`, so each response path must report "length" when max_tokens ended it. +# Found by review: the non-streaming /v1/completions path hardcoded "stop" (#227). +# The non-streaming /v1/completions case is asserted alongside that fix. +SSE_FINISH=' +import json,sys +reason=None +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 ch in d.get("choices",[]): + if ch.get("finish_reason"): reason=ch["finish_reason"] +print(reason)' +LONG_ASK="Write a very long story about a dragon, at least five hundred words." + +log "Test 11: finish_reason is \"length\" when max_tokens ends a chat completion" +R=$(chat "{\"model\":\"x\",\"messages\":[{\"role\":\"user\",\"content\":\"$LONG_ASK\"}],\"max_tokens\":5,\"temperature\":0}" \ + | python3 -c 'import json,sys; print(json.load(sys.stdin)["choices"][0]["finish_reason"])') +[ "$R" = "length" ] && pass "chat, non-streaming: finish_reason=length" || fail "chat, non-streaming: finish_reason=$R (expected length)" + +R=$(curl -sf -N "$URL/v1/chat/completions" -H 'Content-Type: application/json' \ + -d "{\"model\":\"x\",\"messages\":[{\"role\":\"user\",\"content\":\"$LONG_ASK\"}],\"max_tokens\":5,\"temperature\":0,\"stream\":true}" | python3 -c "$SSE_FINISH") +[ "$R" = "length" ] && pass "chat, streaming: finish_reason=length" || fail "chat, streaming: finish_reason=$R (expected length)" + +R=$(curl -sf -N "$URL/v1/completions" -H 'Content-Type: application/json' \ + -d "{\"model\":\"x\",\"prompt\":\"$LONG_ASK\",\"max_tokens\":5,\"temperature\":0,\"stream\":true}" | python3 -c "$SSE_FINISH") +[ "$R" = "length" ] && pass "completions, streaming: finish_reason=length" || fail "completions, streaming: finish_reason=$R (expected length)" + +# ── 12. A stop sequence ends generation, and usage only counts what was generated ── +# The streaming paths stop at the match. The non-streaming ones used to keep generating +# to max_tokens or EOS and cut the text afterwards, so the same request billed far more +# completion_tokens (and spent the compute) when it was not streamed. +log "Test 12: completion_tokens after a stop sequence agree between streaming and non-streaming" +STOP_ASK="Write the numbers one to twenty as English words separated by commas, nothing else." +STOP_WORD="Five" +BASE_REQ="\"model\":\"x\",\"messages\":[{\"role\":\"user\",\"content\":\"$STOP_ASK\"}],\"max_tokens\":120,\"temperature\":0" +# The assertion is only meaningful if the stop sequence would actually fire, so check +# the unconstrained reply contains it first; that is a model capability, not a contract. +BASE_TEXT=$(chat "{$BASE_REQ}" | python3 -c 'import json,sys; print(json.load(sys.stdin)["choices"][0]["message"].get("content") or "")') +if ! echo "$BASE_TEXT" | grep -q "$STOP_WORD"; then + skip "model's reply never contained '$STOP_WORD', so the stop sequence cannot fire" +else + NS_TOKENS=$(chat "{$BASE_REQ,\"stop\":[\"$STOP_WORD\"]}" | python3 -c 'import json,sys; print(json.load(sys.stdin)["usage"]["completion_tokens"])') + ST_TOKENS=$(curl -sf -N "$URL/v1/chat/completions" -H 'Content-Type: application/json' \ + -d "{$BASE_REQ,\"stop\":[\"$STOP_WORD\"],\"stream\":true,\"stream_options\":{\"include_usage\":true}}" | python3 -c ' +import json,sys +n=None +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 + if d.get("usage"): n=d["usage"]["completion_tokens"] +print(n)') + # The streamed count is a chunk count and can trail the real one by a few buffered tokens. + if [ -n "$NS_TOKENS" ] && [ "$ST_TOKENS" != "None" ] && [ "$NS_TOKENS" -le $((ST_TOKENS + 4)) ]; then + pass "stop at '$STOP_WORD': non-streaming $NS_TOKENS vs streaming $ST_TOKENS completion tokens" + else + fail "stop at '$STOP_WORD': non-streaming reports $NS_TOKENS completion tokens, streaming $ST_TOKENS — the non-streaming path kept generating past the stop sequence" + fi +fi + +log "Test 12b: /v1/completions stops generating at a stop sequence (non-streaming)" +COMP_PROMPT="Numbers in words: One, Two, Three, Four, Five, Six, Seven, Eight, Nine, Ten, Eleven, Twelve. Again:" +comp_req() { curl -sf "$URL/v1/completions" -H 'Content-Type: application/json' -d "$1"; } +COMP_BASE=$(comp_req "{\"model\":\"x\",\"prompt\":\"$COMP_PROMPT\",\"max_tokens\":60,\"temperature\":0}") +COMP_BASE_TEXT=$(echo "$COMP_BASE" | python3 -c 'import json,sys; print(json.load(sys.stdin)["choices"][0]["text"])') +if ! echo "$COMP_BASE_TEXT" | grep -q "Five"; then + skip "the completion never contained 'Five', so the stop sequence cannot fire" +else + COMP_STOP=$(comp_req "{\"model\":\"x\",\"prompt\":\"$COMP_PROMPT\",\"max_tokens\":60,\"temperature\":0,\"stop\":[\"Five\"]}") + COMP_BASE_N=$(echo "$COMP_BASE" | python3 -c 'import json,sys; print(json.load(sys.stdin)["usage"]["completion_tokens"])') + COMP_STOP_N=$(echo "$COMP_STOP" | python3 -c 'import json,sys; print(json.load(sys.stdin)["usage"]["completion_tokens"])') + COMP_STOP_TEXT=$(echo "$COMP_STOP" | python3 -c 'import json,sys; print(json.load(sys.stdin)["choices"][0]["text"])') + if [ "$COMP_STOP_N" -lt "$COMP_BASE_N" ] && ! echo "$COMP_STOP_TEXT" | grep -q "Five"; then + pass "completions stop at 'Five': $COMP_STOP_N tokens vs $COMP_BASE_N without the stop" + else + fail "completions stop at 'Five': $COMP_STOP_N tokens vs $COMP_BASE_N without the stop (text: $(echo "$COMP_STOP_TEXT" | head -c 60))" + fi +fi + log "═══════════════════════════════════════" log "Results: $PASS passed, $FAIL failed, $SKIP skipped, $TOTAL total" log "═══════════════════════════════════════" From 0cfa3357974e3a510b432b00f98c35f3bf28647f Mon Sep 17 00:00:00 2001 From: Hamed Shams Date: Sun, 11 Oct 2026 22:42:06 +0300 Subject: [PATCH 2/2] fix: report finish_reason "length" from non-streaming /v1/completions handleTextNonStreaming took generationTokenCount from the final .info event but ignored info.stopReason and hardcoded finish_reason to "stop", so a completion cut off by max_tokens looked like a natural end. It now uses the same mapping as handleChatNonStreaming; the streaming text path already had it. A client stop sequence still reports "stop" whatever ended generation. Contract Test 11 gains the non-streaming /v1/completions case it listed as pending, and TextCompletionFinishReasonTests drives the handler with a synthetic stream. Closes #227. --- Sources/SwiftLM/Server.swift | 14 ++- .../TextCompletionFinishReasonTests.swift | 98 +++++++++++++++++++ tests/test-contract.sh | 6 +- 3 files changed, 116 insertions(+), 2 deletions(-) create mode 100644 tests/SwiftLMTests/TextCompletionFinishReasonTests.swift diff --git a/Sources/SwiftLM/Server.swift b/Sources/SwiftLM/Server.swift index 77fbd8e..ba0e42f 100644 --- a/Sources/SwiftLM/Server.swift +++ b/Sources/SwiftLM/Server.swift @@ -3866,6 +3866,7 @@ func handleTextNonStreaming( ) async throws -> Response { var fullText = "" var completionTokenCount = 0 + var generationStopReason: GenerateStopReason = .stop let stopWindow = stopScanWindow(stopSequences) generationLoop: for await generation in stream { switch generation { @@ -3883,6 +3884,7 @@ func handleTextNonStreaming( try? await Task.sleep(for: .microseconds(50)) } case .info(let info): + generationStopReason = info.stopReason completionTokenCount = info.generationTokenCount // real count; see handleChatStreaming case .toolCall, .rejectedToolCall: break @@ -3892,7 +3894,17 @@ func handleTextNonStreaming( await stats.requestFinished(tokens: completionTokenCount, duration: duration) slot.release() - var finishReason = "stop" + // Same mapping as handleChatNonStreaming (#227). This path used to hardcode "stop", + // so a completion cut off by max_tokens was reported as a natural end; the + // streaming text path and both chat paths already said "length". A stop sequence + // still wins below, whatever ended generation. + var finishReason: String + switch generationStopReason { + case .length: + finishReason = "length" + default: + finishReason = "stop" + } if let (trimmedText, _) = checkStopSequences(fullText, stopSequences: stopSequences) { fullText = trimmedText finishReason = "stop" diff --git a/tests/SwiftLMTests/TextCompletionFinishReasonTests.swift b/tests/SwiftLMTests/TextCompletionFinishReasonTests.swift new file mode 100644 index 0000000..1069bc3 --- /dev/null +++ b/tests/SwiftLMTests/TextCompletionFinishReasonTests.swift @@ -0,0 +1,98 @@ +import XCTest +import Foundation +import Hummingbird +import HTTPTypes +import MLXLMCommon +import NIOCore +@testable import SwiftLM + +// MARK: - `finish_reason` on the non-streaming /v1/completions path +// +// The handler used to hardcode "stop": it read `info.generationTokenCount` and ignored +// `info.stopReason`, so a completion cut off by `max_tokens` was reported as a natural +// stop. The chat handler and the *streaming* text handler both already mapped `.length`. +// These drive the handler with a synthetic generation stream, no model needed. +final class TextCompletionFinishReasonTests: XCTestCase { + + 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 {} + } + + /// Runs the non-streaming text handler over `events` and decodes the JSON it answers with. + private func respond(to events: [Generation], stop: [String] = []) async throws -> [String: Any] { + 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 = try await handleTextNonStreaming( + stream: stream, modelId: "test-model", stopSequences: stop, + promptTokenCount: 3, slot: slot, stats: ServerStats(), genStart: Date()) + + let sink = BodySink() + try await response.body.write(CollectingWriter(sink: sink)) + return try XCTUnwrap(JSONSerialization.jsonObject(with: sink.data) as? [String: Any]) + } + + private func info(stopReason: GenerateStopReason, tokens: Int) -> Generation { + .info(GenerateCompletionInfo( + promptTokenCount: 3, generationTokenCount: tokens, + promptTime: 0, generationTime: 0, stopReason: stopReason)) + } + + private func finishReason(_ body: [String: Any]) throws -> String { + let choices = try XCTUnwrap(body["choices"] as? [[String: Any]]) + return try XCTUnwrap(choices.first?["finish_reason"] as? String) + } + + func testMaxTokensReportsLength() async throws { + let body = try await respond(to: [ + .chunk("one", tokenId: 1), .chunk(" two", tokenId: 2), + info(stopReason: .length, tokens: 2), + ]) + XCTAssertEqual(try finishReason(body), "length") + let choices = try XCTUnwrap(body["choices"] as? [[String: Any]]) + XCTAssertEqual(choices.first?["text"] as? String, "one two", "the text itself is untouched") + } + + func testEndOfSequenceReportsStop() async throws { + let body = try await respond(to: [ + .chunk("done", tokenId: 1), info(stopReason: .stop, tokens: 1), + ]) + XCTAssertEqual(try finishReason(body), "stop") + } + + /// A client stop sequence is a stop even when the generator ran into its token limit + /// on the same chunk: the text was cut at the stop, so that is what happened. + func testClientStopSequenceReportsStop() async throws { + let body = try await respond(to: [ + .chunk("alpha", tokenId: 1), .chunk("<|end|>beta", tokenId: 2), + info(stopReason: .length, tokens: 2), + ], stop: ["<|end|>"]) + XCTAssertEqual(try finishReason(body), "stop") + let choices = try XCTUnwrap(body["choices"] as? [[String: Any]]) + XCTAssertEqual(choices.first?["text"] as? String, "alpha") + } + + /// The `.info` event carries the real token count (the chunk count undercounts when + /// the decoder buffers), and that must survive the finish_reason change. + func testUsageStillComesFromTheInfoEvent() async throws { + let body = try await respond(to: [ + .chunk("x", tokenId: 1), info(stopReason: .length, tokens: 7), + ]) + let usage = try XCTUnwrap(body["usage"] as? [String: Any]) + XCTAssertEqual(usage["completion_tokens"] as? Int, 7) + XCTAssertEqual(usage["prompt_tokens"] as? Int, 3) + XCTAssertEqual(usage["total_tokens"] as? Int, 10) + } +} diff --git a/tests/test-contract.sh b/tests/test-contract.sh index ce614f4..1f207a1 100755 --- a/tests/test-contract.sh +++ b/tests/test-contract.sh @@ -295,7 +295,6 @@ fi # A client cannot tell a truncated completion from a finished one except through # `finish_reason`, so each response path must report "length" when max_tokens ended it. # Found by review: the non-streaming /v1/completions path hardcoded "stop" (#227). -# The non-streaming /v1/completions case is asserted alongside that fix. SSE_FINISH=' import json,sys reason=None @@ -322,6 +321,11 @@ R=$(curl -sf -N "$URL/v1/completions" -H 'Content-Type: application/json' \ -d "{\"model\":\"x\",\"prompt\":\"$LONG_ASK\",\"max_tokens\":5,\"temperature\":0,\"stream\":true}" | python3 -c "$SSE_FINISH") [ "$R" = "length" ] && pass "completions, streaming: finish_reason=length" || fail "completions, streaming: finish_reason=$R (expected length)" +R=$(curl -sf "$URL/v1/completions" -H 'Content-Type: application/json' \ + -d "{\"model\":\"x\",\"prompt\":\"$LONG_ASK\",\"max_tokens\":5,\"temperature\":0}" \ + | python3 -c 'import json,sys; print(json.load(sys.stdin)["choices"][0]["finish_reason"])') +[ "$R" = "length" ] && pass "completions, non-streaming: finish_reason=length" || fail "completions, non-streaming: finish_reason=$R (expected length; #227)" + # ── 12. A stop sequence ends generation, and usage only counts what was generated ── # The streaming paths stop at the match. The non-streaming ones used to keep generating # to max_tokens or EOS and cut the text afterwards, so the same request billed far more