1use serde_json::Value;
23
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
27pub enum NativeSchema {
28 ClaudeJsonSchema,
32 }
38
39pub struct Schema {
42 raw: Value,
43 text: String,
45 validator: jsonschema::Validator,
46}
47
48impl Schema {
49 pub fn compile(text: &str) -> Result<Schema, String> {
53 let raw: Value =
54 serde_json::from_str(text).map_err(|e| format!("schema is not valid JSON: {e}"))?;
55 let validator =
56 jsonschema::validator_for(&raw).map_err(|e| format!("not a valid JSON Schema: {e}"))?;
57 let text = serde_json::to_string(&raw).unwrap_or_else(|_| text.to_string());
61 Ok(Schema {
62 raw,
63 text,
64 validator,
65 })
66 }
67
68 pub fn as_value(&self) -> &Value {
70 &self.raw
71 }
72
73 pub fn as_text(&self) -> &str {
75 &self.text
76 }
77
78 pub fn validate(&self, instance: &Value) -> Result<(), Vec<String>> {
81 let errors: Vec<String> = self
82 .validator
83 .iter_errors(instance)
84 .map(|e| e.to_string())
85 .collect();
86 if errors.is_empty() {
87 Ok(())
88 } else {
89 Err(errors)
90 }
91 }
92}
93
94#[derive(Debug, Clone, PartialEq)]
97pub struct Check {
98 pub value: Option<Value>,
99 pub errors: Vec<String>,
100}
101
102impl Check {
103 pub fn is_valid(&self) -> bool {
105 self.value.is_some() && self.errors.is_empty()
106 }
107}
108
109pub fn check(schema: &Schema, native: Option<NativeSchema>, text: &str, stdout: &str) -> Check {
113 match extract_value(native, text, stdout) {
114 Some(value) => match schema.validate(&value) {
115 Ok(()) => Check {
116 value: Some(value),
117 errors: Vec::new(),
118 },
119 Err(errors) => Check {
120 value: Some(value),
121 errors,
122 },
123 },
124 None => Check {
125 value: None,
126 errors: vec!["no JSON value could be extracted from the response".to_string()],
127 },
128 }
129}
130
131pub fn extract_value(native: Option<NativeSchema>, text: &str, stdout: &str) -> Option<Value> {
134 if native == Some(NativeSchema::ClaudeJsonSchema) {
135 if let Ok(doc) = serde_json::from_str::<Value>(stdout.trim()) {
139 if let Some(value) = doc.get("structured_output") {
140 if !value.is_null() {
141 return Some(value.clone());
142 }
143 }
144 }
145 }
146 extract_json(text)
147}
148
149pub fn extract_json(text: &str) -> Option<Value> {
154 let trimmed = text.trim();
155 if trimmed.is_empty() {
156 return None;
157 }
158 if let Ok(value) = serde_json::from_str::<Value>(trimmed) {
159 return Some(value);
160 }
161 if let Some(inner) = fenced_block(trimmed) {
162 let body = inner.trim();
163 if let Ok(value) = serde_json::from_str::<Value>(body) {
164 return Some(value);
165 }
166 if let Some(value) = first_json_value(body) {
167 return Some(value);
168 }
169 }
170 first_json_value(trimmed)
171}
172
173fn fenced_block(s: &str) -> Option<&str> {
176 let start = s.find("```")?;
177 let after = &s[start + 3..];
178 let body_start = after.find('\n').map(|i| i + 1)?;
180 let body = &after[body_start..];
181 let end = body.find("```")?;
182 Some(&body[..end])
183}
184
185fn first_json_value(s: &str) -> Option<Value> {
189 let start = s.find(['{', '['])?;
190 serde_json::Deserializer::from_str(&s[start..])
191 .into_iter::<Value>()
192 .next()
193 .and_then(Result::ok)
194}
195
196pub fn prompt_instruction(schema_text: &str) -> String {
206 format!(
207 "You must respond with a single JSON value that strictly conforms to the \
208 following JSON Schema. Output ONLY that JSON value — no prose, no \
209 explanation, and no Markdown code fences. JSON Schema: {schema_text}"
210 )
211}
212
213pub fn retry_instruction(schema_text: &str, previous: &str, errors: &[String]) -> String {
219 let errors = if errors.is_empty() {
220 "(no JSON value could be extracted from the response)".to_string()
221 } else {
222 errors.join("; ")
223 };
224 let previous = flatten_whitespace(previous);
225 format!(
226 "Your previous response did not conform to the required JSON Schema. \
227 Previous response: {previous} -- Validation errors: {errors}. \
228 Respond again with ONLY a single JSON value that strictly conforms to \
229 this JSON Schema (no prose, no code fences): {schema_text}"
230 )
231}
232
233fn flatten_whitespace(s: &str) -> String {
237 s.split_whitespace().collect::<Vec<_>>().join(" ")
238}
239
240#[cfg(test)]
241mod tests {
242 use super::*;
243 use serde_json::json;
244
245 const PERSON: &str = r#"{"type":"object","properties":{"name":{"type":"string"},
246 "age":{"type":"integer"}},"required":["name","age"],"additionalProperties":false}"#;
247
248 fn person() -> Schema {
249 Schema::compile(PERSON).expect("schema compiles")
250 }
251
252 #[test]
253 fn compile_rejects_non_json_and_invalid_schema() {
254 let err = Schema::compile("not json")
255 .err()
256 .expect("non-json rejected");
257 assert!(err.contains("not valid JSON"), "{err}");
258 let err = Schema::compile(r#"{"type": 5}"#)
260 .err()
261 .expect("invalid schema rejected");
262 assert!(err.contains("not a valid JSON Schema"), "{err}");
263 }
264
265 #[test]
266 fn validate_accepts_conforming_and_lists_errors_otherwise() {
267 let s = person();
268 assert!(s.validate(&json!({"name": "Ada", "age": 36})).is_ok());
269 let errs = s
270 .validate(&json!({"name": "Ada"}))
271 .expect_err("missing required age");
272 assert!(!errs.is_empty());
273 assert!(s
275 .validate(&json!({"name": "Ada", "age": 1, "x": true}))
276 .is_err());
277 }
278
279 #[test]
280 fn as_text_is_canonical_compact_json() {
281 let s = Schema::compile(r#"{ "type" : "string" }"#).unwrap();
282 assert_eq!(s.as_text(), r#"{"type":"string"}"#);
283 assert_eq!(s.as_value(), &json!({"type": "string"}));
284 }
285
286 #[test]
287 fn extract_json_parses_plain_document() {
288 let v = extract_json(r#" {"a":1} "#).unwrap();
289 assert_eq!(v, json!({"a": 1}));
290 }
291
292 #[test]
293 fn extract_json_unwraps_fenced_block() {
294 let text = "Here you go:\n```json\n{\"name\":\"Ada\",\"age\":36}\n```\nDone.";
295 assert_eq!(extract_json(text).unwrap(), json!({"name":"Ada","age":36}));
296 let text = "```\n[1, 2, 3]\n```";
298 assert_eq!(extract_json(text).unwrap(), json!([1, 2, 3]));
299 }
300
301 #[test]
302 fn extract_json_recovers_value_inside_a_noisy_fence() {
303 let text = "```json\n{\"name\":\"Ada\",\"age\":36} (that's her)\n```";
307 assert_eq!(extract_json(text).unwrap(), json!({"name":"Ada","age":36}));
308 }
309
310 #[test]
311 fn extract_json_finds_object_embedded_in_prose() {
312 let text = "Sure! {\"name\": \"Ada\", \"age\": 36} — hope that helps.";
313 assert_eq!(extract_json(text).unwrap(), json!({"name":"Ada","age":36}));
314 }
315
316 #[test]
317 fn extract_json_returns_none_when_absent() {
318 assert!(extract_json("no json here at all").is_none());
319 assert!(extract_json(" ").is_none());
320 assert_eq!(extract_json("```json\n{\"a\":1}").unwrap(), json!({"a": 1}));
323 }
324
325 #[test]
326 fn check_validates_prompt_based_text() {
327 let s = person();
328 let ok = check(&s, None, r#"{"name":"Ada","age":36}"#, "");
329 assert!(ok.is_valid());
330 assert_eq!(ok.value.unwrap(), json!({"name":"Ada","age":36}));
331
332 let bad = check(&s, None, r#"{"name":"Ada"}"#, "");
333 assert!(!bad.is_valid());
334 assert!(bad.value.is_some());
335 assert!(!bad.errors.is_empty());
336
337 let none = check(&s, None, "I can't do that", "");
338 assert!(!none.is_valid());
339 assert!(none.value.is_none());
340 assert_eq!(none.errors.len(), 1);
341 }
342
343 #[test]
344 fn check_native_reads_structured_output_field() {
345 let s = person();
346 let stdout = r#"{"type":"result","result":"Here is Ada.",
349 "structured_output":{"name":"Ada","age":36}}"#;
350 let c = check(
351 &s,
352 Some(NativeSchema::ClaudeJsonSchema),
353 "Here is Ada.",
354 stdout,
355 );
356 assert!(c.is_valid());
357 assert_eq!(c.value.unwrap(), json!({"name":"Ada","age":36}));
358 }
359
360 #[test]
361 fn native_falls_back_to_text_when_field_absent() {
362 let s = person();
363 let c = check(
366 &s,
367 Some(NativeSchema::ClaudeJsonSchema),
368 r#"{"name":"Ada","age":36}"#,
369 r#"{"type":"result","result":"{\"name\":\"Ada\",\"age\":36}"}"#,
370 );
371 assert!(c.is_valid());
372 }
373
374 #[test]
375 fn instructions_carry_the_schema_and_errors() {
376 let instr = prompt_instruction(r#"{"type":"string"}"#);
377 assert!(instr.contains(r#"{"type":"string"}"#));
378 assert!(instr.contains("ONLY"));
379
380 let retry = retry_instruction(
381 r#"{"type":"object"}"#,
382 "oops not json",
383 &["missing required property 'name'".to_string()],
384 );
385 assert!(retry.contains("oops not json"));
386 assert!(retry.contains("missing required property 'name'"));
387 assert!(retry.contains(r#"{"type":"object"}"#));
388
389 let retry = retry_instruction("{}", "prose", &[]);
392 assert!(retry.contains("no JSON value could be extracted"));
393 }
394
395 #[test]
396 fn instructions_are_single_line_for_cmd_shim_safety() {
397 let instr = prompt_instruction("{\"type\":\"object\"}");
402 assert!(!instr.contains('\n'), "prompt_instruction must be one line");
403 let retry = retry_instruction(
404 "{\"type\":\"object\"}",
405 "line one\nline two\r\nline three",
406 &["err a".to_string(), "err b".to_string()],
407 );
408 assert!(!retry.contains('\n'), "retry_instruction must be one line");
409 assert!(!retry.contains('\r'), "retry_instruction must be one line");
410 assert!(retry.contains("line one line two line three"), "{retry}");
412 assert!(retry.contains("err a; err b"));
413 }
414
415 #[test]
416 fn validator_is_shareable_across_threads() {
417 fn assert_sync<T: Sync>() {}
420 assert_sync::<Schema>();
421 }
422}