diff --git a/Sources/SwiftLM/Server.swift b/Sources/SwiftLM/Server.swift
index e5cd29a..927f75e 100644
--- a/Sources/SwiftLM/Server.swift
+++ b/Sources/SwiftLM/Server.swift
@@ -564,6 +564,9 @@ struct MLXServer: AsyncParsableCommand {
@Flag(name: .long, help: "Enable thinking/reasoning mode (Qwen3.5 etc). Default: disabled")
var thinking: Bool = false
+ @Option(name: .long, help: "Hard cap on reasoning tokens per request (overridable per request with thinking_budget). After N thinking tokens the server closes the thinking block and the model answers. 0 closes it immediately. Needs thinking enabled and --max-tokens (or the request's max_tokens) large enough to hold the budget plus the answer. Default: no cap")
+ var thinkingBudget: Int?
+
@Flag(name: .long, help: "Do not echo generated tokens to stdout as they stream. Request log lines are unaffected.")
var noTokenEcho: Bool = false
@@ -628,6 +631,9 @@ struct MLXServer: AsyncParsableCommand {
var mtpAssistantModel: String?
func validate() throws {
+ if let thinkingBudget, thinkingBudget < 0 {
+ throw ValidationError("--thinking-budget must be 0 or more.")
+ }
if vision && noVision {
throw ValidationError("--vision and --no-vision are mutually exclusive.")
}
@@ -1382,6 +1388,7 @@ struct MLXServer: AsyncParsableCommand {
minP: self.minP,
repeatPenalty: self.repeatPenalty,
thinking: self.thinking,
+ thinkingBudget: self.thinkingBudget,
maxPromptTokens: self.maxPromptTokens,
promptCacheEntries: self.promptCacheEntries,
tokenEcho: !self.noTokenEcho,
@@ -1699,6 +1706,8 @@ struct ServerConfig: Sendable {
let minP: Float?
let repeatPenalty: Float?
let thinking: Bool
+ /// `--thinking-budget`: default cap on reasoning tokens. nil = no cap.
+ let thinkingBudget: Int?
/// `--max-prompt-tokens`: longer prompts are rejected before prefill. nil = no limit.
let maxPromptTokens: Int?
/// `--prompt-cache-entries`: LRU capacity of the prompt cache.
@@ -2555,6 +2564,26 @@ func handleChatCompletion(
}
let promptTokens = lmInput.text.tokens.asArray(Int.self)
+ // ── Reasoning budget: per-request `thinking_budget` wins over `--thinking-budget` ──
+ let budgetComponents: GenerationComponents?
+ do {
+ budgetComponents = try thinkingBudgetComponents(
+ budget: chatReq.thinkingBudget ?? config.thinkingBudget, enableThinking: enableThinking,
+ reasoning: await container.configuration.reasoningConfig,
+ tokenizer: await container.tokenizer, parameters: params)
+ } catch let error as ThinkingBudgetError {
+ print("srv slot_reject: id 0 | thinking budget: \(error.errorDescription ?? "invalid")")
+ slot.release()
+ await stats.requestFinished(tokens: 0, duration: 0)
+ return Response(
+ status: .badRequest, headers: jsonHeaders(),
+ body: .init(byteBuffer: ByteBuffer(string: thinkingBudgetErrorBody(error))))
+ }
+ let genComponents = budgetComponents ?? GenerationComponents()
+ if budgetComponents != nil, draftModelRef != nil || config.mtp || dflashModel != nil {
+ print("[SwiftLM] Warning: the thinking budget is not applied with speculative, MTP or DFlash decoding.")
+ }
+
// ── Issue #108: does the template leave the thinking block open? ──
// Qwen3/3.5/3.6 templates append `\n` to the generation prompt when
// enable_thinking is true, so the model emits reasoning with no opening tag of
@@ -2729,7 +2758,7 @@ func handleChatCompletion(
}
stream = try MLXLMCommon.generate(
input: LMInput(text: sliceText(lmInput.text, from: boundary)),
- cache: cache, parameters: params, context: context)
+ cache: cache, parameters: params, context: context, components: genComponents)
} else if let boundary = rotatingBoundary {
var start = await promptCache.restore(newTokens: promptTokens, into: cache) ?? 0
if start >= promptLength {
@@ -2753,7 +2782,7 @@ func handleChatCompletion(
}
stream = try MLXLMCommon.generate(
input: LMInput(text: sliceText(lmInput.text, from: start)),
- cache: cache, parameters: params, context: context)
+ cache: cache, parameters: params, context: context, components: genComponents)
} else if let draftRef = draftModelRef {
// Speculative decoding path: draft model generates candidates, main model verifies.
// Bypass prompt cache to avoid draft/main KV drift on partial-match restores.
@@ -2781,7 +2810,8 @@ func handleChatCompletion(
)
} else {
stream = try MLXLMCommon.generate(
- input: trimmedInput, cache: cache, parameters: params, context: context
+ input: trimmedInput, cache: cache, parameters: params, context: context,
+ components: genComponents
)
}
} else {
@@ -2792,7 +2822,8 @@ func handleChatCompletion(
)
} else {
stream = try MLXLMCommon.generate(
- input: lmInput, cache: cache, parameters: params, context: context
+ input: lmInput, cache: cache, parameters: params, context: context,
+ components: genComponents
)
}
}
@@ -3479,10 +3510,10 @@ func handleChatNonStreaming(
print("srv debug: pre-extract fullText=\(fullText.prefix(40).debugDescription)")
let (extracted, remaining) = extractThinkingBlock(from: fullText, alreadyOpen: thinkingPreOpened)
print("srv debug: extracted=\(extracted != nil ? "true" : "false"), remaining_len=\(remaining.count)")
- if let extracted {
- reasoningContent = extracted
- responseContent = remaining
- }
+ // `remaining` is already `fullText` when no tag was found, and it drops the closing
+ // tag when the reasoning was empty (`` first, e.g. a zero thinking budget).
+ reasoningContent = extracted
+ responseContent = remaining
}
// ── JSON mode validation ──
@@ -4142,6 +4173,41 @@ func promptTooLongBody(promptTokens: Int, limit: Int?) -> String? {
return "{\"error\":{\"message\":\"This model's maximum context length is \(limit) tokens. However, your messages resulted in \(promptTokens) tokens.\",\"type\":\"invalid_request_error\",\"code\":\"context_length_exceeded\"}}"
}
+/// Components that cap reasoning at `budget` tokens, or nil when no cap applies (no budget,
+/// thinking off, or a model that declares no reasoning protocol). Throws `ThinkingBudgetError`
+/// for a budget the request cannot honour, e.g. a `max_tokens` too small for the budget plus
+/// the answer.
+///
+/// The transition is `.immediate`: the model's own `` closes the span with no extra
+/// text. Qwen3.5 declares no validated hard-budget transition, so this is an opt-in.
+func thinkingBudgetComponents(
+ budget: Int?, enableThinking: Bool, reasoning: ReasoningConfig?,
+ tokenizer: any MLXLMCommon.Tokenizer, parameters: GenerateParameters
+) throws -> GenerationComponents? {
+ guard let budget, enableThinking, let reasoning else { return nil }
+ let configuration = try ThinkingBudgetConfiguration(
+ maximumTokenCount: budget, transitionOverride: .immediate)
+ let components = try GenerationComponents().applyingThinkingBudget(
+ configuration, reasoning: reasoning, tokenizer: tokenizer)
+ try components.validate(parameters: parameters)
+ return components
+}
+
+/// OpenAI-shaped 400 body for a reasoning budget the request cannot honour.
+func thinkingBudgetErrorBody(_ error: ThinkingBudgetError) -> String {
+ let payload: [String: Any] = ["error": [
+ "message": error.errorDescription ?? "Invalid thinking budget.",
+ "type": "invalid_request_error",
+ "code": "invalid_thinking_budget",
+ ]]
+ guard let data = try? JSONSerialization.data(withJSONObject: payload),
+ let json = String(data: data, encoding: .utf8)
+ else {
+ return "{\"error\":{\"message\":\"Invalid thinking budget.\",\"type\":\"invalid_request_error\",\"code\":\"invalid_thinking_budget\"}}"
+ }
+ return json
+}
+
/// Releases the slot and returns a 400 when the prompt is over `--max-prompt-tokens`,
/// else nil. Runs before prefill, so a rejected prompt costs no GPU work.
func rejectLongPrompt(promptTokens: Int, limit: Int?, slot: GenerationSlot, stats: ServerStats) async -> Response? {
@@ -4556,6 +4622,8 @@ struct ChatCompletionRequest: Decodable {
let responseFormat: ResponseFormat?
/// Per-request Jinja template kwargs (e.g. {"enable_thinking": false} for Qwen3/Qwen3.5)
let chatTemplateKwargs: [String: Bool]?
+ /// Per-request reasoning-token cap; overrides `--thinking-budget`.
+ let thinkingBudget: Int?
/// Top-level thinking override emitted by Aegis-AI gateway
let enableThinking: Bool?
/// Number of bits for native MLX quantized KV cache (nil = no quantization).
@@ -4577,6 +4645,7 @@ struct ChatCompletionRequest: Decodable {
case responseFormat = "response_format"
case chatTemplateKwargs = "chat_template_kwargs"
case enableThinking = "enable_thinking"
+ case thinkingBudget = "thinking_budget"
case kvBits = "kv_bits"
}
}
diff --git a/tests/SwiftLMTests/ThinkingBudgetTests.swift b/tests/SwiftLMTests/ThinkingBudgetTests.swift
new file mode 100644
index 0000000..82669d0
--- /dev/null
+++ b/tests/SwiftLMTests/ThinkingBudgetTests.swift
@@ -0,0 +1,136 @@
+import XCTest
+import Foundation
+import MLXLMCommon
+@testable import SwiftLM
+
+/// `--thinking-budget` / `thinking_budget` cap a request's reasoning tokens. A budget the
+/// request cannot honour is a 400 before any GPU work, and a request that does not think is
+/// left alone.
+final class ThinkingBudgetTests: XCTestCase {
+
+ private let tokenizer = ByteTokenizer()
+ /// Qwen-style tags that the byte tokenizer can spell; Qwen3.5 declares no hard-budget
+ /// transition of its own, so the server supplies `.immediate`.
+ private let reasoning = ReasoningConfig(
+ startDelimiter: "", endDelimiter: "",
+ promptStrategy: .templateFlag(key: "enable_thinking", defaultOn: true),
+ isSpecialToken: false)
+
+ private func components(
+ budget: Int?, thinking: Bool = true, reasoning: ReasoningConfig?? = .none, maxTokens: Int? = 4096
+ ) throws -> GenerationComponents? {
+ try thinkingBudgetComponents(
+ budget: budget, enableThinking: thinking, reasoning: reasoning ?? self.reasoning,
+ tokenizer: tokenizer, parameters: GenerateParameters(maxTokens: maxTokens))
+ }
+
+ // MARK: - When a budget applies
+
+ func testNoBudgetMeansNoComponents() throws {
+ XCTAssertNil(try components(budget: nil))
+ }
+
+ func testThinkingOffLeavesTheRequestAlone() throws {
+ XCTAssertNil(try components(budget: 64, thinking: false))
+ }
+
+ func testModelWithoutAReasoningProtocolIsLeftAlone() throws {
+ XCTAssertNil(try components(budget: 64, reasoning: .some(nil)))
+ }
+
+ func testBudgetBuildsComponents() throws {
+ XCTAssertNotNil(try components(budget: 64))
+ }
+
+ func testZeroBudgetIsAllowed() throws {
+ XCTAssertNotNil(try components(budget: 0))
+ }
+
+ // MARK: - Budgets the request cannot honour
+
+ func testMaxTokensTooSmallForBudgetAndAnswerIsRejected() {
+ XCTAssertThrowsError(try components(budget: 1000, maxTokens: 500)) { error in
+ guard case .insufficientGenerationTokenLimit(_, let actual) = error as? ThinkingBudgetError else {
+ return XCTFail("unexpected error \(error)")
+ }
+ XCTAssertEqual(actual, 500)
+ }
+ }
+
+ func testUnlimitedMaxTokensAlwaysFits() throws {
+ XCTAssertNotNil(try components(budget: 100_000, maxTokens: nil))
+ }
+
+ func testNegativeBudgetIsRejected() {
+ XCTAssertThrowsError(try components(budget: -1)) { error in
+ XCTAssertEqual(error as? ThinkingBudgetError, .invalidMaximumTokenCount(-1))
+ }
+ }
+
+ // MARK: - 400 body
+
+ func testErrorBodyIsOpenAIShaped() throws {
+ let body = thinkingBudgetErrorBody(.insufficientGenerationTokenLimit(required: 1100, actual: 500))
+ let json = try XCTUnwrap(JSONSerialization.jsonObject(with: Data(body.utf8)) as? [String: Any])
+ let error = try XCTUnwrap(json["error"] as? [String: Any])
+ XCTAssertEqual(error["type"] as? String, "invalid_request_error")
+ XCTAssertEqual(error["code"] as? String, "invalid_thinking_budget")
+ XCTAssertFalse((error["message"] as? String ?? "").isEmpty)
+ }
+
+ // MARK: - Empty reasoning (budget 0)
+
+ /// A zero budget makes the model emit `` first. The handler keeps `remaining`
+ /// whether or not any reasoning was extracted, so this contract must hold: no reasoning,
+ /// and no closing tag left in the content.
+ func testEmptyReasoningDropsTheClosingTag() {
+ let (reasoning, content) = extractThinkingBlock(from: "\n\n9 sheep are left.", alreadyOpen: true)
+ XCTAssertNil(reasoning)
+ XCTAssertEqual(content, "9 sheep are left.")
+ }
+
+ func testNoTagsLeavesTheTextUntouched() {
+ let (reasoning, content) = extractThinkingBlock(from: "plain answer", alreadyOpen: false)
+ XCTAssertNil(reasoning)
+ XCTAssertEqual(content, "plain answer")
+ }
+
+ // MARK: - Flag and request parsing
+
+ func testFlagParsesAndDefaultsToNil() throws {
+ XCTAssertEqual(try MLXServer.parse(["--model", "m", "--thinking-budget", "2048"]).thinkingBudget, 2048)
+ XCTAssertNil(try MLXServer.parse(["--model", "m"]).thinkingBudget)
+ }
+
+ func testNegativeFlagIsRejected() {
+ XCTAssertThrowsError(try MLXServer.parse(["--model", "m", "--thinking-budget", "-5"]))
+ }
+
+ func testRequestFieldDecodes() throws {
+ let json = #"{"model":"m","messages":[{"role":"user","content":"hi"}],"thinking_budget":128}"#
+ let request = try JSONDecoder().decode(ChatCompletionRequest.self, from: Data(json.utf8))
+ XCTAssertEqual(request.thinkingBudget, 128)
+ }
+}
+
+private struct ByteTokenizer: MLXLMCommon.Tokenizer {
+ func encode(text: String, addSpecialTokens: Bool) -> [Int] { text.utf8.map(Int.init) }
+ func decode(tokenIds: [Int], skipSpecialTokens: Bool) -> String {
+ String(decoding: tokenIds.map(UInt8.init), as: UTF8.self)
+ }
+ func convertTokenToId(_ token: String) -> Int? {
+ guard token.utf8.count == 1 else { return nil }
+ return token.utf8.first.map(Int.init)
+ }
+ func convertIdToToken(_ id: Int) -> String? {
+ guard (0 ... 255).contains(id) else { return nil }
+ return String(decoding: [UInt8(id)], as: UTF8.self)
+ }
+ var bosToken: String? { nil }
+ var eosToken: String? { nil }
+ var unknownToken: String? { nil }
+ func applyChatTemplate(
+ messages: [[String: any Sendable]], tools: [[String: any Sendable]]?,
+ additionalContext: [String: any Sendable]?
+ ) throws -> [Int] { [] }
+}