1use std::pin::Pin;
32
33use agent_sdk_foundation::llm::{
34 ChatOutcome, ChatRequest, ChatResponse, ContentBlock, Message, ResponseFormat, Tool,
35 ToolChoice, Usage,
36};
37use agent_sdk_foundation::types::ToolTier;
38use futures::{Stream, StreamExt};
39
40use crate::provider::{LlmProvider, StructuredOutputSupport};
41use crate::streaming::{StreamAccumulator, StreamDelta, StreamErrorKind};
42
43#[cfg(feature = "openai")]
44use crate::impls::openai_schema::normalized_strict_schema;
45
46const RESPOND_TOOL_NAME: &str = "respond";
48
49#[derive(Debug, Clone, Copy)]
51pub struct StructuredConfig {
52 pub max_retries: u32,
58}
59
60impl Default for StructuredConfig {
61 fn default() -> Self {
62 Self { max_retries: 2 }
63 }
64}
65
66#[derive(Debug, Clone)]
68pub struct StructuredOutput {
69 pub value: serde_json::Value,
71 pub response: ChatResponse,
74 pub retries: u32,
77}
78
79#[derive(Debug, thiserror::Error)]
84pub enum StructuredOutputError {
85 #[error("structured output requested without a response_format on the request")]
88 MissingResponseFormat,
89
90 #[error("invalid output JSON schema: {0}")]
92 InvalidSchema(String),
93
94 #[error("output JSON schema is incompatible with OpenAI strict mode: {0}")]
97 IncompatibleSchema(String),
98
99 #[error("model produced no structured output to validate")]
102 NoStructuredOutput,
103
104 #[error("provider returned a non-success outcome: {0}")]
107 ProviderOutcome(String),
108
109 #[error(
113 "structured output failed schema validation after {attempts} attempt(s); last errors: {errors}"
114 )]
115 RetriesExhausted {
116 attempts: u32,
118 errors: String,
120 last_value: Option<serde_json::Value>,
122 },
123
124 #[error(transparent)]
126 Transport(#[from] anyhow::Error),
127}
128
129pub async fn run_structured(
142 provider: &dyn LlmProvider,
143 mut request: ChatRequest,
144 config: StructuredConfig,
145) -> Result<StructuredOutput, StructuredOutputError> {
146 let response_format = request
147 .response_format
148 .clone()
149 .ok_or(StructuredOutputError::MissingResponseFormat)?;
150
151 let validator = validator_for_provider(provider, &response_format)?;
153
154 let support = provider.structured_output_support();
155 if matches!(support, StructuredOutputSupport::ToolForcing) {
156 apply_tool_forcing(&mut request, &response_format);
157 }
158
159 let max_attempts = config.max_retries.saturating_add(1);
160 let mut last_value: Option<serde_json::Value> = None;
161 let mut last_errors = String::new();
162
163 for attempt in 0..max_attempts {
164 let attempt_request = if attempt + 1 == max_attempts {
167 std::mem::replace(&mut request, ChatRequest::new(String::new(), Vec::new()))
168 } else {
169 request.clone()
170 };
171 let outcome = provider.chat(attempt_request).await?;
172 let response = match outcome {
173 ChatOutcome::Success(response) => response,
174 ChatOutcome::RateLimited(_) => {
175 return Err(StructuredOutputError::ProviderOutcome(
176 "rate limited".to_owned(),
177 ));
178 }
179 ChatOutcome::InvalidRequest(msg) => {
180 return Err(StructuredOutputError::ProviderOutcome(format!(
181 "invalid request: {msg}"
182 )));
183 }
184 ChatOutcome::ServerError(msg) => {
185 return Err(StructuredOutputError::ProviderOutcome(format!(
186 "server error: {msg}"
187 )));
188 }
189 _ => {
192 return Err(StructuredOutputError::ProviderOutcome(
193 "unrecognized provider outcome".to_owned(),
194 ));
195 }
196 };
197
198 let candidate = extract_candidate(&response, support);
199 let Some(value) = candidate else {
200 if attempt + 1 >= max_attempts {
203 return Err(StructuredOutputError::NoStructuredOutput);
204 }
205 append_correction(
206 &mut request,
207 &response,
208 support,
209 "Your previous reply did not contain a structured answer. \
210 Respond with a single JSON value that satisfies the requested schema.",
211 );
212 "missing structured output".clone_into(&mut last_errors);
213 continue;
214 };
215
216 let errors = collect_schema_errors(&validator, &value);
217
218 if errors.is_empty() {
219 return Ok(StructuredOutput {
220 value,
221 response,
222 retries: attempt,
223 });
224 }
225
226 last_errors = errors.join("; ");
227 last_value = Some(value);
228
229 if attempt + 1 < max_attempts {
230 let correction = format!(
231 "Your previous JSON output did not satisfy the schema. \
232 Fix these validation errors and resend the full JSON value: {last_errors}"
233 );
234 append_correction(&mut request, &response, support, &correction);
235 }
236 }
237
238 Err(StructuredOutputError::RetriesExhausted {
239 attempts: max_attempts,
240 errors: last_errors,
241 last_value,
242 })
243}
244
245#[derive(Debug, Clone)]
247pub enum StructuredStreamUpdate {
248 Partial(serde_json::Value),
254 Final(StructuredOutput),
257}
258
259pub type StructuredStream<'a> =
261 Pin<Box<dyn Stream<Item = Result<StructuredStreamUpdate, StructuredOutputError>> + Send + 'a>>;
262
263pub fn run_structured_stream(
281 provider: &dyn LlmProvider,
282 request: ChatRequest,
283 config: StructuredConfig,
284) -> StructuredStream<'_> {
285 Box::pin(async_stream::stream! {
286 let mut request = request;
287 let Some(response_format) = request.response_format.clone() else {
288 yield Err(StructuredOutputError::MissingResponseFormat);
289 return;
290 };
291 let validator = match validator_for_provider(provider, &response_format) {
292 Ok(validator) => validator,
293 Err(error) => {
294 yield Err(error);
295 return;
296 }
297 };
298
299 let support = provider.structured_output_support();
300 if matches!(support, StructuredOutputSupport::ToolForcing) {
301 apply_tool_forcing(&mut request, &response_format);
302 }
303
304 let max_attempts = config.max_retries.saturating_add(1);
305 let model = provider.model().to_owned();
306 let mut last_value: Option<serde_json::Value> = None;
307 let mut last_errors = String::new();
308
309 for attempt in 0..max_attempts {
310 let response = if attempt == 0 {
313 let mut attempt_stream =
314 Box::pin(stream_first_attempt(provider, request.clone(), support, model.clone()));
315 let mut completed: Option<ChatResponse> = None;
316 while let Some(item) = attempt_stream.next().await {
317 match item {
318 StreamAttemptItem::Partial(value) => {
319 yield Ok(StructuredStreamUpdate::Partial(value));
320 }
321 StreamAttemptItem::Complete(response) => completed = Some(response),
322 StreamAttemptItem::Failed(error) => {
323 yield Err(error);
324 return;
325 }
326 }
327 }
328 match completed {
331 Some(response) => response,
332 None => return,
333 }
334 } else {
335 match provider.chat(request.clone()).await {
336 Ok(ChatOutcome::Success(response)) => response,
337 Ok(other) => {
338 yield Err(non_success_outcome_error(&other));
339 return;
340 }
341 Err(e) => {
342 yield Err(StructuredOutputError::Transport(e));
343 return;
344 }
345 }
346 };
347
348 let Some(value) = extract_candidate(&response, support) else {
349 if attempt + 1 >= max_attempts {
350 yield Err(StructuredOutputError::NoStructuredOutput);
351 return;
352 }
353 append_correction(
354 &mut request,
355 &response,
356 support,
357 "Your previous reply did not contain a structured answer. \
358 Respond with a single JSON value that satisfies the requested schema.",
359 );
360 "missing structured output".clone_into(&mut last_errors);
361 continue;
362 };
363
364 let errors = collect_schema_errors(&validator, &value);
365
366 if errors.is_empty() {
367 yield Ok(StructuredStreamUpdate::Final(StructuredOutput {
368 value,
369 response,
370 retries: attempt,
371 }));
372 return;
373 }
374
375 last_errors = errors.join("; ");
376 last_value = Some(value);
377
378 if attempt + 1 < max_attempts {
379 let correction = format!(
380 "Your previous JSON output did not satisfy the schema. \
381 Fix these validation errors and resend the full JSON value: {last_errors}"
382 );
383 append_correction(&mut request, &response, support, &correction);
384 }
385 }
386
387 yield Err(StructuredOutputError::RetriesExhausted {
388 attempts: max_attempts,
389 errors: last_errors,
390 last_value,
391 });
392 })
393}
394
395enum StreamAttemptItem {
398 Partial(serde_json::Value),
399 Complete(ChatResponse),
400 Failed(StructuredOutputError),
401}
402
403fn stream_first_attempt(
408 provider: &dyn LlmProvider,
409 request: ChatRequest,
410 support: StructuredOutputSupport,
411 model: String,
412) -> impl Stream<Item = StreamAttemptItem> + Send + '_ {
413 async_stream::stream! {
414 let mut accumulator = StreamAccumulator::new();
415 let mut partial_buf = String::new();
416 let mut respond_tool_ids: std::collections::HashSet<String> =
417 std::collections::HashSet::new();
418 let mut last_partial: Option<serde_json::Value> = None;
419 let mut stream_error: Option<(String, StreamErrorKind)> = None;
420
421 let mut stream = provider.chat_stream(request);
422 while let Some(item) = stream.next().await {
423 let delta = match item {
424 Ok(delta) => delta,
425 Err(e) => {
426 yield StreamAttemptItem::Failed(StructuredOutputError::Transport(e));
427 return;
428 }
429 };
430
431 accumulate_partial_buffer(&delta, support, &mut partial_buf, &mut respond_tool_ids);
432 if let StreamDelta::Error { message, kind } = &delta {
433 stream_error = Some((message.clone(), *kind));
434 }
435 accumulator.apply(&delta);
436
437 if let Some(value) = partial_from_buffer(&partial_buf)
438 && last_partial.as_ref() != Some(&value)
439 {
440 last_partial = Some(value.clone());
441 yield StreamAttemptItem::Partial(value);
442 }
443 }
444
445 if let Some((message, kind)) = stream_error {
446 yield StreamAttemptItem::Failed(stream_error_to_outcome(&message, kind));
447 return;
448 }
449
450 yield StreamAttemptItem::Complete(build_streamed_response(accumulator, model));
451 }
452}
453
454fn accumulate_partial_buffer(
459 delta: &StreamDelta,
460 support: StructuredOutputSupport,
461 buffer: &mut String,
462 respond_tool_ids: &mut std::collections::HashSet<String>,
463) {
464 match (support, delta) {
465 (StructuredOutputSupport::Native, StreamDelta::TextDelta { delta, .. }) => {
466 buffer.push_str(delta);
467 }
468 (StructuredOutputSupport::ToolForcing, StreamDelta::ToolUseStart { id, name, .. })
469 if name == RESPOND_TOOL_NAME =>
470 {
471 respond_tool_ids.insert(id.clone());
472 }
473 (StructuredOutputSupport::ToolForcing, StreamDelta::ToolInputDelta { id, delta, .. })
474 if respond_tool_ids.contains(id) =>
475 {
476 buffer.push_str(delta);
477 }
478 _ => {}
479 }
480}
481
482fn stream_error_to_outcome(message: &str, kind: StreamErrorKind) -> StructuredOutputError {
484 let label = match kind {
485 StreamErrorKind::RateLimited(_) => "rate limited".to_owned(),
486 StreamErrorKind::InvalidRequest => format!("invalid request: {message}"),
487 _ => format!("server error: {message}"),
488 };
489 StructuredOutputError::ProviderOutcome(label)
490}
491
492fn non_success_outcome_error(outcome: &ChatOutcome) -> StructuredOutputError {
494 let label = match outcome {
495 ChatOutcome::RateLimited(_) => "rate limited".to_owned(),
496 ChatOutcome::InvalidRequest(msg) => format!("invalid request: {msg}"),
497 ChatOutcome::ServerError(msg) => format!("server error: {msg}"),
498 _ => "unrecognized provider outcome".to_owned(),
499 };
500 StructuredOutputError::ProviderOutcome(label)
501}
502
503fn build_streamed_response(mut accumulator: StreamAccumulator, model: String) -> ChatResponse {
505 let usage = accumulator.take_usage().unwrap_or(Usage {
506 served_speed: None,
507 input_tokens: 0,
508 output_tokens: 0,
509 cached_input_tokens: 0,
510 cache_creation_input_tokens: 0,
511 });
512 let stop_reason = accumulator.take_stop_reason();
513 ChatResponse {
514 id: String::new(),
515 content: accumulator.into_content_blocks(),
516 model,
517 stop_reason,
518 usage,
519 }
520}
521
522fn partial_from_buffer(buffer: &str) -> Option<serde_json::Value> {
528 let trimmed = buffer.trim_start();
529 let body = trimmed
531 .strip_prefix("```")
532 .and_then(|rest| rest.split_once('\n').map(|(_, body)| body))
533 .unwrap_or(trimmed)
534 .trim();
535 if body.is_empty() {
536 return None;
537 }
538 let repaired = repair_partial_json(body);
539 serde_json::from_str::<serde_json::Value>(&repaired)
540 .ok()
541 .filter(|value| value.is_object() || value.is_array())
542}
543
544fn repair_partial_json(buffer: &str) -> String {
551 let mut in_string = false;
552 let mut escape = false;
553 let mut stack: Vec<char> = Vec::new();
554
555 for ch in buffer.chars() {
556 if in_string {
557 if escape {
558 escape = false;
559 } else if ch == '\\' {
560 escape = true;
561 } else if ch == '"' {
562 in_string = false;
563 }
564 continue;
565 }
566 match ch {
567 '"' => in_string = true,
568 '{' => stack.push('}'),
569 '[' => stack.push(']'),
570 '}' | ']' => {
571 stack.pop();
572 }
573 _ => {}
574 }
575 }
576
577 let mut out = buffer.to_owned();
578 if escape {
579 out.pop();
580 }
581 if in_string {
582 out.push('"');
583 }
584 out.truncate(out.trim_end().len());
585 if out.ends_with(',') {
586 out.pop();
587 out.truncate(out.trim_end().len());
588 } else if out.ends_with(':') {
589 out.push_str(" null");
590 }
591 for closer in stack.iter().rev() {
592 out.push(*closer);
593 }
594 out
595}
596
597fn collect_schema_errors(
600 validator: &jsonschema::Validator,
601 value: &serde_json::Value,
602) -> Vec<String> {
603 validator
604 .iter_errors(value)
605 .map(|error| format!("at `{}`: {error}", error.instance_path()))
606 .collect()
607}
608
609#[cfg(feature = "openai")]
618fn validation_schema_for_provider(
619 provider: &dyn LlmProvider,
620 response_format: &ResponseFormat,
621) -> Result<serde_json::Value, StructuredOutputError> {
622 if response_format.strict
623 && matches!(
624 provider.structured_output_support(),
625 StructuredOutputSupport::Native
626 )
627 && matches!(provider.provider(), "openai" | "openai-responses")
628 {
629 return normalized_strict_schema(&response_format.schema)
630 .map_err(|error| StructuredOutputError::IncompatibleSchema(error.to_string()));
631 }
632
633 Ok(response_format.schema.clone())
634}
635
636#[cfg(feature = "openai")]
637fn validator_for_provider(
638 provider: &dyn LlmProvider,
639 response_format: &ResponseFormat,
640) -> Result<jsonschema::Validator, StructuredOutputError> {
641 let validation_schema = validation_schema_for_provider(provider, response_format)?;
642 jsonschema::validator_for(&validation_schema)
643 .map_err(|error| StructuredOutputError::InvalidSchema(error.to_string()))
644}
645
646#[cfg(not(feature = "openai"))]
647fn validator_for_provider(
648 _provider: &dyn LlmProvider,
649 response_format: &ResponseFormat,
650) -> Result<jsonschema::Validator, StructuredOutputError> {
651 jsonschema::validator_for(&response_format.schema)
652 .map_err(|error| StructuredOutputError::InvalidSchema(error.to_string()))
653}
654
655fn apply_tool_forcing(request: &mut ChatRequest, response_format: &ResponseFormat) {
657 let respond_tool = Tool {
658 name: RESPOND_TOOL_NAME.to_owned(),
659 description: format!(
660 "Return the final answer as structured data named `{}`. \
661 You MUST call this tool exactly once with arguments matching the schema.",
662 response_format.name
663 ),
664 input_schema: response_format.schema.clone(),
665 display_name: "Structured response".to_owned(),
666 tier: ToolTier::Observe,
667 };
668
669 match request.tools {
670 Some(ref mut tools) => {
671 tools.retain(|t| t.name != RESPOND_TOOL_NAME);
672 tools.push(respond_tool);
673 }
674 None => request.tools = Some(vec![respond_tool]),
675 }
676 request.tool_choice = Some(ToolChoice::Tool(RESPOND_TOOL_NAME.to_owned()));
677 }
688
689fn extract_candidate(
692 response: &ChatResponse,
693 support: StructuredOutputSupport,
694) -> Option<serde_json::Value> {
695 match support {
696 StructuredOutputSupport::ToolForcing => {
697 response.content.iter().find_map(|block| match block {
698 ContentBlock::ToolUse { name, input, .. } if name == RESPOND_TOOL_NAME => {
699 Some(input.clone())
700 }
701 _ => None,
702 })
703 }
704 StructuredOutputSupport::Native => {
705 let text = response.first_text()?;
706 parse_json_text(text)
707 }
708 }
709}
710
711fn parse_json_text(text: &str) -> Option<serde_json::Value> {
717 let trimmed = text.trim();
718 let unfenced = strip_code_fence(trimmed);
719 serde_json::from_str(unfenced).ok()
720}
721
722fn strip_code_fence(text: &str) -> &str {
724 let Some(rest) = text.strip_prefix("```") else {
725 return text;
726 };
727 let rest = rest.split_once('\n').map_or(rest, |(_, body)| body);
729 rest.strip_suffix("```")
730 .map_or(text, |inner| inner.trim_end_matches('`').trim())
731}
732
733fn append_correction(
744 request: &mut ChatRequest,
745 previous: &ChatResponse,
746 support: StructuredOutputSupport,
747 correction: &str,
748) {
749 request
750 .messages
751 .push(Message::assistant_with_content(previous.content.clone()));
752
753 let respond_tool_use_id = if matches!(support, StructuredOutputSupport::ToolForcing) {
754 previous.content.iter().find_map(|block| match block {
755 ContentBlock::ToolUse { id, name, .. } if name == RESPOND_TOOL_NAME => Some(id.clone()),
756 _ => None,
757 })
758 } else {
759 None
760 };
761
762 match respond_tool_use_id {
763 Some(tool_use_id) => {
764 request
765 .messages
766 .push(Message::tool_result(tool_use_id, correction, true));
767 }
768 None => request.messages.push(Message::user(correction)),
769 }
770}
771
772#[cfg(test)]
773mod tests {
774 use super::*;
775
776 use std::sync::Mutex;
777 use std::sync::atomic::{AtomicUsize, Ordering};
778
779 use agent_sdk_foundation::llm::{StopReason, Usage};
780 use anyhow::Result;
781 use async_trait::async_trait;
782
783 use crate::streaming::StreamBox;
784
785 struct ScriptedProvider {
790 provider_name: &'static str,
791 model: String,
792 support: StructuredOutputSupport,
793 outcomes: Mutex<std::collections::VecDeque<ChatOutcome>>,
794 seen_requests: Mutex<Vec<ChatRequest>>,
795 calls: AtomicUsize,
796 }
797
798 impl ScriptedProvider {
799 fn new(
800 provider_name: &'static str,
801 support: StructuredOutputSupport,
802 outcomes: Vec<ChatOutcome>,
803 ) -> Self {
804 Self {
805 provider_name,
806 model: "scripted-model".to_owned(),
807 support,
808 outcomes: Mutex::new(outcomes.into()),
809 seen_requests: Mutex::new(Vec::new()),
810 calls: AtomicUsize::new(0),
811 }
812 }
813
814 fn call_count(&self) -> usize {
815 self.calls.load(Ordering::SeqCst)
816 }
817 }
818
819 #[async_trait]
820 impl LlmProvider for ScriptedProvider {
821 async fn chat(&self, request: ChatRequest) -> Result<ChatOutcome> {
822 self.calls.fetch_add(1, Ordering::SeqCst);
823 self.seen_requests
824 .lock()
825 .expect("seen_requests lock")
826 .push(request);
827 let outcome = self
828 .outcomes
829 .lock()
830 .expect("outcomes lock")
831 .pop_front()
832 .expect("ScriptedProvider: ran out of scripted outcomes");
833 Ok(outcome)
834 }
835
836 fn chat_stream(&self, _request: ChatRequest) -> StreamBox<'_> {
837 Box::pin(async_stream::stream! {
838 yield Err(anyhow::anyhow!("streaming not used in structured tests"));
839 })
840 }
841
842 fn model(&self) -> &str {
843 &self.model
844 }
845
846 fn provider(&self) -> &'static str {
847 self.provider_name
848 }
849
850 fn structured_output_support(&self) -> StructuredOutputSupport {
851 self.support
852 }
853 }
854
855 fn person_schema() -> serde_json::Value {
856 serde_json::json!({
857 "type": "object",
858 "properties": {
859 "name": { "type": "string" },
860 "age": { "type": "integer", "minimum": 0 }
861 },
862 "required": ["name", "age"],
863 "additionalProperties": false
864 })
865 }
866
867 fn request_with_format() -> ChatRequest {
868 ChatRequest {
869 system: String::new(),
870 messages: vec![Message::user("Describe a person.")],
871 tools: None,
872 max_tokens: 256,
873 max_tokens_explicit: true,
874 session_id: None,
875 cached_content: None,
876 thinking: None,
877 tool_choice: None,
878 response_format: Some(ResponseFormat::new("person", person_schema())),
879 cache: None,
880 }
881 }
882
883 fn request_with_optional_property_format() -> ChatRequest {
884 let mut request = request_with_format();
885 request.response_format = Some(ResponseFormat::new(
886 "optional-person",
887 serde_json::json!({
888 "type": "object",
889 "properties": {
890 "name": {"type": "string"},
891 "nickname": {"type": "string"}
892 },
893 "required": ["name"]
894 }),
895 ));
896 request
897 }
898
899 fn success(content: Vec<ContentBlock>) -> ChatOutcome {
900 ChatOutcome::Success(ChatResponse {
901 id: "resp".to_owned(),
902 content,
903 model: "scripted-model".to_owned(),
904 stop_reason: Some(StopReason::EndTurn),
905 usage: Usage {
906 served_speed: None,
907 input_tokens: 1,
908 output_tokens: 1,
909 cached_input_tokens: 0,
910 cache_creation_input_tokens: 0,
911 },
912 })
913 }
914
915 fn text_block(text: &str) -> Vec<ContentBlock> {
916 vec![ContentBlock::Text {
917 text: text.to_owned(),
918 }]
919 }
920
921 fn respond_tool_block(input: serde_json::Value) -> Vec<ContentBlock> {
922 vec![ContentBlock::ToolUse {
923 id: "call_1".to_owned(),
924 name: RESPOND_TOOL_NAME.to_owned(),
925 input,
926 thought_signature: None,
927 }]
928 }
929
930 #[tokio::test]
933 async fn native_happy_path_validates_json_text() -> Result<()> {
934 let provider = ScriptedProvider::new(
935 "openai",
936 StructuredOutputSupport::Native,
937 vec![success(text_block(r#"{"name": "Ada", "age": 36}"#))],
938 );
939
940 let out = run_structured(
941 &provider,
942 request_with_format(),
943 StructuredConfig::default(),
944 )
945 .await?;
946
947 assert_eq!(out.value["name"], "Ada");
948 assert_eq!(out.value["age"], 36);
949 assert_eq!(out.retries, 0);
950 assert_eq!(provider.call_count(), 1);
951 Ok(())
952 }
953
954 #[tokio::test]
955 async fn native_happy_path_strips_markdown_fence() -> Result<()> {
956 let provider = ScriptedProvider::new(
957 "gemini",
958 StructuredOutputSupport::Native,
959 vec![success(text_block(
960 "```json\n{\"name\": \"Grace\", \"age\": 45}\n```",
961 ))],
962 );
963
964 let out = run_structured(
965 &provider,
966 request_with_format(),
967 StructuredConfig::default(),
968 )
969 .await?;
970
971 assert_eq!(out.value["name"], "Grace");
972 Ok(())
973 }
974
975 #[cfg(feature = "openai")]
976 #[tokio::test]
977 async fn native_openai_validates_against_the_normalized_strict_schema() -> Result<()> {
978 for provider_name in ["openai", "openai-responses"] {
979 let provider = ScriptedProvider::new(
980 provider_name,
981 StructuredOutputSupport::Native,
982 vec![success(text_block(r#"{"name": "Ada", "nickname": null}"#))],
983 );
984
985 let output = run_structured(
986 &provider,
987 request_with_optional_property_format(),
988 StructuredConfig::default(),
989 )
990 .await?;
991
992 assert_eq!(output.value["nickname"], serde_json::Value::Null);
993 assert_eq!(provider.call_count(), 1);
994 }
995 Ok(())
996 }
997
998 #[tokio::test]
999 async fn native_non_openai_uses_the_caller_schema() -> Result<()> {
1000 let provider = ScriptedProvider::new(
1001 "gemini",
1002 StructuredOutputSupport::Native,
1003 vec![success(text_block(r#"{"name": "Ada"}"#))],
1004 );
1005
1006 let output = run_structured(
1007 &provider,
1008 request_with_optional_property_format(),
1009 StructuredConfig::default(),
1010 )
1011 .await?;
1012
1013 assert_eq!(output.value["name"], "Ada");
1014 assert!(output.value.get("nickname").is_none());
1015 Ok(())
1016 }
1017
1018 #[cfg(feature = "openai")]
1019 #[tokio::test]
1020 async fn native_openai_rejects_dynamic_property_schemas_before_calling_provider() {
1021 let provider = ScriptedProvider::new("openai", StructuredOutputSupport::Native, Vec::new());
1022 let mut request = request_with_format();
1023 request.response_format = Some(ResponseFormat::new(
1024 "dynamic-properties",
1025 serde_json::json!({
1026 "type": "object",
1027 "properties": {"fixed": {"type": "string"}},
1028 "additionalProperties": {"type": "string"}
1029 }),
1030 ));
1031
1032 let result = run_structured(&provider, request, StructuredConfig::default()).await;
1033
1034 assert!(matches!(
1035 result,
1036 Err(StructuredOutputError::IncompatibleSchema(_))
1037 ));
1038 assert_eq!(provider.call_count(), 0);
1039 }
1040
1041 #[tokio::test]
1044 async fn tool_forcing_happy_path_reads_tool_input() -> Result<()> {
1045 let provider = ScriptedProvider::new(
1046 "anthropic",
1047 StructuredOutputSupport::ToolForcing,
1048 vec![success(respond_tool_block(
1049 serde_json::json!({"name": "Linus", "age": 54}),
1050 ))],
1051 );
1052
1053 let out = run_structured(
1054 &provider,
1055 request_with_format(),
1056 StructuredConfig::default(),
1057 )
1058 .await?;
1059
1060 assert_eq!(out.value["name"], "Linus");
1061 assert_eq!(out.retries, 0);
1062
1063 let (has_respond_tool, forces_respond) = {
1065 let seen = provider.seen_requests.lock().expect("seen lock");
1066 let tools = seen[0].tools.as_ref().expect("tools injected");
1067 (
1068 tools.iter().any(|t| t.name == RESPOND_TOOL_NAME),
1069 matches!(
1070 seen[0].tool_choice,
1071 Some(ToolChoice::Tool(ref n)) if n == RESPOND_TOOL_NAME
1072 ),
1073 )
1074 };
1075 assert!(has_respond_tool);
1076 assert!(forces_respond);
1077 Ok(())
1078 }
1079
1080 #[tokio::test]
1083 async fn mismatch_then_retry_succeeds() -> Result<()> {
1084 let provider = ScriptedProvider::new(
1085 "openai",
1086 StructuredOutputSupport::Native,
1087 vec![
1088 success(text_block(r#"{"name": "Ada", "age": "old"}"#)),
1090 success(text_block(r#"{"name": "Ada", "age": 36}"#)),
1092 ],
1093 );
1094
1095 let out = run_structured(
1096 &provider,
1097 request_with_format(),
1098 StructuredConfig { max_retries: 2 },
1099 )
1100 .await?;
1101
1102 assert_eq!(out.value["age"], 36);
1103 assert_eq!(out.retries, 1);
1104 assert_eq!(provider.call_count(), 2);
1105
1106 let grew = {
1109 let seen = provider.seen_requests.lock().expect("seen lock");
1110 seen[1].messages.len() > seen[0].messages.len()
1111 };
1112 assert!(grew);
1113 Ok(())
1114 }
1115
1116 #[tokio::test]
1117 async fn tool_forcing_retry_appends_tool_result_for_forced_tool_use() -> Result<()> {
1118 use agent_sdk_foundation::llm::Content;
1119
1120 let provider = ScriptedProvider::new(
1121 "anthropic",
1122 StructuredOutputSupport::ToolForcing,
1123 vec![
1124 success(respond_tool_block(serde_json::json!({"name": "x"}))),
1126 success(respond_tool_block(
1128 serde_json::json!({"name": "x", "age": 1}),
1129 )),
1130 ],
1131 );
1132
1133 let out = run_structured(
1134 &provider,
1135 request_with_format(),
1136 StructuredConfig { max_retries: 1 },
1137 )
1138 .await?;
1139 assert_eq!(out.retries, 1);
1140
1141 let seen = provider.seen_requests.lock().expect("seen lock");
1145 let retry = &seen[1];
1146
1147 let assistant_tool_use_id = retry
1148 .messages
1149 .iter()
1150 .find_map(|m| match &m.content {
1151 Content::Blocks(blocks) => blocks.iter().find_map(|b| match b {
1152 ContentBlock::ToolUse { id, name, .. } if name == RESPOND_TOOL_NAME => {
1153 Some(id.clone())
1154 }
1155 _ => None,
1156 }),
1157 Content::Text(_) => None,
1158 })
1159 .expect("assistant respond tool_use present in retry");
1160
1161 let has_matching_result = retry.messages.iter().any(|m| match &m.content {
1162 Content::Blocks(blocks) => blocks.iter().any(|b| {
1163 matches!(
1164 b,
1165 ContentBlock::ToolResult { tool_use_id, .. }
1166 if *tool_use_id == assistant_tool_use_id
1167 )
1168 }),
1169 Content::Text(_) => false,
1170 });
1171 drop(seen);
1172 assert!(
1173 has_matching_result,
1174 "retry must carry a tool_result for the forced tool_use id"
1175 );
1176 Ok(())
1177 }
1178
1179 #[tokio::test]
1182 async fn retry_exhaustion_yields_typed_error() -> Result<()> {
1183 let provider = ScriptedProvider::new(
1184 "anthropic",
1185 StructuredOutputSupport::ToolForcing,
1186 vec![
1187 success(respond_tool_block(serde_json::json!({"name": "x"}))),
1188 success(respond_tool_block(serde_json::json!({"name": "y"}))),
1189 success(respond_tool_block(serde_json::json!({"name": "z"}))),
1190 ],
1191 );
1192
1193 let err = run_structured(
1194 &provider,
1195 request_with_format(),
1196 StructuredConfig { max_retries: 2 },
1197 )
1198 .await
1199 .expect_err("schema never satisfied");
1200
1201 match err {
1202 StructuredOutputError::RetriesExhausted {
1203 attempts,
1204 last_value,
1205 ..
1206 } => {
1207 assert_eq!(attempts, 3, "1 initial + 2 retries");
1208 assert_eq!(
1209 last_value.as_ref().and_then(|v| v["name"].as_str()),
1210 Some("z")
1211 );
1212 }
1213 other => panic!("expected RetriesExhausted, got {other:?}"),
1214 }
1215 assert_eq!(provider.call_count(), 3);
1217 Ok(())
1218 }
1219
1220 #[tokio::test]
1221 async fn zero_retries_fails_after_single_attempt() -> Result<()> {
1222 let provider = ScriptedProvider::new(
1223 "openai",
1224 StructuredOutputSupport::Native,
1225 vec![success(text_block(r#"{"name": "Ada"}"#))],
1226 );
1227
1228 let err = run_structured(
1229 &provider,
1230 request_with_format(),
1231 StructuredConfig { max_retries: 0 },
1232 )
1233 .await
1234 .expect_err("missing required `age`");
1235
1236 assert!(matches!(
1237 err,
1238 StructuredOutputError::RetriesExhausted { attempts: 1, .. }
1239 ));
1240 assert_eq!(provider.call_count(), 1);
1241 Ok(())
1242 }
1243
1244 #[tokio::test]
1247 async fn missing_response_format_is_typed_error() {
1248 let provider = ScriptedProvider::new(
1249 "openai",
1250 StructuredOutputSupport::Native,
1251 vec![success(text_block("{}"))],
1252 );
1253 let mut req = request_with_format();
1254 req.response_format = None;
1255
1256 let err = run_structured(&provider, req, StructuredConfig::default())
1257 .await
1258 .expect_err("no response format");
1259 assert!(matches!(err, StructuredOutputError::MissingResponseFormat));
1260 }
1261
1262 #[tokio::test]
1263 async fn invalid_schema_is_typed_error() {
1264 let provider = ScriptedProvider::new(
1265 "openai",
1266 StructuredOutputSupport::Native,
1267 vec![success(text_block("{}"))],
1268 );
1269 let mut req = request_with_format();
1270 req.response_format = Some(ResponseFormat::new("bad", serde_json::json!({"type": 123})));
1272
1273 let err = run_structured(&provider, req, StructuredConfig::default())
1274 .await
1275 .expect_err("invalid schema");
1276 assert!(matches!(err, StructuredOutputError::InvalidSchema(_)));
1277 }
1278
1279 #[tokio::test]
1280 async fn provider_rate_limit_surfaces_as_typed_error() {
1281 let provider = ScriptedProvider::new(
1282 "openai",
1283 StructuredOutputSupport::Native,
1284 vec![ChatOutcome::RateLimited(None)],
1285 );
1286
1287 let err = run_structured(
1288 &provider,
1289 request_with_format(),
1290 StructuredConfig::default(),
1291 )
1292 .await
1293 .expect_err("rate limited");
1294 assert!(matches!(err, StructuredOutputError::ProviderOutcome(_)));
1295 }
1296
1297 #[tokio::test]
1298 async fn no_structured_output_on_final_attempt_errors() {
1299 let provider = ScriptedProvider::new(
1301 "openai",
1302 StructuredOutputSupport::Native,
1303 vec![
1304 success(text_block("I cannot do that.")),
1305 success(text_block("Still prose, sorry.")),
1306 ],
1307 );
1308
1309 let err = run_structured(
1310 &provider,
1311 request_with_format(),
1312 StructuredConfig { max_retries: 1 },
1313 )
1314 .await
1315 .expect_err("never produced JSON");
1316 assert!(matches!(err, StructuredOutputError::NoStructuredOutput));
1317 assert_eq!(provider.call_count(), 2);
1318 }
1319
1320 struct StreamingProvider {
1326 provider_name: &'static str,
1327 model: String,
1328 support: StructuredOutputSupport,
1329 deltas: Mutex<Vec<StreamDelta>>,
1330 }
1331
1332 impl StreamingProvider {
1333 fn new(
1334 provider_name: &'static str,
1335 support: StructuredOutputSupport,
1336 deltas: Vec<StreamDelta>,
1337 ) -> Self {
1338 Self {
1339 provider_name,
1340 model: "scripted-model".to_owned(),
1341 support,
1342 deltas: Mutex::new(deltas),
1343 }
1344 }
1345 }
1346
1347 #[async_trait]
1348 impl LlmProvider for StreamingProvider {
1349 async fn chat(&self, _request: ChatRequest) -> Result<ChatOutcome> {
1350 Ok(ChatOutcome::ServerError("chat() not used".to_owned()))
1351 }
1352
1353 fn chat_stream(&self, _request: ChatRequest) -> StreamBox<'_> {
1354 let deltas = self.deltas.lock().map(|d| d.clone()).unwrap_or_default();
1355 Box::pin(async_stream::stream! {
1356 for delta in deltas {
1357 yield Ok(delta);
1358 }
1359 })
1360 }
1361
1362 fn model(&self) -> &str {
1363 &self.model
1364 }
1365
1366 fn provider(&self) -> &'static str {
1367 self.provider_name
1368 }
1369
1370 fn structured_output_support(&self) -> StructuredOutputSupport {
1371 self.support
1372 }
1373 }
1374
1375 async fn drive_stream(
1376 mut stream: StructuredStream<'_>,
1377 ) -> Result<(Vec<serde_json::Value>, Option<StructuredOutput>)> {
1378 let mut partials = Vec::new();
1379 let mut final_out = None;
1380 while let Some(update) = stream.next().await {
1381 match update? {
1382 StructuredStreamUpdate::Partial(value) => partials.push(value),
1383 StructuredStreamUpdate::Final(out) => final_out = Some(out),
1384 }
1385 }
1386 Ok((partials, final_out))
1387 }
1388
1389 #[tokio::test]
1390 async fn streaming_native_emits_partials_then_validated_final() -> Result<()> {
1391 let provider = StreamingProvider::new(
1392 "openai",
1393 StructuredOutputSupport::Native,
1394 vec![
1395 StreamDelta::TextDelta {
1396 delta: r#"{"name": "Ada""#.to_owned(),
1397 block_index: 0,
1398 },
1399 StreamDelta::TextDelta {
1400 delta: r#", "age": 36}"#.to_owned(),
1401 block_index: 0,
1402 },
1403 StreamDelta::Done {
1404 stop_reason: Some(StopReason::EndTurn),
1405 served_route: None,
1406 },
1407 ],
1408 );
1409
1410 let stream = run_structured_stream(
1411 &provider,
1412 request_with_format(),
1413 StructuredConfig::default(),
1414 );
1415 let (partials, final_out) = drive_stream(stream).await?;
1416
1417 assert!(!partials.is_empty(), "expected at least one partial");
1418 assert_eq!(partials[0]["name"], "Ada");
1420 let final_out = final_out.expect("a validated final value");
1421 assert_eq!(final_out.value["name"], "Ada");
1422 assert_eq!(final_out.value["age"], 36);
1423 assert_eq!(final_out.retries, 0);
1424 Ok(())
1425 }
1426
1427 #[cfg(feature = "openai")]
1428 #[tokio::test]
1429 async fn streaming_native_openai_uses_the_normalized_strict_schema() -> Result<()> {
1430 let provider = StreamingProvider::new(
1431 "openai-responses",
1432 StructuredOutputSupport::Native,
1433 vec![
1434 StreamDelta::TextDelta {
1435 delta: r#"{"name": "Ada", "nickname": null}"#.to_owned(),
1436 block_index: 0,
1437 },
1438 StreamDelta::Done {
1439 stop_reason: Some(StopReason::EndTurn),
1440 served_route: None,
1441 },
1442 ],
1443 );
1444
1445 let stream = run_structured_stream(
1446 &provider,
1447 request_with_optional_property_format(),
1448 StructuredConfig::default(),
1449 );
1450 let (_, final_output) = drive_stream(stream).await?;
1451 let final_output = final_output.ok_or_else(|| anyhow::anyhow!("missing final output"))?;
1452
1453 assert_eq!(final_output.value["nickname"], serde_json::Value::Null);
1454 Ok(())
1455 }
1456
1457 #[tokio::test]
1458 async fn streaming_tool_forcing_reads_partial_tool_input() -> Result<()> {
1459 let provider = StreamingProvider::new(
1460 "anthropic",
1461 StructuredOutputSupport::ToolForcing,
1462 vec![
1463 StreamDelta::ToolUseStart {
1464 id: "call_1".to_owned(),
1465 name: RESPOND_TOOL_NAME.to_owned(),
1466 block_index: 0,
1467 thought_signature: None,
1468 },
1469 StreamDelta::ToolInputDelta {
1470 id: "call_1".to_owned(),
1471 delta: r#"{"name": "Linus""#.to_owned(),
1472 block_index: 0,
1473 },
1474 StreamDelta::ToolInputDelta {
1475 id: "call_1".to_owned(),
1476 delta: r#", "age": 54}"#.to_owned(),
1477 block_index: 0,
1478 },
1479 StreamDelta::Done {
1480 stop_reason: Some(StopReason::ToolUse),
1481 served_route: None,
1482 },
1483 ],
1484 );
1485
1486 let stream = run_structured_stream(
1487 &provider,
1488 request_with_format(),
1489 StructuredConfig::default(),
1490 );
1491 let (partials, final_out) = drive_stream(stream).await?;
1492
1493 assert_eq!(partials[0]["name"], "Linus");
1494 let final_out = final_out.expect("a validated final value");
1495 assert_eq!(final_out.value["age"], 54);
1496 Ok(())
1497 }
1498
1499 #[tokio::test]
1500 async fn streaming_missing_response_format_errors() {
1501 let provider =
1502 StreamingProvider::new("openai", StructuredOutputSupport::Native, Vec::new());
1503 let mut req = request_with_format();
1504 req.response_format = None;
1505
1506 let mut stream = run_structured_stream(&provider, req, StructuredConfig::default());
1507 let first = stream.next().await.expect("one item");
1508 assert!(matches!(
1509 first,
1510 Err(StructuredOutputError::MissingResponseFormat)
1511 ));
1512 }
1513
1514 #[test]
1515 fn partial_from_buffer_repairs_incomplete_json() {
1516 assert_eq!(
1517 partial_from_buffer(r#"{"name": "Ada""#).map(|v| v["name"].clone()),
1518 Some(serde_json::json!("Ada"))
1519 );
1520 assert_eq!(
1521 partial_from_buffer(r#"{"a": 1,"#),
1522 Some(serde_json::json!({"a": 1}))
1523 );
1524 assert_eq!(
1525 partial_from_buffer(r#"{"a":"#),
1526 Some(serde_json::json!({"a": null}))
1527 );
1528 assert!(partial_from_buffer("").is_none());
1529 assert!(partial_from_buffer("not json").is_none());
1530 }
1531
1532 #[test]
1533 fn apply_tool_forcing_forces_the_respond_tool() {
1534 let mut request = request_with_format();
1543 request.thinking = Some(agent_sdk_foundation::llm::ThinkingConfig::new(10_000));
1544 let response_format = request
1545 .response_format
1546 .clone()
1547 .expect("request_with_format sets a response format");
1548
1549 apply_tool_forcing(&mut request, &response_format);
1550
1551 assert!(matches!(
1552 request.tool_choice,
1553 Some(ToolChoice::Tool(ref name)) if name == RESPOND_TOOL_NAME
1554 ));
1555 assert!(
1556 request
1557 .tools
1558 .as_ref()
1559 .is_some_and(|tools| tools.iter().any(|t| t.name == RESPOND_TOOL_NAME)),
1560 "the respond tool must be present"
1561 );
1562 assert!(request.thinking.is_some());
1565 }
1566}