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("task_key"));
366 assert!(subtask_properties.contains_key("title"));
367 assert!(subtask_properties.contains_key("touched_areas"));
368 }
369
370 #[test]
371 fn test_agent_response_output_schema_all_properties_are_required() {
374 let schema = agent_response_output_schema(SchemaRequiredPolicy::AllProperties);
376
377 assert!(
379 all_properties_in_required(&schema),
380 "every object with `properties` should list all keys in `required`"
381 );
382 }
383
384 #[test]
385 fn test_agent_response_output_schema_minimum_policy_requires_only_answer() {
389 let schema = agent_response_output_schema(SchemaRequiredPolicy::MinimumProtocolKeys);
391 let required_fields = schema
392 .get("required")
393 .and_then(Value::as_array)
394 .expect("schema required fields should exist");
395
396 assert_eq!(
398 required_fields,
399 &vec![Value::String("answer".to_string())],
400 "only `answer` should be required; demanding optional response fields rejects \
401 ordinary replies that omit them"
402 );
403 }
404
405 #[test]
406 fn test_agent_response_output_schema_does_not_contain_one_of() {
409 let schema = agent_response_output_schema(SchemaRequiredPolicy::AllProperties);
411
412 assert!(!contains_schema_key(&schema, "oneOf"));
414 }
415
416 #[test]
417 fn test_agent_response_output_schema_does_not_contain_schema_metadata() {
420 let schema = agent_response_output_schema(SchemaRequiredPolicy::MinimumProtocolKeys);
422
423 assert!(!contains_schema_key(&schema, "$schema"));
425 }
426
427 #[test]
428 fn test_agent_response_json_schema_requires_answer_key() {
431 let schema = agent_response_json_schema();
433 let required_fields = schema
434 .get("required")
435 .and_then(Value::as_array)
436 .expect("schema required fields should exist");
437
438 assert!(
440 required_fields
441 .iter()
442 .any(|value| value.as_str() == Some("answer")),
443 "prompt schema should require `answer` to align with parser key-presence check"
444 );
445 }
446
447 #[test]
448 fn test_agent_response_json_schema_sets_additional_properties_false() {
452 let schema = agent_response_json_schema();
454
455 assert!(
457 all_properties_objects_deny_additional(&schema),
458 "every object with `properties` should set `additionalProperties: false`"
459 );
460 }
461
462 #[test]
463 fn test_inject_additional_properties_false_preserves_existing_value() {
466 let mut schema = serde_json::json!({
468 "type": "object",
469 "properties": {
470 "extra": { "type": "object", "additionalProperties": { "type": "string" } }
471 }
472 });
473
474 inject_additional_properties_false(&mut schema);
476
477 assert_eq!(schema["additionalProperties"], Value::Bool(false));
479 assert_eq!(
481 schema["properties"]["extra"]["additionalProperties"],
482 serde_json::json!({ "type": "string" })
483 );
484 }
485
486 #[test]
487 fn test_agent_response_output_schema_ref_objects_have_no_sibling_keywords() {
489 let schema = agent_response_output_schema(SchemaRequiredPolicy::AllProperties);
491
492 assert!(!contains_ref_with_sibling_keywords(&schema));
494 }
495
496 #[test]
497 fn test_agent_response_json_schema_json_is_parseable_value() {
499 let schema_json = agent_response_json_schema_json();
501 let parsed_schema: Value =
502 serde_json::from_str(&schema_json).expect("schema string should parse as JSON");
503 let schema_value = agent_response_json_schema();
504
505 assert_eq!(parsed_schema, schema_value);
507 }
508
509 #[test]
510 fn test_agent_response_json_schema_preserves_explicit_payload_metadata() {
513 let schema = agent_response_json_schema();
515
516 assert_eq!(
518 schema.get("title").and_then(Value::as_str),
519 Some("AgentResponse")
520 );
521 assert_eq!(
522 schema.get("description").and_then(Value::as_str),
523 Some(
524 "Wire-format protocol payload used for schema-driven provider output. Return this \
525 object as the entire assistant response payload. Providers that support output \
526 schemas (for example, Codex app-server) are asked to emit this object directly."
527 )
528 );
529 }
530
531 #[test]
532 fn test_agent_response_json_schema_preserves_nested_metadata() {
534 let schema = agent_response_json_schema();
536 let question_definition = schema
537 .get("$defs")
538 .and_then(|value| value.get("QuestionItem"))
539 .and_then(Value::as_object)
540 .expect("question definition should exist");
541 let review_comment_outcome_definition = schema
542 .get("$defs")
543 .and_then(|value| value.get("ReviewCommentOutcome"))
544 .and_then(Value::as_object)
545 .expect("review comment outcome definition should exist");
546 let review_comment_resolution_definition = schema
547 .get("$defs")
548 .and_then(|value| value.get("ReviewCommentResolution"))
549 .and_then(Value::as_object)
550 .expect("review comment resolution definition should exist");
551 let summary_definition = schema
552 .get("$defs")
553 .and_then(|value| value.get("AgentResponseSummary"))
554 .and_then(Value::as_object)
555 .expect("summary definition should exist");
556
557 assert_eq!(
559 question_definition.get("title").and_then(Value::as_str),
560 Some("QuestionItem")
561 );
562 assert_eq!(
563 review_comment_outcome_definition
564 .get("title")
565 .and_then(Value::as_str),
566 Some("ReviewCommentOutcome")
567 );
568 assert_eq!(
569 review_comment_resolution_definition
570 .get("title")
571 .and_then(Value::as_str),
572 Some("ReviewCommentResolution")
573 );
574 assert_eq!(
575 summary_definition.get("title").and_then(Value::as_str),
576 Some("AgentResponseSummary")
577 );
578 }
579
580 #[test]
581 fn test_agent_response_json_schema_preserves_field_metadata() {
583 let schema = agent_response_json_schema();
585 let response_properties = schema
586 .get("properties")
587 .and_then(Value::as_object)
588 .expect("response properties should exist");
589 let question_properties = schema_definition_properties(&schema, "QuestionItem");
590 let review_comment_outcome_properties =
591 schema_definition_properties(&schema, "ReviewCommentOutcome");
592 let summary_properties = schema_definition_properties(&schema, "AgentResponseSummary");
593 let expected_questions_description = questions_field_description();
594
595 assert_schema_property_title_and_description(
597 response_properties,
598 "answer",
599 "answer",
600 "Markdown answer text for delivered work, status updates, or concise completion \
601 notes. Keep clarification requests out of this field and emit them through \
602 `questions` instead.",
603 );
604 assert_eq!(
605 response_properties
606 .get("questions")
607 .and_then(|value| value.get("description"))
608 .and_then(Value::as_str),
609 Some(expected_questions_description.as_str())
610 );
611 assert_schema_property_title_and_description(
612 response_properties,
613 "review_comment_outcomes",
614 "review_comment_outcomes",
615 "Per-thread outcomes for an agent-driven forge comment-resolution turn. Emit an empty \
616 array unless the prompt explicitly supplies forge thread IDs. Copy each reported \
617 `thread_id` exactly from the prompt.",
618 );
619 assert_schema_property_title_and_description(
620 response_properties,
621 "summary",
622 "summary",
623 "Structured summary for session-discussion turns, kept outside `answer` markdown. Use \
624 `null` for one-shot prompts and legacy payloads.",
625 );
626 assert_schema_property_title_and_description(
627 question_properties,
628 "text",
629 "text",
630 "Human-readable markdown text for this question. Ask one specific actionable question \
631 instead of bundling multiple decisions into one item.",
632 );
633 assert_schema_property_title(question_properties, "options", "options");
634 assert_schema_property_title_and_description(
635 review_comment_outcome_properties,
636 "reply",
637 "reply",
638 "Concise reply suitable for posting to the forge review thread.",
639 );
640 assert_schema_property_title_and_description(
641 review_comment_outcome_properties,
642 "resolution",
643 "resolution",
644 "Whether the targeted thread was fixed or required no change.",
645 );
646 assert_schema_property_title_and_description(
647 review_comment_outcome_properties,
648 "thread_id",
649 "thread_id",
650 "Opaque forge thread identifier copied exactly from the turn prompt.",
651 );
652 assert_schema_property_title_and_description(
653 summary_properties,
654 "turn",
655 "turn",
656 "Concise summary of only the work completed in the current turn.",
657 );
658 assert_schema_property_title_and_description(
659 summary_properties,
660 "session",
661 "session",
662 "Cumulative summary of active changes on the current session branch.",
663 );
664 }
665
666 #[test]
667 fn test_agent_response_json_schema_keeps_optional_summary_field() {
670 let schema = agent_response_json_schema();
672 let response_required_fields = schema
673 .get("required")
674 .and_then(Value::as_array)
675 .cloned()
676 .unwrap_or_default();
677 let question_definition = schema
678 .get("$defs")
679 .and_then(|value| value.get("QuestionItem"))
680 .and_then(Value::as_object)
681 .expect("question definition should exist");
682 let question_required_fields = question_definition
683 .get("required")
684 .and_then(Value::as_array)
685 .cloned()
686 .unwrap_or_default();
687
688 assert!(
690 response_required_fields
691 .iter()
692 .all(|field| field.as_str() != Some("summary")),
693 "raw prompt schema should keep optional summary fields optional"
694 );
695 assert!(
696 question_required_fields
697 .iter()
698 .all(|field| field.as_str() != Some("options")),
699 "question schema should keep `options` optional for omitted empty lists"
700 );
701 }
702
703 #[test]
704 fn test_agent_response_output_schema_json_is_parseable_value() {
707 let schema_json =
709 agent_response_output_schema_json(SchemaRequiredPolicy::MinimumProtocolKeys);
710 let parsed_schema: Value =
711 serde_json::from_str(&schema_json).expect("schema string should parse as JSON");
712 let schema_value = agent_response_output_schema(SchemaRequiredPolicy::MinimumProtocolKeys);
713
714 assert_eq!(parsed_schema, schema_value);
716 }
717
718 fn contains_schema_key(value: &Value, key: &str) -> bool {
720 match value {
721 Value::Object(object) => {
722 if object.contains_key(key) {
723 return true;
724 }
725
726 object
727 .values()
728 .any(|nested_value| contains_schema_key(nested_value, key))
729 }
730 Value::Array(array) => array
731 .iter()
732 .any(|nested_value| contains_schema_key(nested_value, key)),
733 _ => false,
734 }
735 }
736
737 fn contains_ref_with_sibling_keywords(value: &Value) -> bool {
739 match value {
740 Value::Object(object) => {
741 if object.contains_key("$ref") && object.len() > 1 {
742 return true;
743 }
744
745 object.values().any(contains_ref_with_sibling_keywords)
746 }
747 Value::Array(array) => array.iter().any(contains_ref_with_sibling_keywords),
748 _ => false,
749 }
750 }
751
752 fn all_properties_in_required(value: &Value) -> bool {
755 match value {
756 Value::Object(object) => {
757 if let Some(properties) = object.get("properties").and_then(Value::as_object) {
758 let required_keys: Vec<&str> = object
759 .get("required")
760 .and_then(Value::as_array)
761 .map(|array| array.iter().filter_map(Value::as_str).collect())
762 .unwrap_or_default();
763
764 for key in properties.keys() {
765 if !required_keys.contains(&key.as_str()) {
766 return false;
767 }
768 }
769 }
770
771 object.values().all(all_properties_in_required)
772 }
773 Value::Array(array) => array.iter().all(all_properties_in_required),
774 _ => true,
775 }
776 }
777
778 fn all_properties_objects_deny_additional(value: &Value) -> bool {
781 match value {
782 Value::Object(object) => {
783 if object.contains_key("properties")
784 && object.get("additionalProperties") != Some(&Value::Bool(false))
785 {
786 return false;
787 }
788
789 object.values().all(all_properties_objects_deny_additional)
790 }
791 Value::Array(array) => array.iter().all(all_properties_objects_deny_additional),
792 _ => true,
793 }
794 }
795
796 fn schema_definition_properties<'a>(
798 schema: &'a Value,
799 definition_name: &str,
800 ) -> &'a serde_json::Map<String, Value> {
801 schema
802 .get("$defs")
803 .and_then(|value| value.get(definition_name))
804 .and_then(|value| value.get("properties"))
805 .and_then(Value::as_object)
806 .expect("schema definition properties should exist")
807 }
808
809 fn assert_schema_property_title(
811 properties: &serde_json::Map<String, Value>,
812 property_name: &str,
813 expected_title: &str,
814 ) {
815 assert_eq!(
816 properties
817 .get(property_name)
818 .and_then(|value| value.get("title"))
819 .and_then(Value::as_str),
820 Some(expected_title)
821 );
822 }
823
824 fn assert_schema_property_title_and_description(
827 properties: &serde_json::Map<String, Value>,
828 property_name: &str,
829 expected_title: &str,
830 expected_description: &str,
831 ) {
832 assert_schema_property_title(properties, property_name, expected_title);
833 assert_eq!(
834 properties
835 .get(property_name)
836 .and_then(|value| value.get("description"))
837 .and_then(Value::as_str),
838 Some(expected_description)
839 );
840 }
841}