fix: strip NSNull from chat templates before Jinja render (#168)

Sources/SwiftLM/Server.swift

@@ -182,8 +182,15 @@ private struct TransformersTokenizerBridge: MLXLMCommon.Tokenizer, Sendable {
         additionalContext: [String: any Sendable]?
     ) throws -> [Int] {
         do {
+            // Issue #168: JSON `null` in tool schemas decodes to NSNull via AnyCodable.
+            // swift-jinja's Value.init(any:) has no NSNull case, so a single null anywhere
+            // in a tool spec aborted the whole render with a misleading "Optional<Any>"
+            // conversion error that was then mislabeled as a broken template. Strip nulls
+            // (and unwrap nested optionals) before handing anything to the template engine.
             return try upstream.applyChatTemplate(
-                messages: messages, tools: tools, additionalContext: additionalContext)
+                messages: messages.map { $0.mapValuesDeep(sanitizeForJinja) },
+                tools: tools?.map { $0.mapValuesDeep(sanitizeForJinja) },
+                additionalContext: additionalContext.map { $0.mapValuesDeep(sanitizeForJinja) })
         } catch Tokenizers.TokenizerError.missingChatTemplate {
             throw MLXLMCommon.TokenizerError.missingChatTemplate
         } catch {
@@ -196,6 +203,42 @@ private struct TransformersTokenizerBridge: MLXLMCommon.Tokenizer, Sendable {
     }
 }
 
+/// Returns `nil` when the value must be dropped (JSON `null` / NSNull), otherwise a
+/// structure with every nested null removed. See `TransformersTokenizerBridge.applyChatTemplate`.
+func sanitizeForJinja(_ value: any Sendable) -> any Sendable? {
+    if value is NSNull { return nil }
+    let mirror = Mirror(reflecting: value)
+    if mirror.displayStyle == .optional {
+        guard let child = mirror.children.first else { return nil }
+        return sanitizeForJinja(child.value as any Sendable)
+    }
+    if let dict = value as? [String: Any] {
+        var out: [String: any Sendable] = [:]
+        for (key, val) in dict {
+            if let cleaned = sanitizeForJinja(val as any Sendable) {
+                out[key] = cleaned
+            }
+        }
+        return out
+    }
+    if let arr = value as? [Any] {
+        return arr.compactMap { sanitizeForJinja($0 as any Sendable) }
+    }
+    return value
+}
+
+extension Dictionary where Key == String, Value == any Sendable {
+    func mapValuesDeep(_ transform: (any Sendable) -> any Sendable?) -> [String: any Sendable] {
+        var out: [String: any Sendable] = [:]
+        for (key, val) in self {
+            if let cleaned = transform(val) {
+                out[key] = cleaned
+            }
+        }
+        return out
+    }
+}
+
 // ── CLI ──────────────────────────────────────────────────────────────────────
 
 final class ProgressTracker {
@@ -1115,6 +1158,37 @@ struct MLXServer: AsyncParsableCommand {
             // supply. Neither is a reason to refuse to start.
         }
 
+        // Issue #168: render once more with a minimal non-empty tools array. The probe
+        // above uses `tools: nil`, so a checkpoint whose template (or tool payload)
+        // breaks only under the tools branch loaded clean and then failed every real
+        // agentic request. Warn rather than abort: a tools-broken model still serves
+        // plain chat and /v1/completions.
+        do {
+            let probeTokenizer = await container.tokenizer
+            let probeTool: [String: any Sendable] = [
+                "type": "function",
+                "function": [
+                    "name": "probe",
+                    "description": "startup chat-template tools probe",
+                    "parameters": [
+                        "type": "object",
+                        "properties": [
+                            "query": ["type": "string", "default": NSNull() as any Sendable]
+                        ],
+                    ] as [String: any Sendable],
+                ] as [String: any Sendable],
+            ]
+            _ = try probeTokenizer.applyChatTemplate(
+                messages: [["role": "user", "content": "ping"]],
+                tools: [probeTool],
+                additionalContext: ["add_generation_prompt": true]
+            )
+        } catch let error as MalformedChatTemplate {
+            print("[SwiftLM] ⚠️  chat-template tools probe failed (plain chat may still work): \(error.description)")
+        } catch {
+            // Same lenient pass as above: no template, or context the probe lacks.
+        }
+
         print("[SwiftLM] Model loaded. Starting HTTP server on \(host):\(port)")
 
         // ── Capture CLI defaults into a shared config ──
@@ -3161,23 +3235,26 @@ func sseChunk(modelId: String, reasoningContent: String?, content: String?, fini
 
 /// Prefill-progress heartbeat chunk — emitted every 2s while the server is processing the prompt
 /// when explicitly enabled via `X-SwiftLM-Prefill-Progress: true`.
-/// It is sent as a named SSE event (`event: prefill_progress`) to avoid breaking strict
-/// OpenAI-compatible clients (e.g. OpenCode), which reject unknown `data:` objects.
+/// It is sent as a named SSE event (`event: prefill_progress`).
 /// Format mirrors llama-server's slot_update event:
 ///   n_past          : tokens evaluated so far (real value from chunked prefill, or 0 for single-chunk)
 ///   n_prompt_tokens : total prompt token count
 ///   fraction        : n_past / n_prompt_tokens (0.0–1.0), useful for progress bars
 ///   elapsed_seconds : wall-clock time since the request started
 /// Note: `model` is intentionally omitted — clients can correlate from preceding stream chunks.
 /// Note: `on` is accepted as a truthy header value for parity with common reverse proxy conventions.
+/// Issue #168: `choices: []` is present so strict OpenAI chunk validators (opencode's
+/// ChatCompletionChunk union) accept the payload even when they parse every `data:` line
+/// regardless of `event:`. The named event alone was not enough.
 func ssePrefillChunk(nPast: Int = 0, promptTokens: Int, elapsedSeconds: Int) -> String {
     let fraction = promptTokens > 0 ? Double(nPast) / Double(promptTokens) : 0.0
     let chunk: [String: Any] = [
         "status": "processing",
         "n_past": nPast,
         "n_prompt_tokens": promptTokens,
         "fraction": fraction,
-        "elapsed_seconds": elapsedSeconds
+        "elapsed_seconds": elapsedSeconds,
+        "choices": [Any]()
     ]
     let data = try! JSONSerialization.data(withJSONObject: chunk)
     return "event: prefill_progress\r\ndata: \(String(data: data, encoding: .utf8)!)\r\n\r\n"

tests/SwiftLMTests/JinjaSanitizerTests.swift

@@ -0,0 +1,83 @@
+import XCTest
+import Foundation
+@testable import SwiftLM
+
+/// Issue #168: JSON `null` (NSNull) in tool schemas used to abort Jinja.Value conversion
+/// with a misleading "Optional<Any>" error, mislabeled as a broken chat template.
+final class JinjaSanitizerTests: XCTestCase {
+
+    func testDropsTopLevelNSNull() {
+        XCTAssertNil(sanitizeForJinja(NSNull()))
+    }
+
+    func testDropsNestedObjectNulls() throws {
+        let tool: [String: any Sendable] = [
+            "type": "function",
+            "function": [
+                "name": "probe",
+                "parameters": [
+                    "type": "object",
+                    "properties": [
+                        "query": [
+                            "type": "string",
+                            "default": NSNull() as any Sendable,
+                        ] as [String: any Sendable],
+                    ] as [String: any Sendable],
+                ] as [String: any Sendable],
+            ] as [String: any Sendable],
+        ]
+
+        let cleaned = try XCTUnwrap(sanitizeForJinja(tool) as? [String: any Sendable])
+        let fn = try XCTUnwrap(cleaned["function"] as? [String: any Sendable])
+        let params = try XCTUnwrap(fn["parameters"] as? [String: any Sendable])
+        let props = try XCTUnwrap(params["properties"] as? [String: any Sendable])
+        let query = try XCTUnwrap(props["query"] as? [String: any Sendable])
+
+        XCTAssertNil(query["default"], "null default must be stripped")
+        XCTAssertEqual(query["type"] as? String, "string")
+        XCTAssertEqual(params["type"] as? String, "object")
+    }
+
+    func testDropsNullArrayElements() throws {
+        let value: [String: any Sendable] = [
+            "enum": [1, NSNull(), 2] as [any Sendable]
+        ]
+        let cleaned = try XCTUnwrap(sanitizeForJinja(value) as? [String: any Sendable])
+        let enumValues = try XCTUnwrap(cleaned["enum"] as? [Any])
+        XCTAssertEqual(enumValues.count, 2)
+        XCTAssertFalse(enumValues.contains { $0 is NSNull })
+    }
+
+    func testPreservesScalarsAndStructure() throws {
+        let value: [String: any Sendable] = [
+            "type": "string",
+            "minimum": 0,
+            "required": ["command"] as [any Sendable],
+            "flag": true,
+        ]
+        let cleaned = try XCTUnwrap(sanitizeForJinja(value) as? [String: any Sendable])
+        XCTAssertEqual(cleaned["type"] as? String, "string")
+        XCTAssertEqual(cleaned["minimum"] as? Int, 0)
+        XCTAssertEqual(cleaned["flag"] as? Bool, true)
+        XCTAssertEqual((cleaned["required"] as? [Any])?.count, 1)
+    }
+
+    func testUnwrapsNestedOptional() throws {
+        let wrapped: any Sendable = Optional<String>.some("hi")
+        let cleaned = sanitizeForJinja(wrapped)
+        XCTAssertEqual(cleaned as? String, "hi")
+
+        let empty: any Sendable = Optional<String>.none
+        XCTAssertNil(sanitizeForJinja(empty))
+    }
+
+    func testMapValuesDeepOnToolDict() throws {
+        let dict: [String: any Sendable] = [
+            "keep": "x",
+            "drop": NSNull(),
+        ]
+        let cleaned = dict.mapValuesDeep(sanitizeForJinja)
+        XCTAssertEqual(cleaned["keep"] as? String, "x")
+        XCTAssertNil(cleaned["drop"])
+    }
+}

tests/SwiftLMTests/ServerSSETests.swift

@@ -46,7 +46,9 @@ final class ServerSSETests: XCTestCase {
         XCTAssertEqual(json["n_prompt_tokens"] as? Int, 128)
         XCTAssertEqual(json["elapsed_seconds"] as? Int, 4)
         XCTAssertNil(json["object"])
-        XCTAssertNil(json["choices"])
+        // Issue #168: empty choices keeps strict OpenAI chunk validators (opencode) happy
+        // even when they validate every data: line regardless of event name.
+        XCTAssertEqual((json["choices"] as? [Any])?.count, 0)
     }
 
     // MARK: - 1b: Zero-token boundary (no divide-by-zero crash)
@@ -93,7 +95,24 @@ final class ServerSSETests: XCTestCase {
         XCTAssertNil(json["id"],      "prefill chunk must not carry an id field")
         XCTAssertNil(json["object"],  "prefill chunk must not carry an object field")
         XCTAssertNil(json["model"],   "prefill chunk must not carry a model field")
-        XCTAssertNil(json["choices"], "prefill chunk must not carry a choices field")
+        // Issue #168: choices must be present but empty — opencode's chunk union
+        // requires `choices` (or `error`); omitting it fails type validation.
+        XCTAssertEqual((json["choices"] as? [Any])?.count, 0,
+                       "prefill chunk must carry an empty choices array")
+    }
+
+    // MARK: - Issue #168: empty choices is what strict validators require
+
+    func testPrefillChunk_ChoicesIsEmptyArrayForStrictValidators() throws {
+        let chunk = ssePrefillChunk(nPast: 1, promptTokens: 4, elapsedSeconds: 1)
+        let prefix = "event: prefill_progress\r\ndata: "
+        let suffix = "\r\n\r\n"
+        let payload = String(chunk.dropFirst(prefix.count).dropLast(suffix.count))
+        let data = try XCTUnwrap(payload.data(using: .utf8))
+        let json = try XCTUnwrap(JSONSerialization.jsonObject(with: data) as? [String: Any])
+
+        let choices = try XCTUnwrap(json["choices"] as? [Any], "choices must be present")
+        XCTAssertTrue(choices.isEmpty)
     }
 
     // MARK: - 1e: PrefillState.finish() is idempotent (Issue #2 guard)

tests/test-opencode.sh

@@ -222,6 +222,81 @@ else
     fail "opencode-shaped request failed: $AGENT_OUT"
 fi
 
+# ── Test 3: tools + heartbeat combined (Issue #168) ────────────────
+# Test 1 covers heartbeat without tools; Test 2 covers tools without the heartbeat
+# header. Issue #168 reported both gaps at once: a tools-bearing request with
+# X-SwiftLM-Prefill-Progress enabled failed on the null-bearing tool schema (Jinja
+# NSNull conversion) *and* would have hit opencode's strict validation of the
+# prefill_progress payload. Exercise the intersection here.
+log "Test 3: tools + prefill-progress heartbeat (Issue #168)"
+
+cat << 'PYEOF' > /tmp/opencode_tools_heartbeat_test.py
+import json, os, sys
+import openai
+
+client = openai.OpenAI(base_url=os.environ["OPENAI_BASE_URL"], api_key="sk-test", max_retries=0)
+
+# A tool schema shaped like opencode's zod/effect output — includes JSON nulls
+# (`default: null`) that previously crashed swift-jinja Value.init(any:) with
+# "Cannot convert value of type Optional<Any> to Jinja Value" (#168).
+TOOLS = [
+    {"type": "function", "function": {
+        "name": "bash",
+        "description": "Execute a shell command",
+        "parameters": {"type": "object",
+                       "properties": {
+                           "command": {"type": "string"},
+                           "timeout": {"type": "integer", "default": None},
+                       },
+                       "required": ["command"]}}},
+]
+MESSAGES = [
+    {"role": "system", "content": "You are a coding agent."},
+    {"role": "user", "content": "Say hi."},
+]
+
+try:
+    stream = client.chat.completions.create(
+        model=os.environ["MODEL"], messages=MESSAGES, tools=TOOLS,
+        stream=True, max_tokens=64, temperature=0,
+        stream_options={"include_usage": True},
+        # Enables the named `event: prefill_progress` heartbeat payloads.
+        extra_headers={"X-SwiftLM-Prefill-Progress": "true"},
+    )
+except Exception as e:
+    print(f"Error: request rejected: {e}")
+    sys.exit(1)
+
+chunks = 0
+finish = None
+try:
+    for chunk in stream:
+        chunks += 1
+        for choice in chunk.choices:
+            if choice.finish_reason:
+                finish = choice.finish_reason
+except Exception as e:
+    print(f"Error: SSE stream failed to parse (heartbeat or tools payload rejected): {e}")
+    sys.exit(1)
+
+if chunks == 0:
+    print("Error: stream produced no chunks")
+    sys.exit(1)
+
+print(f"Success: {chunks} chunks, finish_reason={finish}")
+PYEOF
+
+set +e
+HB_OUT=$("$VENV_DIR/bin/python" /tmp/opencode_tools_heartbeat_test.py 2>&1)
+HB_EXIT=$?
+set -e
+
+if [ $HB_EXIT -eq 0 ]; then
+    pass "tools + heartbeat stream accepted — $HB_OUT"
+else
+    fail "tools + heartbeat stream rejected: $HB_OUT"
+fi
+
 # ── Results ──────────────────────────────────────────────────────────
 echo ""
 log "═══════════════════════════════════════"