fix: unwrap JSON/tag-wrapped model output before it reaches the user

Makefile

@@ -104,14 +104,14 @@ endif
 test: $(TEST_RUNNER)
 	@$(TEST_RUNNER)
 
-$(TEST_RUNNER): Sources/AppContextService.swift Sources/AppleFoundationModelsPostProcessor.swift Sources/DictionaryStore.swift Sources/TranscriptTidier.swift Sources/TransformStore.swift Sources/WakePhraseMatcher.swift Tests/AppContextServiceTests.swift Tests/DictionaryStoreTests.swift Tests/TranscriptTidierTests.swift Tests/TransformStoreTests.swift Tests/WakePhraseMatcherTests.swift Sources/ScratchCommandMatcher.swift Tests/ScratchCommandMatcherTests.swift Sources/RawRevertEligibility.swift Tests/RawRevertEligibilityTests.swift Sources/ShortcutCore/ShortcutModels.swift Tests/ShortcutCancelBindingTests.swift Sources/ShortcutCore/MouseDictationButton.swift Tests/MouseDictationButtonTests.swift Tests/SmartCleanupValidationTests.swift
+$(TEST_RUNNER): Sources/AppContextService.swift Sources/AppleFoundationModelsPostProcessor.swift Sources/DictionaryStore.swift Sources/TranscriptTidier.swift Sources/TransformStore.swift Sources/WakePhraseMatcher.swift Tests/AppContextServiceTests.swift Tests/DictionaryStoreTests.swift Tests/TranscriptTidierTests.swift Tests/TransformStoreTests.swift Tests/WakePhraseMatcherTests.swift Sources/ScratchCommandMatcher.swift Tests/ScratchCommandMatcherTests.swift Sources/RawRevertEligibility.swift Tests/RawRevertEligibilityTests.swift Sources/ShortcutCore/ShortcutModels.swift Tests/ShortcutCancelBindingTests.swift Sources/ShortcutCore/MouseDictationButton.swift Tests/MouseDictationButtonTests.swift Tests/SmartCleanupValidationTests.swift Tests/StructuredOutputUnwrapTests.swift
 	@mkdir -p "$(BUILD_DIR)"
 	swiftc \
 		-parse-as-library \
 		-o "$(TEST_RUNNER)" \
 		-sdk $(shell xcrun --show-sdk-path) \
 		-target $(ARCH)-apple-macosx26.0 \
-		Sources/AppContextService.swift Sources/AppleFoundationModelsPostProcessor.swift Sources/DictionaryStore.swift Sources/TranscriptTidier.swift Sources/TransformStore.swift Sources/WakePhraseMatcher.swift Tests/AppContextServiceTests.swift Tests/DictionaryStoreTests.swift Tests/TranscriptTidierTests.swift Tests/TransformStoreTests.swift Tests/WakePhraseMatcherTests.swift Sources/ScratchCommandMatcher.swift Tests/ScratchCommandMatcherTests.swift Sources/RawRevertEligibility.swift Tests/RawRevertEligibilityTests.swift Sources/ShortcutCore/ShortcutModels.swift Tests/ShortcutCancelBindingTests.swift Sources/ShortcutCore/MouseDictationButton.swift Tests/MouseDictationButtonTests.swift Tests/SmartCleanupValidationTests.swift
+		Sources/AppContextService.swift Sources/AppleFoundationModelsPostProcessor.swift Sources/DictionaryStore.swift Sources/TranscriptTidier.swift Sources/TransformStore.swift Sources/WakePhraseMatcher.swift Tests/AppContextServiceTests.swift Tests/DictionaryStoreTests.swift Tests/TranscriptTidierTests.swift Tests/TransformStoreTests.swift Tests/WakePhraseMatcherTests.swift Sources/ScratchCommandMatcher.swift Tests/ScratchCommandMatcherTests.swift Sources/RawRevertEligibility.swift Tests/RawRevertEligibilityTests.swift Sources/ShortcutCore/ShortcutModels.swift Tests/ShortcutCancelBindingTests.swift Sources/ShortcutCore/MouseDictationButton.swift Tests/MouseDictationButtonTests.swift Tests/SmartCleanupValidationTests.swift Tests/StructuredOutputUnwrapTests.swift
 
 icon: $(ICON_ICNS)
 

