From 08ee817d544110278e0320ccfb150c8307c34b5f Mon Sep 17 00:00:00 2001 From: idevlab Date: Mon, 31 Aug 2026 21:19:03 +0800 Subject: [PATCH 1/7] fix: harden ANE prompt and context limits --- Sources/LLM/EspressoLLMEngine.swift | 28 ++++++++++++++++++++++++++-- 1 file changed, 26 insertions(+), 2 deletions(-) diff --git a/Sources/LLM/EspressoLLMEngine.swift b/Sources/LLM/EspressoLLMEngine.swift index a63008e..1549a30 100644 --- a/Sources/LLM/EspressoLLMEngine.swift +++ b/Sources/LLM/EspressoLLMEngine.swift @@ -16,6 +16,8 @@ private func aneLMTokenCallback(_ token: Int32, _ context: UnsafeMutableRawPoint } actor EspressoLLMEngine { + static let maximumContextTokens = 2_048 + private final class LoadedModel { let path: String let runtime: OpaquePointer @@ -90,6 +92,18 @@ actor EspressoLLMEngine { .map { Int32(clamping: $0) } guard !promptTokens.isEmpty else { throw EspressoLLMError.runtimeFailure } + let nativeMaxTokens = max(1, maxTokens) + guard Self.requestFitsContextWindow( + promptTokenCount: promptTokens.count, + maxTokens: nativeMaxTokens + ) else { + let requiredCacheSlots = promptTokens.count + max(0, nativeMaxTokens - 1) + throw recordFailure(ANELMNativeError( + "ANE-LM request needs \(requiredCacheSlots) KV-cache slots; " + + "the packaged runtime supports \(Self.maximumContextTokens)" + )) + } + let context = ANELMGenerationContext() context.tokens.reserveCapacity(max(0, maxTokens)) let contextPointer = Unmanaged.passUnretained(context).toOpaque() @@ -103,7 +117,7 @@ actor EspressoLLMEngine { model.runtime, tokens.baseAddress, tokens.count, - Int32(clamping: max(1, maxTokens)), + Int32(clamping: nativeMaxTokens), Float(temperature), 1.2, Int32(clamping: model.samplerVocabularySize), @@ -153,6 +167,16 @@ actor EspressoLLMEngine { _ = try await makeValidatedTokenizer(at: url) } + static func requestFitsContextWindow( + promptTokenCount: Int, + maxTokens: Int + ) -> Bool { + guard promptTokenCount > 0, maxTokens > 0 else { return false } + let generatedCacheSlots = max(0, maxTokens - 1) + guard generatedCacheSlots <= maximumContextTokens else { return false } + return promptTokenCount <= maximumContextTokens - generatedCacheSlots + } + private struct ValidatedTokenizer { let tokenizer: any Tokenizers.Tokenizer let samplerVocabularySize: Int @@ -211,7 +235,7 @@ actor EspressoLLMEngine { if modelName.lowercased().contains("qwen") { return "<|im_start|>system\n\(system)<|im_end|>\n" + "<|im_start|>user\n\(user)<|im_end|>\n" - + "<|im_start|>assistant\n" + + "<|im_start|>assistant\n\n\n\n\n" } return "System:\n\(system)\n\nUser:\n\(user)\n\nAssistant:\n" } From 9a41d01bf7e8893a251dc9d9e0cb23aab9a92041 Mon Sep 17 00:00:00 2001 From: idevlab Date: Mon, 31 Aug 2026 21:19:22 +0800 Subject: [PATCH 2/7] fix: unload ANE before MLX fallback --- Sources/Processing/TextProcessor+Generation.swift | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/Sources/Processing/TextProcessor+Generation.swift b/Sources/Processing/TextProcessor+Generation.swift index 0a3d4ae..f952712 100644 --- a/Sources/Processing/TextProcessor+Generation.swift +++ b/Sources/Processing/TextProcessor+Generation.swift @@ -48,6 +48,9 @@ extension TextProcessor { temperature: temperature ) }, + prepareForMLXFallback: { + await self.espressoLLM.unload() + }, mlx: { try await self.llm.loadModel(id: options.llmModel) return try await self.llm.generate( @@ -59,9 +62,8 @@ extension TextProcessor { } ) if result.usedMLX { - _ = await espressoLLM.consumeLastFailureMessage() await Self.recordEspressoOutcome(.fallback) - Log.info("[TextProcessor] ANE-LM failed; used the selected MLX model") + Log.info("[TextProcessor] ANE-LM failed; unloaded it and used the selected MLX model") } else { await Self.clearEspressoOutcome() } @@ -69,7 +71,6 @@ extension TextProcessor { } catch is CancellationError { throw CancellationError() } catch let error as EspressoMLXFallbackError { - _ = await espressoLLM.consumeLastFailureMessage() await Self.recordEspressoOutcome(.unavailable) Log.sensitive("[TextProcessor] ANE-LM and MLX fallback failed: \(error.details)") Log.error("[TextProcessor] MLX fallback unavailable") @@ -89,6 +90,7 @@ extension TextProcessor { static func runEspressoWithMLXFallback( fallbackEnabled: Bool = true, espresso: () async throws -> Value, + prepareForMLXFallback: () async -> Void = {}, mlx: () async throws -> Value ) async throws -> (value: Value, usedMLX: Bool) { do { @@ -99,6 +101,8 @@ extension TextProcessor { try Task.checkCancellation() guard fallbackEnabled else { throw error } let espressoFailure = error.localizedDescription + await prepareForMLXFallback() + try Task.checkCancellation() do { let value = try await mlx() try Task.checkCancellation() From 4dd9f8193eb6c1fe829affea2a0d0b530deb99b3 Mon Sep 17 00:00:00 2001 From: idevlab Date: Mon, 31 Aug 2026 21:19:45 +0800 Subject: [PATCH 3/7] fix: release ANE before warmup fallback --- Sources/Processing/TextProcessor+Models.swift | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/Sources/Processing/TextProcessor+Models.swift b/Sources/Processing/TextProcessor+Models.swift index 29820e2..44a5620 100644 --- a/Sources/Processing/TextProcessor+Models.swift +++ b/Sources/Processing/TextProcessor+Models.swift @@ -163,10 +163,10 @@ extension TextProcessor { let result = try await Self.runEspressoWithMLXFallback( fallbackEnabled: fallbackToMLXOnEspressoFailure, espresso: { try await self.espressoLLM.loadModel(path: espressoModelPath) }, + prepareForMLXFallback: { await self.espressoLLM.unload() }, mlx: { try await self.llm.loadModel(id: model) } ) if result.usedMLX { - _ = await espressoLLM.consumeLastFailureMessage() return (true, nil, .fallback) } } @@ -176,7 +176,6 @@ extension TextProcessor { } catch let error as EspressoMLXFallbackError { Log.sensitive("[TextProcessor] ANE-LM and MLX warmup failed: \(error.details)") Log.error("[TextProcessor] MLX fallback unavailable during warmup") - _ = await espressoLLM.consumeLastFailureMessage() return (false, EspressoGenerationOutcome.unavailable.message, .unavailable) } catch { Log.error("[TextProcessor] LLM warmup failed: \(error.localizedDescription)") From 9e1b8c2da0fa45b5c76b40af6f6b5a604c5f9ab6 Mon Sep 17 00:00:00 2001 From: idevlab Date: Mon, 31 Aug 2026 21:20:15 +0800 Subject: [PATCH 4/7] test: cover ANE prompt and context guards --- Tests/OpenTypeTests/ANELMRuntimeTests.swift | 55 +++++++++++++++++++-- 1 file changed, 50 insertions(+), 5 deletions(-) diff --git a/Tests/OpenTypeTests/ANELMRuntimeTests.swift b/Tests/OpenTypeTests/ANELMRuntimeTests.swift index 40dce3e..8afd849 100644 --- a/Tests/OpenTypeTests/ANELMRuntimeTests.swift +++ b/Tests/OpenTypeTests/ANELMRuntimeTests.swift @@ -120,6 +120,41 @@ final class ANELMRuntimeTests: XCTestCase { await XCTAssertThrowsErrorAsync(try await EspressoLLMEngine.validateModelDirectory(at: invalidHeads)) } + func testQwenPromptDisablesThinkingBeforeGeneration() { + let prompt = EspressoLLMEngine.formatPrompt( + user: "Return JSON.", + system: "Do not explain.", + modelName: "Qwen3" + ) + + XCTAssertTrue(prompt.hasSuffix( + "<|im_start|>assistant\n\n\n\n\n" + )) + } + + func testContextWindowGuardReservesGeneratedCacheSlots() { + XCTAssertTrue(EspressoLLMEngine.requestFitsContextWindow( + promptTokenCount: 2_048, + maxTokens: 1 + )) + XCTAssertTrue(EspressoLLMEngine.requestFitsContextWindow( + promptTokenCount: 2_047, + maxTokens: 2 + )) + XCTAssertFalse(EspressoLLMEngine.requestFitsContextWindow( + promptTokenCount: 2_049, + maxTokens: 1 + )) + XCTAssertFalse(EspressoLLMEngine.requestFitsContextWindow( + promptTokenCount: 2_048, + maxTokens: 2 + )) + XCTAssertFalse(EspressoLLMEngine.requestFitsContextWindow( + promptTokenCount: 1, + maxTokens: 2_049 + )) + } + func testRealGenerationLifecycleWhenModelIsProvided() async throws { guard let modelPath = ProcessInfo.processInfo.environment["UTTER_ANE_TEST_MODEL"], !modelPath.isEmpty else { @@ -142,14 +177,24 @@ final class ANELMRuntimeTests: XCTestCase { let requestCount = iterations / lifecycleCount + (lifecycle < iterations % lifecycleCount ? 1 : 0) var lifecycleSamples: [Int] = [] - for _ in 0..")) + if verifiesStructuredCommandOutput { + XCTAssertNotNil(SpokenEditCommandLLMResolver.resolution(from: output)) + } let sample = try residentSizeKB() residentSamples.append(sample) lifecycleSamples.append(sample) From d6d6ba401aaa7680561f3a7f0ab9864d0f879db8 Mon Sep 17 00:00:00 2001 From: idevlab Date: Mon, 31 Aug 2026 21:20:47 +0800 Subject: [PATCH 5/7] test: cover ANE fallback cleanup and tracking --- .../OpenTypeTests/EspressoFallbackTests.swift | 30 ++++++++++++++++++- 1 file changed, 29 insertions(+), 1 deletion(-) diff --git a/Tests/OpenTypeTests/EspressoFallbackTests.swift b/Tests/OpenTypeTests/EspressoFallbackTests.swift index f13dca4..452a4cd 100644 --- a/Tests/OpenTypeTests/EspressoFallbackTests.swift +++ b/Tests/OpenTypeTests/EspressoFallbackTests.swift @@ -56,13 +56,38 @@ final class EspressoFallbackTests: XCTestCase { XCTAssertTrue(result.usedMLX) } + func testEspressoFailureReleasesANEStateBeforeMLXFallback() async throws { + var espressoIsLoaded = true + var espressoWasLoadedWhenMLXStarted = true + + let result = try await TextProcessor.runEspressoWithMLXFallback( + espresso: { throw StubError.espresso }, + prepareForMLXFallback: { + espressoIsLoaded = false + }, + mlx: { + espressoWasLoadedWhenMLXStarted = espressoIsLoaded + return "mlx output" + } + ) + + XCTAssertEqual(result.value, "mlx output") + XCTAssertTrue(result.usedMLX) + XCTAssertFalse(espressoIsLoaded) + XCTAssertFalse(espressoWasLoadedWhenMLXStarted) + } + func testDisabledFallbackDoesNotRunMLX() async { + var preparedForFallback = false var ranMLX = false do { _ = try await TextProcessor.runEspressoWithMLXFallback( fallbackEnabled: false, espresso: { throw StubError.espresso }, + prepareForMLXFallback: { + preparedForFallback = true + }, mlx: { ranMLX = true return "mlx output" @@ -71,6 +96,7 @@ final class EspressoFallbackTests: XCTestCase { XCTFail("Expected the Espresso failure") } catch { XCTAssertEqual(error.localizedDescription, "espresso failed") + XCTAssertFalse(preparedForFallback) XCTAssertFalse(ranMLX) } } @@ -214,11 +240,13 @@ final class EspressoFallbackTests: XCTestCase { temperature: 0 ) if index == 0 { + outcome = await processor.consumeEspressoOutcome() + let espressoIsLoaded = await processor.espressoLLM.isLoaded + XCTAssertFalse(espressoIsLoaded) baselineFootprint = currentMemoryFootprint() options.localLLMBackend = .mlx } } - outcome = await processor.consumeEspressoOutcome() } XCTAssertFalse(output.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty) From 095d3e261549efef751efe8370623b1778beb18d Mon Sep 17 00:00:00 2001 From: idevlab Date: Mon, 31 Aug 2026 21:24:59 +0800 Subject: [PATCH 6/7] test: make fallback cleanup result type explicit --- Tests/OpenTypeTests/EspressoFallbackTests.swift | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Tests/OpenTypeTests/EspressoFallbackTests.swift b/Tests/OpenTypeTests/EspressoFallbackTests.swift index 452a4cd..fd3bbe1 100644 --- a/Tests/OpenTypeTests/EspressoFallbackTests.swift +++ b/Tests/OpenTypeTests/EspressoFallbackTests.swift @@ -60,7 +60,7 @@ final class EspressoFallbackTests: XCTestCase { var espressoIsLoaded = true var espressoWasLoadedWhenMLXStarted = true - let result = try await TextProcessor.runEspressoWithMLXFallback( + let result: (value: String, usedMLX: Bool) = try await TextProcessor.runEspressoWithMLXFallback( espresso: { throw StubError.espresso }, prepareForMLXFallback: { espressoIsLoaded = false From 5e9a6b809108646476a7d5245cc8115b346511b4 Mon Sep 17 00:00:00 2001 From: idevlab Date: Mon, 31 Aug 2026 22:23:38 +0800 Subject: [PATCH 7/7] Fix Qwen prompt version gating --- Sources/LLM/EspressoLLMEngine.swift | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/Sources/LLM/EspressoLLMEngine.swift b/Sources/LLM/EspressoLLMEngine.swift index 1549a30..c413859 100644 --- a/Sources/LLM/EspressoLLMEngine.swift +++ b/Sources/LLM/EspressoLLMEngine.swift @@ -232,10 +232,14 @@ actor EspressoLLMEngine { } static func formatPrompt(user: String, system: String, modelName: String) -> String { - if modelName.lowercased().contains("qwen") { + let normalizedModelName = modelName.lowercased() + if normalizedModelName.contains("qwen") { + let assistantPrefix = normalizedModelName.contains("qwen3") + ? "<|im_start|>assistant\n\n\n\n\n" + : "<|im_start|>assistant\n" return "<|im_start|>system\n\(system)<|im_end|>\n" + "<|im_start|>user\n\(user)<|im_end|>\n" - + "<|im_start|>assistant\n\n\n\n\n" + + assistantPrefix } return "System:\n\(system)\n\nUser:\n\(user)\n\nAssistant:\n" }