1use serde_json::Value;
5
6use super::model::{AgentResponse, questions_field_description, subtasks_field_description};
7
8#[derive(Clone, Copy, Debug, Eq, PartialEq)]
17pub enum SchemaRequiredPolicy {
18 AllProperties,
20 MinimumProtocolKeys,
22}
23
24pub fn agent_response_output_schema(required_policy: SchemaRequiredPolicy) -> Value {
31 let mut value = agent_response_json_schema();
32 normalize_schema_for_transport(&mut value, required_policy);
33
34 value
35}
36
37pub fn agent_response_json_schema_json() -> String {
43 let schema = agent_response_json_schema();
44
45 stringify_schema_json(&schema)
46}
47
48pub fn agent_response_output_schema_json(required_policy: SchemaRequiredPolicy) -> String {
54 let schema = agent_response_output_schema(required_policy);
55
56 stringify_schema_json(&schema)
57}
58
59fn agent_response_json_schema() -> Value {
65 let schema = schemars::schema_for!(AgentResponse);
66 let mut schema_value = serde_json::to_value(schema).unwrap_or(Value::Null);
67
68 inject_dynamic_schema_guidance(&mut schema_value);
69 inject_additional_properties_false(&mut schema_value);
70 inject_minimum_required_protocol_key(&mut schema_value);
71
72 schema_value
73}
74
75fn inject_dynamic_schema_guidance(schema: &mut Value) {
78 let Some(properties) = schema.get_mut("properties").and_then(Value::as_object_mut) else {
79 return;
80 };
81
82 for (property_name, description) in [
83 ("questions", questions_field_description()),
84 ("subtasks", subtasks_field_description()),
85 ] {
86 let Some(property) = properties
87 .get_mut(property_name)
88 .and_then(Value::as_object_mut)
89 else {
90 continue;
91 };
92
93 property.insert("description".to_string(), Value::String(description));
94 }
95}
96
97fn inject_additional_properties_false(value: &mut Value) {
107 match value {
108 Value::Object(object) => {
109 if object.contains_key("properties") && !object.contains_key("additionalProperties") {
110 object.insert("additionalProperties".to_string(), Value::Bool(false));
111 }
112
113 for nested_value in object.values_mut() {
114 inject_additional_properties_false(nested_value);
115 }
116 }
117 Value::Array(array) => {
118 for nested_value in array {
119 inject_additional_properties_false(nested_value);
120 }
121 }
122 _ => {}
123 }
124}
125
126fn inject_minimum_required_protocol_key(schema: &mut Value) {
135 let Some(object) = schema.as_object_mut() else {
136 return;
137 };
138
139 let required = object
140 .entry("required")
141 .or_insert_with(|| Value::Array(Vec::new()));
142
143 let Some(required_array) = required.as_array_mut() else {
144 return;
145 };
146
147 let already_listed = required_array
148 .iter()
149 .any(|value| value.as_str() == Some("answer"));
150
151 if !already_listed {
152 required_array.push(Value::String("answer".to_string()));
153 }
154}
155
156fn normalize_schema_for_transport(value: &mut Value, required_policy: SchemaRequiredPolicy) {
165 match value {
166 Value::Object(object) => {
167 object.remove("$schema");
168
169 for nested_value in object.values_mut() {
170 normalize_schema_for_transport(nested_value, required_policy);
171 }
172
173 normalize_ref_object_for_codex(object);
174 if required_policy == SchemaRequiredPolicy::AllProperties {
175 normalize_required_for_codex(object);
176 }
177
178 let one_of_values = object
179 .get("oneOf")
180 .and_then(Value::as_array)
181 .map(|items| {
182 items
183 .iter()
184 .filter_map(Value::as_object)
185 .map(|item| item.get("const").and_then(Value::as_str))
186 .collect::<Option<Vec<_>>>()
187 })
188 .map(|option| {
189 option.map(|values| {
190 values
191 .into_iter()
192 .map(ToString::to_string)
193 .collect::<Vec<_>>()
194 })
195 });
196
197 if let Some(Some(enum_variants)) = one_of_values {
198 object.remove("oneOf");
199 object.insert("type".to_string(), Value::String("string".to_string()));
200 object.insert(
201 "enum".to_string(),
202 Value::Array(enum_variants.into_iter().map(Value::String).collect()),
203 );
204 }
205 }
206 Value::Array(array) => {
207 for nested_value in array {
208 normalize_schema_for_transport(nested_value, required_policy);
209 }
210 }
211 _ => {}
212 }
213}
214
215fn normalize_ref_object_for_codex(object: &mut serde_json::Map<String, Value>) {
221 let Some(reference) = object.get("$ref").cloned() else {
222 return;
223 };
224
225 object.clear();
226 object.insert("$ref".to_string(), reference);
227}
228
229fn normalize_required_for_codex(object: &mut serde_json::Map<String, Value>) {
235 let Some(properties) = object.get("properties").and_then(Value::as_object) else {
236 return;
237 };
238
239 let property_keys: Vec<String> = properties.keys().cloned().collect();
240 if property_keys.is_empty() {
241 return;
242 }
243
244 let required = object
245 .entry("required")
246 .or_insert_with(|| Value::Array(Vec::new()));
247
248 let Some(required_array) = required.as_array_mut() else {
249 return;
250 };
251
252 for key in &property_keys {
253 let already_listed = required_array
254 .iter()
255 .any(|value| value.as_str() == Some(key));
256
257 if !already_listed {
258 required_array.push(Value::String(key.clone()));
259 }
260 }
261}
262
263fn stringify_schema_json(schema: &Value) -> String {
265 serde_json::to_string_pretty(schema).unwrap_or("null".to_string())
266}
267
268#[cfg(test)]
269mod tests {
270 use super::*;
271
272 #[test]
273 fn test_agent_response_output_schema_contains_required_fields() {
275 let schema = agent_response_output_schema(SchemaRequiredPolicy::AllProperties);
277 let required_fields = schema
278 .get("required")
279 .and_then(Value::as_array)
280 .expect("schema required fields should exist");
281 let properties = schema
282 .get("properties")
283 .and_then(Value::as_object)
284 .expect("schema properties should exist");
285
286 assert!(
288 required_fields
289 .iter()
290 .any(|value| value.as_str() == Some("answer"))
291 );
292 assert!(
293 required_fields
294 .iter()
295 .any(|value| value.as_str() == Some("questions"))
296 );
297 assert!(
298 required_fields
299 .iter()
300 .any(|value| value.as_str() == Some("review_comment_outcomes"))
301 );
302 assert!(
303 required_fields
304 .iter()
305 .any(|value| value.as_str() == Some("subtasks"))
306 );
307 assert!(
308 required_fields
309 .iter()
310 .any(|value| value.as_str() == Some("summary"))
311 );
312 assert!(properties.contains_key("answer"));
313 assert!(properties.contains_key("questions"));
314 assert!(properties.contains_key("review_comment_outcomes"));
315 assert!(properties.contains_key("subtasks"));
316 assert!(properties.contains_key("summary"));
317 }
318
319 #[test]
320 fn test_inject_dynamic_schema_guidance_skips_absent_properties() {
323 let mut schema = serde_json::json!({
325 "properties": {
326 "answer": { "type": "string" }
327 }
328 });
329
330 inject_dynamic_schema_guidance(&mut schema);
332
333 assert_eq!(
335 schema,
336 serde_json::json!({
337 "properties": {
338 "answer": { "type": "string" }
339 }
340 })
341 );
342 }
343
344 #[test]
345 fn test_agent_response_json_schema_injects_subtasks_description() {
348 let schema = agent_response_json_schema();
350 let response_properties = schema
351 .get("properties")
352 .and_then(Value::as_object)
353 .expect("response properties should exist");
354 let subtask_properties = schema_definition_properties(&schema, "SubtaskItem");
355
356 assert_eq!(
358 response_properties
359 .get("subtasks")
360 .and_then(|value| value.get("description"))
361 .and_then(Value::as_str),
362 Some(subtasks_field_description().as_str())
363 );
364 assert!(subtask_properties.contains_key("prompt"));
365 assert!(subtask_properties.contains_key("kind"));
366 assert!(subtask_properties.contains_key("task_key"));
367 assert!(subtask_properties.contains_key("title"));
368 assert!(subtask_properties.contains_key("touched_areas"));
369 }
370
371 #[test]
372 fn test_agent_response_output_schema_all_properties_are_required() {
375 let schema = agent_response_output_schema(SchemaRequiredPolicy::AllProperties);
377
378 assert!(
380 all_properties_in_required(&schema),
381 "every object with `properties` should list all keys in `required`"
382 );
383 }
384
385 #[test]
386 fn test_agent_response_output_schema_minimum_policy_requires_only_answer() {
390 let schema = agent_response_output_schema(SchemaRequiredPolicy::MinimumProtocolKeys);
392 let required_fields = schema
393 .get("required")
394 .and_then(Value::as_array)
395 .expect("schema required fields should exist");
396
397 assert_eq!(
399 required_fields,
400 &vec![Value::String("answer".to_string())],
401 "only `answer` should be required; demanding optional response fields rejects \
402 ordinary replies that omit them"
403 );
404 }
405
406 #[test]
407 fn test_agent_response_output_schema_does_not_contain_one_of() {
410 let schema = agent_response_output_schema(SchemaRequiredPolicy::AllProperties);
412
413 assert!(!contains_schema_key(&schema, "oneOf"));
415 }
416
417 #[test]
418 fn test_agent_response_output_schema_does_not_contain_schema_metadata() {
421 let schema = agent_response_output_schema(SchemaRequiredPolicy::MinimumProtocolKeys);
423
424 assert!(!contains_schema_key(&schema, "$schema"));
426 }
427
428 #[test]
429 fn test_agent_response_json_schema_requires_answer_key() {
432 let schema = agent_response_json_schema();
434 let required_fields = schema
435 .get("required")
436 .and_then(Value::as_array)
437 .expect("schema required fields should exist");
438
439 assert!(
441 required_fields
442 .iter()
443 .any(|value| value.as_str() == Some("answer")),
444 "prompt schema should require `answer` to align with parser key-presence check"
445 );
446 }
447
448 #[test]
449 fn test_agent_response_json_schema_sets_additional_properties_false() {
453 let schema = agent_response_json_schema();
455
456 assert!(
458 all_properties_objects_deny_additional(&schema),
459 "every object with `properties` should set `additionalProperties: false`"
460 );
461 }
462
463 #[test]
464 fn test_inject_additional_properties_false_preserves_existing_value() {
467 let mut schema = serde_json::json!({
469 "type": "object",
470 "properties": {
471 "extra": { "type": "object", "additionalProperties": { "type": "string" } }
472 }
473 });
474
475 inject_additional_properties_false(&mut schema);
477
478 assert_eq!(schema["additionalProperties"], Value::Bool(false));
480 assert_eq!(
482 schema["properties"]["extra"]["additionalProperties"],
483 serde_json::json!({ "type": "string" })
484 );
485 }
486
487 #[test]
488 fn test_agent_response_output_schema_ref_objects_have_no_sibling_keywords() {
490 let schema = agent_response_output_schema(SchemaRequiredPolicy::AllProperties);
492
493 assert!(!contains_ref_with_sibling_keywords(&schema));
495 }
496
497 #[test]
498 fn test_agent_response_json_schema_json_is_parseable_value() {
500 let schema_json = agent_response_json_schema_json();
502 let parsed_schema: Value =
503 serde_json::from_str(&schema_json).expect("schema string should parse as JSON");
504 let schema_value = agent_response_json_schema();
505
506 assert_eq!(parsed_schema, schema_value);
508 }
509
510 #[test]
511 fn test_agent_response_json_schema_preserves_explicit_payload_metadata() {
514 let schema = agent_response_json_schema();
516
517 assert_eq!(
519 schema.get("title").and_then(Value::as_str),
520 Some("AgentResponse")
521 );
522 assert_eq!(
523 schema.get("description").and_then(Value::as_str),
524 Some(
525 "Wire-format protocol payload used for schema-driven provider output. Return this \
526 object as the entire assistant response payload. Providers that support output \
527 schemas (for example, Codex app-server) are asked to emit this object directly."
528 )
529 );
530 }
531
532 #[test]
533 fn test_agent_response_json_schema_preserves_nested_metadata() {
535 let schema = agent_response_json_schema();
537 let question_definition = schema
538 .get("$defs")
539 .and_then(|value| value.get("QuestionItem"))
540 .and_then(Value::as_object)
541 .expect("question definition should exist");
542 let review_comment_outcome_definition = schema
543 .get("$defs")
544 .and_then(|value| value.get("ReviewCommentOutcome"))
545 .and_then(Value::as_object)
546 .expect("review comment outcome definition should exist");
547 let review_comment_resolution_definition = schema
548 .get("$defs")
549 .and_then(|value| value.get("ReviewCommentResolution"))
550 .and_then(Value::as_object)
551 .expect("review comment resolution definition should exist");
552 let summary_definition = schema
553 .get("$defs")
554 .and_then(|value| value.get("AgentResponseSummary"))
555 .and_then(Value::as_object)
556 .expect("summary definition should exist");
557
558 assert_eq!(
560 question_definition.get("title").and_then(Value::as_str),
561 Some("QuestionItem")
562 );
563 assert_eq!(
564 review_comment_outcome_definition
565 .get("title")
566 .and_then(Value::as_str),
567 Some("ReviewCommentOutcome")
568 );
569 assert_eq!(
570 review_comment_resolution_definition
571 .get("title")
572 .and_then(Value::as_str),
573 Some("ReviewCommentResolution")
574 );
575 assert_eq!(
576 summary_definition.get("title").and_then(Value::as_str),
577 Some("AgentResponseSummary")
578 );
579 }
580
581 #[test]
582 fn test_agent_response_json_schema_preserves_field_metadata() {
584 let schema = agent_response_json_schema();
586 let response_properties = schema
587 .get("properties")
588 .and_then(Value::as_object)
589 .expect("response properties should exist");
590 let question_properties = schema_definition_properties(&schema, "QuestionItem");
591 let review_comment_outcome_properties =
592 schema_definition_properties(&schema, "ReviewCommentOutcome");
593 let summary_properties = schema_definition_properties(&schema, "AgentResponseSummary");
594 let expected_questions_description = questions_field_description();
595
596 assert_schema_property_title_and_description(
598 response_properties,
599 "answer",
600 "answer",
601 "Markdown answer text for delivered work, status updates, or concise completion \
602 notes. Keep clarification requests out of this field and emit them through \
603 `questions` instead.",
604 );
605 assert_eq!(
606 response_properties
607 .get("questions")
608 .and_then(|value| value.get("description"))
609 .and_then(Value::as_str),
610 Some(expected_questions_description.as_str())
611 );
612 assert_schema_property_title_and_description(
613 response_properties,
614 "review_comment_outcomes",
615 "review_comment_outcomes",
616 "Per-thread outcomes for an agent-driven forge comment-resolution turn. Emit an empty \
617 array unless the prompt explicitly supplies forge thread IDs. Copy each reported \
618 `thread_id` exactly from the prompt.",
619 );
620 assert_schema_property_title_and_description(
621 response_properties,
622 "summary",
623 "summary",
624 "Structured summary for session-discussion turns, kept outside `answer` markdown. Use \
625 `null` for one-shot prompts and legacy payloads.",
626 );
627 assert_schema_property_title_and_description(
628 question_properties,
629 "text",
630 "text",
631 "Human-readable markdown text for this question. Ask one specific actionable question \
632 instead of bundling multiple decisions into one item.",
633 );
634 assert_schema_property_title(question_properties, "options", "options");
635 assert_schema_property_title_and_description(
636 review_comment_outcome_properties,
637 "reply",
638 "reply",
639 "Concise reply suitable for posting to the forge review thread.",
640 );
641 assert_schema_property_title_and_description(
642 review_comment_outcome_properties,
643 "resolution",
644 "resolution",
645 "Whether the targeted thread was fixed or required no change.",
646 );
647 assert_schema_property_title_and_description(
648 review_comment_outcome_properties,
649 "thread_id",
650 "thread_id",
651 "Opaque forge thread identifier copied exactly from the turn prompt.",
652 );
653 assert_schema_property_title_and_description(
654 summary_properties,
655 "turn",
656 "turn",
657 "Concise summary of only the work completed in the current turn.",
658 );
659 assert_schema_property_title_and_description(
660 summary_properties,
661 "session",
662 "session",
663 "Cumulative summary of active changes on the current session branch.",
664 );
665 }
666
667 #[test]
668 fn test_agent_response_json_schema_keeps_optional_summary_field() {
671 let schema = agent_response_json_schema();
673 let response_required_fields = schema
674 .get("required")
675 .and_then(Value::as_array)
676 .cloned()
677 .unwrap_or_default();
678 let question_definition = schema
679 .get("$defs")
680 .and_then(|value| value.get("QuestionItem"))
681 .and_then(Value::as_object)
682 .expect("question definition should exist");
683 let question_required_fields = question_definition
684 .get("required")
685 .and_then(Value::as_array)
686 .cloned()
687 .unwrap_or_default();
688
689 assert!(
691 response_required_fields
692 .iter()
693 .all(|field| field.as_str() != Some("summary")),
694 "raw prompt schema should keep optional summary fields optional"
695 );
696 assert!(
697 question_required_fields
698 .iter()
699 .all(|field| field.as_str() != Some("options")),
700 "question schema should keep `options` optional for omitted empty lists"
701 );
702 }
703
704 #[test]
705 fn test_agent_response_output_schema_json_is_parseable_value() {
708 let schema_json =
710 agent_response_output_schema_json(SchemaRequiredPolicy::MinimumProtocolKeys);
711 let parsed_schema: Value =
712 serde_json::from_str(&schema_json).expect("schema string should parse as JSON");
713 let schema_value = agent_response_output_schema(SchemaRequiredPolicy::MinimumProtocolKeys);
714
715 assert_eq!(parsed_schema, schema_value);
717 }
718
719 fn contains_schema_key(value: &Value, key: &str) -> bool {
721 match value {
722 Value::Object(object) => {
723 if object.contains_key(key) {
724 return true;
725 }
726
727 object
728 .values()
729 .any(|nested_value| contains_schema_key(nested_value, key))
730 }
731 Value::Array(array) => array
732 .iter()
733 .any(|nested_value| contains_schema_key(nested_value, key)),
734 _ => false,
735 }
736 }
737
738 fn contains_ref_with_sibling_keywords(value: &Value) -> bool {
740 match value {
741 Value::Object(object) => {
742 if object.contains_key("$ref") && object.len() > 1 {
743 return true;
744 }
745
746 object.values().any(contains_ref_with_sibling_keywords)
747 }
748 Value::Array(array) => array.iter().any(contains_ref_with_sibling_keywords),
749 _ => false,
750 }
751 }
752
753 fn all_properties_in_required(value: &Value) -> bool {
756 match value {
757 Value::Object(object) => {
758 if let Some(properties) = object.get("properties").and_then(Value::as_object) {
759 let required_keys: Vec<&str> = object
760 .get("required")
761 .and_then(Value::as_array)
762 .map(|array| array.iter().filter_map(Value::as_str).collect())
763 .unwrap_or_default();
764
765 for key in properties.keys() {
766 if !required_keys.contains(&key.as_str()) {
767 return false;
768 }
769 }
770 }
771
772 object.values().all(all_properties_in_required)
773 }
774 Value::Array(array) => array.iter().all(all_properties_in_required),
775 _ => true,
776 }
777 }
778
779 fn all_properties_objects_deny_additional(value: &Value) -> bool {
782 match value {
783 Value::Object(object) => {
784 if object.contains_key("properties")
785 && object.get("additionalProperties") != Some(&Value::Bool(false))
786 {
787 return false;
788 }
789
790 object.values().all(all_properties_objects_deny_additional)
791 }
792 Value::Array(array) => array.iter().all(all_properties_objects_deny_additional),
793 _ => true,
794 }
795 }
796
797 fn schema_definition_properties<'a>(
799 schema: &'a Value,
800 definition_name: &str,
801 ) -> &'a serde_json::Map<String, Value> {
802 schema
803 .get("$defs")
804 .and_then(|value| value.get(definition_name))
805 .and_then(|value| value.get("properties"))
806 .and_then(Value::as_object)
807 .expect("schema definition properties should exist")
808 }
809
810 fn assert_schema_property_title(
812 properties: &serde_json::Map<String, Value>,
813 property_name: &str,
814 expected_title: &str,
815 ) {
816 assert_eq!(
817 properties
818 .get(property_name)
819 .and_then(|value| value.get("title"))
820 .and_then(Value::as_str),
821 Some(expected_title)
822 );
823 }
824
825 fn assert_schema_property_title_and_description(
828 properties: &serde_json::Map<String, Value>,
829 property_name: &str,
830 expected_title: &str,
831 expected_description: &str,
832 ) {
833 assert_schema_property_title(properties, property_name, expected_title);
834 assert_eq!(
835 properties
836 .get(property_name)
837 .and_then(|value| value.get("description"))
838 .and_then(Value::as_str),
839 Some(expected_description)
840 );
841 }
842}