fix: pass chat_template_kwargs to the chat template as sent

README.md

@@ -687,6 +687,7 @@ In addition to the standard OpenAI fields (`temperature`, `top_p`, `max_tokens`,
 |---|---|---|
 | `kv_bits` | `int` (4 or 8) | Enable **MLX-native quantized KV cache** for this request. Uses `QuantizedKVCache` (standard group quantization) instead of `KVCacheSimple`. Separate from `--turbo-kv`. Reduces KV memory ~2–4× at mild quality cost. |
 | `enable_thinking` | `bool` | Force-enable or disable chain-of-thought thinking blocks for Gemma-4 / Qwen3. |
+| `chat_template_kwargs` | `object` | Passed to the chat template as sent, any JSON value: `{"thinking": true}` for DeepSeek V3.1, `{"reasoning_effort": "high"}` for gpt-oss. An `enable_thinking` inside it is overridden by the top-level field above. |
 | `kv_group_size` | `int` | Group size for `kv_bits` quantization (default: `64`). |
 | `top_k` | `int` | Per-request top-k sampling override (0 = disabled). |
 | `min_p` | `float` | Per-request min-p sampling threshold (0 = disabled). |

Sources/MLXInferenceCore/HFModelSearch.swift

@@ -160,6 +160,12 @@ public final class HFModelSearchService: ObservableObject {
     private var currentSort = HFSortOption.trending
     private var currentSizeFilter = HFSizeFilter.all
     private var debounceTask: Task<Void, Never>? = nil
+    /// Counts searches. A page fetch remembers the count it started under and publishes
+    /// nothing once a newer search has started: cancelling the debounce task does not stop
+    /// a fetch that is already awaiting the network, and a `loadMore()` fetch is not in
+    /// that task at all. Before this, a superseded fetch wrote "Search failed: cancelled"
+    /// and `isSearching = false` over the search that replaced it.
+    private var generation = 0
 
     private init() {}
 
@@ -168,6 +174,7 @@ public final class HFModelSearchService: ObservableObject {
     /// Debounced search — safe to call on every keystroke.
     public func search(query: String, sort: HFSortOption = .trending, sizeFilter: HFSizeFilter = .all) {
         debounceTask?.cancel()
+        generation += 1
         debounceTask = Task {
             // 300ms debounce
             try? await Task.sleep(nanoseconds: 300_000_000)
@@ -225,6 +232,10 @@ public final class HFModelSearchService: ObservableObject {
 
     private func fetchPage() async {
         print("HFSearch: fetchPage started. Query: '\(currentQuery)' Sort: \(currentSort.rawValue)")
+        let generation = self.generation
+        // True once a newer search has started (or this fetch's task was cancelled with
+        // it): the published state belongs to that search, so this one stops here.
+        func superseded() -> Bool { Task.isCancelled || generation != self.generation }
         isSearching = true
         errorMessage = nil
 
@@ -233,7 +244,9 @@ public final class HFModelSearchService: ObservableObject {
         // and hit the model detail endpoint. This is needed because the HF
         // search API filters on pipeline_tag which many MLX uploads don't match.
         if isDirectRepoQuery(currentQuery) {
-            if let result = await fetchDirectRepo(currentQuery) {
+            let result = await fetchDirectRepo(currentQuery)
+            if superseded() { return }
+            if let result {
                 results = [result]
             } else {
                 // Also try with prefix variants (e.g. user typed partial name without quant suffix)
@@ -283,6 +296,7 @@ public final class HFModelSearchService: ObservableObject {
 
             do {
                 let (data, response) = try await URLSession.shared.data(from: urlToFetch)
+                if superseded() { return }
                 guard let http = response as? HTTPURLResponse, http.statusCode == 200 else {
                     errorMessage = "HuggingFace API unavailable"
                     break
@@ -335,6 +349,7 @@ public final class HFModelSearchService: ObservableObject {
                             if let size = size { page[index].usedStorage = size }
                         }
                     }
+                    if superseded() { return }
                     localResults.append(contentsOf: page)
                 }
 
@@ -345,17 +360,22 @@ public final class HFModelSearchService: ObservableObject {
                     hasMore = true
                 }
             } catch is CancellationError {
-                break
+                return
+            } catch let error as URLError where error.code == .cancelled {
+                // The request was cancelled with this fetch's task: a newer search's doing.
+                return
             } catch {
+                if superseded() { return }
                 errorMessage = "Search failed: \(error.localizedDescription)"
                 break
             }
         }
 
+        if superseded() { return }
         if !localResults.isEmpty {
             results.append(contentsOf: localResults)
         }
-        
+
         isSearching = false
         print("HFSearch: fetchPage finished")
     }

Sources/SwiftLM/Server.swift

@@ -2516,6 +2516,18 @@ func handleChatCompletion(
         )
     }
 
+    // ── chat_template_kwargs is free-form JSON, but enable_thinking stays a boolean ──
+    // It decides the precedence below, so a string or a number there is a client error,
+    // answered the way the whole field used to be answered before it took other types.
+    if let raw = chatReq.chatTemplateKwargs?["enable_thinking"], !(raw.value is Bool) {
+        let errBody = "{\"error\":{\"message\":\"Invalid type for 'chat_template_kwargs.enable_thinking': expected a boolean.\",\"type\":\"invalid_request_error\",\"param\":\"chat_template_kwargs.enable_thinking\"}}"
+        return Response(
+            status: .badRequest,
+            headers: jsonHeaders(),
+            body: .init(byteBuffer: ByteBuffer(string: errBody))
+        )
+    }
+
     var generateParams = GenerateParameters(
         maxTokens: tokenLimit,
         maxKVSize: config.ctxSize,
@@ -2583,7 +2595,7 @@ func handleChatCompletion(
     var enableThinking: Bool
     if let explicitTopLevel = chatReq.enableThinking {
         enableThinking = explicitTopLevel
-    } else if let kwargs = chatReq.chatTemplateKwargs, let perRequest = kwargs["enable_thinking"] {
+    } else if let perRequest = chatReq.chatTemplateKwargs?["enable_thinking"]?.value as? Bool {
         enableThinking = perRequest  // per-request override wins
     } else {
         enableThinking = config.thinking  // fall back to server --thinking flag
@@ -2617,7 +2629,8 @@ func handleChatCompletion(
 
     // The Jinja template evaluates `not enable_thinking | default(false)`. If we pass nil instead of
     // true, it evaluates to false and still breaks. We MUST explicitly pass the boolean.
-    let templateContext: [String: any Sendable] = ["enable_thinking": enableThinking]
+    let templateContext = chatTemplateContext(
+        kwargs: chatReq.chatTemplateKwargs, enableThinking: enableThinking)
     let userInput = UserInput(chat: chatMessages, tools: toolSpecs, additionalContext: templateContext)
     print("[Server Debug] Created UserInput with \(userInput.images.count) images and \(userInput.audios.count) audio inputs.")
     let lmInput = try await container.prepare(input: userInput)
@@ -4419,6 +4432,18 @@ func checkStopSequences(_ text: String, stopSequences: [String], lookback: Int?
     return (String(text[text.startIndex..<earliest.index]), earliest.stop)
 }
 
+/// The `additionalContext` a chat request hands to its chat template: every
+/// `chat_template_kwargs` entry as the client sent it, with `enable_thinking` set to the
+/// value the request's precedence rules decided (top-level field, then the kwarg, then the
+/// server flag). Templates read their own names here (`thinking`, `reasoning_effort`,
+/// `model_identity`, …), so nothing but that one key is interpreted by the server.
+/// Nulls inside the values are stripped later by `sanitizeForJinja`, as for tools.
+func chatTemplateContext(kwargs: [String: AnyCodable]?, enableThinking: Bool) -> [String: any Sendable] {
+    var context = AnyCodable.toSendable(kwargs) ?? [:]
+    context["enable_thinking"] = enableThinking
+    return context
+}
+
 /// The sampler seed for a request `seed`. OpenAI's `seed` is a plain integer, so negative
 /// values do arrive, and `UInt64(seed)` trapped on them — one such request took the whole
 /// server down (#219). Taking the bit pattern keeps every distinct Int a distinct seed.
@@ -4920,8 +4945,11 @@ struct ChatCompletionRequest: Decodable {
     let seed: Int?
     let streamOptions: StreamOptions?
     let responseFormat: ResponseFormat?
-    /// Per-request Jinja template kwargs (e.g. {"enable_thinking": false} for Qwen3/Qwen3.5)
-    let chatTemplateKwargs: [String: Bool]?
+    /// Per-request Jinja template kwargs, passed to the chat template as sent: any JSON
+    /// value, since templates read more than booleans (`{"enable_thinking": false}` for
+    /// Qwen3, `{"thinking": true}` for DeepSeek V3.1, `{"reasoning_effort": "high"}` for
+    /// gpt-oss). `enable_thinking` keeps its own precedence rules in `handleChatCompletion`.
+    let chatTemplateKwargs: [String: AnyCodable]?
     /// Per-request reasoning-token cap; overrides `--thinking-budget`.
     let thinkingBudget: Int?
     /// Top-level thinking override emitted by Aegis-AI gateway

docs/api-response-formats.md

@@ -185,6 +185,7 @@ for await (const chunk of stream) {
 | `response_format` | object | none | `{"type": "json_object"}` for JSON mode |
 | `tools` | array | none | Tool/function definitions |
 | `enable_thinking` | bool | false | Enable `<think>` reasoning tokens |
+| `chat_template_kwargs` | object | none | Passed to the chat template as sent (e.g. `{"thinking": true}`, `{"reasoning_effort": "high"}`) |
 
 ---
 

tests/SwiftBuddyTests/HFModelSearchTests.swift

@@ -50,6 +50,32 @@ final class HFModelSearchTests: XCTestCase {
         XCTAssertNil(service.errorMessage)
     }
 
+    /// A search that replaces one still awaiting the network must not inherit its state.
+    /// The replaced fetch's request is cancelled, which it used to report as
+    /// "Search failed: cancelled" over the new search, and it also cleared `isSearching`
+    /// while the new search was still running. That is the flake CI hit in
+    /// testStrictMLXFilterDisabled whenever an earlier test left a fetch in flight.
+    @MainActor
+    func testSupersededSearchPublishesNothing() async throws {
+        let service = HFModelSearchService.shared
+        service.errorMessage = nil
+        service.results = []
+        service.strictMLX = true
+
+        service.search(query: "gpt", sort: .trending)
+        // Past the 300 ms debounce, so the first fetch is on the wire when it is replaced.
+        try? await Task.sleep(nanoseconds: 450_000_000)
+        service.search(query: "mistral", sort: .trending)
+
+        for _ in 0..<40 {
+            try? await Task.sleep(nanoseconds: 100_000_000)
+            if !service.isSearching && service.results.count > 0 { break }
+        }
+
+        try XCTSkipIf(service.isSearching || service.results.count == 0, "Skipping due to transient HF rate limiting on GitHub Actions IP bounds")
+        XCTAssertNil(service.errorMessage, "the replaced search's cancellation leaked into the new search")
+    }
+
     // Feature 3: Empty query trending
     @MainActor
     func testFeature3_EmptyQueryTrending() async throws {

tests/SwiftLMTests/ChatRequestParsingTests.swift

@@ -322,4 +322,62 @@ final class ChatRequestParsingTests: XCTestCase {
         XCTAssertEqual(req.maxCompletionTokens, 64)
         XCTAssertNil(req.maxTokens)
     }
+
+    // MARK: - chat_template_kwargs is passed to the template as sent (#259)
+
+    /// The field decoded as `[String: Bool]`, so a string or a number anywhere in it failed
+    /// the whole request, and only `enable_thinking` was read. Templates take other names
+    /// and types: DeepSeek-V3.1 `thinking`, gpt-oss `reasoning_effort` and `model_identity`.
+    func testChatTemplateKwargsDecodeAnyJSONValue() throws {
+        let req = try decode("""
+        {
+            "messages": [{"role": "user", "content": "hi"}],
+            "chat_template_kwargs": {
+                "enable_thinking": false,
+                "thinking": true,
+                "reasoning_effort": "high",
+                "budget": 3,
+                "nested": {"a": [1, "b", null]}
+            }
+        }
+        """)
+        let kwargs = try XCTUnwrap(req.chatTemplateKwargs)
+        XCTAssertEqual(kwargs["enable_thinking"]?.value as? Bool, false)
+        XCTAssertEqual(kwargs["thinking"]?.value as? Bool, true)
+        XCTAssertEqual(kwargs["reasoning_effort"]?.value as? String, "high")
+        XCTAssertEqual(kwargs["budget"]?.value as? Int, 3)
+        let nested = try XCTUnwrap(kwargs["nested"]?.value as? [String: Any])
+        let list = try XCTUnwrap(nested["a"] as? [Any])
+        XCTAssertEqual(list.count, 3)
+        XCTAssertEqual(list[0] as? Int, 1)
+        XCTAssertEqual(list[1] as? String, "b")
+        XCTAssertTrue(list[2] is NSNull)
+    }
+
+    func testChatTemplateKwargsAbsentDecodesToNil() throws {
+        let req = try decode(#"{"messages":[{"role":"user","content":"hi"}]}"#)
+        XCTAssertNil(req.chatTemplateKwargs)
+    }
+
+    /// Every kwarg reaches the template context; `enable_thinking` carries the verdict of
+    /// the precedence rules in `handleChatCompletion`, not the raw kwarg.
+    func testChatTemplateContextForwardsEveryKwargAndSetsEnableThinking() throws {
+        let req = try decode("""
+        {
+            "messages": [{"role": "user", "content": "hi"}],
+            "chat_template_kwargs": {"enable_thinking": true, "thinking": true, "reasoning_effort": "low"}
+        }
+        """)
+        let context = chatTemplateContext(kwargs: req.chatTemplateKwargs, enableThinking: false)
+        XCTAssertEqual(context["enable_thinking"] as? Bool, false)
+        XCTAssertEqual(context["thinking"] as? Bool, true)
+        XCTAssertEqual(context["reasoning_effort"] as? String, "low")
+        XCTAssertEqual(context.count, 3)
+    }
+
+    func testChatTemplateContextWithoutKwargsHasOnlyEnableThinking() {
+        let context = chatTemplateContext(kwargs: nil, enableThinking: true)
+        XCTAssertEqual(context["enable_thinking"] as? Bool, true)
+        XCTAssertEqual(context.count, 1)
+    }
 }

tests/test-contract.sh

@@ -553,6 +553,31 @@ PY
 early_exit_then_follow_up "stop sequence" "stop-seq" "the" ""
 early_exit_then_follow_up "client disconnect" "disconnect" "" "leave"
 
+# ── 17. chat_template_kwargs is passed to the template as sent ──────────────
+# #259: the field decoded as [String: Bool], so a string or a number anywhere in it was a
+# 400, and every key but enable_thinking was dropped before the template saw it. Templates
+# read other names (DeepSeek-V3.1 `thinking`, gpt-oss `reasoning_effort`). The CI model's
+# template ignores names it does not know, so this checks the wire contract only: any JSON
+# value is accepted, and a non-boolean enable_thinking is still refused by name.
+log "Test 17: chat_template_kwargs takes any JSON value"
+KW=$(chat '{"model":"x","messages":[{"role":"user","content":"Say hello in one short sentence."}],"max_tokens":16,"temperature":0,"chat_template_kwargs":{"reasoning_effort":"low","budget":3,"thinking":true,"nested":{"a":[1,"b"]}}}' || true)
+KW_CONTENT=$(echo "$KW" | python3 -c 'import json,sys; print(json.load(sys.stdin)["choices"][0]["message"].get("content") or "")' 2>/dev/null || echo "__no_body__")
+if [ "$KW_CONTENT" = "__no_body__" ] || [ -z "$KW_CONTENT" ]; then
+    fail "string, number and nested kwargs were rejected: $(echo "$KW" | head -c 160)"
+else
+    pass "string, number, boolean and nested kwargs are accepted"
+fi
+
+log "Test 17b: a non-boolean enable_thinking is a 400 that names the param"
+CODE=$(curl -s -o /tmp/SwiftLM-contract-kwargs.json -w '%{http_code}' "$URL/v1/chat/completions" \
+    -H 'Content-Type: application/json' -d '{"model":"x","messages":[{"role":"user","content":"hi"}],"max_tokens":4,"chat_template_kwargs":{"enable_thinking":"yes"}}')
+ERR_TYPE=$(python3 -c 'import json; e=json.load(open("/tmp/SwiftLM-contract-kwargs.json"))["error"]; print(e.get("type"), e.get("param"))' 2>/dev/null || echo "unparseable")
+if [ "$CODE" = "400" ] && [ "$ERR_TYPE" = "invalid_request_error chat_template_kwargs.enable_thinking" ]; then
+    pass "400 invalid_request_error naming param 'chat_template_kwargs.enable_thinking'"
+else
+    fail "got HTTP $CODE with error '$ERR_TYPE' (expected 400 invalid_request_error chat_template_kwargs.enable_thinking)"
+fi
+
 log "═══════════════════════════════════════"
 log "Results: $PASS passed, $FAIL failed, $SKIP skipped, $TOTAL total"
 log "═══════════════════════════════════════"