Skip to main content

a3s_code_core/llm/
structured.rs

1//! Structured object generation from LLM output.
2//!
3//! Provides reliable JSON object generation with schema validation, automatic
4//! repair, and streaming partial object support. Works across all providers by
5//! selecting the best available mode (strict JSON schema, json_mode, tool-call,
6//! or prompt-only).
7
8use super::{LlmClient, Message, StreamEvent, TokenUsage, ToolDefinition};
9use anyhow::{bail, Context, Result};
10use serde::{Deserialize, Serialize};
11use serde_json::Value;
12use tokio_util::sync::CancellationToken;
13
14// ---------------------------------------------------------------------------
15// Public types
16// ---------------------------------------------------------------------------
17
18/// Mode selection for structured output generation.
19#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
20#[serde(rename_all = "snake_case")]
21pub enum StructuredMode {
22    /// Auto-select best mode based on provider capabilities.
23    Auto,
24    /// OpenAI native strict JSON schema (response_format.type = json_schema).
25    Strict,
26    /// OpenAI json_object mode (guarantees valid JSON, not schema-conformant).
27    Json,
28    /// Use tool-calling: inject a synthetic tool whose parameters IS the schema.
29    /// Works on all providers that support tool use (Anthropic, OpenAI, etc).
30    Tool,
31    /// Prompt-only: append schema instructions to the prompt. Least reliable.
32    Prompt,
33}
34
35/// Request specification for structured object generation.
36#[derive(Debug, Clone)]
37pub struct StructuredRequest {
38    pub prompt: String,
39    pub system: Option<String>,
40    pub schema: Value,
41    pub schema_name: String,
42    pub schema_description: Option<String>,
43    pub mode: StructuredMode,
44    pub max_repair_attempts: u8,
45}
46
47/// Result of a successful structured generation.
48#[derive(Debug, Clone, Serialize)]
49pub struct StructuredResult {
50    pub object: Value,
51    pub raw_text: Option<String>,
52    pub usage: TokenUsage,
53    pub repair_rounds: u8,
54    pub mode_used: StructuredMode,
55}
56
57/// Provider-native structured-output capability.
58///
59/// Each [`LlmClient`] reports this so the structured engine can request the
60/// strongest enforcement the provider actually supports. Defaults to
61/// [`NativeStructuredSupport::None`] for clients that don't override it.
62#[derive(Debug, Clone, Copy, PartialEq, Eq)]
63pub enum NativeStructuredSupport {
64    /// No native enforcement — rely on prompt instructions + lenient extraction.
65    None,
66    /// Can force a specific tool call (Anthropic `tool_choice`, OpenAI function
67    /// `tool_choice`). Guarantees the model emits the structured tool call
68    /// instead of free-form prose.
69    ForcedTool,
70    /// Supports OpenAI-style `response_format` (`json_object` and
71    /// `json_schema` + `strict`) in addition to forced tool calls.
72    JsonSchema,
73}
74
75/// A native `response_format` request for OpenAI-compatible providers.
76#[derive(Debug, Clone, PartialEq)]
77pub enum ResponseFormat {
78    /// `{"type":"json_object"}` — guarantees syntactically valid JSON, but not
79    /// schema conformance.
80    JsonObject,
81    /// `{"type":"json_schema","json_schema":{name,schema,strict:true}}` —
82    /// parser-enforced schema conformance.
83    JsonSchema { name: String, schema: Value },
84}
85
86/// Instruction telling a provider how to enforce structured output for a call.
87///
88/// Carries the union of intents; each provider honors what it supports and
89/// ignores the rest (e.g. Anthropic has no `response_format`, so it only acts
90/// on `force_tool`). The default (`force_tool: None, response_format: None`)
91/// reproduces an ordinary completion, which is why the trait's default
92/// `complete_structured` impl is behavior-preserving.
93#[derive(Debug, Clone, Default, PartialEq)]
94pub struct StructuredDirective {
95    /// Force the model to call exactly this tool (provider `tool_choice`).
96    pub force_tool: Option<String>,
97    /// Request a provider-native `response_format` (OpenAI-compatible only).
98    pub response_format: Option<ResponseFormat>,
99}
100
101/// Callback for streaming partial object snapshots.
102pub type PartialObjectCallback = Box<dyn Fn(&Value) + Send>;
103
104/// Provider-facing schema envelope.
105///
106/// Function/tool parameters are most reliable when the top-level schema is an
107/// object. Inspired by Vercel AI SDK's `Output.array` / `Output.choice`
108/// wrappers, A3S sends top-level arrays and scalar schemas inside a small object
109/// envelope, then unwraps the validated value before returning it to callers.
110#[derive(Debug, Clone, Copy, PartialEq, Eq)]
111enum SchemaEnvelope {
112    Direct,
113    Elements,
114    Value,
115}
116
117impl SchemaEnvelope {
118    fn for_schema(schema: &Value) -> Self {
119        if schema_is_object_like(schema) {
120            Self::Direct
121        } else if schema.get("type").and_then(Value::as_str) == Some("array") {
122            Self::Elements
123        } else {
124            Self::Value
125        }
126    }
127
128    fn response_schema(self, schema: &Value) -> Value {
129        match self {
130            Self::Direct => schema.clone(),
131            Self::Elements => serde_json::json!({
132                "type": "object",
133                "required": ["elements"],
134                "additionalProperties": false,
135                "properties": {
136                    "elements": schema
137                }
138            }),
139            Self::Value => serde_json::json!({
140                "type": "object",
141                "required": ["value"],
142                "additionalProperties": false,
143                "properties": {
144                    "value": schema
145                }
146            }),
147        }
148    }
149
150    fn unwrap_final(self, value: &Value) -> Option<Value> {
151        match self {
152            Self::Direct => Some(value.clone()),
153            Self::Elements => value.get("elements").cloned(),
154            Self::Value => value.get("value").cloned(),
155        }
156    }
157
158    fn project_partial(self, value: &Value, repaired: bool) -> Option<Value> {
159        match self {
160            Self::Direct => Some(value.clone()),
161            Self::Elements => {
162                let mut elements = value.get("elements")?.as_array()?.clone();
163                // A repaired parse may include a synthetic last element that was
164                // closed only so the partial JSON can parse. Match Vercel's
165                // array streaming behavior: publish only completed elements.
166                if repaired && !elements.is_empty() {
167                    elements.pop();
168                }
169                Some(Value::Array(elements))
170            }
171            Self::Value => value.get("value").cloned(),
172        }
173    }
174
175    fn instruction(self) -> &'static str {
176        match self {
177            Self::Direct => "",
178            Self::Elements => {
179                "The provider-facing response schema wraps the requested array in an `elements` field. Follow that schema exactly; callers receive the unwrapped array."
180            }
181            Self::Value => {
182                "The provider-facing response schema wraps the requested scalar/enum value in a `value` field. Follow that schema exactly; callers receive the unwrapped value."
183            }
184        }
185    }
186}
187
188fn schema_is_object_like(schema: &Value) -> bool {
189    schema.get("type").and_then(Value::as_str) == Some("object")
190        || schema.get("properties").is_some()
191        || schema.get("required").is_some()
192        || schema.get("additionalProperties").is_some()
193}
194
195// ---------------------------------------------------------------------------
196// Core generation: blocking (non-streaming)
197// ---------------------------------------------------------------------------
198
199/// Generate a structured JSON object using the given LLM client.
200///
201/// Selects the best mode based on `req.mode`, calls the LLM, validates against
202/// the schema, and retries with repair prompts if validation fails.
203pub async fn generate_blocking(
204    client: &dyn LlmClient,
205    req: &StructuredRequest,
206) -> Result<StructuredResult> {
207    let mode = resolve_mode(req.mode, client.native_structured_support());
208    let envelope = SchemaEnvelope::for_schema(&req.schema);
209    let mut messages = build_initial_messages(req, mode);
210    let system = build_system_prompt(req, mode);
211    let tools = build_tools(req, mode);
212    let directive = build_directive(req, mode);
213
214    let mut total_usage = TokenUsage::default();
215    let mut repair_rounds: u8 = 0;
216
217    loop {
218        let resp = client
219            .complete_structured(&messages, Some(&system), &tools, &directive)
220            .await
221            .context("LLM call failed during structured generation")?;
222
223        accumulate_usage(&mut total_usage, &resp.usage);
224
225        // Mine the object from every place a model might have parked it (tool call,
226        // text content, AND the reasoning channel), trying each balanced JSON
227        // candidate against the schema. Reasoning models routinely leave `content`
228        // empty and emit the object inside `reasoning`, so without the reasoning
229        // fallback generate_object failed with "no structured output" across models.
230        let candidates = extract_raw_candidates(&resp.message, mode);
231        let resolution = resolve_structured(&candidates, &req.schema, envelope);
232
233        if let Some((value, raw)) = resolution.valid {
234            return Ok(StructuredResult {
235                object: value,
236                raw_text: Some(raw),
237                usage: total_usage,
238                repair_rounds,
239                mode_used: mode,
240            });
241        }
242
243        if repair_rounds >= req.max_repair_attempts {
244            return Err(match resolution.invalid {
245                Some((_, errors)) => anyhow::anyhow!(
246                    "Structured output failed schema validation after {} repair attempts. Errors: {}",
247                    repair_rounds,
248                    errors.join("; ")
249                ),
250                None => anyhow::anyhow!(
251                    "Structured output parsing failed after {} repair attempts: no JSON object found in tool call, text content, or reasoning channel",
252                    repair_rounds
253                ),
254            });
255        }
256
257        repair_rounds += 1;
258        let (repair_msg, raw_for_ctx) = match resolution.invalid {
259            Some((raw, errors)) => (build_repair_message(&raw, &errors), raw),
260            None => {
261                let raw = resolution.raw_seen.unwrap_or_default();
262                (build_parse_failure_repair(&raw), raw)
263            }
264        };
265        append_repair_context(
266            &mut messages,
267            &resp.message,
268            &repair_msg,
269            mode,
270            &raw_for_ctx,
271        );
272    }
273}
274
275// ---------------------------------------------------------------------------
276// Core generation: streaming
277// ---------------------------------------------------------------------------
278
279/// Generate a structured JSON object with streaming partial updates.
280///
281/// Calls `on_partial` with progressively more complete partial objects as tokens
282/// arrive. Returns the final validated object.
283///
284/// In streaming mode, `max_repair_attempts` defaults to 0 because a repair
285/// would reset the partial object stream (confusing for consumers).
286pub async fn generate_streaming(
287    client: &dyn LlmClient,
288    req: &StructuredRequest,
289    on_partial: PartialObjectCallback,
290) -> Result<StructuredResult> {
291    let mode = resolve_mode(req.mode, client.native_structured_support());
292    let envelope = SchemaEnvelope::for_schema(&req.schema);
293    let messages = build_initial_messages(req, mode);
294    let system = build_system_prompt(req, mode);
295    let tools = build_tools(req, mode);
296    let directive = build_directive(req, mode);
297
298    let cancel_token = CancellationToken::new();
299    let mut rx = client
300        .complete_streaming_structured(&messages, Some(&system), &tools, &directive, cancel_token)
301        .await
302        .context("LLM streaming call failed during structured generation")?;
303
304    let mut json_buffer = String::new();
305    let mut last_valid_partial: Option<Value> = None;
306    let mut final_response: Option<super::LlmResponse> = None;
307    let mut last_parse_len: usize = 0;
308    // Minimum bytes of new data before attempting a partial parse (reduces CPU)
309    const PARSE_THRESHOLD: usize = 8;
310
311    while let Some(event) = rx.recv().await {
312        match event {
313            StreamEvent::ToolUseInputDelta(delta) if mode == StructuredMode::Tool => {
314                if final_response.is_some() {
315                    continue;
316                }
317                json_buffer.push_str(&delta);
318                if json_buffer.len() - last_parse_len >= PARSE_THRESHOLD {
319                    if let Some(partial) = parse_partial_json(&json_buffer) {
320                        if let Some(projected) =
321                            envelope.project_partial(&partial.value, partial.repaired)
322                        {
323                            if last_valid_partial.as_ref() != Some(&projected) {
324                                on_partial(&projected);
325                                last_valid_partial = Some(projected);
326                            }
327                        }
328                    }
329                    last_parse_len = json_buffer.len();
330                }
331            }
332            StreamEvent::TextDelta(delta) if mode != StructuredMode::Tool => {
333                if final_response.is_some() {
334                    continue;
335                }
336                json_buffer.push_str(&delta);
337                if json_buffer.len() - last_parse_len >= PARSE_THRESHOLD {
338                    if let Some(json_start) = find_json_start(&json_buffer) {
339                        let candidate = &json_buffer[json_start..];
340                        if let Some(partial) = parse_partial_json(candidate) {
341                            if let Some(projected) =
342                                envelope.project_partial(&partial.value, partial.repaired)
343                            {
344                                if last_valid_partial.as_ref() != Some(&projected) {
345                                    on_partial(&projected);
346                                    last_valid_partial = Some(projected);
347                                }
348                            }
349                        }
350                    }
351                    last_parse_len = json_buffer.len();
352                }
353            }
354            StreamEvent::Done(resp) => {
355                final_response = Some(resp);
356            }
357            _ => {}
358        }
359    }
360
361    let resp = final_response.context("Stream ended without Done event")?;
362    // Same multi-source resolution as the blocking path: the final message may carry
363    // the object in the tool call, the text content, or the reasoning channel.
364    let candidates = extract_raw_candidates(&resp.message, mode);
365    let resolution = resolve_structured(&candidates, &req.schema, envelope);
366    let (value, raw_text) = match resolution.valid {
367        Some(vr) => vr,
368        None => {
369            return Err(match resolution.invalid {
370                Some((_, errors)) => anyhow::anyhow!(
371                    "Streamed structured output failed schema validation: {}",
372                    errors.join("; ")
373                ),
374                None => anyhow::anyhow!(
375                    "Streamed output produced no parseable JSON object (checked tool call, text content, and reasoning channel)"
376                ),
377            });
378        }
379    };
380
381    // Emit final complete object
382    on_partial(&value);
383
384    Ok(StructuredResult {
385        object: value,
386        raw_text: Some(raw_text),
387        usage: resp.usage,
388        repair_rounds: 0,
389        mode_used: mode,
390    })
391}
392
393// ---------------------------------------------------------------------------
394// JSON extraction and parsing
395// ---------------------------------------------------------------------------
396
397/// Extract a JSON value from potentially dirty LLM output.
398///
399/// Handles: raw JSON, markdown code fences, leading/trailing prose.
400pub fn extract_json_value(text: &str) -> Result<Value> {
401    let trimmed = text.trim();
402
403    // 1. Direct parse
404    if let Ok(v) = serde_json::from_str::<Value>(trimmed) {
405        if v.is_object() || v.is_array() {
406            return Ok(v);
407        }
408    }
409
410    // 2. Strip markdown code fence
411    if let Some(inner) = strip_code_fence(trimmed) {
412        if let Ok(v) = serde_json::from_str::<Value>(inner.trim()) {
413            if v.is_object() || v.is_array() {
414                return Ok(v);
415            }
416        }
417    }
418
419    // 3. Find balanced JSON substring (first { to matching })
420    if let Some(candidate) = find_balanced_json_object(trimmed) {
421        if let Ok(v) = serde_json::from_str::<Value>(candidate) {
422            return Ok(v);
423        }
424    }
425
426    // 4. Try array
427    if let Some(candidate) = find_balanced_json_array(trimmed) {
428        if let Ok(v) = serde_json::from_str::<Value>(candidate) {
429            return Ok(v);
430        }
431    }
432
433    bail!("No valid JSON object found in LLM output")
434}
435
436/// Strip ```json ... ``` or ``` ... ``` fences.
437fn strip_code_fence(text: &str) -> Option<&str> {
438    let start_patterns = ["```json\n", "```json\r\n", "```\n", "```\r\n"];
439    for pat in &start_patterns {
440        if let Some(rest) = text.strip_prefix(pat) {
441            // Find closing fence
442            if let Some(end) = rest.rfind("```") {
443                return Some(&rest[..end]);
444            }
445        }
446    }
447    // Also handle inline: ```json{...}```
448    if let Some(inner) = text.strip_prefix("```json") {
449        if let Some(end) = inner.rfind("```") {
450            return Some(inner[..end].trim());
451        }
452    }
453    if let Some(inner) = text.strip_prefix("```") {
454        if let Some(end) = inner.rfind("```") {
455            return Some(inner[..end].trim());
456        }
457    }
458    None
459}
460
461/// Find the first balanced `{...}` substring using bracket counting.
462fn find_balanced_json_object(text: &str) -> Option<&str> {
463    find_balanced(text, '{', '}')
464}
465
466/// Find the first balanced `[...]` substring.
467fn find_balanced_json_array(text: &str) -> Option<&str> {
468    find_balanced(text, '[', ']')
469}
470
471fn find_balanced(text: &str, open: char, close: char) -> Option<&str> {
472    find_balanced_range(text, open, close).map(|(start, end)| &text[start..end])
473}
474
475/// Byte range `[start, end)` of the first balanced `open..close` substring (quote-aware).
476fn find_balanced_range(text: &str, open: char, close: char) -> Option<(usize, usize)> {
477    let bytes = text.as_bytes();
478    let open_byte = open as u8;
479    let close_byte = close as u8;
480
481    // Find the first unquoted occurrence of `open`
482    let mut in_string = false;
483    let mut escape_next = false;
484    let mut start = None;
485
486    for (i, &b) in bytes.iter().enumerate() {
487        if escape_next {
488            escape_next = false;
489            continue;
490        }
491        match b {
492            b'\\' if in_string => escape_next = true,
493            b'"' => in_string = !in_string,
494            _ if in_string => {}
495            _ if b == open_byte => {
496                start = Some(i);
497                break;
498            }
499            _ => {}
500        }
501    }
502
503    let start = start?;
504    let mut depth = 0i32;
505    in_string = false;
506    escape_next = false;
507
508    for (i, &b) in bytes[start..].iter().enumerate() {
509        if escape_next {
510            escape_next = false;
511            continue;
512        }
513        match b {
514            b'\\' if in_string => escape_next = true,
515            b'"' => in_string = !in_string,
516            _ if in_string => {}
517            _ if b == open_byte => depth += 1,
518            _ if b == close_byte => {
519                depth -= 1;
520                if depth == 0 {
521                    return Some((start, start + i + 1));
522                }
523            }
524            _ => {}
525        }
526    }
527    None
528}
529
530/// Every top-level balanced `open..close` substring, in document order.
531///
532/// Reasoning traces often contain several objects (worked examples, partial drafts)
533/// before the final answer, so callers validate each against the schema and keep the
534/// one that fits rather than blindly trusting the first `{...}`.
535fn find_all_balanced(text: &str, open: char, close: char) -> Vec<String> {
536    let mut out = Vec::new();
537    let mut base = 0usize;
538    while base < text.len() {
539        match find_balanced_range(&text[base..], open, close) {
540            Some((start, end)) => {
541                out.push(text[base + start..base + end].to_string());
542                base += end;
543            }
544            None => break,
545        }
546    }
547    out
548}
549
550/// Find the byte offset where JSON content starts in a text stream.
551/// Skips leading prose/whitespace to find `{` or `[` that isn't inside a string.
552fn find_json_start(text: &str) -> Option<usize> {
553    // Skip past code fence markers if present
554    let (search_text, offset) = if let Some(rest) = text.strip_prefix("```json") {
555        (rest, 7)
556    } else if let Some(rest) = text.strip_prefix("```") {
557        (rest, 3)
558    } else {
559        (text, 0)
560    };
561
562    let mut in_string = false;
563    let mut escape_next = false;
564    for (i, &b) in search_text.as_bytes().iter().enumerate() {
565        if escape_next {
566            escape_next = false;
567            continue;
568        }
569        match b {
570            b'\\' if in_string => {
571                escape_next = true;
572            }
573            b'"' => {
574                in_string = !in_string;
575            }
576            b'{' | b'[' if !in_string => {
577                return Some(offset + i);
578            }
579            _ => {}
580        }
581    }
582    None
583}
584
585// ---------------------------------------------------------------------------
586// Partial JSON parsing (for streaming)
587// ---------------------------------------------------------------------------
588
589/// Attempt to parse a potentially incomplete JSON string into the most complete
590/// valid partial object possible.
591///
592/// Strategy: try parsing as-is first. If that fails, run a single-pass JSON
593/// scanner that trims incomplete trailing tokens and closes open containers.
594/// This mirrors the practical behavior of Vercel AI SDK's partial JSON parser
595/// while keeping A3S's final schema validation strict.
596#[cfg(test)]
597fn try_parse_partial_json(text: &str) -> Option<Value> {
598    parse_partial_json(text).map(|parsed| parsed.value)
599}
600
601#[derive(Debug, Clone, PartialEq)]
602struct PartialJsonValue {
603    value: Value,
604    repaired: bool,
605}
606
607fn parse_partial_json(text: &str) -> Option<PartialJsonValue> {
608    let trimmed = text.trim();
609    if trimmed.is_empty() {
610        return None;
611    }
612
613    // Fast path: already valid
614    if let Ok(v) = serde_json::from_str::<Value>(trimmed) {
615        if v.is_object() || v.is_array() {
616            return Some(PartialJsonValue {
617                value: v,
618                repaired: false,
619            });
620        }
621    }
622
623    let repaired = fix_partial_json(trimmed);
624    if repaired.trim().is_empty() || repaired == trimmed {
625        return None;
626    }
627
628    serde_json::from_str::<Value>(&repaired)
629        .ok()
630        .filter(|v| v.is_object() || v.is_array())
631        .map(|value| PartialJsonValue {
632            value,
633            repaired: true,
634        })
635}
636
637#[derive(Debug, Clone, Copy, PartialEq, Eq)]
638enum PartialJsonState {
639    Root,
640    Finish,
641    InsideString,
642    InsideStringEscape,
643    InsideStringUnicodeEscape,
644    InsideLiteral,
645    InsideNumber,
646    InsideObjectStart,
647    InsideObjectKey,
648    InsideObjectAfterKey,
649    InsideObjectBeforeValue,
650    InsideObjectAfterValue,
651    InsideObjectAfterComma,
652    InsideArrayStart,
653    InsideArrayAfterValue,
654    InsideArrayAfterComma,
655}
656
657fn is_json_hex_digit(ch: char) -> bool {
658    ch.is_ascii_hexdigit()
659}
660
661fn process_partial_value_start(
662    ch: char,
663    end: usize,
664    swap_state: PartialJsonState,
665    stack: &mut Vec<PartialJsonState>,
666    last_valid_end: &mut usize,
667    literal_start: &mut Option<usize>,
668    start: usize,
669) {
670    match ch {
671        '"' => {
672            *last_valid_end = end;
673            stack.pop();
674            stack.push(swap_state);
675            stack.push(PartialJsonState::InsideString);
676        }
677        'f' | 't' | 'n' => {
678            *last_valid_end = end;
679            *literal_start = Some(start);
680            stack.pop();
681            stack.push(swap_state);
682            stack.push(PartialJsonState::InsideLiteral);
683        }
684        '-' => {
685            stack.pop();
686            stack.push(swap_state);
687            stack.push(PartialJsonState::InsideNumber);
688        }
689        '0'..='9' => {
690            *last_valid_end = end;
691            stack.pop();
692            stack.push(swap_state);
693            stack.push(PartialJsonState::InsideNumber);
694        }
695        '{' => {
696            *last_valid_end = end;
697            stack.pop();
698            stack.push(swap_state);
699            stack.push(PartialJsonState::InsideObjectStart);
700        }
701        '[' => {
702            *last_valid_end = end;
703            stack.pop();
704            stack.push(swap_state);
705            stack.push(PartialJsonState::InsideArrayStart);
706        }
707        _ => {}
708    }
709}
710
711fn process_after_partial_object_value(
712    ch: char,
713    end: usize,
714    stack: &mut Vec<PartialJsonState>,
715    last_valid_end: &mut usize,
716) {
717    match ch {
718        ',' => {
719            stack.pop();
720            stack.push(PartialJsonState::InsideObjectAfterComma);
721        }
722        '}' => {
723            *last_valid_end = end;
724            stack.pop();
725        }
726        _ => {}
727    }
728}
729
730fn process_after_partial_array_value(
731    ch: char,
732    end: usize,
733    stack: &mut Vec<PartialJsonState>,
734    last_valid_end: &mut usize,
735) {
736    match ch {
737        ',' => {
738            stack.pop();
739            stack.push(PartialJsonState::InsideArrayAfterComma);
740        }
741        ']' => {
742            *last_valid_end = end;
743            stack.pop();
744        }
745        _ => {}
746    }
747}
748
749fn fix_partial_json(input: &str) -> String {
750    use PartialJsonState::*;
751
752    let mut stack = vec![Root];
753    let mut last_valid_end = 0usize;
754    let mut literal_start: Option<usize> = None;
755    let mut unicode_escape_digits = 0usize;
756
757    for (start, ch) in input.char_indices() {
758        let end = start + ch.len_utf8();
759        let current_state = *stack.last().unwrap_or(&Finish);
760
761        match current_state {
762            Root => {
763                process_partial_value_start(
764                    ch,
765                    end,
766                    Finish,
767                    &mut stack,
768                    &mut last_valid_end,
769                    &mut literal_start,
770                    start,
771                );
772            }
773            InsideObjectStart => match ch {
774                '"' => {
775                    stack.pop();
776                    stack.push(InsideObjectKey);
777                }
778                '}' => {
779                    last_valid_end = end;
780                    stack.pop();
781                }
782                _ => {}
783            },
784            InsideObjectAfterComma => {
785                if ch == '"' {
786                    stack.pop();
787                    stack.push(InsideObjectKey);
788                }
789            }
790            InsideObjectKey => {
791                if ch == '"' {
792                    stack.pop();
793                    stack.push(InsideObjectAfterKey);
794                }
795            }
796            InsideObjectAfterKey => {
797                if ch == ':' {
798                    stack.pop();
799                    stack.push(InsideObjectBeforeValue);
800                }
801            }
802            InsideObjectBeforeValue => {
803                process_partial_value_start(
804                    ch,
805                    end,
806                    InsideObjectAfterValue,
807                    &mut stack,
808                    &mut last_valid_end,
809                    &mut literal_start,
810                    start,
811                );
812            }
813            InsideObjectAfterValue => {
814                process_after_partial_object_value(ch, end, &mut stack, &mut last_valid_end);
815            }
816            InsideString => match ch {
817                '"' => {
818                    stack.pop();
819                    last_valid_end = end;
820                }
821                '\\' => {
822                    stack.push(InsideStringEscape);
823                }
824                _ => {
825                    last_valid_end = end;
826                }
827            },
828            InsideArrayStart => match ch {
829                ']' => {
830                    last_valid_end = end;
831                    stack.pop();
832                }
833                _ => {
834                    last_valid_end = end;
835                    process_partial_value_start(
836                        ch,
837                        end,
838                        InsideArrayAfterValue,
839                        &mut stack,
840                        &mut last_valid_end,
841                        &mut literal_start,
842                        start,
843                    );
844                }
845            },
846            InsideArrayAfterValue => match ch {
847                ',' => {
848                    stack.pop();
849                    stack.push(InsideArrayAfterComma);
850                }
851                ']' => {
852                    last_valid_end = end;
853                    stack.pop();
854                }
855                _ => {
856                    last_valid_end = end;
857                }
858            },
859            InsideArrayAfterComma => {
860                process_partial_value_start(
861                    ch,
862                    end,
863                    InsideArrayAfterValue,
864                    &mut stack,
865                    &mut last_valid_end,
866                    &mut literal_start,
867                    start,
868                );
869            }
870            InsideStringEscape => {
871                stack.pop();
872                if ch == 'u' {
873                    unicode_escape_digits = 0;
874                    stack.push(InsideStringUnicodeEscape);
875                } else {
876                    last_valid_end = end;
877                }
878            }
879            InsideStringUnicodeEscape => {
880                if is_json_hex_digit(ch) {
881                    unicode_escape_digits += 1;
882                    if unicode_escape_digits == 4 {
883                        stack.pop();
884                        last_valid_end = end;
885                    }
886                }
887            }
888            InsideNumber => match ch {
889                '0'..='9' => {
890                    last_valid_end = end;
891                }
892                'e' | 'E' | '-' | '.' => {}
893                ',' => {
894                    stack.pop();
895                    if stack.last() == Some(&InsideArrayAfterValue) {
896                        process_after_partial_array_value(ch, end, &mut stack, &mut last_valid_end);
897                    }
898                    if stack.last() == Some(&InsideObjectAfterValue) {
899                        process_after_partial_object_value(
900                            ch,
901                            end,
902                            &mut stack,
903                            &mut last_valid_end,
904                        );
905                    }
906                }
907                '}' => {
908                    stack.pop();
909                    if stack.last() == Some(&InsideObjectAfterValue) {
910                        process_after_partial_object_value(
911                            ch,
912                            end,
913                            &mut stack,
914                            &mut last_valid_end,
915                        );
916                    }
917                }
918                ']' => {
919                    stack.pop();
920                    if stack.last() == Some(&InsideArrayAfterValue) {
921                        process_after_partial_array_value(ch, end, &mut stack, &mut last_valid_end);
922                    }
923                }
924                _ => {
925                    stack.pop();
926                }
927            },
928            InsideLiteral => {
929                let partial_literal = literal_start
930                    .and_then(|s| input.get(s..end))
931                    .unwrap_or_default();
932                if !("false".starts_with(partial_literal)
933                    || "true".starts_with(partial_literal)
934                    || "null".starts_with(partial_literal))
935                {
936                    stack.pop();
937                    if stack.last() == Some(&InsideObjectAfterValue) {
938                        process_after_partial_object_value(
939                            ch,
940                            end,
941                            &mut stack,
942                            &mut last_valid_end,
943                        );
944                    } else if stack.last() == Some(&InsideArrayAfterValue) {
945                        process_after_partial_array_value(ch, end, &mut stack, &mut last_valid_end);
946                    }
947                } else {
948                    last_valid_end = end;
949                }
950            }
951            Finish => {}
952        }
953    }
954
955    let mut result = input[..last_valid_end].to_string();
956    for state in stack.iter().rev() {
957        match state {
958            InsideString => result.push('"'),
959            InsideObjectKey
960            | InsideObjectAfterKey
961            | InsideObjectAfterComma
962            | InsideObjectStart
963            | InsideObjectBeforeValue
964            | InsideObjectAfterValue => result.push('}'),
965            InsideArrayStart | InsideArrayAfterComma | InsideArrayAfterValue => result.push(']'),
966            InsideLiteral => {
967                let partial_literal = literal_start
968                    .and_then(|s| input.get(s..input.len()))
969                    .unwrap_or_default();
970                if "true".starts_with(partial_literal) {
971                    result.push_str(&"true"[partial_literal.len()..]);
972                } else if "false".starts_with(partial_literal) {
973                    result.push_str(&"false"[partial_literal.len()..]);
974                } else if "null".starts_with(partial_literal) {
975                    result.push_str(&"null"[partial_literal.len()..]);
976                }
977            }
978            _ => {}
979        }
980    }
981
982    result
983}
984
985// ---------------------------------------------------------------------------
986// Schema validation
987// ---------------------------------------------------------------------------
988
989/// Validate a JSON value against a JSON Schema.
990/// Returns Ok(()) on success, or a list of human-readable error strings.
991fn validate_against_schema(value: &Value, schema: &Value) -> Result<(), Vec<String>> {
992    // We do a basic recursive validation here. For production, consider using
993    // the `jsonschema` crate, but to avoid adding a heavy dependency we implement
994    // the subset of JSON Schema that matters for structured output.
995    let errors = basic_schema_validate(value, schema, "");
996    if errors.is_empty() {
997        Ok(())
998    } else {
999        Err(errors)
1000    }
1001}
1002
1003/// Basic JSON Schema validator covering the most common constraints.
1004fn basic_schema_validate(value: &Value, schema: &Value, path: &str) -> Vec<String> {
1005    let mut errors = Vec::new();
1006
1007    // Handle $ref — not supported in basic validator, skip
1008    if schema.get("$ref").is_some() {
1009        return errors;
1010    }
1011
1012    // Handle anyOf / oneOf: value must match at least one sub-schema
1013    if let Some(any_of) = schema
1014        .get("anyOf")
1015        .or_else(|| schema.get("oneOf"))
1016        .and_then(|v| v.as_array())
1017    {
1018        let matched = any_of
1019            .iter()
1020            .any(|sub| basic_schema_validate(value, sub, path).is_empty());
1021        if !matched {
1022            errors.push(format!(
1023                "{}: value does not match any variant in anyOf/oneOf",
1024                path_or_root(path),
1025            ));
1026        }
1027        return errors;
1028    }
1029
1030    // Handle enum
1031    if let Some(enum_values) = schema.get("enum").and_then(|v| v.as_array()) {
1032        if !enum_values.contains(value) {
1033            errors.push(format!(
1034                "{}: value {:?} not in enum {:?}",
1035                path_or_root(path),
1036                value,
1037                enum_values
1038            ));
1039        }
1040        return errors;
1041    }
1042
1043    // Handle const
1044    if let Some(const_val) = schema.get("const") {
1045        if value != const_val {
1046            errors.push(format!(
1047                "{}: expected const {:?}, got {:?}",
1048                path_or_root(path),
1049                const_val,
1050                value
1051            ));
1052        }
1053        return errors;
1054    }
1055
1056    // Type checking (supports nullable via type array: ["string", "null"])
1057    if let Some(type_val) = schema.get("type") {
1058        let type_ok = if let Some(type_str) = type_val.as_str() {
1059            check_type(value, type_str)
1060        } else if let Some(type_arr) = type_val.as_array() {
1061            type_arr
1062                .iter()
1063                .filter_map(|t| t.as_str())
1064                .any(|t| check_type(value, t))
1065        } else {
1066            true
1067        };
1068        if !type_ok {
1069            errors.push(format!(
1070                "{}: expected type {:?}, got {:?}",
1071                path_or_root(path),
1072                type_val,
1073                value_type_name(value)
1074            ));
1075            return errors;
1076        }
1077    }
1078
1079    // Object validation
1080    if let Some(obj) = value.as_object() {
1081        if let Some(properties) = schema.get("properties").and_then(|v| v.as_object()) {
1082            for (key, prop_schema) in properties {
1083                if let Some(child_value) = obj.get(key) {
1084                    let child_path = if path.is_empty() {
1085                        format!(".{}", key)
1086                    } else {
1087                        format!("{}.{}", path, key)
1088                    };
1089                    errors.extend(basic_schema_validate(child_value, prop_schema, &child_path));
1090                }
1091            }
1092        }
1093
1094        if let Some(required) = schema.get("required").and_then(|v| v.as_array()) {
1095            for req_field in required {
1096                if let Some(field_name) = req_field.as_str() {
1097                    if !obj.contains_key(field_name) {
1098                        errors.push(format!(
1099                            "{}: missing required field '{}'",
1100                            path_or_root(path),
1101                            field_name
1102                        ));
1103                    }
1104                }
1105            }
1106        }
1107
1108        // additionalProperties: false
1109        if schema.get("additionalProperties") == Some(&Value::Bool(false)) {
1110            if let Some(properties) = schema.get("properties").and_then(|v| v.as_object()) {
1111                for key in obj.keys() {
1112                    if !properties.contains_key(key) {
1113                        errors.push(format!(
1114                            "{}: unexpected additional property '{}'",
1115                            path_or_root(path),
1116                            key
1117                        ));
1118                    }
1119                }
1120            }
1121        }
1122    }
1123
1124    // Array validation
1125    if let Some(arr) = value.as_array() {
1126        if let Some(items_schema) = schema.get("items") {
1127            for (i, item) in arr.iter().enumerate() {
1128                let child_path = format!("{}[{}]", path, i);
1129                errors.extend(basic_schema_validate(item, items_schema, &child_path));
1130            }
1131        }
1132        if let Some(min) = schema.get("minItems").and_then(|v| v.as_u64()) {
1133            if (arr.len() as u64) < min {
1134                errors.push(format!(
1135                    "{}: array has {} items, minimum is {}",
1136                    path_or_root(path),
1137                    arr.len(),
1138                    min
1139                ));
1140            }
1141        }
1142        if let Some(max) = schema.get("maxItems").and_then(|v| v.as_u64()) {
1143            if (arr.len() as u64) > max {
1144                errors.push(format!(
1145                    "{}: array has {} items, maximum is {}",
1146                    path_or_root(path),
1147                    arr.len(),
1148                    max
1149                ));
1150            }
1151        }
1152    }
1153
1154    // String validation
1155    if let Some(s) = value.as_str() {
1156        if let Some(min_len) = schema.get("minLength").and_then(|v| v.as_u64()) {
1157            if (s.chars().count() as u64) < min_len {
1158                errors.push(format!(
1159                    "{}: string length {} < minLength {}",
1160                    path_or_root(path),
1161                    s.chars().count(),
1162                    min_len
1163                ));
1164            }
1165        }
1166        if let Some(max_len) = schema.get("maxLength").and_then(|v| v.as_u64()) {
1167            if (s.chars().count() as u64) > max_len {
1168                errors.push(format!(
1169                    "{}: string length {} > maxLength {}",
1170                    path_or_root(path),
1171                    s.chars().count(),
1172                    max_len
1173                ));
1174            }
1175        }
1176        if let Some(pattern) = schema.get("pattern").and_then(|v| v.as_str()) {
1177            if let Ok(re) = regex::Regex::new(pattern) {
1178                if !re.is_match(s) {
1179                    errors.push(format!(
1180                        "{}: string does not match pattern '{}'",
1181                        path_or_root(path),
1182                        pattern
1183                    ));
1184                }
1185            }
1186        }
1187    }
1188
1189    // Number validation
1190    if let Some(n) = value.as_f64() {
1191        if let Some(min) = schema.get("minimum").and_then(|v| v.as_f64()) {
1192            if n < min {
1193                errors.push(format!(
1194                    "{}: value {} < minimum {}",
1195                    path_or_root(path),
1196                    n,
1197                    min
1198                ));
1199            }
1200        }
1201        if let Some(max) = schema.get("maximum").and_then(|v| v.as_f64()) {
1202            if n > max {
1203                errors.push(format!(
1204                    "{}: value {} > maximum {}",
1205                    path_or_root(path),
1206                    n,
1207                    max
1208                ));
1209            }
1210        }
1211        if let Some(exc_min) = schema.get("exclusiveMinimum").and_then(|v| v.as_f64()) {
1212            if n <= exc_min {
1213                errors.push(format!(
1214                    "{}: value {} <= exclusiveMinimum {}",
1215                    path_or_root(path),
1216                    n,
1217                    exc_min
1218                ));
1219            }
1220        }
1221        if let Some(exc_max) = schema.get("exclusiveMaximum").and_then(|v| v.as_f64()) {
1222            if n >= exc_max {
1223                errors.push(format!(
1224                    "{}: value {} >= exclusiveMaximum {}",
1225                    path_or_root(path),
1226                    n,
1227                    exc_max
1228                ));
1229            }
1230        }
1231    }
1232
1233    errors
1234}
1235
1236fn check_type(value: &Value, type_str: &str) -> bool {
1237    match type_str {
1238        "object" => value.is_object(),
1239        "array" => value.is_array(),
1240        "string" => value.is_string(),
1241        "number" => value.is_number(),
1242        "integer" => {
1243            value.is_i64()
1244                || value.is_u64()
1245                || value
1246                    .as_f64()
1247                    .map(|f| f.fract() == 0.0 && f.is_finite())
1248                    .unwrap_or(false)
1249        }
1250        "boolean" => value.is_boolean(),
1251        "null" => value.is_null(),
1252        _ => true,
1253    }
1254}
1255
1256fn path_or_root(path: &str) -> &str {
1257    if path.is_empty() {
1258        "$"
1259    } else {
1260        path
1261    }
1262}
1263
1264fn value_type_name(value: &Value) -> &'static str {
1265    match value {
1266        Value::Null => "null",
1267        Value::Bool(_) => "boolean",
1268        Value::Number(_) => "number",
1269        Value::String(_) => "string",
1270        Value::Array(_) => "array",
1271        Value::Object(_) => "object",
1272    }
1273}
1274
1275// ---------------------------------------------------------------------------
1276// Message/prompt construction helpers
1277// ---------------------------------------------------------------------------
1278
1279/// Resolve the requested mode against the provider's native capability.
1280///
1281/// Prefer native enforcement only when the client explicitly reports support.
1282/// Unknown OpenAI-compatible endpoints can hang when sent `tool_choice` or
1283/// `response_format`, so unsupported requests degrade to prompt+schema parsing
1284/// instead of optimistic native parameters.
1285fn resolve_mode(requested: StructuredMode, support: NativeStructuredSupport) -> StructuredMode {
1286    match (requested, support) {
1287        (StructuredMode::Prompt, _) => StructuredMode::Prompt,
1288        (StructuredMode::Strict, NativeStructuredSupport::JsonSchema) => StructuredMode::Strict,
1289        (StructuredMode::Json, NativeStructuredSupport::JsonSchema) => StructuredMode::Json,
1290        (StructuredMode::Auto | StructuredMode::Tool, NativeStructuredSupport::JsonSchema) => {
1291            StructuredMode::Tool
1292        }
1293        (
1294            StructuredMode::Auto
1295            | StructuredMode::Tool
1296            | StructuredMode::Strict
1297            | StructuredMode::Json,
1298            NativeStructuredSupport::ForcedTool,
1299        ) => StructuredMode::Tool,
1300        (
1301            StructuredMode::Auto
1302            | StructuredMode::Tool
1303            | StructuredMode::Strict
1304            | StructuredMode::Json,
1305            NativeStructuredSupport::None,
1306        ) => StructuredMode::Prompt,
1307    }
1308}
1309
1310/// Build the provider directive for an already-resolved mode.
1311fn build_directive(req: &StructuredRequest, mode: StructuredMode) -> StructuredDirective {
1312    match mode {
1313        StructuredMode::Tool => StructuredDirective {
1314            force_tool: Some(format!("emit_{}", req.schema_name)),
1315            response_format: None,
1316        },
1317        StructuredMode::Strict => StructuredDirective {
1318            force_tool: None,
1319            response_format: Some(ResponseFormat::JsonSchema {
1320                name: req.schema_name.clone(),
1321                schema: SchemaEnvelope::for_schema(&req.schema).response_schema(&req.schema),
1322            }),
1323        },
1324        StructuredMode::Json => StructuredDirective {
1325            force_tool: None,
1326            response_format: Some(ResponseFormat::JsonObject),
1327        },
1328        StructuredMode::Auto | StructuredMode::Prompt => StructuredDirective::default(),
1329    }
1330}
1331
1332fn build_initial_messages(req: &StructuredRequest, mode: StructuredMode) -> Vec<Message> {
1333    let envelope = SchemaEnvelope::for_schema(&req.schema);
1334    let response_schema = envelope.response_schema(&req.schema);
1335    let envelope_instruction = envelope.instruction();
1336    match mode {
1337        StructuredMode::Tool => {
1338            // For tool mode, the prompt is the user message; the LLM will respond
1339            // with a tool call whose input is the structured object.
1340            vec![Message::user(&req.prompt)]
1341        }
1342        StructuredMode::Prompt | StructuredMode::Json => {
1343            // Prompt mode and json_object mode both need the schema in the prompt:
1344            // json_object only guarantees *syntactic* validity, so the model still
1345            // has to be told the shape it should produce.
1346            let augmented = format!(
1347                "{}\n\n{}{}\n\nYou MUST respond with ONLY a valid JSON object (no markdown, no explanation) that conforms to this JSON Schema:\n\n```json\n{}\n```",
1348                req.prompt,
1349                envelope_instruction,
1350                if envelope_instruction.is_empty() { "" } else { "\n" },
1351                serde_json::to_string_pretty(&response_schema).unwrap_or_default()
1352            );
1353            vec![Message::user(&augmented)]
1354        }
1355        _ => {
1356            // Strict mode: the schema constraint is enforced by the provider via
1357            // response_format.json_schema, so the user message is just the prompt.
1358            vec![Message::user(&req.prompt)]
1359        }
1360    }
1361}
1362
1363fn build_system_prompt(req: &StructuredRequest, mode: StructuredMode) -> String {
1364    let base = req.system.as_deref().unwrap_or("");
1365    let envelope_instruction = SchemaEnvelope::for_schema(&req.schema).instruction();
1366
1367    match mode {
1368        StructuredMode::Tool => {
1369            format!(
1370                "{}{}You MUST respond by calling the `emit_{}` tool exactly once with a valid argument matching the schema. Do not output any text outside the tool call.{}{}",
1371                base,
1372                if base.is_empty() { "" } else { "\n\n" },
1373                req.schema_name,
1374                if envelope_instruction.is_empty() { "" } else { "\n\n" },
1375                envelope_instruction
1376            )
1377        }
1378        StructuredMode::Prompt | StructuredMode::Json => {
1379            format!(
1380                "{}{}You are a structured data extraction assistant. Always respond with valid JSON only, no markdown fences, no explanation text.{}{}",
1381                base,
1382                if base.is_empty() { "" } else { "\n\n" },
1383                if envelope_instruction.is_empty() { "" } else { "\n\n" },
1384                envelope_instruction,
1385            )
1386        }
1387        _ => base.to_string(),
1388    }
1389}
1390
1391fn build_tools(req: &StructuredRequest, mode: StructuredMode) -> Vec<ToolDefinition> {
1392    match mode {
1393        StructuredMode::Tool => {
1394            vec![ToolDefinition {
1395                name: format!("emit_{}", req.schema_name),
1396                description: req
1397                    .schema_description
1398                    .clone()
1399                    .unwrap_or_else(|| format!("Emit a structured {} object", req.schema_name)),
1400                parameters: SchemaEnvelope::for_schema(&req.schema).response_schema(&req.schema),
1401            }]
1402        }
1403        _ => vec![],
1404    }
1405}
1406
1407/// Outcome of mining a response for the structured object across all candidate sources.
1408struct StructuredResolution {
1409    /// A schema-valid object plus the raw source string it came from.
1410    valid: Option<(Value, String)>,
1411    /// First parseable-but-schema-invalid object source + its validation errors,
1412    /// used to build a targeted repair prompt.
1413    invalid: Option<(String, Vec<String>)>,
1414    /// First non-empty raw candidate, shown verbatim in a parse-failure repair prompt.
1415    raw_seen: Option<String>,
1416}
1417
1418/// Append `s` to `out` if it is non-empty and not already present (trimmed, deduped).
1419fn push_candidate(out: &mut Vec<String>, s: String) {
1420    let trimmed = s.trim();
1421    if !trimmed.is_empty() && !out.iter().any(|c| c == trimmed) {
1422        out.push(trimmed.to_string());
1423    }
1424}
1425
1426/// Ordered raw strings to mine for the structured object, most authoritative first:
1427/// tool-call arguments, then text content, then the reasoning channel.
1428///
1429/// The reasoning fallback is the crux of the cross-model fix: reasoning models
1430/// (GLM/zhipu, DeepSeek-R1, kimi…) frequently emit the final object inside
1431/// `reasoning` with `content` empty and no tool call. Earlier extraction only looked
1432/// at the tool call / text, so those models yielded an empty string and the whole
1433/// generate_object failed even though a perfectly good object was produced.
1434fn extract_raw_candidates(message: &super::Message, mode: StructuredMode) -> Vec<String> {
1435    let mut out: Vec<String> = Vec::new();
1436    if mode == StructuredMode::Tool {
1437        if let Some(call) = message.tool_calls().first() {
1438            push_candidate(
1439                &mut out,
1440                serde_json::to_string(&call.args).unwrap_or_default(),
1441            );
1442        }
1443    }
1444    push_candidate(&mut out, message.text());
1445    if let Some(reasoning) = message.reasoning_content.as_deref() {
1446        push_candidate(&mut out, reasoning.to_string());
1447    }
1448    out
1449}
1450
1451/// Every JSON object/array value mineable from possibly-dirty text, in document order
1452/// (direct parse, code fences, then all balanced `{...}` / `[...]`). Deduped.
1453#[cfg(test)]
1454fn extract_all_json_values(text: &str) -> Vec<Value> {
1455    extract_json_candidates(text, false)
1456}
1457
1458/// Every JSON value mineable from possibly-dirty text for schema-aware structured
1459/// resolution. When `include_direct_scalars` is true, direct raw/fenced scalar JSON
1460/// is retained so top-level scalar schemas can recover non-enveloped model output.
1461fn extract_json_candidates(text: &str, include_direct_scalars: bool) -> Vec<Value> {
1462    let trimmed = text.trim();
1463    let mut values: Vec<Value> = Vec::new();
1464    let consider = |candidate: &str, values: &mut Vec<Value>, allow_scalar: bool| {
1465        if let Ok(v) = serde_json::from_str::<Value>(candidate.trim()) {
1466            if (v.is_object() || v.is_array() || allow_scalar) && !values.contains(&v) {
1467                values.push(v);
1468            }
1469        }
1470    };
1471    consider(trimmed, &mut values, include_direct_scalars);
1472    if let Some(inner) = strip_code_fence(trimmed) {
1473        consider(inner, &mut values, include_direct_scalars);
1474    }
1475    for candidate in find_all_balanced(trimmed, '{', '}') {
1476        consider(&candidate, &mut values, false);
1477    }
1478    for candidate in find_all_balanced(trimmed, '[', ']') {
1479        consider(&candidate, &mut values, false);
1480    }
1481    values
1482}
1483
1484/// Try every raw candidate × every JSON value it yields against the schema; return the
1485/// first schema-valid value, else the best parseable-but-invalid value (for repair).
1486fn resolve_structured(
1487    candidates: &[String],
1488    schema: &Value,
1489    envelope: SchemaEnvelope,
1490) -> StructuredResolution {
1491    let mut invalid: Option<(String, Vec<String>)> = None;
1492    let mut raw_seen: Option<String> = None;
1493    let response_schema = envelope.response_schema(schema);
1494    for raw in candidates {
1495        if raw_seen.is_none() && !raw.trim().is_empty() {
1496            raw_seen = Some(raw.clone());
1497        }
1498        for value in extract_json_candidates(raw, envelope == SchemaEnvelope::Value) {
1499            match validate_against_schema(&value, schema) {
1500                Ok(()) => {
1501                    return StructuredResolution {
1502                        valid: Some((value, raw.clone())),
1503                        invalid,
1504                        raw_seen,
1505                    };
1506                }
1507                Err(errors) => {
1508                    if invalid.is_none() {
1509                        invalid = Some((raw.clone(), errors));
1510                    }
1511                }
1512            }
1513
1514            if envelope != SchemaEnvelope::Direct {
1515                match validate_against_schema(&value, &response_schema) {
1516                    Ok(()) => {
1517                        if let Some(unwrapped) = envelope.unwrap_final(&value) {
1518                            match validate_against_schema(&unwrapped, schema) {
1519                                Ok(()) => {
1520                                    return StructuredResolution {
1521                                        valid: Some((unwrapped, raw.clone())),
1522                                        invalid,
1523                                        raw_seen,
1524                                    };
1525                                }
1526                                Err(errors) => {
1527                                    if invalid.is_none() {
1528                                        invalid = Some((raw.clone(), errors));
1529                                    }
1530                                }
1531                            }
1532                        } else if invalid.is_none() {
1533                            invalid = Some((
1534                                raw.clone(),
1535                                vec!["$: response envelope was missing the expected value field"
1536                                    .to_string()],
1537                            ));
1538                        }
1539                    }
1540                    Err(errors) => {
1541                        if invalid.is_none() {
1542                            invalid = Some((raw.clone(), errors));
1543                        }
1544                    }
1545                }
1546            }
1547        }
1548    }
1549    StructuredResolution {
1550        valid: None,
1551        invalid,
1552        raw_seen,
1553    }
1554}
1555
1556/// UTF-8-safe truncation to at most `max` bytes (never splits a multibyte char —
1557/// repair prompts echo arbitrary model output, including CJK).
1558fn truncate_utf8(s: &str, max: usize) -> &str {
1559    if s.len() <= max {
1560        return s;
1561    }
1562    let mut end = max;
1563    while end > 0 && !s.is_char_boundary(end) {
1564        end -= 1;
1565    }
1566    &s[..end]
1567}
1568
1569/// Repair prompt for when nothing parseable was produced at all.
1570fn build_parse_failure_repair(raw_text: &str) -> String {
1571    if raw_text.trim().is_empty() {
1572        return "Your previous response contained no JSON. Respond with ONLY a single valid JSON object that matches the schema — no prose, no markdown, no analysis, and put the object in your reply content (not in a thinking/reasoning aside).".to_string();
1573    }
1574    format!(
1575        "Your previous output could not be parsed as a JSON object:\n\n{}\n\nReturn ONLY a single valid JSON object matching the schema — no prose, no markdown.",
1576        truncate_utf8(raw_text, 2000)
1577    )
1578}
1579
1580fn build_repair_message(raw_text: &str, errors: &[String]) -> String {
1581    // Truncate raw output in repair message to avoid blowing context
1582    let truncated_raw = if raw_text.len() > 2000 {
1583        format!(
1584            "{}...[truncated, {} bytes total]",
1585            truncate_utf8(raw_text, 2000),
1586            raw_text.len()
1587        )
1588    } else {
1589        raw_text.to_string()
1590    };
1591    format!(
1592        "Your previous output failed schema validation:\n\n{}\n\nValidation errors:\n{}\n\nPlease return ONLY a corrected JSON object that fixes these errors. No explanation, no markdown.",
1593        truncated_raw,
1594        errors.iter().map(|e| format!("- {}", e)).collect::<Vec<_>>().join("\n")
1595    )
1596}
1597
1598fn accumulate_usage(total: &mut TokenUsage, delta: &TokenUsage) {
1599    total.prompt_tokens += delta.prompt_tokens;
1600    total.completion_tokens += delta.completion_tokens;
1601    total.total_tokens += delta.total_tokens;
1602}
1603
1604/// Append repair context to the message history, respecting conversation structure.
1605///
1606/// In tool mode, the LLM returned a tool_use block. The correct follow-up is:
1607///   assistant (tool_use) → user (tool_result with error) → assistant (retry)
1608/// In text modes, it's simply:
1609///   assistant (text) → user (repair request) → assistant (retry)
1610fn append_repair_context(
1611    messages: &mut Vec<Message>,
1612    assistant_msg: &Message,
1613    repair_text: &str,
1614    mode: StructuredMode,
1615    _raw_text: &str,
1616) {
1617    if mode == StructuredMode::Tool {
1618        // Push the original assistant message (with tool_use block intact)
1619        messages.push(assistant_msg.clone());
1620        // Find the tool_use ID to construct a proper tool_result
1621        let tool_use_id = assistant_msg
1622            .tool_calls()
1623            .first()
1624            .map(|tc| tc.id.clone())
1625            .unwrap_or_else(|| "unknown".to_string());
1626        // Return the error as a tool_result so the conversation stays valid
1627        messages.push(Message::tool_result(&tool_use_id, repair_text, true));
1628    } else {
1629        // Text modes: push assistant text then user repair request
1630        messages.push(assistant_msg.clone());
1631        messages.push(Message::user(repair_text));
1632    }
1633}
1634
1635// ---------------------------------------------------------------------------
1636// Tests
1637// ---------------------------------------------------------------------------
1638
1639#[cfg(test)]
1640#[path = "structured_tests.rs"]
1641mod structured_tests;