feat: --max-prompt-tokens rejects over-long prompts before prefill

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` |
 | `--no-token-echo` | `false` | Stop echoing generated tokens to stdout as they stream. Request log lines (`srv ...`) are unaffected |
 
 ## 🔧 Per-Request API Parameters

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
 
@@ -1358,6 +1361,7 @@ struct MLXServer: AsyncParsableCommand {
             minP: self.minP,
             repeatPenalty: self.repeatPenalty,
             thinking: self.thinking,
+            maxPromptTokens: self.maxPromptTokens,
             tokenEcho: !self.noTokenEcho,
             isVision: loadedAsVision,
             prefillSize: self.prefillSize,
@@ -1652,6 +1656,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?
     /// Echo each generated chunk to stdout (`--no-token-echo` turns it off).
     let tokenEcho: Bool
     let isVision: Bool
@@ -2152,6 +2158,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? ──
@@ -3226,6 +3237,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
 
@@ -3720,6 +3736,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 {

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)
+    }
+}