Sources/AppleFoundationModelsPostProcessor.swift

@@ -595,7 +595,8 @@ actor AppleFoundationModelsPostProcessor {
     /// (e.g. "wrap this in a div") is never stripped.
     private static let wrapperTags = [
         "response", "result", "output", "answer", "reply", "message",
-        "bulleted_list", "numbered_list", "list", "rewritten_text"
+        "bulleted_list", "numbered_list", "list", "rewritten_text",
+        "cleaned_text", "clean_text"
     ]
     /// Prompt sections the model sometimes replays before its actual answer.
     private static let echoedPromptTags = [
@@ -607,6 +608,12 @@ actor AppleFoundationModelsPostProcessor {
         var value = raw.trimmingCharacters(in: .whitespacesAndNewlines)
         let options: String.CompareOptions = [.regularExpression, .caseInsensitive]
 
+        // The small model sometimes returns its answer wrapped in a markdown
+        // code fence and/or a JSON object like {"cleaned_text": "…"} despite the
+        // instructions. Peel those structured wrappers before the tag loop so
+        // the text they contain — not the scaffolding — reaches the user.
+        value = Self.unwrapStructuredOutput(value)
+
         while !value.isEmpty {
             let before = value
             for tag in Self.echoedPromptTags {
@@ -640,6 +647,68 @@ actor AppleFoundationModelsPostProcessor {
         return value
     }
 
+    /// JSON keys the model uses when it wraps a plain answer in an object, e.g.
+    /// `{"cleaned_text": "…"}`. Ordered by preference for extraction.
+    private static let jsonTextKeys = [
+        "cleaned_text", "clean_text", "cleaned", "corrected_text",
+        "rewritten_text", "text", "output", "result", "response", "answer"
+    ]
+
+    /// Peels a single-purpose JSON object the model emits around its answer, so
+    /// `{"cleaned_text": "Hi."}` — bare or wrapped in a ```` ```json ```` fence —
+    /// collapses to `Hi.`. A code fence is removed *only* when it wraps such an
+    /// object, so a code block the user genuinely dictated survives untouched.
+    private static func unwrapStructuredOutput(_ input: String) -> String {
+        let value = input.trimmingCharacters(in: .whitespacesAndNewlines)
+        if let unwrapped = Self.jsonTextValue(in: value) {
+            return unwrapped.trimmingCharacters(in: .whitespacesAndNewlines)
+        }
+        let defenced = Self.stripCodeFence(value)
+        if defenced != value, let unwrapped = Self.jsonTextValue(in: defenced) {
+            return unwrapped.trimmingCharacters(in: .whitespacesAndNewlines)
+        }
+        return value
+    }
+
+    /// Removes a surrounding ```` ``` ```` fence (with optional language tag)
+    /// only when the whole string is one fenced block, so inline prose that
+    /// merely mentions backticks is untouched.
+    private static func stripCodeFence(_ text: String) -> String {
+        let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines)
+        guard trimmed.hasPrefix("```") else { return trimmed }
+        var lines = trimmed.components(separatedBy: .newlines)
+        guard lines.count >= 2,
+              let opener = lines.first?.trimmingCharacters(in: .whitespaces),
+              opener.range(of: "^```[a-zA-Z0-9+#-]*$", options: .regularExpression) != nil,
+              lines.last?.trimmingCharacters(in: .whitespaces) == "```" else {
+            return trimmed
+        }
+        lines.removeFirst()
+        lines.removeLast()
+        return lines.joined(separator: "\n").trimmingCharacters(in: .whitespacesAndNewlines)
+    }
+
+    /// Returns the wrapped text when `input` is a JSON object whose keys are all
+    /// known answer-wrapper keys. Requiring *every* key to be known leaves a
+    /// genuinely dictated object such as `{"name": "Ada"}` untouched.
+    private static func jsonTextValue(in input: String) -> String? {
+        let trimmed = input.trimmingCharacters(in: .whitespacesAndNewlines)
+        guard trimmed.hasPrefix("{"), trimmed.hasSuffix("}"),
+              let data = trimmed.data(using: .utf8),
+              let object = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
+              !object.isEmpty else {
+            return nil
+        }
+        var lowered: [String: Any] = [:]
+        for (key, value) in object { lowered[key.lowercased()] = value }
+        let known = Set(Self.jsonTextKeys)
+        guard lowered.keys.allSatisfy({ known.contains($0) }) else { return nil }
+        for key in Self.jsonTextKeys {
+            if let value = lowered[key] as? String { return value }
+        }
+        return nil
+    }
+
     private func makeSession(instructions: String) -> LanguageModelSession {
         LanguageModelSession(model: model, tools: [], instructions: instructions)
     }

Tests/AppContextServiceTests.swift

@@ -29,6 +29,7 @@ struct AppContextServiceTests {
         ShortcutCancelBindingTests.run()
         MouseDictationButtonTests.run()
         SmartCleanupValidationTests.run()
+        StructuredOutputUnwrapTests.run()
         print("MegaphoneTests passed")
     }
 

Tests/StructuredOutputUnwrapTests.swift

@@ -0,0 +1,66 @@
+import Foundation
+
+/// Covers issue #14: the on-device model occasionally wraps its answer in a JSON
+/// object or `<cleaned_text>` tags, which `normalizeCommandOutput` must peel so
+/// the user gets their text and not the scaffolding.
+enum StructuredOutputUnwrapTests {
+    static func run() {
+        testBareJSONObjectIsUnwrapped()
+        testFencedJSONObjectIsUnwrapped()
+        testCleanedTextTagsAreStripped()
+        testGenuineDataObjectIsUntouched()
+        testDictatedCodeBlockSurvives()
+        testOrdinaryProseIsUnchanged()
+    }
+
+    private static func testBareJSONObjectIsUnwrapped() {
+        expect("{\n  \"cleaned_text\": \"Testing, testing.\"\n}", becomes: "Testing, testing.")
+        expect("{\"text\": \"- YOLO.\"}", becomes: "- YOLO.")
+        expect("{\"clean_text\":\"What's up?\"}", becomes: "What's up?")
+    }
+
+    private static func testFencedJSONObjectIsUnwrapped() {
+        expect("```json\n{\n  \"cleaned_text\": \"Testing, testing.\"\n}\n```", becomes: "Testing, testing.")
+        expect("```\n{\"text\": \"- YOLO.\"}\n```", becomes: "- YOLO.")
+    }
+
+    private static func testCleanedTextTagsAreStripped() {
+        expect("<cleaned_text>\nWhat's up, my fellow humans?\n</cleaned_text>", becomes: "What's up, my fellow humans?")
+        expect("<clean_text>How do you do, fellow kids?</clean_text>", becomes: "How do you do, fellow kids?")
+    }
+
+    /// An object whose keys are not all answer-wrappers is real dictated content,
+    /// not scaffolding, so it must pass through verbatim.
+    private static func testGenuineDataObjectIsUntouched() {
+        let object = "{\"name\": \"Ada\", \"text\": \"hello\"}"
+        expect(object, becomes: object)
+    }
+
+    /// De-fencing is scoped to JSON wrappers, so a code block the user actually
+    /// dictated (e.g. via a wake command) keeps its fence.
+    private static func testDictatedCodeBlockSurvives() {
+        let block = "```swift\nprint(\"hi\")\n```"
+        expect(block, becomes: block)
+    }
+
+    private static func testOrdinaryProseIsUnchanged() {
+        expect("Let's ship the release on Wednesday.", becomes: "Let's ship the release on Wednesday.")
+        // A stray brace in prose is not a JSON object and must not be touched.
+        expect("Use the {placeholder} token here.", becomes: "Use the {placeholder} token here.")
+    }
+
+    private static func expect(
+        _ input: String,
+        becomes expected: String,
+        file: StaticString = #file,
+        line: UInt = #line
+    ) {
+        let actual = AppleFoundationModelsPostProcessor.normalizeCommandOutput(input)
+        if actual != expected {
+            fatalError(
+                "\(file):\(line): normalizeCommandOutput(\(input.debugDescription)) == "
+                    + "\(actual.debugDescription), expected \(expected.debugDescription)"
+            )
+        }
+    }
+}