Let the model choose the variant of an anyOf

[2570165831]

Fixes #289.

Constrained generation (MLX, Core ML, llama.cpp) always generated the first variant of an anyOf, so anyOf: [person, .null] could never produce null, and a union of object types always produced the first type.

As discussed in #289:

• References are resolved and nested unions flattened, keeping declaration order (flattenedVariants).
• When variants start differently, the generator samples once among the tokens that can start each variant, without decoding the sample (like sampleWhetherToCloseEmptyArray(items:)), and continues with the variants that can start with it. This covers [T, null] in either order and unions of different JSON types.
• Object variants all start with {, so generateObject(oneOf:) offers the keys of every variant still in play and each emitted key keeps only the variants that declare it. The object can close once some live variant has all its required keys.
• A key shared by several variants whose values are string choices narrows by the chosen value (generateValue(forKey:narrowing:)). Since Support constant and anyOf string guides #304 a .constant(_:) guide is a one-choice enum, so constants narrow the same way, and Generate enums with associated values as Foundation Models does #301's type keys will too.
• Variants that start alike but aren't all objects (two strings, integer and number, two arrays) fall back to the first of them, with a comment at both spots, as agreed.

Tests (AnyOfGenerationTests, MockTokenBackend with scripted sampling): nullable string → null; null-first union → string; nullable object from a DynamicGenerationSchema → either branch; object union follows the first key; discriminated union by the shared key's value, with the payload key first, and the payloadless case; a .constant-guided discriminator built with DynamicGenerationSchema; and the documented fallback for two strings.

Testing on macOS 27, Xcode 27 beta 4 (Swift 6.4):

• swift format lint --strict --recursive .: clean.
• swift build --build-tests --build-system native: OK. With the default build system the test bundle fails to link on this beta (undefined SwiftSyntax symbols in AnyLanguageModelMacros), unrelated to this change, so I used the native one.
• CI=1 swift test --skip-build --build-system native: 682 tests in 71 suites passed. Without CI the OllamaLanguageModel suite runs and fails here because no Ollama server is running locally.

Not run: the MLX,Llama,CoreML traits, Linux, and generation with a real local model.

[Copilot]

🟡 Changes recommended

Equivalent field schemas with different descriptions can prematurely discard viable object variants.

1 open finding

• Ignore descriptions when comparing Node variants for union generation · New

What changed in this PR

Enables constrained generators to select among anyOf variants instead of always using the first.

Changes:

• Samples among distinct variant start tokens.
• Narrows object variants by emitted keys and discriminator values.
• Adds scripted tests for nullable and object unions.

File	Description
StructuredGeneration.swift	Implements variant selection and object narrowing.
AnyOfGenerationTests.swift	Tests nullable, object, and discriminated unions.

🧠 Review effort: Balanced

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[mattt]

Thanks, @2570165831. This follows the plan from #289 closely, and the tests make the new behavior easy to follow.

A few things turned up while I tried edge cases:

1. In generateValue(forKey:narrowing:), if the value from generateChoice isn't exactly one of the choices, live ends up empty, and the object closes on the next pass without the remaining required keys. That can happen when a token's text differs from its source text, like a piece with a leading space. The key path already has guard !live.isEmpty. Could the value path either throw or leave live as it was when nothing matches?

2. When a shared key's values are all objects, the first variant still wins, since every value starts with {. For example, with anyOf: [{x: {y: …}}, {x: {z: …}}], the model can't produce {"x": {"z": …}}. This will matter after #301 for two cases with struct payloads, if the model writes the payload key before type. If handling it here is more than you want to take on, I'm happy to leave it for a follow-up. A comment at that spot naming the object case would help in the meantime.

3. withoutDescriptions doesn't resolve .ref, so x: ref A and x: ref B that differ only in their descriptions still narrow to the first variant.

4. Two small ones: the comment above the loop in generateObject(oneOf:) still says } is allowed once every required property has been emitted, and discriminatedUnionCanPickThePayloadlessCase passes without this change because its case comes first. Listing it second would make the test cover the fix.

[Copilot]

🔵 Needs a closer look

Reference aliases can narrow incorrectly, deep branches can disappear, and object unions repeatedly scan the vocabulary.

0 open findings

1 resolved since last review

• Ignore descriptions when comparing Node variants for union generation

Previously missed (3)

In code that hasn't changed since last review

Avoid repeated vocabulary scans for all-object unions

Sources/​AnyLanguageModel/​Shared/​StructuredGeneration.swift:428

For the all-object unions this PR primarily targets, this call computes itemStartTokens once per variant before discovering that every set is the same. Each object computation scans the entire tokenizer vocabulary, so an enum with N object cases incurs N full vocabulary scans on every generation. Detect the all-object case before probing start tokens; mixed unions can retain the post-sampling object check.

Do not discard valid branches at recursion depth limit

Sources/​AnyLanguageModel/​Shared/​StructuredGeneration.swift:437

A valid branch nested 64 levels deep is silently converted to an empty list here. Consequently, a deeply nested sole branch now throws emptyAnyOf, while a deep branch beside another candidate simply disappears and can never be chosen; the previous recursive generation did not impose this limit. Track active references to detect cycles and report a dedicated depth/recursion error rather than treating the branch as absent.

