diff --git a/README.md b/README.md index b7ec12b..d259764 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` | | `--no-token-echo` | `false` | Stop echoing generated tokens to stdout as they stream. Request log lines (`srv ...`) are unaffected | ## 🔧 Per-Request API Parameters diff --git a/Sources/SwiftLM/Server.swift b/Sources/SwiftLM/Server.swift index 9823593..e608b56 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 @@ -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 { 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) + } +}