Fix guided generation refused after engine load on static-shape engines

swift/Sources/CoreAILanguageModels/InferenceEngines/InferenceEngine.swift

@@ -202,6 +202,15 @@ extension InferenceEngine {
     public var supportsLogits: Bool { false }
 }
 
+extension InferenceEngine {
+    /// Guided/structured generation needs either per-step logits (CPU-side
+    /// constrained decoding) or GPU-side constrained sampling
+    /// (`ConstrainedGenerationCapable`).
+    public var supportsGuidedGeneration: Bool {
+        supportsLogits || self is any ConstrainedGenerationCapable
+    }
+}
+
 extension InferenceEngine {
     /// Default: no prefix hits (engine doesn't track history).
     public var lastPrefixHitCount: Int { 0 }

swift/Sources/CoreAILanguageModels/LanguageModel/CoreAILanguageModel.swift

@@ -199,15 +199,10 @@ public struct CoreAILanguageModel: LanguageModel {
         resources.unloadResources()
     }
 
-    /// Whether guided generation is available for this model.
+    /// Whether guided generation is available for this model. Assumes yes
+    /// before an engine is loaded.
     private var isGuidedGenerationSupported: Bool {
-        if let isConstrainedCapable = resources.loadedEngineIsConstrainedCapable {
-            return isConstrainedCapable
-        }
-        if let supportsLogits = resources.loadedEngineSupportsLogits {
-            return supportsLogits
-        }
-        return true
+        resources.loadedEngineSupportsGuidedGeneration ?? true
     }
 
     // MARK: - Executor
@@ -283,7 +278,7 @@ public struct CoreAILanguageModel: LanguageModel {
 
                 // Check if guided generation is requested
                 if let schema = request.schema {
-                    guard engine.supportsLogits || engine is any ConstrainedGenerationCapable else {
+                    guard engine.supportsGuidedGeneration else {
                         throw LanguageModelError.unsupportedCapability(
                             .init(
                                 capability: .guidedGeneration,

swift/Sources/CoreAILanguageModels/LanguageModel/ModelResources.swift

@@ -36,18 +36,10 @@ final class ModelResources: ResourceManaging {
 
     var isLoaded: Bool { state.withLock { $0.loaded != nil } }
 
-    /// `supportsLogits` of the resident engine, or `nil` when nothing is loaded.
-    /// Used only for best-effort capability reporting before a load.
-    var loadedEngineSupportsLogits: Bool? {
-        state.withLock { $0.loaded?.supportsLogits }
-    }
-
-    /// Whether the loaded engine supports GPU-side constrained generation, or `nil` when unloaded.
-    var loadedEngineIsConstrainedCapable: Bool? {
-        state.withLock { engine in
-            guard let loaded = engine.loaded else { return nil }
-            return loaded is any ConstrainedGenerationCapable
-        }
+    /// Guided-generation capability of the resident engine, or `nil` when nothing
+    /// is loaded. Mirrors the check the executor enforces per-request.
+    var loadedEngineSupportsGuidedGeneration: Bool? {
+        state.withLock { $0.loaded?.supportsGuidedGeneration }
     }
 
     /// Returns the engine, loading it on first use. Concurrent callers share one

swift/Tests/LanguageModelsTests/ModelResourcesTests.swift

@@ -107,6 +107,42 @@ struct ModelResourcesTests {
         #expect(resources.isLoaded)
     }
 
+    @Test(
+        "loadedEngineSupportsGuidedGeneration stays true after loading a logits-only engine"
+    )
+    func guidedGenerationSupportedForLogitsOnlyEngineAfterLoad() async throws {
+        // MockEngine(vocabSize: 100) mirrors StaticShapeEngine: supportsLogits == true,
+        // but it does not conform to ConstrainedGenerationCapable.
+        let resources = ModelResources { MockEngine(vocabSize: 100) }
+
+        // Nothing loaded yet — best-effort reporting defaults to "supported".
+        #expect(resources.loadedEngineSupportsGuidedGeneration == nil)
+
+        _ = try await resources.engine()
+
+        // Regression check: before the fix this incorrectly returned `false` once
+        // an engine was loaded, because `loadedEngineIsConstrainedCapable` (false)
+        // shadowed `loadedEngineSupportsLogits` (true).
+        #expect(resources.loadedEngineSupportsGuidedGeneration == true)
+    }
+
+    @Test(
+        "loadedEngineSupportsGuidedGeneration is true for a constrained-capable engine without logits"
+    )
+    func guidedGenerationSupportedForConstrainedOnlyEngineAfterLoad() async throws {
+        // MockConstrainedEngine conforms to ConstrainedGenerationCapable but has
+        // supportsLogits == false — the other branch of the capability OR. Guided
+        // generation must still be reported supported via GPU-side constrained
+        // sampling.
+        let resources = ModelResources { MockConstrainedEngine(scriptedTokens: []) }
+
+        #expect(resources.loadedEngineSupportsGuidedGeneration == nil)
+
+        _ = try await resources.engine()
+
+        #expect(resources.loadedEngineSupportsGuidedGeneration == true)
+    }
+
     @Test("unloadResources during an active borrow defers teardown until it finishes")
     func unloadDeferredDuringActiveBorrow() async throws {
         let resources = ModelResources { MockEngine() }