From 4e63909d7753c31909da300ed2ab86e7a6c8578e Mon Sep 17 00:00:00 2001 From: CodeAndCanvas728 Date: Sun, 4 Oct 2026 15:07:06 +0200 Subject: [PATCH] feat: --max-prompt-tokens rejects over-long prompts before prefill An opt-in limit: a prompt longer than N tokens gets an OpenAI-shaped 400 (code context_length_exceeded) on /v1/chat/completions and /v1/completions, streaming or not, before any prefill work. The slot and stats are released. Agent clients recognise that code and compact their history instead of waiting minutes on a prefill and retrying it. It is a separate option because --ctx-size is documented as a sliding window; rejecting at --ctx-size would change what that flag does. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_014J4GsqznbKNX8tcKzNRxEr --- README.md | 1 + Sources/SwiftLM/Server.swift | 36 ++++++++++++++++++ .../SwiftLMTests/PromptLengthGuardTests.swift | 38 +++++++++++++++++++ 3 files changed, 75 insertions(+) create mode 100644 tests/SwiftLMTests/PromptLengthGuardTests.swift diff --git a/README.md b/README.md index 6bf11ab..5d226c5 100644 --- a/README.md +++ b/README.md @@ -641,6 +641,7 @@ curl http://localhost:5413/v1/chat/completions \ | `--num-draft-tokens` | `4` | Tokens per speculation round. Auto-capped to 1 when combined with `--stream-experts`. | | `--dflash` | `false` | Enable DFlash block-diffusion speculative decoding. Requires a compatible DFlash draft model | | `--dflash-block-size`| (auto) | Number of tokens per DFlash draft block. Defaults to draft model config | +| `--max-prompt-tokens` | (none) | Reject prompts longer than this many tokens with an OpenAI-style 400 (`context_length_exceeded`) before prefill, so agent clients compact instead of waiting on a prefill the machine can't hold. Independent of `--ctx-size` | ## 🔧 Per-Request API Parameters diff --git a/Sources/SwiftLM/Server.swift b/Sources/SwiftLM/Server.swift index bd7225d..75cf9dc 100644 --- a/Sources/SwiftLM/Server.swift +++ b/Sources/SwiftLM/Server.swift @@ -555,6 +555,9 @@ struct MLXServer: AsyncParsableCommand { @Option(name: .long, help: "Number of parallel request slots") var parallel: Int = 1 + @Option(name: .long, help: "Reject prompts longer than this many tokens with an OpenAI-style 400 (context_length_exceeded) before prefill. Default: no limit") + var maxPromptTokens: Int? + @Flag(name: .long, help: "Enable thinking/reasoning mode (Qwen3.5 etc). Default: disabled") var thinking: Bool = false @@ -1355,6 +1358,7 @@ struct MLXServer: AsyncParsableCommand { minP: self.minP, repeatPenalty: self.repeatPenalty, thinking: self.thinking, + maxPromptTokens: self.maxPromptTokens, isVision: loadedAsVision, prefillSize: self.prefillSize, turboKV: self.turboKV, @@ -1648,6 +1652,8 @@ struct ServerConfig: Sendable { let minP: Float? let repeatPenalty: Float? let thinking: Bool + /// `--max-prompt-tokens`: longer prompts are rejected before prefill. nil = no limit. + let maxPromptTokens: Int? let isVision: Bool let prefillSize: Int /// When true, each KVCacheSimple layer compresses history > 8192 tokens to 3-bit PolarQuant. @@ -2095,6 +2101,11 @@ func handleChatCompletion( // ── Prompt caching: full token sequence for prefix matching ── let promptTokenCount = lmInput.text.tokens.size + if let rejection = await rejectLongPrompt( + promptTokens: promptTokenCount, limit: config.maxPromptTokens, slot: slot, stats: stats) + { + return rejection + } let promptTokens = lmInput.text.tokens.asArray(Int.self) // ── Issue #108: does the template leave the thinking block open? ── @@ -3106,6 +3117,11 @@ func handleTextCompletion( // ── Get actual prompt token count before generate() to avoid data race ── let promptTokenCount = lmInput.text.tokens.size + if let rejection = await rejectLongPrompt( + promptTokens: promptTokenCount, limit: config.maxPromptTokens, slot: slot, stats: stats) + { + return rejection + } let modelId = config.modelId @@ -3575,6 +3591,26 @@ func sseHeaders() -> HTTPFields { ]) } +/// OpenAI-shaped `context_length_exceeded` body when `promptTokens` is over the +/// `--max-prompt-tokens` limit, else nil (always nil when no limit is set). +/// +/// It uses OpenAI's wording and code so agent clients recognise it and compact +/// their history instead of retrying the same prompt. +func promptTooLongBody(promptTokens: Int, limit: Int?) -> String? { + guard let limit, promptTokens > limit else { return nil } + 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\"}}" +} + +/// 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? { + guard let limit, let body = promptTooLongBody(promptTokens: promptTokens, limit: limit) else { return nil } + print("srv slot_reject: id 0 | prompt=\(promptTokens)t exceeds max_prompt_tokens=\(limit)") + slot.release() + await stats.requestFinished(tokens: 0, duration: 0) + return Response(status: .badRequest, headers: jsonHeaders(), body: .init(byteBuffer: ByteBuffer(string: body))) +} + /// Build an OpenAI-style `{"error":{...}}` body for a server error. The message /// is JSON-encoded, so quotes, backslashes and newlines in it stay valid JSON. func errorJSON(_ error: Error) -> String { diff --git a/tests/SwiftLMTests/PromptLengthGuardTests.swift b/tests/SwiftLMTests/PromptLengthGuardTests.swift new file mode 100644 index 0000000..089e5f1 --- /dev/null +++ b/tests/SwiftLMTests/PromptLengthGuardTests.swift @@ -0,0 +1,38 @@ +import XCTest +import Foundation +@testable import SwiftLM + +/// `--max-prompt-tokens` rejects over-long prompts before prefill, so a prompt the +/// machine can't hold fails fast with a 400 instead of driving the server into swap. +final class PromptLengthGuardTests: XCTestCase { + + func testNoLimitAllowsAnything() { + XCTAssertNil(promptTooLongBody(promptTokens: 1_000_000, limit: nil)) + } + + func testPromptAtTheLimitIsAllowed() { + XCTAssertNil(promptTooLongBody(promptTokens: 32_768, limit: 32_768)) + } + + func testPromptOverTheLimitIsRejected() { + XCTAssertNotNil(promptTooLongBody(promptTokens: 32_769, limit: 32_768)) + } + + /// Agent clients key on OpenAI's code and wording to decide to compact. + func testRejectionIsOpenAIShaped() throws { + let body = try XCTUnwrap(promptTooLongBody(promptTokens: 34_412, limit: 32_768)) + let json = try XCTUnwrap(JSONSerialization.jsonObject(with: Data(body.utf8)) as? [String: Any]) + let error = try XCTUnwrap(json["error"] as? [String: Any]) + XCTAssertEqual(error["code"] as? String, "context_length_exceeded") + XCTAssertEqual(error["type"] as? String, "invalid_request_error") + let message = try XCTUnwrap(error["message"] as? String) + XCTAssertTrue(message.contains("maximum context length is 32768 tokens"), message) + XCTAssertTrue(message.contains("34412 tokens"), message) + } + + func testFlagParses() throws { + let server = try MLXServer.parse(["--model", "m", "--max-prompt-tokens", "20000"]) + XCTAssertEqual(server.maxPromptTokens, 20_000) + XCTAssertNil(try MLXServer.parse(["--model", "m"]).maxPromptTokens) + } +}