Step 40: JSON Schema — schema-compiled constrained decoding (#78)

[kkokosa]

Closes #78

Summary

• Schema-compiled constrained decoding: JsonSchemaConstraint layers a SchemaTracker struct on the existing JsonCharParser to enforce JSON Schema semantics during generation — type constraints, required properties, enum values, nested structures
• Schema compilation: SchemaCompiler converts JSON Schema (System.Text.Json) into an immutable CompiledSchema flat node array with property-name tries for efficient key masking. Resolves $ref/$defs at compile time
• Zero-alloc Clone(): Both JsonCharParser and SchemaTracker are value types with InlineArray stacks (~1.3KB + ~580B). Clone is a struct copy — zero heap allocations on the hot path
• CLI + API integration: --response-format json_schema --schema <json|@file>, ResponseFormat.JsonSchema record, TextGenerator pattern-match dispatch
• 43 unit tests covering compiler, trie, tracker, and full constraint integration. Smoke test script test_models_json_schema.py for cross-architecture validation

Technical Details

Architecture

JsonSchemaConstraint : IDecodingConstraint
  ├── _parser:   JsonCharParser   (struct — JSON syntax validation)
  ├── _tracker:  SchemaTracker    (struct — schema semantics)
  └── _maskCache: Dictionary<long, TokenMask>  (shared across clones)

The tracker observes structural events (object open, key complete, value start/end, array item, close) by comparing parser state transitions. At each decode step, IsCharAllowedBySchema() intersects JSON-valid characters with schema-valid characters before the parser advances.

Schema Features (MVP)

type (object/array/string/number/integer/boolean/null), properties, required, additionalProperties: false, items, enum, const, $ref/$defs (non-recursive), anyOf, nested objects/arrays

Key Design Decisions

• Layered on JsonCharParser rather than a combined PDA — reuses battle-tested JSON syntax validation, keeps schema concerns separate and testable
• PropertyNameTrie for O(1)-per-char key restriction — avoids testing every vocab token against every property name
• ulong bitmask (64 props max) for required property tracking — zero-alloc, covers 99.9% of real schemas
• Bounded mask cache (4096 entries, clear-on-overflow) — working set during generation is typically much smaller
• Integer enforcement at tracker level — rejects ./e/E in number states when schema type is integer

Key Files

File	Changes
src/DotLLM.Engine/Constraints/Schema/SchemaNode.cs	Compiled schema node: types, properties, required bitmask, enum/const
src/DotLLM.Engine/Constraints/Schema/SchemaCompiler.cs	JSON Schema → CompiledSchema with $ref resolution
src/DotLLM.Engine/Constraints/Schema/PropertyNameTrie.cs	Char trie for property name + enum value validation
src/DotLLM.Engine/Constraints/Schema/SchemaTracker.cs	~1.3KB struct tracking schema position via InlineArray stacks
src/DotLLM.Engine/Constraints/Schema/CompiledSchema.cs	Immutable container (flat node array + tries)
src/DotLLM.Engine/Constraints/JsonSchemaConstraint.cs	IDecodingConstraint with layered parser+tracker
src/DotLLM.Engine/Constraints/JsonCharParser.cs	5 internal readonly accessors for tracker observation
src/DotLLM.Core/Configuration/ResponseFormat.cs	JsonSchema sealed record variant
src/DotLLM.Engine/TextGenerator.cs	Constraint dispatch via pattern match (2 sites)
src/DotLLM.Cli/Commands/RunCommand.cs	--schema option, json_schema format
src/DotLLM.Cli/Commands/ChatCommand.cs	--schema option, json_schema format
scripts/test-schema.json	Test schema (name:string + age:integer)
scripts/test_models_json_schema.py	Cross-architecture smoke test

Test Plan

• Unit tests pass: dotnet test --filter "FullyQualifiedName~Constraint" — 134 tests (9 existing + 43 new schema + 82 other)
• Existing JsonConstraintTests pass unchanged (no regression)
• Build succeeds with zero errors
• Smoke test: python scripts/test_models_json_schema.py across cached models
• Manual test: dotnet run --project src/DotLLM.Cli -- run <model> --response-format json_schema --schema @scripts/test-schema.json -p "Output a JSON with name and age" -n 64 --temp 0

🤖 Generated with Claude Code

[kkokosa]

@codex please review this PR thoroughly - performance issues and possible improvements, potential bugs, code quality and architecture

[kkokosa]

@gemini please review this PR thoroughly - performance issues and possible improvements, potential bugs, code quality and architecture

[github-actions]

Gemini ✦

This is a very solid implementation. Layering the schema semantics on top of the existing JsonCharParser PDA is a great architectural choice—it keeps the syntax validation isolated and highly optimized while allowing the schema tracker to observe state transitions. The zero-alloc struct cloning and trie-based property masking are also well aligned with the engine's performance goals.

I have reviewed the PR and found one critical correctness issue regarding cache key collisions, along with a few performance improvement suggestions for the hot path.

1. Critical Bug: Cache Key Collisions in GetSchemaStateKey

In SchemaTracker.cs, the state key used for the _maskCache compresses multiple fields into a 64-bit long using HashCode.Combine:

return ((long)parserKey << 32) | (uint)HashCode.Combine(nodeIdx, emitted.GetHashCode(), triePos);

HashCode.Combine is not guaranteed to be collision-free. Because emitted (64-bit), nodeIdx, and triePos are compressed into 32 bits, hash collisions are highly probable over long generation sequences. A collision here will cause the constraint to return an incorrect TokenMask from the cache, which could allow invalid tokens (breaking the schema) or block valid ones (halting generation).

Fix: Use a custom struct as the dictionary key instead of a long. This guarantees 100% correctness and is still allocation-free for lookups.

internal readonly record struct SchemaStateKey(
    int ParserKey, 
    int NodeIdx, 
    ulong EmittedProps, 
    int TriePos
);

// In SchemaTracker.cs
public readonly SchemaStateKey GetSchemaStateKey(in JsonCharParser parser)
{
    int parserKey = parser.GetEffectiveStateKey();
    ulong emitted = _stackDepth > 0 ? _emittedProps[_stackDepth - 1] : 0;
    int triePos = _inKeyString ? _trieNodeIndex : (_inEnumString ? _enumTrieNodeIndex : 0);
    return new SchemaStateKey(parserKey, _currentNodeIndex, emitted, triePos);
}

// In JsonSchemaConstraint.cs
private readonly Dictionary<SchemaStateKey, TokenMask> _maskCache;

2. Performance: PropertyNameTrie Node Children

In PropertyNameTrie.cs, TrieNode uses a Dictionary<char, int> for its children.
While the trie itself is shared and built once, dictionary lookups during TryGetChild on the inference hot path (inside IsTokenValid for every token) introduce pointer chasing and overhead.

Improvement: Since the trie is immutable after compilation and nodes typically have very few children (often just 1 for linear string segments), replacing the Dictionary with a small inline array or a flat (char, int)[] searched via linear/binary search will significantly improve L1 cache locality.

3. Performance: Struct Copying Overhead in BuildAndCacheMask

In JsonSchemaConstraint.IsTokenValid, you are copying _parser and _tracker for every token in the vocabulary:

var parserClone = _parser;
var trackerClone = _tracker;

Because SchemaTracker contains several InlineArray stacks, its size is ~1.3KB. Copying 1.3KB for 128,000 tokens (e.g., Llama 3 vocab) results in ~160MB of memory copied per cache miss.

While memcpy in the L1 cache is fast and this only happens on cache misses, it's a noticeable CPU overhead.
Future optimization (not a blocker for this PR): You could optimize this by grouping tokens by their first character, or by walking a Trie representation of the tokenizer vocabulary so that the tracker state is only cloned when branching, rather than from the root state for every single token.

4. Minor: Integer Type Validation

In test_models_json_schema.py, you have a special case for Python's bool being a subclass of int.
In SchemaTracker.IsCharAllowedBySchema, you correctly reject . and e/E for integer types:

JsonParserState.InNumberDot => true, // Wait, does this reject dots for integers?

Make sure that if node.AllowedTypes is strictly JsonSchemaType.Integer, IsCharAllowedBySchema returns false when c == '.' or c == 'e'. The current diff shows JsonParserState.InNumberDot => true; which might rely on the parser, but the parser doesn't know about the schema type. Ensure InNumberDot checks the schema type!

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 5a4975cfe1

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Enforce enum/const constraints for non-string values

StartValueString only activates value checking when EnumTrieIndex is present, which is built only for string enums/const. As a result, schemas like {"const": 1} or {"enum": [true, false]} are effectively treated as plain type constraints and can emit values outside the declared set. This breaks schema correctness for numeric/boolean/null/object enums and consts.

Useful? React with 👍 / 👎.

[dotllm-claude-code-bot]

Acknowledged — this is a known MVP limitation. Non-string enum/const values (e.g. {"const": 1}, {"enum": [true, false]}) are only type-restricted, not character-level constrained. Full enforcement requires character-sequence matching for literals and numbers, which is a fundamentally different mechanism from the string trie approach.

Added a TODO comment at StartValueString documenting this limitation for a follow-up.

[chatgpt-codex-connector]

Preserve anyOf branch constraints after selecting a type

For anyOf, the tracker currently unions only AllowedTypes and never records which alternative is being followed. That means once a value starts, nested constraints from each branch (required keys, enums, object shapes, etc.) are not enforced, so outputs can satisfy neither branch while still being accepted by the mask.

Useful? React with 👍 / 👎.

[dotllm-claude-code-bot]

Acknowledged — this is a known MVP limitation (same approach as Outlines and XGrammar). The current overapproximation (union of types from all branches) is safe: output is always valid JSON, but may not satisfy every branch perfectly.

Full branch narrowing requires maintaining parallel tracker states (one per viable branch), which is significant complexity. Added a TODO comment at IsValueStartCharAllowed documenting this for a follow-up.

[chatgpt-codex-connector]

Honor required keys that are absent from properties

Required entries are only encoded when their name exists in propertyNames; missing names are silently ignored. In JSON Schema, required still mandates presence even if the key has no schema in properties (especially when additionalProperties is true), so this allows objects to close while omitting required keys.

Useful? React with 👍 / 👎.

[dotllm-claude-code-bot]

Fixed in 5b4b179. SchemaCompiler now checks for required keys absent from properties and adds unconstrained schema nodes for them at compile time. The required bitmask and property name trie are rebuilt to include the added keys, so the constraint correctly blocks } until all required properties (even those without explicit schemas) have been emitted.

[dotllm-claude-code-bot]

Thanks for the thorough review. All items addressed in 5b4b179:

1. Cache key collisions (critical) — Fixed. Replaced long hash key with SchemaStateKey readonly record struct (int ParserKey, int NodeIdx, ulong EmittedProps, int TriePos). All fields compared exactly — no hash compression, zero collision risk. readonly record struct auto-generates correct GetHashCode/Equals.

2. PropertyNameTrie children (perf) — Fixed. Replaced Dictionary<char, int> with sorted (char Key, int Child)[] flat array + linear search. Typical nodes have 1–5 children where linear search beats dictionary, and the flat array is L1-friendly.

3. Struct copying overhead (perf) — Added TODO comment on IsTokenValid documenting the ~160MB/cache-miss cost and the tokenizer-trie optimization path. Not a blocker for this PR.

4. Integer InNumberDot (minor) — No change needed. The . is blocked earlier at InNumberIntDigits/InNumberZero via IsNumberContinuationAllowed, so the parser never reaches InNumberDot for integer-only types. The InNumberDot => true line is only reachable for number types where fractional digits are valid.