From eb756b54fadd56d20fa3c0768b79cc54448bfd8b Mon Sep 17 00:00:00 2001 From: Mattt Zmuda Date: Thu, 17 Sep 2026 04:43:20 -0700 Subject: [PATCH 1/4] Port JSONCompleter and drop the PartialJSONDecoder dependency GeneratedContent(json:) completed truncated JSON by appending a single closing brace, bracket, or empty string. That failed on every realistic streaming cut, including nested objects, open arrays, partial literals, and the example in its own doc comment, and turned a trailing backslash into a literal quote. The system model separately pulled in PartialJSONDecoder to get real completion for the same job. Port JSONCompleter from PartialJSONDecoder as an internal type and use it in GeneratedContent(json:). It walks the text as a JSON value and appends exactly the closing characters the input is missing, with a depth limit. Text that ends inside an escape sequence is cut back to the escape start so the result stays valid JSON. The system model now calls GeneratedContent(json:) directly, and the package dependency is removed. --- Package.resolved | 11 +- Package.swift | 2 - .../AnyLanguageModel/GeneratedContent.swift | 32 +- .../Models/SystemLanguageModel.swift | 19 +- .../Shared/JSONCompleter.swift | 411 ++++++++++++++++++ .../GeneratedContentJSONTests.swift | 28 +- .../JSONCompleterTests.swift | 139 ++++++ 7 files changed, 589 insertions(+), 53 deletions(-) create mode 100644 Sources/AnyLanguageModel/Shared/JSONCompleter.swift create mode 100644 Tests/AnyLanguageModelTests/JSONCompleterTests.swift diff --git a/Package.resolved b/Package.resolved index e486dbf2..bd83c028 100644 --- a/Package.resolved +++ b/Package.resolved @@ -1,5 +1,5 @@ { - "originHash" : "b46e46d156bb5bfeea0c5c4e9b59d51ebea883a40627ee34aed14b97f142276d", + "originHash" : "afcaf7d64f1529491a2f05d864dc9f4984f3ed6c77a583eb96b474a6a001e3ea", "pins" : [ { "identity" : "eventsource", @@ -28,15 +28,6 @@ "version" : "2.10549.0" } }, - { - "identity" : "partialjsondecoder", - "kind" : "remoteSourceControl", - "location" : "https://github.com/mattt/PartialJSONDecoder", - "state" : { - "revision" : "e4d389e6bcc6771bb988d1a8a17695d8bfa97172", - "version" : "1.0.0" - } - }, { "identity" : "swift-asn1", "kind" : "remoteSourceControl", diff --git a/Package.swift b/Package.swift index 2c4a8d63..e7bdfa5f 100644 --- a/Package.swift +++ b/Package.swift @@ -41,7 +41,6 @@ let package = Package( ), .package(url: "https://github.com/mattt/JSONSchema", from: "1.3.0"), .package(url: "https://github.com/mattt/llama.swift", .upToNextMajor(from: "2.10549.0")), - .package(url: "https://github.com/mattt/PartialJSONDecoder", from: "1.0.0"), .package(url: "https://github.com/ml-explore/mlx-swift-lm", from: "3.31.4"), .package(url: "https://github.com/swiftlang/swift-syntax", from: "602.0.0"), .package(url: "https://github.com/swift-server/async-http-client.git", from: "1.24.0"), @@ -53,7 +52,6 @@ let package = Package( .target(name: "AnyLanguageModelMacros"), .product(name: "EventSource", package: "EventSource"), .product(name: "JSONSchema", package: "JSONSchema"), - .product(name: "PartialJSONDecoder", package: "PartialJSONDecoder"), .product( name: "MLXLLM", package: "mlx-swift-lm", diff --git a/Sources/AnyLanguageModel/GeneratedContent.swift b/Sources/AnyLanguageModel/GeneratedContent.swift index 82f460b3..7b8cc15b 100644 --- a/Sources/AnyLanguageModel/GeneratedContent.swift +++ b/Sources/AnyLanguageModel/GeneratedContent.swift @@ -179,33 +179,17 @@ public struct GeneratedContent: Sendable, Equatable, Generable, CustomDebugStrin return } - // Handle incomplete JSON by attempting to complete it - let completedJSON = String(decoding: data, as: UTF8.self) - .trimmingCharacters(in: .whitespacesAndNewlines) - - // Try adding closing braces/brackets to make it valid - var attempts: [String] = [completedJSON] - - // If it looks like an incomplete object, try closing it - if completedJSON.hasPrefix("{") && !completedJSON.hasSuffix("}") { - attempts.append(completedJSON + "}") - attempts.append(completedJSON + "\"\"}") // incomplete string value - } - - // If it looks like an incomplete array, try closing it - if completedJSON.hasPrefix("[") && !completedJSON.hasSuffix("]") { - attempts.append(completedJSON + "]") - } - - for attempt in attempts { - if let parsed = try? JSONSerialization.jsonObject(with: Data(attempt.utf8), options: [.fragmentsAllowed]) { - self = try Self.fromJSONObject(parsed) - return - } + // Handle incomplete JSON by completing it and parsing again + let json = String(decoding: data, as: UTF8.self) + if let completed = try? JSONCompleter().complete(json), + let parsed = try? JSONSerialization.jsonObject(with: Data(completed.utf8), options: [.fragmentsAllowed]) + { + self = try Self.fromJSONObject(parsed) + return } // If all else fails, treat it as a string - self.init(kind: .string(completedJSON)) + self.init(kind: .string(json.trimmingCharacters(in: .whitespacesAndNewlines))) } private static func fromJSONObject(_ value: Any) throws -> GeneratedContent { diff --git a/Sources/AnyLanguageModel/Models/SystemLanguageModel.swift b/Sources/AnyLanguageModel/Models/SystemLanguageModel.swift index 834ec577..c8ff63b3 100644 --- a/Sources/AnyLanguageModel/Models/SystemLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/SystemLanguageModel.swift @@ -1,7 +1,6 @@ #if canImport(FoundationModels) import FoundationModels import Foundation - import PartialJSONDecoder import JSONSchema @@ -788,9 +787,7 @@ return finalize(content: content) } catch { // Attempt partial JSON decoding before surfacing an error. - let decoder = PartialJSONDecoder() - let jsonString = fmResponse.content.jsonString - if let partialContent = try? decoder.decode(GeneratedContent.self, from: jsonString).value, + if let partialContent = try? GeneratedContent(json: fmResponse.content.jsonString), let content = try? type.init(partialContent) { return finalize(content: content) @@ -870,7 +867,6 @@ func processStructuredStream(_ fmSession: FoundationModels.LanguageModelSession) async { let fmSchema = FoundationModels.GenerationSchema(schema) - let partialDecoder = PartialJSONDecoder() let fmStream = fmSession.streamResponse( to: fmPrompt, schema: fmSchema, @@ -898,12 +894,7 @@ lastLength: &lastLength ) - let jsonString = accumulatedText - if let partialContent = try? partialDecoder.decode( - GeneratedContent.self, - from: jsonString - ) - .value { + if let partialContent = try? GeneratedContent(json: accumulatedText) { let partial: Content.PartiallyGenerated? = try? .init(partialContent) if let partial { continuation.yield(.init(content: partial, rawContent: partialContent)) @@ -937,11 +928,7 @@ ?? GeneratedContent(jsonString) // Prefer partial decoding so we can surface intermediate snapshots. - if let partialContent = try? partialDecoder.decode( - GeneratedContent.self, - from: jsonString - ) - .value { + if let partialContent = try? GeneratedContent(json: jsonString) { let partial: Content.PartiallyGenerated? = try? .init(partialContent) if let partial { continuation.yield(.init(content: partial, rawContent: partialContent)) diff --git a/Sources/AnyLanguageModel/Shared/JSONCompleter.swift b/Sources/AnyLanguageModel/Shared/JSONCompleter.swift new file mode 100644 index 00000000..faf728b0 --- /dev/null +++ b/Sources/AnyLanguageModel/Shared/JSONCompleter.swift @@ -0,0 +1,411 @@ +import Foundation + +// Ported from PartialJSONDecoder (https://github.com/mattt/PartialJSONDecoder), MIT licensed. + +/// An error that occurs while completing partial JSON. +enum JSONCompletionError: Error, Equatable { + /// The input contains a value that JSON does not allow, such as `NaN` or `Infinity`. + case invalidValue(String) + + /// The input nests deeper than the completer's limit. + case depthLimitExceeded(Int) +} + +/// Completes partial JSON text by appending the closing characters it is missing. +/// +/// The completer scans the text as a JSON value and computes the suffix needed to +/// close every open string, array, and object at the point where the text ends. +/// Unfinished literals such as `tr` or `1.` are completed to `true` and `1.0`, +/// a key with no value receives `null`, and a trailing comma is dropped. +/// +/// The result is intended for a JSON parser, not for display: +/// text that is already complete is returned unchanged, +/// and text that is not JSON at all is returned unchanged too. +struct JSONCompleter: Sendable { + /// The completion for a partial JSON value. + /// + /// `string` holds the characters to append, + /// and `endIndex` is the index in the original text immediately after the portion to keep. + /// The kept portion can be shorter than the input when the text ends inside an escape sequence. + typealias Completion = (string: String, endIndex: String.Index) + + /// The maximum nesting depth the completer accepts before it throws. + /// + /// This bounds recursion on adversarial or malformed input. + var maximumDepth: Int = 64 + + /// Creates a completer with the default depth limit. + init() {} + + /// Returns the input with any missing closing characters appended. + /// + /// - Parameter json: Partial JSON text. + /// - Returns: Text that closes every structure the input left open. + /// - Throws: ``JSONCompletionError`` if the input contains a non-JSON literal + /// or nests deeper than ``maximumDepth``. + func complete(_ json: String) throws -> String { + guard !json.isEmpty else { return "" } + + if let completion = try completion(for: json, from: json.startIndex) { + return json[.. Completion? { + let start = skipWhitespace(json, from: startIndex) + guard start < json.endIndex else { return nil } + + return try completeValue(json, from: start, depth: 0) + } + + // MARK: - + + private func skipWhitespace(_ json: String, from index: String.Index) -> String.Index { + var current = index + while current < json.endIndex && json[current].isWhitespace { + current = json.index(after: current) + } + return current + } + + private func completeValue(_ json: String, from startIndex: String.Index, depth: Int) throws -> Completion? { + guard depth < maximumDepth else { + throw JSONCompletionError.depthLimitExceeded(maximumDepth) + } + + let start = skipWhitespace(json, from: startIndex) + guard start < json.endIndex else { return nil } + + switch json[start] { + case "{": + return try completeObject(json, from: start, depth: depth + 1) + case "[": + return try completeArray(json, from: start, depth: depth + 1) + case "\"": + return completeString(json, from: start) + case "-": + let next = json.index(after: start) + if next < json.endIndex, json[next] == "I" { + throw JSONCompletionError.invalidValue("-Infinity") + } + return completeNumber(json, from: start) + case "0" ... "9": + return completeNumber(json, from: start) + case "t": + return completeLiteral(json, from: start, literal: "true") + case "f": + return completeLiteral(json, from: start, literal: "false") + case "n": + return completeLiteral(json, from: start, literal: "null") + case "I": + throw JSONCompletionError.invalidValue("Infinity") + case "N": + throw JSONCompletionError.invalidValue("NaN") + default: + return nil + } + } + + /// Completes a string that starts at the given index. + /// + /// If the text ends inside an escape sequence, the completion keeps only the + /// text before the backslash so that the result is valid JSON. + private func completeString(_ json: String, from startIndex: String.Index) -> Completion? { + guard startIndex < json.endIndex, json[startIndex] == "\"" else { return nil } + + var current = json.index(after: startIndex) + while current < json.endIndex { + let char = json[current] + if char == "\\" { + let escapeStart = current + current = json.index(after: current) + guard current < json.endIndex else { + return (string: "\"", endIndex: escapeStart) + } + if json[current] == "u" { + var remaining = 4 + current = json.index(after: current) + while remaining > 0 && current < json.endIndex && json[current].isHexDigit { + current = json.index(after: current) + remaining -= 1 + } + if remaining > 0 { + return (string: "\"", endIndex: escapeStart) + } + continue + } + } else if char == "\"" { + return nil + } + current = json.index(after: current) + } + + return (string: "\"", endIndex: current) + } + + private func completeArray(_ json: String, from startIndex: String.Index, depth: Int) throws -> Completion? { + guard startIndex < json.endIndex, json[startIndex] == "[" else { return nil } + + var current = skipWhitespace(json, from: json.index(after: startIndex)) + var requiresComma = false + var lastValidIndex = current + + if current >= json.endIndex || json[current] == "]" { + return (string: "]", endIndex: current) + } + + while current < json.endIndex { + if json[current] == "]" { + return nil + } + + if requiresComma { + guard json[current] == "," else { + return (string: "]", endIndex: lastValidIndex) + } + requiresComma = false + current = skipWhitespace(json, from: json.index(after: current)) + if current >= json.endIndex { break } + lastValidIndex = current + } + + if json[current] == "]" { + return nil + } + + if let elementCompletion = try completeValue(json, from: current, depth: depth + 1) { + return (string: elementCompletion.string + "]", endIndex: elementCompletion.endIndex) + } + + current = findEndOfCompleteValue(json, from: current) + lastValidIndex = current + requiresComma = true + } + + return (string: "]", endIndex: lastValidIndex) + } + + private func completeObject(_ json: String, from startIndex: String.Index, depth: Int) throws -> Completion? { + guard startIndex < json.endIndex, json[startIndex] == "{" else { return nil } + + var current = skipWhitespace(json, from: json.index(after: startIndex)) + var requiresComma = false + var lastValidIndex = current + + if current >= json.endIndex || json[current] == "}" { + return (string: "}", endIndex: current) + } + + while current < json.endIndex { + if json[current] == "}" { + return nil + } + + if requiresComma { + guard json[current] == "," else { + return (string: "}", endIndex: lastValidIndex) + } + requiresComma = false + current = skipWhitespace(json, from: json.index(after: current)) + if current >= json.endIndex { break } + lastValidIndex = current + } + + if json[current] == "}" { + return nil + } + + // Key + if let keyCompletion = completeString(json, from: current) { + return (string: keyCompletion.string + ": null}", endIndex: keyCompletion.endIndex) + } + let keyEnd = findEndOfCompleteValue(json, from: current) + guard keyEnd > current else { + return (string: "}", endIndex: lastValidIndex) + } + current = keyEnd + lastValidIndex = current + + // Colon + current = skipWhitespace(json, from: current) + guard current < json.endIndex, json[current] == ":" else { + return (string: ": null}", endIndex: lastValidIndex) + } + current = json.index(after: current) + lastValidIndex = current + + // Value + current = skipWhitespace(json, from: current) + guard current < json.endIndex else { + return (string: "null}", endIndex: lastValidIndex) + } + + if let valueCompletion = try completeValue(json, from: current, depth: depth + 1) { + return (string: valueCompletion.string + "}", endIndex: valueCompletion.endIndex) + } + + current = findEndOfCompleteValue(json, from: current) + lastValidIndex = current + requiresComma = true + } + + return (string: "}", endIndex: lastValidIndex) + } + + private func completeNumber(_ json: String, from startIndex: String.Index) -> Completion? { + var current = startIndex + + if current < json.endIndex && json[current] == "-" { + current = json.index(after: current) + } + + guard current < json.endIndex else { + return (string: "0", endIndex: current) + } + + if json[current] == "." { + return (string: "0.0", endIndex: current) + } + + while current < json.endIndex && json[current].isNumber { + current = json.index(after: current) + } + + if current < json.endIndex && json[current] == "." { + current = json.index(after: current) + let fractionStart = current + while current < json.endIndex && json[current].isNumber { + current = json.index(after: current) + } + if current == fractionStart { + return (string: "0", endIndex: current) + } + } + + if current < json.endIndex && (json[current] == "e" || json[current] == "E") { + current = json.index(after: current) + if current < json.endIndex && (json[current] == "+" || json[current] == "-") { + current = json.index(after: current) + } + if current >= json.endIndex || !json[current].isNumber { + return (string: "0", endIndex: current) + } + while current < json.endIndex && json[current].isNumber { + current = json.index(after: current) + } + } + + return nil + } + + private func completeLiteral(_ json: String, from startIndex: String.Index, literal: String) -> Completion? { + var current = startIndex + var remaining = literal[...] + + while current < json.endIndex, let expected = remaining.first { + guard json[current] == expected else { return nil } + current = json.index(after: current) + remaining = remaining.dropFirst() + } + + guard !remaining.isEmpty else { return nil } + return (string: String(remaining), endIndex: current) + } + + /// Returns the index immediately after the complete value that starts at the given index. + private func findEndOfCompleteValue(_ json: String, from startIndex: String.Index) -> String.Index { + let start = skipWhitespace(json, from: startIndex) + guard start < json.endIndex else { return start } + + if let completion = try? completeValue(json, from: start, depth: 0) { + return completion.endIndex + } + + switch json[start] { + case "\"": + var current = json.index(after: start) + var isEscaped = false + while current < json.endIndex { + let char = json[current] + if char == "\\" { + isEscaped.toggle() + } else if char == "\"" && !isEscaped { + return json.index(after: current) + } else { + isEscaped = false + } + current = json.index(after: current) + } + return current + case "{": + return findMatchingBrace(json, from: start, open: "{", close: "}") + case "[": + return findMatchingBrace(json, from: start, open: "[", close: "]") + case "t" where json[start...].hasPrefix("true"): + return json.index(start, offsetBy: 4) + case "f" where json[start...].hasPrefix("false"): + return json.index(start, offsetBy: 5) + case "n" where json[start...].hasPrefix("null"): + return json.index(start, offsetBy: 4) + case "-", "0" ... "9": + var current = start + while current < json.endIndex && "0123456789.-+eE".contains(json[current]) { + current = json.index(after: current) + } + return current + default: + return start + } + } + + /// Returns the index immediately after the bracket that closes the one at the given index. + private func findMatchingBrace( + _ json: String, + from startIndex: String.Index, + open: Character, + close: Character + ) -> String.Index { + var level = 0 + var current = startIndex + var inString = false + var isEscaped = false + + while current < json.endIndex { + let char = json[current] + + if inString { + if char == "\\" { + isEscaped.toggle() + } else if char == "\"" && !isEscaped { + inString = false + } else { + isEscaped = false + } + } else if char == "\"" { + inString = true + isEscaped = false + } else if char == open { + level += 1 + } else if char == close { + level -= 1 + if level == 0 { + return json.index(after: current) + } + } + current = json.index(after: current) + } + + return current + } +} diff --git a/Tests/AnyLanguageModelTests/GeneratedContentJSONTests.swift b/Tests/AnyLanguageModelTests/GeneratedContentJSONTests.swift index 6aae424e..c2d94824 100644 --- a/Tests/AnyLanguageModelTests/GeneratedContentJSONTests.swift +++ b/Tests/AnyLanguageModelTests/GeneratedContentJSONTests.swift @@ -40,7 +40,7 @@ struct GeneratedContentJSONTests { } @Test func dataInitializerCompletesPartialJSON() throws { - let partial = Data(#"{"title": "A story of""#.utf8) + let partial = Data(#"{"title": "A story of"#.utf8) let content = try GeneratedContent(json: partial) #expect(try content.value(String.self, forProperty: "title") == "A story of") @@ -48,6 +48,32 @@ struct GeneratedContentJSONTests { #expect(try GeneratedContent(json: partialArray).kind == .array([1, 2, 3].map { GeneratedContent($0) })) } + @Test(arguments: [ + (#"{"a": {"b": "x"#, #"{"a": {"b": "x"}}"#), + (#"{"a": [1, 2"#, #"{"a": [1, 2]}"#), + (#"[{"a": 1}, {"b": "#, #"[{"a": 1}, {"b": null}]"#), + (#"{"a": 1, "b": tr"#, #"{"a": 1, "b": true}"#), + (#"{"a": 12."#, #"{"a": 12.0}"#), + (#"{"a": "esc\"#, #"{"a": "esc"}"#), + (#"{"a": "q\" more\u00"#, #"{"a": "q\" more"}"#), + (#"{"na"#, #"{"na": null}"#), + ]) + func completesTruncatedStreamingJSON(partial: String, expected: String) throws { + let fromString = try GeneratedContent(json: partial) + let fromData = try GeneratedContent(json: Data(partial.utf8)) + let expectedContent = try GeneratedContent(json: expected) + #expect(fromString.jsonValue == expectedContent.jsonValue) + #expect(fromData.jsonValue == expectedContent.jsonValue) + } + + @Test func partialGenerableDecodesFromTruncatedJSON() throws { + let partial = #"{"title": "Dune", "pages": 41"# + let idea = try NovelIdea.PartiallyGenerated(GeneratedContent(json: partial)) + #expect(idea.title == "Dune") + #expect(idea.pages == 41) + #expect(idea.tags == nil) + } + @Test func dataInitializerFallsBackToString() throws { let content = try GeneratedContent(json: Data(" not json ".utf8)) #expect(content.kind == .string("not json")) diff --git a/Tests/AnyLanguageModelTests/JSONCompleterTests.swift b/Tests/AnyLanguageModelTests/JSONCompleterTests.swift new file mode 100644 index 00000000..d87f3ca4 --- /dev/null +++ b/Tests/AnyLanguageModelTests/JSONCompleterTests.swift @@ -0,0 +1,139 @@ +import Foundation +import Testing + +@testable import AnyLanguageModel + +@Suite("JSONCompleter") +struct JSONCompleterTests { + let completer = JSONCompleter() + + @Test func leavesCompleteJSONUnchanged() throws { + #expect(try completer.complete("42") == "42") + #expect(try completer.complete("\"hello\"") == "\"hello\"") + #expect(try completer.complete("[1, 2, 3]") == "[1, 2, 3]") + #expect(try completer.complete("{\"a\": 1}") == "{\"a\": 1}") + #expect(try completer.complete("true") == "true") + #expect(try completer.complete("null") == "null") + } + + @Test func closesOpenStructures() throws { + #expect(try completer.complete("[1, 2, 3") == "[1, 2, 3]") + #expect(try completer.complete("{\"a\": 1") == "{\"a\": 1}") + #expect(try completer.complete("\"hello") == "\"hello\"") + #expect( + try completer.complete(#"{"name": "Alice", "age": 30, "hobbies": ["reading", "hiking"#) + == #"{"name": "Alice", "age": 30, "hobbies": ["reading", "hiking"]}"# + ) + } + + @Test func reportsCompletionSuffix() throws { + #expect(try completer.completion(for: "[1, 2, 3]", from: "[1, 2, 3]".startIndex) == nil) + #expect(try completer.completion(for: "{\"a\": 1}", from: "{\"a\": 1}".startIndex) == nil) + + #expect(try completer.completion(for: "[1, 2, 3", from: "[1, 2, 3".startIndex)?.string == "]") + #expect(try completer.completion(for: "{\"a\": 1", from: "{\"a\": 1".startIndex)?.string == "}") + #expect(try completer.completion(for: "\"hello", from: "\"hello".startIndex)?.string == "\"") + + let nested = "{\"obj\": {\"arr\": [1, 2," + #expect(try completer.completion(for: nested, from: nested.startIndex)?.string == "]}}") + + let partial = "{\"complete\": true, \"partial\": {\"arr\": [1, 2," + let midIndex = partial.lastIndex(of: "{")! + #expect(try completer.completion(for: partial, from: midIndex)?.string == "]}") + } + + @Test func enforcesDepthLimit() throws { + var limited = JSONCompleter() + limited.maximumDepth = 10 + + #expect(throws: JSONCompletionError.depthLimitExceeded(10)) { + try limited.complete(String(repeating: "[", count: 20)) + } + #expect(try limited.complete(String(repeating: "[", count: 5)) == "[[[[[]]]]]") + #expect(JSONCompleter().maximumDepth >= 32) + } + + @Test func rejectsNonJSONNumbers() { + #expect(throws: JSONCompletionError.invalidValue("NaN")) { try completer.complete("[NaN") } + #expect(throws: JSONCompletionError.invalidValue("Infinity")) { try completer.complete("Infinity") } + #expect(throws: JSONCompletionError.invalidValue("-Infinity")) { try completer.complete("{\"a\": -Inf") } + } + + @Test func closesComplexNestedStructure() throws { + let partial = """ + { + "name": "Complex Test", + "data": { + "numbers": [1, 2, 3, 4, 5], + "boolean": true, + "nested": { + "array": [ + {"id": 1, "value": "first"}, + {"id": 2, "value": "second"}, + {"id": 3, "value": "third" + """ + let completed = try completer.complete(partial) + #expect(completed.hasSuffix("}]}}}")) + #expect(isValidJSON(completed)) + } + + @Test func handlesEscapes() throws { + let escaped = #""Special \"quoted\" and \n newline and \t tab and ♥ unicode"# + #expect(try completer.complete(escaped) == escaped + "\"") + + // Text that ends inside an escape sequence is cut back to the escape start. + #expect(try completer.complete(#""Partial escape: \"#) == #""Partial escape: ""#) + #expect(try completer.complete(#""Unicode escape: \u26"#) == #""Unicode escape: ""#) + #expect(try completer.complete(#"{"a": "esc\"#) == #"{"a": "esc"}"#) + #expect(try completer.complete(#"{"a": "done\" next"#) == #"{"a": "done\" next"}"#) + #expect(isValidJSON(try completer.complete(#"{"a": "x\\"#))) + } + + @Test func handlesEmptyAndWhitespaceInput() throws { + #expect(try completer.complete("") == "") + #expect(try completer.complete(" ") == " ") + #expect(try completer.complete("\n\t\r ") == "\n\t\r ") + #expect(try completer.complete("{ ") == "{ }") + #expect(try completer.complete("[ ") == "[ ]") + } + + @Test func completesPartialNumbers() throws { + #expect(try completer.complete("123") == "123") + #expect(try completer.complete("123.") == "123.0") + #expect(try completer.complete("123.4") == "123.4") + #expect(try completer.complete("-") == "-0") + #expect(try completer.complete("-.") == "-0.0") + #expect(try completer.complete("-123.") == "-123.0") + #expect(try completer.complete("1.23e") == "1.23e0") + #expect(try completer.complete("1.23e+") == "1.23e+0") + #expect(try completer.complete("1.23e-") == "1.23e-0") + } + + @Test func completesPartialLiterals() throws { + #expect(try completer.complete("{\"a\": 1, \"b\": tr") == "{\"a\": 1, \"b\": true}") + #expect(try completer.complete("[fa") == "[false]") + #expect(try completer.complete("{\"a\": n") == "{\"a\": null}") + } + + @Test func completesObjectsWithMissingValues() throws { + #expect(try completer.complete("{\"key\":") == "{\"key\":null}") + #expect(try completer.complete("{\"key\": \"value") == "{\"key\": \"value\"}") + #expect(try completer.complete("{\"key1\": true, \"key2\":") == "{\"key1\": true, \"key2\":null}") + #expect(try completer.complete("{\"key\": 42,") == "{\"key\": 42}") + #expect(try completer.complete("{\"ke") == "{\"ke\": null}") + #expect( + try completer.complete("{\"outer\": {\"inner\": [1, 2, {\"nested\":") + == "{\"outer\": {\"inner\": [1, 2, {\"nested\":null}]}}" + ) + } + + @Test func completesArraysWithMissingValues() throws { + #expect(try completer.complete("[1, 2, 3,") == "[1, 2, 3]") + #expect(try completer.complete("[1, 2,") == "[1, 2]") + #expect(try completer.complete("[[1, 2], [3,") == "[[1, 2], [3]]") + } + + private func isValidJSON(_ json: String) -> Bool { + (try? JSONSerialization.jsonObject(with: Data(json.utf8), options: [.fragmentsAllowed])) != nil + } +} From ba1f8d6379f17a480f4f336b6cf77de96659674d Mon Sep 17 00:00:00 2001 From: Mattt Zmuda Date: Thu, 17 Sep 2026 05:13:47 -0700 Subject: [PATCH 2/4] Fix JSONCompleter dropping values after empty containers or padded commas A complete empty object or array was reported as needing a closing character, so the enclosing container closed right after it and every later value was discarded. Whitespace between a complete value and the following comma was treated as a missing comma with the same effect. Treat an empty container as complete and skip whitespace after a value before looking for the comma. Also cite the exact upstream revision in the attribution: the ported source is MIT licensed at that revision, and the same source ships in the 1.0.0 release under Apache-2.0. --- .../Shared/JSONCompleter.swift | 10 ++++--- .../GeneratedContentJSONTests.swift | 2 ++ .../JSONCompleterTests.swift | 27 +++++++++++++++++++ 3 files changed, 36 insertions(+), 3 deletions(-) diff --git a/Sources/AnyLanguageModel/Shared/JSONCompleter.swift b/Sources/AnyLanguageModel/Shared/JSONCompleter.swift index faf728b0..1d96e3ac 100644 --- a/Sources/AnyLanguageModel/Shared/JSONCompleter.swift +++ b/Sources/AnyLanguageModel/Shared/JSONCompleter.swift @@ -1,6 +1,8 @@ import Foundation -// Ported from PartialJSONDecoder (https://github.com/mattt/PartialJSONDecoder), MIT licensed. +// Ported from PartialJSONDecoder (https://github.com/mattt/PartialJSONDecoder) +// at revision d331b237cafe56c5233557bdf1f5a3415393435c, MIT licensed. +// The same source ships in the 1.0.0 release under the Apache License 2.0. /// An error that occurs while completing partial JSON. enum JSONCompletionError: Error, Equatable { @@ -161,7 +163,7 @@ struct JSONCompleter: Sendable { var requiresComma = false var lastValidIndex = current - if current >= json.endIndex || json[current] == "]" { + if current >= json.endIndex { return (string: "]", endIndex: current) } @@ -190,6 +192,7 @@ struct JSONCompleter: Sendable { current = findEndOfCompleteValue(json, from: current) lastValidIndex = current + current = skipWhitespace(json, from: current) requiresComma = true } @@ -203,7 +206,7 @@ struct JSONCompleter: Sendable { var requiresComma = false var lastValidIndex = current - if current >= json.endIndex || json[current] == "}" { + if current >= json.endIndex { return (string: "}", endIndex: current) } @@ -257,6 +260,7 @@ struct JSONCompleter: Sendable { current = findEndOfCompleteValue(json, from: current) lastValidIndex = current + current = skipWhitespace(json, from: current) requiresComma = true } diff --git a/Tests/AnyLanguageModelTests/GeneratedContentJSONTests.swift b/Tests/AnyLanguageModelTests/GeneratedContentJSONTests.swift index c2d94824..d6bca8ae 100644 --- a/Tests/AnyLanguageModelTests/GeneratedContentJSONTests.swift +++ b/Tests/AnyLanguageModelTests/GeneratedContentJSONTests.swift @@ -57,6 +57,8 @@ struct GeneratedContentJSONTests { (#"{"a": "esc\"#, #"{"a": "esc"}"#), (#"{"a": "q\" more\u00"#, #"{"a": "q\" more"}"#), (#"{"na"#, #"{"na": null}"#), + (#"{"a": {}, "b": 2"#, #"{"a": {}, "b": 2}"#), + (#"{"a": [] , "b": [1 , 2"#, #"{"a": [], "b": [1, 2]}"#), ]) func completesTruncatedStreamingJSON(partial: String, expected: String) throws { let fromString = try GeneratedContent(json: partial) diff --git a/Tests/AnyLanguageModelTests/JSONCompleterTests.swift b/Tests/AnyLanguageModelTests/JSONCompleterTests.swift index d87f3ca4..aba1430a 100644 --- a/Tests/AnyLanguageModelTests/JSONCompleterTests.swift +++ b/Tests/AnyLanguageModelTests/JSONCompleterTests.swift @@ -127,6 +127,33 @@ struct JSONCompleterTests { ) } + @Test func keepsValuesAfterEmptyNestedContainers() throws { + #expect(try completer.complete(#"{"a": {}, "b": 2"#) == #"{"a": {}, "b": 2}"#) + #expect(try completer.complete(#"{"a": [], "b": 2"#) == #"{"a": [], "b": 2}"#) + #expect(try completer.complete(#"[{}, 2"#) == #"[{}, 2]"#) + #expect(try completer.complete(#"[[], [1"#) == #"[[], [1]]"#) + #expect(try completer.complete(#"{"a": { }, "b""#) == #"{"a": { }, "b": null}"#) + #expect(try completer.complete("{}") == "{}") + #expect(try completer.complete("[]") == "[]") + #expect(try completer.completion(for: "{}", from: "{}".startIndex) == nil) + #expect(try completer.completion(for: "[ ]", from: "[ ]".startIndex) == nil) + } + + @Test func keepsValuesAfterWhitespaceBeforeComma() throws { + #expect(try completer.complete("[1 , 2") == "[1 , 2]") + #expect(try completer.complete("[1 ,") == "[1]") + #expect(try completer.complete("[1 ") == "[1]") + #expect(try completer.complete(#"{"a": 1 , "b": 2"#) == #"{"a": 1 , "b": 2}"#) + #expect( + try completer.complete(#"{"a": "x" , "b": [1 , {"c": true , "d""#) + == #"{"a": "x" , "b": [1 , {"c": true , "d": null}]}"# + ) + #expect( + try completer.complete("{\"a\": 1\n,\n\"b\": [\n1\n,\n2\n") + == "{\"a\": 1\n,\n\"b\": [\n1\n,\n2]}" + ) + } + @Test func completesArraysWithMissingValues() throws { #expect(try completer.complete("[1, 2, 3,") == "[1, 2, 3]") #expect(try completer.complete("[1, 2,") == "[1, 2]") From 006a47f6209370f3505b86711cd36e57a7e1dda8 Mon Sep 17 00:00:00 2001 From: Mattt Zmuda Date: Thu, 17 Sep 2026 05:19:39 -0700 Subject: [PATCH 3/4] Simplify JSONCompleter license attribution --- Sources/AnyLanguageModel/Shared/JSONCompleter.swift | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/Sources/AnyLanguageModel/Shared/JSONCompleter.swift b/Sources/AnyLanguageModel/Shared/JSONCompleter.swift index 1d96e3ac..f4b7eb55 100644 --- a/Sources/AnyLanguageModel/Shared/JSONCompleter.swift +++ b/Sources/AnyLanguageModel/Shared/JSONCompleter.swift @@ -1,8 +1,6 @@ import Foundation -// Ported from PartialJSONDecoder (https://github.com/mattt/PartialJSONDecoder) -// at revision d331b237cafe56c5233557bdf1f5a3415393435c, MIT licensed. -// The same source ships in the 1.0.0 release under the Apache License 2.0. +// Ported from PartialJSONDecoder (https://github.com/mattt/PartialJSONDecoder), Apache-2.0 licensed. /// An error that occurs while completing partial JSON. enum JSONCompletionError: Error, Equatable { From 125c7bc4e91b6c27d426bf591c196a1ee685d8e6 Mon Sep 17 00:00:00 2001 From: Mattt Zmuda Date: Thu, 17 Sep 2026 05:42:26 -0700 Subject: [PATCH 4/4] Make JSONCompleter scan complete siblings once findEndOfCompleteValue re-ran the completer on a value its callers had already found complete, then scanned it again to locate its end. Each nested complete sibling doubled the work, so a complete sibling nested 20 deep before an incomplete tail took half a second and 30 deep would take minutes. This runs on every streamed snapshot. Drop the rescan and rely on the linear string, bracket, literal, and number scanners. The depth counter also advanced twice per nesting level, once entering a container and once per element, so the effective limit was half the documented value. Count each level once. --- .../Shared/JSONCompleter.swift | 13 ++++++------- .../JSONCompleterTests.swift | 18 ++++++++++++++++++ 2 files changed, 24 insertions(+), 7 deletions(-) diff --git a/Sources/AnyLanguageModel/Shared/JSONCompleter.swift b/Sources/AnyLanguageModel/Shared/JSONCompleter.swift index f4b7eb55..f52ed0c6 100644 --- a/Sources/AnyLanguageModel/Shared/JSONCompleter.swift +++ b/Sources/AnyLanguageModel/Shared/JSONCompleter.swift @@ -29,7 +29,7 @@ struct JSONCompleter: Sendable { /// The kept portion can be shorter than the input when the text ends inside an escape sequence. typealias Completion = (string: String, endIndex: String.Index) - /// The maximum nesting depth the completer accepts before it throws. + /// The maximum number of nested arrays and objects the completer accepts before it throws. /// /// This bounds recursion on adversarial or malformed input. var maximumDepth: Int = 64 @@ -184,7 +184,7 @@ struct JSONCompleter: Sendable { return nil } - if let elementCompletion = try completeValue(json, from: current, depth: depth + 1) { + if let elementCompletion = try completeValue(json, from: current, depth: depth) { return (string: elementCompletion.string + "]", endIndex: elementCompletion.endIndex) } @@ -252,7 +252,7 @@ struct JSONCompleter: Sendable { return (string: "null}", endIndex: lastValidIndex) } - if let valueCompletion = try completeValue(json, from: current, depth: depth + 1) { + if let valueCompletion = try completeValue(json, from: current, depth: depth) { return (string: valueCompletion.string + "}", endIndex: valueCompletion.endIndex) } @@ -326,14 +326,13 @@ struct JSONCompleter: Sendable { } /// Returns the index immediately after the complete value that starts at the given index. + /// + /// Callers invoke this only after `completeValue` has reported the value complete, + /// so this scans the text once without re-running the completer on it. private func findEndOfCompleteValue(_ json: String, from startIndex: String.Index) -> String.Index { let start = skipWhitespace(json, from: startIndex) guard start < json.endIndex else { return start } - if let completion = try? completeValue(json, from: start, depth: 0) { - return completion.endIndex - } - switch json[start] { case "\"": var current = json.index(after: start) diff --git a/Tests/AnyLanguageModelTests/JSONCompleterTests.swift b/Tests/AnyLanguageModelTests/JSONCompleterTests.swift index aba1430a..23689fef 100644 --- a/Tests/AnyLanguageModelTests/JSONCompleterTests.swift +++ b/Tests/AnyLanguageModelTests/JSONCompleterTests.swift @@ -49,6 +49,10 @@ struct JSONCompleterTests { #expect(throws: JSONCompletionError.depthLimitExceeded(10)) { try limited.complete(String(repeating: "[", count: 20)) } + #expect(throws: JSONCompletionError.depthLimitExceeded(10)) { + try limited.complete(String(repeating: "[", count: 11)) + } + #expect(try limited.complete(String(repeating: "[", count: 10)) == "[[[[[[[[[[]]]]]]]]]]") #expect(try limited.complete(String(repeating: "[", count: 5)) == "[[[[[]]]]]") #expect(JSONCompleter().maximumDepth >= 32) } @@ -154,6 +158,20 @@ struct JSONCompleterTests { ) } + @Test func scansCompleteNestedSiblingsInLinearTime() throws { + // Before the redundant rescan was removed, each nesting level doubled the work, + // so this depth would not finish. + let depth = 60 + let sibling = String(repeating: "[", count: depth) + String(repeating: "]", count: depth) + #expect(try completer.complete("[" + sibling + ", [1") == "[" + sibling + ", [1]]") + + let objectSibling = String(repeating: "{\"a\": ", count: depth) + "1" + String(repeating: "}", count: depth) + #expect( + try completer.complete("{\"x\": " + objectSibling + ", \"y\": \"tail") + == "{\"x\": " + objectSibling + ", \"y\": \"tail\"}" + ) + } + @Test func completesArraysWithMissingValues() throws { #expect(try completer.complete("[1, 2, 3,") == "[1, 2, 3]") #expect(try completer.complete("[1, 2,") == "[1, 2]")