Skip to main content

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::Function { .. } => Err(reject(path, "functions have no JSON form")),
207        Type::BigInt | Type::BigIntLiteral(_) => Err(reject(
208            path,
209            "`bigint` has no JSON number form — use `number` or a `string` encoding",
210        )),
211        Type::Uint8Array => Err(reject(
212            path,
213            "`Uint8Array` has no JSON form — use a `string` encoding such as base64",
214        )),
215        Type::ClassRef { .. } => Err(reject(
216            path,
217            "class types are nominal — use an object type or a data interface",
218        )),
219        Type::InterfaceRef { .. } => Err(reject(
220            path,
221            "this interface has no data shape — method-bearing interfaces are nominal, and recursive ones cannot be inlined",
222        )),
223        Type::NumberEnum { .. } | Type::StringEnum { .. } => Err(reject(
224            path,
225            "enums need their variants resolved — use a union of literals",
226        )),
227        Type::TypeVar(_) | Type::GenericParam { .. } => Err(reject(
228            path,
229            "generic type parameters are erased — the schema is emitted at compile time",
230        )),
231        Type::Never => Err(reject(path, "`never` has no values")),
232        Type::Void => Err(reject(path, "`void` is not a value type")),
233    }
234}
235
236/// The recursion verdict, shared by the back-edge and no-expander cases.
237const RECURSIVE: &str = "recursive types have no finite JSON Schema — \
238                         providers drop `$ref`, so every definition must be inlined";
239
240fn reject(path: &str, reason: &'static str) -> SchemaReject {
241    SchemaReject {
242        path: path.to_string(),
243        reason,
244        cascading: false,
245    }
246}
247
248/// Append `name` to the field path: `"a.b"`, but `"a"` at the root and
249/// `"items[].id"` after an array step (which already ends in a bracket).
250fn push_field(path: &mut String, name: &str) {
251    if !path.is_empty() {
252        path.push('.');
253    }
254    path.push_str(name);
255}
256
257#[cfg(test)]
258mod tests {
259    use std::collections::BTreeMap;
260
261    use super::*;
262    use crate::mangle;
263    use crate::types::{LiteralF64, ObjectField, Package};
264
265    fn obj(fields: &[(&str, Type, bool)]) -> Type {
266        Type::Object {
267            index: None,
268            fields: fields
269                .iter()
270                .map(|(name, ty, optional)| {
271                    (
272                        (*name).to_string(),
273                        if *optional {
274                            ObjectField::optional(ty.clone())
275                        } else {
276                            ObjectField::required(ty.clone())
277                        },
278                    )
279                })
280                .collect::<BTreeMap<_, _>>(),
281        }
282    }
283
284    fn s_lit(s: &str) -> Type {
285        Type::StringLiteral(s.to_string())
286    }
287
288    fn alias_ref(name: &str) -> Type {
289        Type::AliasRef {
290            mangled: mangle::prelude(name),
291            package: Package::prelude(),
292            name: name.to_string(),
293            args: Vec::new(),
294        }
295    }
296
297    #[test]
298    fn flat_object() {
299        let schema = json_schema(&obj(&[
300            ("age", Type::Number, false),
301            ("name", Type::String, false),
302        ]))
303        .expect("emits");
304        assert_eq!(
305            schema,
306            json!({
307                "type": "object",
308                "properties": {
309                    "age": { "type": "number" },
310                    "name": { "type": "string" },
311                },
312                "required": ["age", "name"],
313                "additionalProperties": false,
314            })
315        );
316    }
317
318    #[test]
319    fn nested_object() {
320        let inner = obj(&[("city", Type::String, false)]);
321        let schema = json_schema(&obj(&[("address", inner, false)])).expect("emits");
322        assert_eq!(
323            schema["properties"]["address"]["properties"]["city"],
324            json!({ "type": "string" })
325        );
326        assert_eq!(schema["properties"]["address"]["required"], json!(["city"]));
327    }
328
329    #[test]
330    fn optional_field_is_omitted_from_required_and_is_not_nullable() {
331        let schema = json_schema(&obj(&[
332            ("id", Type::String, false),
333            ("note", Type::String, true),
334        ]))
335        .expect("emits");
336        assert_eq!(schema["required"], json!(["id"]));
337        // Present in `properties`, and as a bare string — the `| undefined` widening
338        // is the read type of the field, not its schema.
339        assert_eq!(schema["properties"]["note"], json!({ "type": "string" }));
340    }
341
342    #[test]
343    fn undefined_projection_preserves_property_requiredness() {
344        let maybe_string = Type::union(vec![Type::String, Type::Undefined]);
345        let schema = json_schema(&obj(&[
346            ("required", maybe_string.clone(), false),
347            ("optional", maybe_string, true),
348            (
349                "nullable",
350                Type::union(vec![Type::String, Type::Null, Type::Undefined]),
351                false,
352            ),
353        ]))
354        .expect("JSON-representable branches remain");
355        assert_eq!(schema["required"], json!(["nullable", "required"]));
356        assert_eq!(schema["properties"]["required"], json!({"type":"string"}));
357        assert_eq!(schema["properties"]["optional"], json!({"type":"string"}));
358        assert_eq!(
359            schema["properties"]["nullable"],
360            json!({"anyOf":[{"type":"string"},{"type":"null"}]})
361        );
362    }
363
364    #[test]
365    fn undefined_only_schema_reports_the_property_path() {
366        let error = json_schema(&obj(&[("missing", Type::Undefined, true)]))
367            .expect_err("undefined is not JSON");
368        assert_eq!(error.path, "missing");
369        assert!(error.reason.contains("undefined has no JSON value"));
370        assert!(json_schema(&Type::Undefined).is_err());
371    }
372
373    #[test]
374    fn optional_tuple_schema_keeps_lengths_and_excludes_null() {
375        let tuple = Type::Tuple(crate::types::TupleType {
376            elements: vec![
377                Type::Number,
378                Type::union(vec![Type::String, Type::Undefined]),
379            ],
380            optional: 1,
381        });
382        let schema = json_schema(&tuple).expect("optional tuple schema");
383        assert_eq!(schema["minItems"], json!(1));
384        assert_eq!(schema["maxItems"], json!(2));
385        assert_eq!(schema["prefixItems"][1], json!({"type":"string"}));
386        assert_eq!(schema["items"], json!(false));
387    }
388
389    #[test]
390    fn array_and_tuple() {
391        let arr = json_schema(&Type::Array(Box::new(Type::Number))).expect("emits");
392        assert_eq!(arr["type"], json!("array"));
393        assert_eq!(arr["items"], json!({ "type": "number" }));
394
395        let tup = json_schema(&Type::tuple(vec![Type::String, Type::Boolean])).expect("emits");
396        assert_eq!(tup["type"], json!("array"));
397        assert_eq!(
398            tup["prefixItems"],
399            json!([{ "type": "string" }, { "type": "boolean" }])
400        );
401        assert_eq!(tup["minItems"], json!(2));
402        assert_eq!(tup["maxItems"], json!(2));
403    }
404
405    #[test]
406    fn union_becomes_any_of() {
407        let ty = Type::union(vec![Type::String, Type::Null]);
408        let schema = json_schema(&ty).expect("emits");
409        let any_of = schema["anyOf"].as_array().expect("anyOf array");
410        assert_eq!(any_of.len(), 2);
411        assert!(any_of.contains(&json!({ "type": "null" })));
412        assert!(any_of.contains(&json!({ "type": "string" })));
413    }
414
415    #[test]
416    fn literals_become_const() {
417        assert_eq!(
418            json_schema(&s_lit("red")).expect("emits"),
419            json!({ "const": "red" })
420        );
421        assert_eq!(
422            json_schema(&Type::NumberLiteral(LiteralF64(3.0))).expect("emits"),
423            json!({ "const": 3.0 })
424        );
425    }
426
427    #[test]
428    fn reduced_interface_matches_the_literal_object_type() {
429        // `reduce_interfaces_to_shapes` runs before the walk, so a data-only
430        // interface arrives here already as a `Type::Object`. Emitting the same
431        // schema for both is what makes that pre-pass sufficient.
432        let shape = obj(&[("id", Type::String, false), ("n", Type::Number, true)]);
433        let literal = obj(&[("id", Type::String, false), ("n", Type::Number, true)]);
434        assert_eq!(
435            json_schema(&shape).expect("emits"),
436            json_schema(&literal).expect("emits")
437        );
438    }
439
440    #[test]
441    fn recursive_alias_is_rejected_rather_than_emitting_a_ref() {
442        // `type Node = { next: Node | null }`. Google's provider strips `$ref`,
443        // so a recursive type must fail at compile time instead.
444        let node = alias_ref("Node");
445        let body = obj(&[("next", Type::union(vec![node.clone(), Type::Null]), false)]);
446        let expand = |ty: &Type| {
447            if matches!(ty.peel(), Type::AliasRef { name, .. } if name == "Node") {
448                body.clone()
449            } else {
450                ty.peel().clone()
451            }
452        };
453        let err = json_schema_with(&node, &expand).expect_err("rejects");
454        assert_eq!(err.path, "next");
455        assert!(err.reason.contains("recursive"), "reason: {}", err.reason);
456    }
457
458    #[test]
459    fn unexpandable_alias_ref_is_rejected() {
460        let err = json_schema(&alias_ref("Missing")).expect_err("rejects");
461        assert!(err.reason.contains("recursive"), "reason: {}", err.reason);
462    }
463
464    #[test]
465    fn function_field_is_rejected_and_the_message_names_the_field() {
466        // This shape compiles today under the cast gate, which returns `None`
467        // for `Function`. The schema gate must name `handler`.
468        let handler = Type::Function {
469            params: vec![Type::Number],
470            ret: Box::new(Type::Number),
471            predicate: None,
472            has_rest: false,
473            optional: 0,
474        };
475        let err = json_schema(&obj(&[
476            ("handler", handler, false),
477            ("id", Type::String, false),
478        ]))
479        .expect_err("rejects");
480        assert_eq!(err.path, "handler");
481        assert!(
482            err.describe().contains("field `handler`"),
483            "{}",
484            err.describe()
485        );
486        assert!(err.reason.contains("function"), "reason: {}", err.reason);
487    }
488
489    #[test]
490    fn nested_field_path_is_built_through_objects_and_arrays() {
491        let handler = Type::Function {
492            params: Vec::new(),
493            ret: Box::new(Type::Void),
494            predicate: None,
495            has_rest: false,
496            optional: 0,
497        };
498        let item = obj(&[("run", handler, false)]);
499        let ty = obj(&[("items", Type::Array(Box::new(item)), false)]);
500        let err = json_schema(&ty).expect_err("rejects");
501        assert_eq!(err.path, "items[].run");
502    }
503
504    #[test]
505    fn tuple_element_path_names_its_index() {
506        let ty = obj(&[("pair", Type::tuple(vec![Type::String, Type::BigInt]), false)]);
507        let err = json_schema(&ty).expect_err("rejects");
508        assert_eq!(err.path, "pair[1]");
509    }
510
511    #[test]
512    fn bigint_and_uint8array_have_distinct_reasons() {
513        let big = json_schema(&obj(&[("n", Type::BigInt, false)])).expect_err("rejects");
514        let bytes = json_schema(&obj(&[("b", Type::Uint8Array, false)])).expect_err("rejects");
515        assert_eq!(big.path, "n");
516        assert_eq!(bytes.path, "b");
517        assert!(big.reason.contains("bigint"), "reason: {}", big.reason);
518        assert!(
519            bytes.reason.contains("Uint8Array"),
520            "reason: {}",
521            bytes.reason
522        );
523        assert_ne!(big.reason, bytes.reason);
524    }
525
526    #[test]
527    fn nested_unknown_is_rejected_though_the_cast_gate_accepts_it() {
528        // The cast gate returns `None` for `Unknown`. If the schema emitter
529        // deferred to it, this would emit `{}` and constrain nothing — the very
530        // degeneration to "any" that inlining exists to avoid.
531        let err = json_schema(&obj(&[("payload", Type::Unknown, false)])).expect_err("rejects");
532        assert_eq!(err.path, "payload");
533        assert!(err.reason.contains("unknown"), "reason: {}", err.reason);
534    }
535
536    #[test]
537    fn an_already_failed_type_rejects_without_asking_for_a_second_diagnostic() {
538        // `Type::Error` means a diagnostic was already reported for this type.
539        // There is still no schema, but the caller must not stack a second
540        // complaint onto the same mistake.
541        let err = json_schema(&obj(&[("bad", Type::Error, false)])).expect_err("rejects");
542        assert_eq!(err.path, "bad");
543        assert!(
544            err.cascading,
545            "an already-failed type is a cascading reject"
546        );
547        // Every other reject wants to be reported.
548        let real = json_schema(&obj(&[("n", Type::BigInt, false)])).expect_err("rejects");
549        assert!(!real.cascading, "a genuine rejection is reportable");
550    }
551
552    #[test]
553    fn emitted_schema_uses_only_the_safe_subset_of_keys() {
554        const SAFE: &[&str] = &[
555            "type",
556            "properties",
557            "required",
558            "additionalProperties",
559            "items",
560            "prefixItems",
561            "minItems",
562            "maxItems",
563            "anyOf",
564            "const",
565        ];
566
567        fn assert_safe(value: &Value, safe: &[&str]) {
568            match value {
569                Value::Object(map) => {
570                    for (key, child) in map {
571                        assert!(safe.contains(&key.as_str()), "unsafe schema key `{key}`");
572                        // `properties` keys are user field names, not keywords.
573                        if key == "properties" {
574                            for nested in map[key].as_object().expect("properties object").values()
575                            {
576                                assert_safe(nested, safe);
577                            }
578                        } else {
579                            assert_safe(child, safe);
580                        }
581                    }
582                }
583                Value::Array(items) => items.iter().for_each(|i| assert_safe(i, safe)),
584                _ => {}
585            }
586        }
587
588        let ty = obj(&[
589            ("flag", Type::Boolean, true),
590            ("kind", Type::union(vec![s_lit("a"), s_lit("b")]), false),
591            (
592                "rows",
593                Type::Array(Box::new(obj(&[("v", Type::Number, false)]))),
594                false,
595            ),
596            ("span", Type::tuple(vec![Type::Number, Type::Number]), false),
597        ]);
598        assert_safe(&json_schema(&ty).expect("emits"), SAFE);
599    }
600
601    #[test]
602    fn emission_is_deterministic() {
603        let ty = obj(&[
604            ("zeta", Type::String, false),
605            ("alpha", Type::Number, true),
606            ("mid", Type::Array(Box::new(Type::Boolean)), false),
607        ]);
608        let first = serde_json::to_string(&json_schema(&ty).expect("emits")).expect("serializes");
609        for _ in 0..8 {
610            let again =
611                serde_json::to_string(&json_schema(&ty).expect("emits")).expect("serializes");
612            assert_eq!(first, again, "schema emission must be byte-identical");
613        }
614    }
615}