submilli_engine/typechecker/
json_schema.rs

1//! JSON Schema emitter for `submilli:llm` typed output.
2//!
3//! `llm.call<T>()` sends a JSON Schema for `T` to the provider so the model is
4//! constrained to the shape the program asked for. The schema is advisory —
5//! every response is still verified by our own structural check — so the job
6//! here is to describe `T` in the narrow dialect every provider agrees on, and
7//! to refuse at compile time anything that dialect cannot carry.
8//!
9//! **Every definition is inlined; `$ref`/`$defs` are never emitted.** Google's
10//! provider strips `$ref`, silently degrading the referenced subschema to "any"
11//! — the response then validates against nothing. A recursive type has no
12//! finite inlining, so it is a compile error here rather than a ref.
13//!
14//! The emitter carries its **own** gate rather than deferring to
15//! `unsupported_cast_target_reason`. The two agree on most variants but differ
16//! where it matters: the cast gate accepts `Unknown` (the runtime widen is a
17//! no-op), while a schema for `unknown` could only be `{}` — the same silent
18//! degeneration to "any" that `$ref` is rejected for. See [`SchemaReject`].
19
20use serde_json::{Value, json};
21
22use crate::types::Type;
23
24/// One-level expander for recursion back-edges, mirroring
25/// [`json_strategy::AliasExpander`](super::json_strategy::AliasExpander).
26/// Given a [`Type::AliasRef`] it returns the alias body (args substituted); for
27/// any other type it returns `ty.peel()` cloned. Lets the emitter resolve a
28/// named alias without owning the type namespace.
29pub type AliasExpander<'a> = &'a dyn Fn(&Type) -> Type;
30
31/// Identity expander — leaves `AliasRef` unexpanded, so any back-edge is
32/// reported as an unresolvable alias. Used by call sites that hold no type
33/// namespace (unit tests, already-expanded types).
34fn no_expand(ty: &Type) -> Type {
35    ty.peel().clone()
36}
37
38/// A type that has no JSON Schema in the safe subset.
39///
40/// `path` names the offending field so the diagnostic points at the thing the
41/// author must change, not merely at the type as a whole (R8). It is a
42/// dotted/bracketed path from the root — `"items[0].handler"` — and is empty
43/// when the root type itself is unrepresentable.
44#[derive(Clone, Debug, PartialEq, Eq)]
45pub struct SchemaReject {
46    pub path: String,
47    pub reason: &'static str,
48    /// The walk hit a [`Type::Error`] — a type that already failed to resolve
49    /// and already reported its own diagnostic. There is still no schema to
50    /// emit, so this is an `Err`, but the caller must stay silent rather than
51    /// cascade a second diagnostic onto the same mistake. This is the one
52    /// reject a caller suppresses; `unsupported_cast_target_reason` reaches the
53    /// same conclusion by returning `None` for `Type::Error`.
54    pub cascading: bool,
55}
56
57impl SchemaReject {
58    /// `"field `items[0].handler`: <reason>"`, or just the reason at the root.
59    /// Callers embed this in the diagnostic message.
60    pub fn describe(&self) -> String {
61        if self.path.is_empty() {
62            self.reason.to_string()
63        } else {
64            format!("field `{}`: {}", self.path, self.reason)
65        }
66    }
67}
68
69/// Emit an inlined JSON Schema for `ty`, or name the first field that has none.
70///
71/// Callers pass a type that has already been through
72/// `reduce_interfaces_to_shapes`, so a surviving [`Type::InterfaceRef`] means
73/// an interface that could not be expanded — method-bearing (nominal) or
74/// recursive — and is rejected.
75pub fn json_schema(ty: &Type) -> Result<Value, SchemaReject> {
76    json_schema_with(ty, &no_expand)
77}
78
79/// Like [`json_schema`] but resolves recursion back-edges ([`Type::AliasRef`])
80/// one level through `expand`. A back-edge that expands to a type already open
81/// on the walk is a genuine cycle and is rejected.
82pub fn json_schema_with(ty: &Type, expand: AliasExpander<'_>) -> Result<Value, SchemaReject> {
83    walk(ty, expand, &mut String::new(), &mut Vec::new())
84}
85
86/// `seen` holds the alias back-edges currently open on this path; re-entering
87/// one is the recursion that cannot be inlined. Same cycle-guard shape as
88/// `unsupported_cast_target_reason` and `reduce_interfaces_to_shapes`, but the
89/// verdict is the opposite: they treat a re-encounter as supported (codegen
90/// emits a recursive validator), while an inlined schema has no such escape.
91fn walk(
92    ty: &Type,
93    expand: AliasExpander<'_>,
94    path: &mut String,
95    seen: &mut Vec<Type>,
96) -> Result<Value, SchemaReject> {
97    let peeled = ty.peel();
98    match peeled {
99        Type::Number => Ok(json!({ "type": "number" })),
100        Type::NumberLiteral(n) => Ok(json!({ "const": n.0 })),
101        Type::String => Ok(json!({ "type": "string" })),
102        Type::StringLiteral(s) => Ok(json!({ "const": s })),
103        Type::Boolean => Ok(json!({ "type": "boolean" })),
104        Type::BooleanLiteral(b) => Ok(json!({ "const": b })),
105        Type::Null => Ok(json!({ "type": "null" })),
106        Type::Object { fields, index } => {
107            // `BTreeMap` iteration is field-name order, so `properties` and
108            // `required` come out canonical without sorting here.
109            let mut properties = serde_json::Map::new();
110            let mut required = Vec::new();
111            for (name, field) in fields {
112                let restore = path.len();
113                push_field(path, name);
114                properties.insert(name.clone(), walk(&field.ty, expand, path, seen)?);
115                path.truncate(restore);
116                // An optional field is absent from `required`. It is *not* also
117                // `| undefined`: that widening is the *read* type of the field, a
118                // different thing from whether the provider must emit the key.
119                if !field.optional {
120                    required.push(Value::String(name.clone()));
121                }
122            }
123            Ok(json!({
124                "type": "object",
125                "properties": Value::Object(properties),
126                "required": Value::Array(required),
127                "additionalProperties": match index { Some(index) => walk(&index.value, expand, path, seen)?, None => Value::Bool(false) },
128            }))
129        }
130        Type::Array(elem) => {
131            let restore = path.len();
132            path.push_str("[]");
133            let items = walk(elem, expand, path, seen)?;
134            path.truncate(restore);
135            Ok(json!({ "type": "array", "items": items }))
136        }
137        Type::Tuple(elems) => {
138            let mut items = Vec::with_capacity(elems.len());
139            for (i, elem) in elems.iter().enumerate() {
140                let restore = path.len();
141                path.push_str(&format!("[{i}]"));
142                items.push(walk(elem, expand, path, seen)?);
143                path.truncate(restore);
144            }
145            let len = elems.len();
146            Ok(json!({
147                "type": "array",
148                "prefixItems": Value::Array(items),
149                "items": false,
150                "minItems": elems.required_len(),
151                "maxItems": len,
152            }))
153        }
154        Type::Union(members) => {
155            let mut any_of = Vec::with_capacity(members.len());
156            for member in members {
157                if matches!(member.peel(), Type::Undefined | Type::Void) {
158                    continue;
159                }
160                any_of.push(walk(member, expand, path, seen)?);
161            }
162            match any_of.len() {
163                0 => Err(reject(
164                    path,
165                    "undefined has no JSON value; include a JSON-representable type",
166                )),
167                1 => any_of
168                    .pop()
169                    .ok_or_else(|| reject(path, "missing JSON union schema")),
170                _ => Ok(json!({ "anyOf": Value::Array(any_of) })),
171            }
172        }
173        Type::Undefined => Err(reject(
174            path,
175            "undefined has no JSON value; include a JSON-representable type",
176        )),
177        // A recursion back-edge: expand once and describe the body inline.
178        // Unlike the cast gate, a re-encounter is fatal — inlining a cycle
179        // does not terminate, and `$ref` is not an option (see module docs).
180        Type::AliasRef { .. } => {
181            if seen.iter().any(|s| s == peeled) {
182                return Err(reject(path, RECURSIVE));
183            }
184            let expanded = expand(peeled);
185            if matches!(expanded.peel(), Type::AliasRef { .. }) {
186                return Err(reject(path, RECURSIVE));
187            }
188            seen.push(peeled.clone());
189            let schema = walk(&expanded, expand, path, seen);
190            seen.pop();
191            schema
192        }
193        // Already-reported failure: still no schema, but the caller stays quiet.
194        Type::Error => Err(SchemaReject {
195            cascading: true,
196            ..reject(path, "this type failed to resolve")
197        }),
198        Type::Alias { ty, .. } | Type::Refined { ty, .. } | Type::Readonly(ty) => {
199            walk(ty, expand, path, seen)
200        }
201        // An unconstrained schema would silently accept anything, unlike an `as unknown` widening.
202        Type::Unknown => Err(reject(
203            path,
204            "`unknown` has no JSON Schema — an unconstrained schema would let the model return anything; name the shape you expect",
205        )),
206        Type::NonPrimitive => Err(reject(
207            path,
208            "`object` has no JSON Schema — it admits any object, array or function; name the shape you expect",
209        )),
210        Type::Function { .. } => Err(reject(path, "functions have no JSON form")),
211        Type::BigInt | Type::BigIntLiteral(_) => Err(reject(
212            path,
213            "`bigint` has no JSON number form — use `number` or a `string` encoding",
214        )),
215        Type::Uint8Array => Err(reject(
216            path,
217            "`Uint8Array` has no JSON form — use a `string` encoding such as base64",
218        )),
219        Type::ClassRef { .. } => Err(reject(
220            path,
221            "class types are nominal — use an object type or a data interface",
222        )),
223        Type::InterfaceRef { .. } => Err(reject(
224            path,
225            "this interface has no data shape — method-bearing interfaces are nominal, and recursive ones cannot be inlined",
226        )),
227        Type::NumberEnum { .. } | Type::StringEnum { .. } => Err(reject(
228            path,
229            "enums need their variants resolved — use a union of literals",
230        )),
231        Type::TypeVar(_) | Type::GenericParam { .. } => Err(reject(
232            path,
233            "generic type parameters are erased — the schema is emitted at compile time",
234        )),
235        Type::Never => Err(reject(path, "`never` has no values")),
236        Type::Void => Err(reject(path, "`void` is not a value type")),
237    }
238}
239
240/// The recursion verdict, shared by the back-edge and no-expander cases.
241const RECURSIVE: &str = "recursive types have no finite JSON Schema — \
242                         providers drop `$ref`, so every definition must be inlined";
243
244fn reject(path: &str, reason: &'static str) -> SchemaReject {
245    SchemaReject {
246        path: path.to_string(),
247        reason,
248        cascading: false,
249    }
250}
251
252/// Append `name` to the field path: `"a.b"`, but `"a"` at the root and
253/// `"items[].id"` after an array step (which already ends in a bracket).
254fn push_field(path: &mut String, name: &str) {
255    if !path.is_empty() {
256        path.push('.');
257    }
258    path.push_str(name);
259}
260
261#[cfg(test)]
262mod tests {
263    use std::collections::BTreeMap;
264
265    use super::*;
266    use crate::mangle;
267    use crate::types::{LiteralF64, ObjectField, Package};
268
269    fn obj(fields: &[(&str, Type, bool)]) -> Type {
270        Type::Object {
271            index: None,
272            fields: fields
273                .iter()
274                .map(|(name, ty, optional)| {
275                    (
276                        (*name).to_string(),
277                        if *optional {
278                            ObjectField::optional(ty.clone())
279                        } else {
280                            ObjectField::required(ty.clone())
281                        },
282                    )
283                })
284                .collect::<BTreeMap<_, _>>(),
285        }
286    }
287
288    fn s_lit(s: &str) -> Type {
289        Type::StringLiteral(s.to_string())
290    }
291
292    fn alias_ref(name: &str) -> Type {
293        Type::AliasRef {
294            mangled: mangle::prelude(name),
295            package: Package::prelude(),
296            name: name.to_string(),
297            args: Vec::new(),
298        }
299    }
300
301    #[test]
302    fn flat_object() {
303        let schema = json_schema(&obj(&[
304            ("age", Type::Number, false),
305            ("name", Type::String, false),
306        ]))
307        .expect("emits");
308        assert_eq!(
309            schema,
310            json!({
311                "type": "object",
312                "properties": {
313                    "age": { "type": "number" },
314                    "name": { "type": "string" },
315                },
316                "required": ["age", "name"],
317                "additionalProperties": false,
318            })
319        );
320    }
321
322    #[test]
323    fn nested_object() {
324        let inner = obj(&[("city", Type::String, false)]);
325        let schema = json_schema(&obj(&[("address", inner, false)])).expect("emits");
326        assert_eq!(
327            schema["properties"]["address"]["properties"]["city"],
328            json!({ "type": "string" })
329        );
330        assert_eq!(schema["properties"]["address"]["required"], json!(["city"]));
331    }
332
333    #[test]
334    fn optional_field_is_omitted_from_required_and_is_not_nullable() {
335        let schema = json_schema(&obj(&[
336            ("id", Type::String, false),
337            ("note", Type::String, true),
338        ]))
339        .expect("emits");
340        assert_eq!(schema["required"], json!(["id"]));
341        // Present in `properties`, and as a bare string — the `| undefined` widening
342        // is the read type of the field, not its schema.
343        assert_eq!(schema["properties"]["note"], json!({ "type": "string" }));
344    }
345
346    #[test]
347    fn undefined_projection_preserves_property_requiredness() {
348        let maybe_string = Type::union(vec![Type::String, Type::Undefined]);
349        let schema = json_schema(&obj(&[
350            ("required", maybe_string.clone(), false),
351            ("optional", maybe_string, true),
352            (
353                "nullable",
354                Type::union(vec![Type::String, Type::Null, Type::Undefined]),
355                false,
356            ),
357        ]))
358        .expect("JSON-representable branches remain");
359        assert_eq!(schema["required"], json!(["nullable", "required"]));
360        assert_eq!(schema["properties"]["required"], json!({"type":"string"}));
361        assert_eq!(schema["properties"]["optional"], json!({"type":"string"}));
362        assert_eq!(
363            schema["properties"]["nullable"],
364            json!({"anyOf":[{"type":"string"},{"type":"null"}]})
365        );
366    }
367
368    #[test]
369    fn undefined_only_schema_reports_the_property_path() {
370        let error = json_schema(&obj(&[("missing", Type::Undefined, true)]))
371            .expect_err("undefined is not JSON");
372        assert_eq!(error.path, "missing");
373        assert!(error.reason.contains("undefined has no JSON value"));
374        assert!(json_schema(&Type::Undefined).is_err());
375    }
376
377    #[test]
378    fn optional_tuple_schema_keeps_lengths_and_excludes_null() {
379        let tuple = Type::Tuple(crate::types::TupleType {
380            elements: vec![
381                Type::Number,
382                Type::union(vec![Type::String, Type::Undefined]),
383            ],
384            optional: 1,
385        });
386        let schema = json_schema(&tuple).expect("optional tuple schema");
387        assert_eq!(schema["minItems"], json!(1));
388        assert_eq!(schema["maxItems"], json!(2));
389        assert_eq!(schema["prefixItems"][1], json!({"type":"string"}));
390        assert_eq!(schema["items"], json!(false));
391    }
392
393    #[test]
394    fn array_and_tuple() {
395        let arr = json_schema(&Type::Array(Box::new(Type::Number))).expect("emits");
396        assert_eq!(arr["type"], json!("array"));
397        assert_eq!(arr["items"], json!({ "type": "number" }));
398
399        let tup = json_schema(&Type::tuple(vec![Type::String, Type::Boolean])).expect("emits");
400        assert_eq!(tup["type"], json!("array"));
401        assert_eq!(
402            tup["prefixItems"],
403            json!([{ "type": "string" }, { "type": "boolean" }])
404        );
405        assert_eq!(tup["minItems"], json!(2));
406        assert_eq!(tup["maxItems"], json!(2));
407    }
408
409    #[test]
410    fn union_becomes_any_of() {
411        let ty = Type::union(vec![Type::String, Type::Null]);
412        let schema = json_schema(&ty).expect("emits");
413        let any_of = schema["anyOf"].as_array().expect("anyOf array");
414        assert_eq!(any_of.len(), 2);
415        assert!(any_of.contains(&json!({ "type": "null" })));
416        assert!(any_of.contains(&json!({ "type": "string" })));
417    }
418
419    #[test]
420    fn literals_become_const() {
421        assert_eq!(
422            json_schema(&s_lit("red")).expect("emits"),
423            json!({ "const": "red" })
424        );
425        assert_eq!(
426            json_schema(&Type::NumberLiteral(LiteralF64(3.0))).expect("emits"),
427            json!({ "const": 3.0 })
428        );
429    }
430
431    #[test]
432    fn reduced_interface_matches_the_literal_object_type() {
433        // `reduce_interfaces_to_shapes` runs before the walk, so a data-only
434        // interface arrives here already as a `Type::Object`. Emitting the same
435        // schema for both is what makes that pre-pass sufficient.
436        let shape = obj(&[("id", Type::String, false), ("n", Type::Number, true)]);
437        let literal = obj(&[("id", Type::String, false), ("n", Type::Number, true)]);
438        assert_eq!(
439            json_schema(&shape).expect("emits"),
440            json_schema(&literal).expect("emits")
441        );
442    }
443
444    #[test]
445    fn recursive_alias_is_rejected_rather_than_emitting_a_ref() {
446        // `type Node = { next: Node | null }`. Google's provider strips `$ref`,
447        // so a recursive type must fail at compile time instead.
448        let node = alias_ref("Node");
449        let body = obj(&[("next", Type::union(vec![node.clone(), Type::Null]), false)]);
450        let expand = |ty: &Type| {
451            if matches!(ty.peel(), Type::AliasRef { name, .. } if name == "Node") {
452                body.clone()
453            } else {
454                ty.peel().clone()
455            }
456        };
457        let err = json_schema_with(&node, &expand).expect_err("rejects");
458        assert_eq!(err.path, "next");
459        assert!(err.reason.contains("recursive"), "reason: {}", err.reason);
460    }
461
462    #[test]
463    fn unexpandable_alias_ref_is_rejected() {
464        let err = json_schema(&alias_ref("Missing")).expect_err("rejects");
465        assert!(err.reason.contains("recursive"), "reason: {}", err.reason);
466    }
467
468    #[test]
469    fn function_field_is_rejected_and_the_message_names_the_field() {
470        // This shape compiles today under the cast gate, which returns `None`
471        // for `Function`. The schema gate must name `handler`.
472        let handler = Type::Function {
473            params: vec![Type::Number],
474            ret: Box::new(Type::Number),
475            predicate: None,
476            has_rest: false,
477            optional: 0,
478        };
479        let err = json_schema(&obj(&[
480            ("handler", handler, false),
481            ("id", Type::String, false),
482        ]))
483        .expect_err("rejects");
484        assert_eq!(err.path, "handler");
485        assert!(
486            err.describe().contains("field `handler`"),
487            "{}",
488            err.describe()
489        );
490        assert!(err.reason.contains("function"), "reason: {}", err.reason);
491    }
492
493    #[test]
494    fn nested_field_path_is_built_through_objects_and_arrays() {
495        let handler = Type::Function {
496            params: Vec::new(),
497            ret: Box::new(Type::Void),
498            predicate: None,
499            has_rest: false,
500            optional: 0,
501        };
502        let item = obj(&[("run", handler, false)]);
503        let ty = obj(&[("items", Type::Array(Box::new(item)), false)]);
504        let err = json_schema(&ty).expect_err("rejects");
505        assert_eq!(err.path, "items[].run");
506    }
507
508    #[test]
509    fn tuple_element_path_names_its_index() {
510        let ty = obj(&[("pair", Type::tuple(vec![Type::String, Type::BigInt]), false)]);
511        let err = json_schema(&ty).expect_err("rejects");
512        assert_eq!(err.path, "pair[1]");
513    }
514
515    #[test]
516    fn bigint_and_uint8array_have_distinct_reasons() {
517        let big = json_schema(&obj(&[("n", Type::BigInt, false)])).expect_err("rejects");
518        let bytes = json_schema(&obj(&[("b", Type::Uint8Array, false)])).expect_err("rejects");
519        assert_eq!(big.path, "n");
520        assert_eq!(bytes.path, "b");
521        assert!(big.reason.contains("bigint"), "reason: {}", big.reason);
522        assert!(
523            bytes.reason.contains("Uint8Array"),
524            "reason: {}",
525            bytes.reason
526        );
527        assert_ne!(big.reason, bytes.reason);
528    }
529
530    #[test]
531    fn nested_unknown_is_rejected_though_the_cast_gate_accepts_it() {
532        // The cast gate returns `None` for `Unknown`. If the schema emitter
533        // deferred to it, this would emit `{}` and constrain nothing — the very
534        // degeneration to "any" that inlining exists to avoid.
535        let err = json_schema(&obj(&[("payload", Type::Unknown, false)])).expect_err("rejects");
536        assert_eq!(err.path, "payload");
537        assert!(err.reason.contains("unknown"), "reason: {}", err.reason);
538    }
539
540    #[test]
541    fn an_already_failed_type_rejects_without_asking_for_a_second_diagnostic() {
542        // `Type::Error` means a diagnostic was already reported for this type.
543        // There is still no schema, but the caller must not stack a second
544        // complaint onto the same mistake.
545        let err = json_schema(&obj(&[("bad", Type::Error, false)])).expect_err("rejects");
546        assert_eq!(err.path, "bad");
547        assert!(
548            err.cascading,
549            "an already-failed type is a cascading reject"
550        );
551        // Every other reject wants to be reported.
552        let real = json_schema(&obj(&[("n", Type::BigInt, false)])).expect_err("rejects");
553        assert!(!real.cascading, "a genuine rejection is reportable");
554    }
555
556    #[test]
557    fn emitted_schema_uses_only_the_safe_subset_of_keys() {
558        const SAFE: &[&str] = &[
559            "type",
560            "properties",
561            "required",
562            "additionalProperties",
563            "items",
564            "prefixItems",
565            "minItems",
566            "maxItems",
567            "anyOf",
568            "const",
569        ];
570
571        fn assert_safe(value: &Value, safe: &[&str]) {
572            match value {
573                Value::Object(map) => {
574                    for (key, child) in map {
575                        assert!(safe.contains(&key.as_str()), "unsafe schema key `{key}`");
576                        // `properties` keys are user field names, not keywords.
577                        if key == "properties" {
578                            for nested in map[key].as_object().expect("properties object").values()
579                            {
580                                assert_safe(nested, safe);
581                            }
582                        } else {
583                            assert_safe(child, safe);
584                        }
585                    }
586                }
587                Value::Array(items) => items.iter().for_each(|i| assert_safe(i, safe)),
588                _ => {}
589            }
590        }
591
592        let ty = obj(&[
593            ("flag", Type::Boolean, true),
594            ("kind", Type::union(vec![s_lit("a"), s_lit("b")]), false),
595            (
596                "rows",
597                Type::Array(Box::new(obj(&[("v", Type::Number, false)]))),
598                false,
599            ),
600            ("span", Type::tuple(vec![Type::Number, Type::Number]), false),
601        ]);
602        assert_safe(&json_schema(&ty).expect("emits"), SAFE);
603    }
604
605    #[test]
606    fn emission_is_deterministic() {
607        let ty = obj(&[
608            ("zeta", Type::String, false),
609            ("alpha", Type::Number, true),
610            ("mid", Type::Array(Box::new(Type::Boolean)), false),
611        ]);
612        let first = serde_json::to_string(&json_schema(&ty).expect("emits")).expect("serializes");
613        for _ in 0..8 {
614            let again =
615                serde_json::to_string(&json_schema(&ty).expect("emits")).expect("serializes");
616            assert_eq!(first, again, "schema emission must be byte-identical");
617        }
618    }
619}
submilli_engine/typechecker/json_schema.rs

submilli_engine/typechecker/
json_schema.rs
