fix: reclaim generation slot on error and stream SSE headers before prefill

tests/SwiftLMTests/GenerationSlotTests.swift

@@ -0,0 +1,120 @@
+import XCTest
+import Foundation
+@testable import SwiftLM
+
+/// GenerationSlot guarantees the AsyncSemaphore slot acquired after `wait()`
+/// is returned exactly once — on explicit `release()` or `deinit` (whichever
+/// comes first). Without it, any throw between `wait()` and the success-path
+/// `signal()` calls permanently leaked the slot; with the default
+/// `--parallel 1` a single failed request wedged the server until restart.
+final class GenerationSlotTests: XCTestCase {
+
+    /// Race `wait()` against a timeout: returns true if the wait completed.
+    private func waitCompletes(_ sem: AsyncSemaphore, timeoutNs: UInt64 = 1_000_000_000) async -> Bool {
+        await withTaskGroup(of: Bool.self) { group in
+            group.addTask {
+                await sem.wait()
+                await sem.signal()  // immediately give it back; only completion matters
+                return true
+            }
+            group.addTask {
+                try? await Task.sleep(nanoseconds: timeoutNs)
+                return false
+            }
+            let first = await group.next() ?? false
+            group.cancelAll()
+            return first
+        }
+    }
+
+    /// Simulates the error path: slot created after `wait()`, then dropped
+    /// without an explicit release (scope unwinds on throw). The deinit must
+    /// reclaim the slot so the next request can proceed.
+    func testDeinitReclaimsSlotWithoutExplicitRelease() async {
+        let sem = AsyncSemaphore(limit: 1)
+        await sem.wait()
+        do {
+            let slot = GenerationSlot(semaphore: sem)
+            _ = slot
+        }  // deinit → release
+
+        let ok = await waitCompletes(sem)
+        XCTAssertTrue(ok, "slot was not reclaimed by deinit — server would wedge after one error at --parallel 1")
+    }
+
+    /// Explicit release on the success path, then deinit when the frame unwinds:
+    /// the once-flag must prevent a second signal (which would over-admit and
+    /// break the parallel limit when other requests are queued).
+    func testExplicitReleaseThenDeinitDoesNotDoubleSignal() async {
+        let sem = AsyncSemaphore(limit: 1)
+        await sem.wait()
+
+        let slot = GenerationSlot(semaphore: sem)
+        slot.release()          // success-path release
+        // slot deinits at end of scope → second release attempt, must no-op
+
+        var bAdmitted = false
+        let bTask = Task {
+            await sem.wait()
+            bAdmitted = true
+            // hold the slot until the test finishes asserting
+            try? await Task.sleep(nanoseconds: 2_000_000_000)
+            await sem.signal()
+        }
+        // Let B queue and be admitted by the single legitimate release.
+        try? await Task.sleep(nanoseconds: 200_000_000)
+        XCTAssertTrue(bAdmitted, "legitimate release should admit exactly one waiter")
+
+        // If double-signalling had happened, the semaphore would have counted
+        // an extra free slot and this second wait would complete while B still
+        // holds the only slot.
+        var cAdmitted = false
+        let cTask = Task {
+            await sem.wait()
+            cAdmitted = true
+            await sem.signal()
+        }
+        try? await Task.sleep(nanoseconds: 300_000_000)
+        XCTAssertFalse(cAdmitted, "double-release over-admitted a second concurrent request (parallel limit broken)")
+
+        await bTask.value
+        await cTask.value
+        _ = slot
+    }
+
+    /// Multiple explicit releases (e.g. a success path plus a deferred cleanup
+    /// that both run) must not signal twice either: exactly one waiter is
+    /// admitted while the holder still holds the slot.
+    func testRepeatedReleaseIsIdempotent() async {
+        let sem = AsyncSemaphore(limit: 1)
+        await sem.wait()
+        let slot = GenerationSlot(semaphore: sem)
+        slot.release()
+        slot.release()
+        slot.release()
+
+        var bAdmitted = false
+        let bTask = Task {
+            await sem.wait()
+            bAdmitted = true
+            // hold the slot so any over-admission would surface as C below
+            try? await Task.sleep(nanoseconds: 2_000_000_000)
+            await sem.signal()
+        }
+        try? await Task.sleep(nanoseconds: 200_000_000)
+        XCTAssertTrue(bAdmitted, "single slot must be reclaimed exactly once")
+
+        var cAdmitted = false
+        let cTask = Task {
+            await sem.wait()
+            cAdmitted = true
+            await sem.signal()
+        }
+        try? await Task.sleep(nanoseconds: 300_000_000)
+        XCTAssertFalse(cAdmitted, "extra signals leaked additional slots beyond the limit")
+
+        await bTask.value
+        await cTask.value
+        withExtendedLifetime(slot) {}  // deinit runs here; flag must make it a no-op
+    }
+}

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)
@@ -120,4 +139,27 @@ final class ServerSSETests: XCTestCase {
         // is irrelevant to correctness. We capture the current contract here.
         // If a post-done guard is added later, add XCTAssertNotEqual(await state.nPast, 999).
     }
+
+    // MARK: - Error event stays valid JSON for any message text
+
+    private struct MessyError: Error, CustomStringConvertible {
+        let description = "bad \"quote\", back\\slash\nnew line"
+    }
+
+    func testErrorChunkEncodesMessageAsValidJSON() throws {
+        let chunk = sseErrorChunk(MessyError())
+
+        let prefix = "data: "
+        let suffix = "\r\n\r\n"
+        XCTAssertTrue(chunk.hasPrefix(prefix))
+        XCTAssertTrue(chunk.hasSuffix(suffix))
+
+        let payload = String(chunk.dropFirst(prefix.count).dropLast(suffix.count))
+        let obj = try XCTUnwrap(
+            JSONSerialization.jsonObject(with: Data(payload.utf8)) as? [String: Any])
+        let err = try XCTUnwrap(obj["error"] as? [String: Any])
+        XCTAssertEqual(err["message"] as? String, MessyError().description)
+        XCTAssertEqual(err["type"] as? String, "server_error")
+        XCTAssertEqual(err["code"] as? String, "internal_error")
+    }
 }

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 "═══════════════════════════════════════"