Resolve references when comparing equivalent union variants

Sources/​AnyLanguageModel/​Shared/​StructuredGeneration.swift:616

Leaving references unresolved means aliases to generation-equivalent schemas still compare unequal. For example, if two object variants share x through refs to separately named but structurally identical object definitions, emitting x reaches the fallback, chooses the first ref, and removes the other variant before a later distinguishing key can be selected. Resolve refs (with cycle protection) while building the description-free comparison shape so equivalent referenced values remain live.

🧠 Review effort: Balanced

[2570165831]

Thanks for the careful review! ca562b6 covers all four points.

1. If the generated value matches none of the choices, generateValue(forKey:of:narrowing:) now leaves live alone, and there's a comment explaining why. I went with the lenient option over throwing because the value is already emitted by then, so throwing would only throw away a usable object. choiceThatDecodesDifferentlyKeepsTheRemainingKeys makes the token for "b" decode as " b" and checks that y is still generated. Before this change, that test closed the object right after x.

2. I handled this one here. When every live variant's value for the shared key resolves to objects (through refs and nested anyOf), the value is now generated with the multi-variant object generator. generateObject(oneOf:) now also returns which of its inputs are still live, and the outer variants whose nested object survived stay live. Identical nested schemas are deduplicated, ignoring descriptions. sharedKeyWithObjectValuesFollowsTheNestedKeys uses anyOf: [{x: {x: …}}, {x: {y: …}, y: …}] with the model writing {"x":{"y":"a"},"y":"a"}. Before this change it failed because only the first variant's nested key was offered. If one of the values can be a non-object, for example {…} | null, it still uses the first-variant fallback described in the existing comment.

3. withoutDescriptions is now an instance method and resolves .ref through schema.defs. A ref that's already being resolved higher up the chain stays as .ref, so recursive definitions stop expanding. sharedKeysWhoseReferencesDifferOnlyInDescriptionKeepBothVariants uses x: ref A vs x: ref B, two integers that differ only in their description. I used integers on purpose: with point 2 in place, object refs already go through the nested-object path, and integers keep this test about point 3 alone. selfReferencingSharedKeysStillGenerate covers the recursion guard. Without the guard it overflows the stack.

4. I fixed the comment above the loop and the one next to where } is appended. Both now say closing is allowed once some live variant has all of its required properties. discriminatedUnionCanPickThePayloadlessCase now lists the payloadless case second, and it fails against the code from before this PR.

For each new or changed test, I reverted just the matching source change and confirmed that test fails while the others still pass. With everything in place, the full suite passes and swift format lint --strict is clean.

[Copilot]

🟡 Changes recommended

Branch narrowing can close on an incorrect variant, and explicit-nil postprocessing can corrupt non-first object variants.

2 open findings

• Explicit-nil postprocessing selects the wrong object variant · New
• Sampled enum branch loses identity, allowing required keys to be skipped · New

🧠 Review effort: Balanced

[2570165831]

Thanks. On the decoded-choice finding: you're right. Keeping every variant live let the x-only case close the object before y. In 02143b4, generateChoice now reports which candidate the sampled tokens spell, and the value path narrows by that candidate, not by the decoded text. The new test choiceThatDecodesDifferentlyCannotCloseBeforeTheRemainingKeys scripts } after a " b" token and checks it's refused; it failed before the change.

On explicit-nil postprocessing: also real, but variant(matching:) already has it on main, for any backend that returns a later variant; this PR makes it reachable for local models too. I'd fix it in a small follow-up that checks string enum and constant values when matching variants, unless you'd rather have it here.

[Copilot]

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

[Copilot]

🟡 Changes recommended

Nested object narrowing can retain invalid variants, and several edge cases can produce schema-invalid output or excessive expansion.

1 open finding

• Depth guard allows exponential flattened schema expansion · New

2 resolved since last review

• Sampled enum branch loses identity, allowing required keys to be skipped
• Explicit-nil postprocessing selects the wrong object variant

🧠 Review effort: Balanced

[2570165831]

Good catch, and this one came in with the PR. With a definition that refers to the next one twice per level, the old code took 16 s and 25 s at 22 levels (4 GB peak) and doubled per level. 829eb07 bounds both walks:

• flattenedVariants expands each definition once per flattening; a repeat reference only repeats variants already in the list. That makes the result linear in the number of definitions and replaces the depth guard.
• withoutDescriptions resolves at most 256 references per call and leaves the rest as references, so the worst case is two schemas comparing unequal, which takes the conservative path. Properties are visited in key order so equal schemas stop at the same place.

unionsThatReferToADefinitionTwiceFlattenOnce and sharedKeysOverADefinitionUsedTwicePerLevelStillCompare use 64 levels and finish in milliseconds; 690 tests pass and swift format lint --strict is clean.

[Copilot]

🟡 Changes recommended

Nested object closure can retain variants with unmet required keys, producing output that matches no union branch.

1 open finding

• Preserve and apply closable variant indices before early exit · New

1 resolved since last review

• Depth guard allows exponential flattened schema expansion

🧠 Review effort: Balanced

[mattt]

Thanks, @2570165831. Copilot found one more case where a nested union could produce JSON that matches none of its variants. I fixed it in #320, which includes your commits, so I'll close this in favor of that one. Thanks again!