nika_engine/runtime/
output.rs

1//! Output Handling - task result processing
2//!
3//! Extracted from runner.rs for cleaner separation:
4//! - `make_task_result`: Convert raw output to TaskResult with format handling
5//! - `validate_schema`: Validate JSON output against JSON Schema (with caching)
6//! - `validate_schema_ref`: Validate against SchemaRef (inline or file)
7//! - `extract_json_from_output`: Extract JSON from markdown code blocks
8//! - `format_validation_errors`: Format errors for retry feedback
9
10use std::sync::{Arc, LazyLock};
11
12use dashmap::DashMap;
13use jsonschema::Validator;
14use serde_json::Value;
15
16use crate::ast::output::SchemaRef;
17use crate::ast::OutputFormat;
18use crate::error::NikaError;
19use crate::store::TaskResult;
20
21/// Maximum number of entries in each cache before eviction.
22const CACHE_MAX_ENTRIES: usize = 512;
23
24/// Global schema cache: path → parsed JSON schema
25/// Avoids re-reading and re-parsing schema files on repeated validations.
26static SCHEMA_CACHE: LazyLock<DashMap<Arc<str>, Arc<Value>>> = LazyLock::new(DashMap::new);
27
28/// Global compiled validator cache: blake3 hash of schema JSON → compiled validator.
29/// Uses blake3 for cryptographic collision resistance (replaces DefaultHasher u64).
30/// Avoids recompiling the same JSON Schema on every validation call (10-50ms each).
31static VALIDATOR_CACHE: LazyLock<DashMap<String, Arc<Validator>>> = LazyLock::new(DashMap::new);
32
33/// Compute a blake3 hash of a JSON Schema value for use as cache key.
34/// Returns a hex-encoded hash string with cryptographic collision resistance.
35fn hash_schema(schema: &Value) -> String {
36    let canonical = serde_json::to_string(schema).unwrap_or_default();
37    blake3::hash(canonical.as_bytes()).to_hex().to_string()
38}
39
40/// Get or compile a JSON Schema validator, using the global cache.
41fn get_or_compile_validator(schema: &Value) -> Result<Arc<Validator>, NikaError> {
42    let key = hash_schema(schema);
43    if let Some(cached) = VALIDATOR_CACHE.get(&key) {
44        return Ok(Arc::clone(cached.value()));
45    }
46    let compiled = jsonschema::validator_for(schema).map_err(|e| NikaError::SchemaFailed {
47        details: format!("Invalid schema: {e}"),
48    })?;
49    let compiled = Arc::new(compiled);
50    // Evict all entries if cache exceeds size cap to prevent unbounded growth
51    if VALIDATOR_CACHE.len() >= CACHE_MAX_ENTRIES {
52        VALIDATOR_CACHE.clear();
53    }
54    VALIDATOR_CACHE.insert(key, Arc::clone(&compiled));
55    Ok(compiled)
56}
57
58/// Extract JSON from LLM output, handling markdown code blocks.
59///
60/// LLMs often wrap JSON in markdown code blocks like:
61/// ```json
62/// {"key": "value"}
63/// ```
64///
65/// This function tries multiple strategies:
66/// 1. Direct JSON parsing (fast path)
67/// 2. Extract from ```json ... ``` blocks
68/// 3. Extract from ``` ... ``` blocks (no language)
69/// 4. Find outermost { } or [ ] brackets
70fn extract_json_from_output(output: &str) -> Result<Value, String> {
71    let trimmed = output.trim();
72
73    // Strategy 1: Direct parse (fast path for well-behaved LLMs)
74    if let Ok(v) = serde_json::from_str::<Value>(trimmed) {
75        return Ok(v);
76    }
77
78    // Strategy 2: Extract from ```json ... ``` blocks
79    if let Some(start) = trimmed.find("```json") {
80        let after_marker = &trimmed[start + 7..];
81        if let Some(end) = after_marker.find("```") {
82            let json_str = after_marker[..end].trim();
83            if let Ok(v) = serde_json::from_str::<Value>(json_str) {
84                return Ok(v);
85            }
86        }
87    }
88
89    // Strategy 3: Extract from ``` ... ``` blocks (no language specifier)
90    if let Some(start) = trimmed.find("```\n") {
91        let after_marker = &trimmed[start + 4..];
92        if let Some(end) = after_marker.find("```") {
93            let json_str = after_marker[..end].trim();
94            if let Ok(v) = serde_json::from_str::<Value>(json_str) {
95                return Ok(v);
96            }
97        }
98    }
99
100    // Strategy 4: Find outermost { } or [ ] brackets
101    let first_brace = trimmed.find('{');
102    let first_bracket = trimmed.find('[');
103
104    let (start_char, end_char, start_pos) = match (first_brace, first_bracket) {
105        (Some(b), Some(k)) if b < k => ('{', '}', b),
106        (Some(_), Some(k)) => ('[', ']', k),
107        (Some(b), None) => ('{', '}', b),
108        (None, Some(k)) => ('[', ']', k),
109        (None, None) => return Err("No JSON object or array found in output".to_string()),
110    };
111
112    // Find matching closing bracket (handle nesting)
113    let substr = &trimmed[start_pos..];
114    let mut depth = 0;
115    let mut end_pos = None;
116
117    for (i, c) in substr.char_indices() {
118        if c == start_char {
119            depth += 1;
120        } else if c == end_char {
121            depth -= 1;
122            if depth == 0 {
123                end_pos = Some(i + 1);
124                break;
125            }
126        }
127    }
128
129    if let Some(end) = end_pos {
130        let json_str = &substr[..end];
131        if let Ok(v) = serde_json::from_str::<Value>(json_str) {
132            return Ok(v);
133        }
134    }
135
136    // Strategy 4b: Greedy fallback — first '{'/[' to last '}'/']'
137    // Handles cases where braces inside JSON string values confuse depth tracking.
138    if let Some(last) = substr.rfind(end_char) {
139        let json_str = &substr[..last + 1];
140        if let Ok(v) = serde_json::from_str::<Value>(json_str) {
141            return Ok(v);
142        }
143    }
144
145    // All strategies failed - return original error
146    Err(format!(
147        "Failed to extract JSON from output. First 200 chars: {}",
148        crate::util::truncate_str(trimmed, 200)
149    ))
150}
151
152/// Convert execution output to TaskResult, parsing as JSON if output format is json.
153/// Also validates against schema if declared.
154///
155/// Empty output with JSON format returns `null` instead of failing.
156/// This enables graceful handling of commands that produce no output.
157pub async fn make_task_result(
158    output: String,
159    policy: Option<&crate::ast::OutputPolicy>,
160    duration: std::time::Duration,
161) -> TaskResult {
162    if let Some(policy) = policy {
163        if policy.format == OutputFormat::Json {
164            // Handle empty output gracefully - return null instead of error
165            if output.trim().is_empty() {
166                tracing::debug!(
167                    target: "nika::output",
168                    "Empty output with JSON format, returning null"
169                );
170                return TaskResult::success(Value::Null, duration);
171            }
172
173            // Parse as JSON, handling markdown code blocks
174            let json_value = match extract_json_from_output(&output) {
175                Ok(v) => v,
176                Err(e) => {
177                    return TaskResult::failed(
178                        format!("NIKA-060: Invalid JSON output: {}", e),
179                        duration,
180                    );
181                }
182            };
183
184            // Validate against schema if declared
185            if let Some(schema_ref) = &policy.schema {
186                if let Err(e) = validate_schema_ref(&json_value, schema_ref).await {
187                    return TaskResult::failed(e.to_string(), duration);
188                }
189            }
190
191            return TaskResult::success(json_value, duration);
192        }
193    }
194    TaskResult::success_str(output, duration)
195}
196
197/// Validate JSON value against a JSON Schema file (with caching)
198///
199/// Schema files are cached after first load to avoid repeated file I/O.
200pub async fn validate_schema(value: &Value, schema_path: &str) -> Result<(), NikaError> {
201    // Try cache first (fast path)
202    let schema = if let Some(cached) = SCHEMA_CACHE.get(schema_path) {
203        Arc::clone(cached.value())
204    } else {
205        // Cache miss: read and parse schema
206        let schema_str =
207            tokio::fs::read_to_string(schema_path)
208                .await
209                .map_err(|e| NikaError::SchemaFailed {
210                    details: format!("Failed to read schema '{}': {}", schema_path, e),
211                })?;
212
213        let schema: Value =
214            serde_json::from_str(&schema_str).map_err(|e| NikaError::SchemaFailed {
215                details: format!("Invalid JSON in schema '{}': {}", schema_path, e),
216            })?;
217
218        // Store in cache (evict all if exceeding size cap)
219        let schema = Arc::new(schema);
220        if SCHEMA_CACHE.len() >= CACHE_MAX_ENTRIES {
221            SCHEMA_CACHE.clear();
222        }
223        SCHEMA_CACHE.insert(Arc::from(schema_path), Arc::clone(&schema));
224        schema
225    };
226
227    // PERF: Use cached compiled validator (H4 — saves 10-50ms per call)
228    let compiled = get_or_compile_validator(&schema).map_err(|_| NikaError::SchemaFailed {
229        details: format!("Invalid schema '{}'", schema_path),
230    })?;
231
232    // Collect all validation errors
233    let errors: Vec<_> = compiled.iter_errors(value).collect();
234    if errors.is_empty() {
235        Ok(())
236    } else {
237        let error_msgs: Vec<String> = errors.iter().map(|e| e.to_string()).collect();
238        Err(NikaError::SchemaFailed {
239            details: error_msgs.join("; "),
240        })
241    }
242}
243
244/// Validate a JSON value against a SchemaRef (inline or file)
245pub async fn validate_schema_ref(value: &Value, schema_ref: &SchemaRef) -> Result<(), NikaError> {
246    match schema_ref {
247        SchemaRef::File(path) => validate_schema(value, path).await,
248        SchemaRef::Inline(schema) => validate_inline_schema(value, schema),
249    }
250}
251
252/// Validate against an inline JSON Schema
253pub fn validate_inline_schema(value: &Value, schema: &Value) -> Result<(), NikaError> {
254    // PERF: Use cached compiled validator (H4)
255    let compiled = get_or_compile_validator(schema)?;
256
257    let errors: Vec<_> = compiled.iter_errors(value).collect();
258    if errors.is_empty() {
259        Ok(())
260    } else {
261        let error_msgs: Vec<String> = errors
262            .iter()
263            .map(|e| format!("- {}: {}", e.instance_path, e))
264            .collect();
265        Err(NikaError::SchemaFailed {
266            details: format!("Output validation failed:\n{}", error_msgs.join("\n")),
267        })
268    }
269}
270
271/// Format validation errors for retry feedback to LLM
272pub fn format_validation_errors(value: &Value, schema: &Value) -> String {
273    // PERF: Use cached compiled validator (H4)
274    let compiled = match get_or_compile_validator(schema) {
275        Ok(c) => c,
276        Err(e) => return format!("{e}"),
277    };
278
279    let errors: Vec<_> = compiled.iter_errors(value).collect();
280    if errors.is_empty() {
281        return "No validation errors".to_string();
282    }
283
284    errors
285        .iter()
286        .map(|e| {
287            format!(
288                "- Path '{}': {} (got: {})",
289                e.instance_path,
290                e,
291                serde_json::to_string(&*e.instance).unwrap_or_default()
292            )
293        })
294        .collect::<Vec<_>>()
295        .join("\n")
296}
297
298/// Extract JSON from LLM output - public version for executor retry loop
299pub fn extract_json(output: &str) -> Result<Value, String> {
300    extract_json_from_output(output)
301}
302
303#[cfg(test)]
304mod tests {
305    use super::*;
306    use std::io::Write;
307    use std::time::Duration;
308    use tempfile::NamedTempFile;
309
310    #[tokio::test]
311    async fn schema_cache_works() {
312        // Create a temp schema file
313        let mut schema_file = NamedTempFile::new().unwrap();
314        writeln!(
315            schema_file,
316            r#"{{"type": "object", "properties": {{"name": {{"type": "string"}}}}}}"#
317        )
318        .unwrap();
319        let schema_path = schema_file.path().to_str().unwrap();
320
321        // First validation - cache miss
322        let value = serde_json::json!({"name": "test"});
323        assert!(validate_schema(&value, schema_path).await.is_ok());
324
325        // Second validation - cache hit (same path)
326        assert!(validate_schema(&value, schema_path).await.is_ok());
327
328        // Cache should have the entry
329        assert!(SCHEMA_CACHE.contains_key(schema_path));
330    }
331
332    #[tokio::test]
333    async fn schema_validation_rejects_invalid() {
334        let mut schema_file = NamedTempFile::new().unwrap();
335        writeln!(schema_file, r#"{{"type": "object", "properties": {{"age": {{"type": "number"}}}}, "required": ["age"]}}"#).unwrap();
336        let schema_path = schema_file.path().to_str().unwrap();
337
338        // Missing required field
339        let value = serde_json::json!({"name": "test"});
340        assert!(validate_schema(&value, schema_path).await.is_err());
341
342        // Correct value
343        let value = serde_json::json!({"age": 25});
344        assert!(validate_schema(&value, schema_path).await.is_ok());
345    }
346
347    #[tokio::test]
348    async fn make_task_result_validates_json_file_schema() {
349        use crate::ast::OutputPolicy;
350
351        let mut schema_file = NamedTempFile::new().unwrap();
352        writeln!(schema_file, r#"{{"type": "object"}}"#).unwrap();
353        let schema_path = schema_file.path().to_string_lossy().to_string();
354
355        let policy = OutputPolicy {
356            format: OutputFormat::Json,
357            schema: Some(SchemaRef::File(schema_path)),
358            from_example: None,
359            max_retries: None,
360            source_structured_spec: None,
361        };
362
363        // Valid JSON object
364        let result = make_task_result(
365            r#"{"key": "value"}"#.to_string(),
366            Some(&policy),
367            Duration::from_millis(100),
368        )
369        .await;
370        assert!(result.is_success());
371
372        // Invalid JSON
373        let result = make_task_result(
374            "not json".to_string(),
375            Some(&policy),
376            Duration::from_millis(100),
377        )
378        .await;
379        assert!(!result.is_success());
380    }
381
382    #[tokio::test]
383    async fn make_task_result_validates_json_inline_schema() {
384        use crate::ast::OutputPolicy;
385
386        let inline_schema = serde_json::json!({
387            "type": "object",
388            "properties": {
389                "name": { "type": "string" }
390            },
391            "required": ["name"]
392        });
393
394        let policy = OutputPolicy {
395            format: OutputFormat::Json,
396            schema: Some(SchemaRef::Inline(inline_schema)),
397            from_example: None,
398            max_retries: None,
399            source_structured_spec: None,
400        };
401
402        // Valid JSON with required field
403        let result = make_task_result(
404            r#"{"name": "test"}"#.to_string(),
405            Some(&policy),
406            Duration::from_millis(100),
407        )
408        .await;
409        assert!(result.is_success());
410
411        // Invalid JSON missing required field
412        let result = make_task_result(
413            r#"{"other": "value"}"#.to_string(),
414            Some(&policy),
415            Duration::from_millis(100),
416        )
417        .await;
418        assert!(!result.is_success());
419    }
420
421    // ══════════════════════════════════════════════════════════════
422    // make_task_result EDGE CASES
423    // ══════════════════════════════════════════════════════════════
424
425    #[tokio::test]
426    async fn make_task_result_no_policy_returns_text() {
427        let result = make_task_result(
428            "plain text output".to_string(),
429            None,
430            Duration::from_millis(50),
431        )
432        .await;
433
434        assert!(result.is_success());
435        // Without policy, output should be stored as string (success_str)
436        assert_eq!(
437            result.output.as_ref(),
438            &serde_json::Value::String("plain text output".to_string())
439        );
440    }
441
442    #[tokio::test]
443    async fn make_task_result_json_no_schema_parses_json() {
444        use crate::ast::OutputPolicy;
445
446        let policy = OutputPolicy {
447            format: OutputFormat::Json,
448            schema: None, // No schema validation
449            from_example: None,
450            max_retries: None,
451            source_structured_spec: None,
452        };
453
454        let result = make_task_result(
455            r#"{"key": "value", "nested": {"a": 1}}"#.to_string(),
456            Some(&policy),
457            Duration::from_millis(50),
458        )
459        .await;
460
461        assert!(result.is_success());
462        // Should be parsed as JSON object, not string
463        assert!(result.output.is_object());
464        assert_eq!(result.output["key"], "value");
465        assert_eq!(result.output["nested"]["a"], 1);
466    }
467
468    #[tokio::test]
469    async fn make_task_result_invalid_json_returns_error_with_code() {
470        use crate::ast::OutputPolicy;
471
472        let policy = OutputPolicy {
473            format: OutputFormat::Json,
474            schema: None,
475            from_example: None,
476            max_retries: None,
477            source_structured_spec: None,
478        };
479
480        let result = make_task_result(
481            "{ invalid json".to_string(),
482            Some(&policy),
483            Duration::from_millis(50),
484        )
485        .await;
486
487        assert!(!result.is_success());
488        // Error should contain NIKA-060 code
489        let error_msg = result.error().expect("Should have error");
490        assert!(
491            error_msg.contains("NIKA-060"),
492            "Error should contain NIKA-060 code: {}",
493            error_msg
494        );
495    }
496
497    #[tokio::test]
498    async fn make_task_result_text_format_returns_raw_string() {
499        use crate::ast::OutputPolicy;
500
501        let policy = OutputPolicy {
502            format: OutputFormat::Text,
503            schema: None,
504            from_example: None,
505            max_retries: None,
506            source_structured_spec: None,
507        };
508
509        // Even valid JSON should be treated as text
510        let result = make_task_result(
511            r#"{"key": "value"}"#.to_string(),
512            Some(&policy),
513            Duration::from_millis(50),
514        )
515        .await;
516
517        assert!(result.is_success());
518        // Should be stored as string, not parsed JSON
519        assert!(result.output.is_string());
520        assert_eq!(
521            result.output.as_ref(),
522            &serde_json::Value::String(r#"{"key": "value"}"#.to_string())
523        );
524    }
525
526    // ══════════════════════════════════════════════════════════════
527    // validate_schema ERROR PATHS
528    // ══════════════════════════════════════════════════════════════
529
530    #[tokio::test]
531    async fn validate_schema_file_not_found_returns_error() {
532        let value = serde_json::json!({"name": "test"});
533        let result = validate_schema(&value, "/nonexistent/path/schema.json").await;
534
535        assert!(result.is_err());
536        let err = result.unwrap_err();
537        let err_string = err.to_string();
538        assert!(
539            err_string.contains("Failed to read schema"),
540            "Error should mention file read failure: {}",
541            err_string
542        );
543    }
544
545    #[tokio::test]
546    async fn validate_schema_invalid_json_in_schema_file() {
547        let mut schema_file = NamedTempFile::new().unwrap();
548        writeln!(schema_file, "{{ not valid json").unwrap();
549        let schema_path = schema_file.path().to_str().unwrap();
550
551        let value = serde_json::json!({"name": "test"});
552        let result = validate_schema(&value, schema_path).await;
553
554        assert!(result.is_err());
555        let err = result.unwrap_err();
556        let err_string = err.to_string();
557        assert!(
558            err_string.contains("Invalid JSON in schema"),
559            "Error should mention invalid JSON: {}",
560            err_string
561        );
562    }
563
564    #[tokio::test]
565    async fn validate_schema_invalid_schema_structure() {
566        let mut schema_file = NamedTempFile::new().unwrap();
567        // Valid JSON but not a valid JSON Schema (type must be a string, not number)
568        writeln!(schema_file, r#"{{"type": 123}}"#).unwrap();
569        let schema_path = schema_file.path().to_str().unwrap();
570
571        let value = serde_json::json!({"name": "test"});
572        let result = validate_schema(&value, schema_path).await;
573
574        assert!(result.is_err());
575        let err = result.unwrap_err();
576        let err_string = err.to_string();
577        assert!(
578            err_string.contains("Invalid schema"),
579            "Error should mention invalid schema: {}",
580            err_string
581        );
582    }
583
584    #[tokio::test]
585    async fn validate_schema_multiple_validation_errors() {
586        let mut schema_file = NamedTempFile::new().unwrap();
587        writeln!(
588            schema_file,
589            r#"{{
590                "type": "object",
591                "properties": {{
592                    "name": {{"type": "string"}},
593                    "age": {{"type": "number"}}
594                }},
595                "required": ["name", "age"]
596            }}"#
597        )
598        .unwrap();
599        let schema_path = schema_file.path().to_str().unwrap();
600
601        // Missing both required fields
602        let value = serde_json::json!({});
603        let result = validate_schema(&value, schema_path).await;
604
605        assert!(result.is_err());
606        let err = result.unwrap_err();
607        let err_string = err.to_string();
608        // Should mention both missing fields
609        assert!(
610            err_string.contains("name") || err_string.contains("required"),
611            "Error should mention validation issues: {}",
612            err_string
613        );
614    }
615
616    // ══════════════════════════════════════════════════════════════
617    // EDGE CASES
618    // ══════════════════════════════════════════════════════════════
619
620    #[tokio::test]
621    async fn make_task_result_large_json_output() {
622        use crate::ast::OutputPolicy;
623
624        let policy = OutputPolicy {
625            format: OutputFormat::Json,
626            schema: None,
627            from_example: None,
628            max_retries: None,
629            source_structured_spec: None,
630        };
631
632        // Generate large JSON array
633        let large_array: Vec<i32> = (0..10000).collect();
634        let json_str = serde_json::to_string(&large_array).unwrap();
635
636        let result = make_task_result(json_str, Some(&policy), Duration::from_millis(100)).await;
637
638        assert!(result.is_success());
639        assert!(result.output.is_array());
640        assert_eq!(result.output.as_array().unwrap().len(), 10000);
641    }
642
643    #[tokio::test]
644    async fn make_task_result_unicode_content() {
645        use crate::ast::OutputPolicy;
646
647        let policy = OutputPolicy {
648            format: OutputFormat::Json,
649            schema: None,
650            from_example: None,
651            max_retries: None,
652            source_structured_spec: None,
653        };
654
655        // JSON with various Unicode characters
656        let json_str = r#"{"greeting": "你好世界", "emoji": "🚀✨", "japanese": "こんにちは"}"#;
657
658        let result = make_task_result(
659            json_str.to_string(),
660            Some(&policy),
661            Duration::from_millis(50),
662        )
663        .await;
664
665        assert!(result.is_success());
666        assert_eq!(result.output["greeting"], "你好世界");
667        assert_eq!(result.output["emoji"], "🚀✨");
668        assert_eq!(result.output["japanese"], "こんにちは");
669    }
670
671    #[tokio::test]
672    async fn schema_cache_concurrent_access() {
673        // Create a temp schema file
674        let mut schema_file = NamedTempFile::new().unwrap();
675        writeln!(schema_file, r#"{{"type": "object"}}"#).unwrap();
676        let schema_path = schema_file.path().to_str().unwrap().to_string();
677
678        // Spawn multiple concurrent validation tasks
679        let handles: Vec<_> = (0..10)
680            .map(|i| {
681                let path = schema_path.clone();
682                tokio::spawn(async move {
683                    let value = serde_json::json!({"id": i});
684                    validate_schema(&value, &path).await
685                })
686            })
687            .collect();
688
689        // All should succeed
690        for handle in handles {
691            let result = handle.await.unwrap();
692            assert!(result.is_ok());
693        }
694    }
695
696    #[tokio::test]
697    async fn make_task_result_preserves_duration() {
698        let duration = Duration::from_secs(5);
699        let result = make_task_result("output".to_string(), None, duration).await;
700
701        assert_eq!(result.duration, duration);
702    }
703
704    #[tokio::test]
705    async fn make_task_result_json_array() {
706        use crate::ast::OutputPolicy;
707
708        let policy = OutputPolicy {
709            format: OutputFormat::Json,
710            schema: None,
711            from_example: None,
712            max_retries: None,
713            source_structured_spec: None,
714        };
715
716        let result = make_task_result(
717            r#"[1, 2, 3, "four"]"#.to_string(),
718            Some(&policy),
719            Duration::from_millis(50),
720        )
721        .await;
722
723        assert!(result.is_success());
724        assert!(result.output.is_array());
725        let arr = result.output.as_array().unwrap();
726        assert_eq!(arr.len(), 4);
727        assert_eq!(arr[3], "four");
728    }
729
730    // ══════════════════════════════════════════════════════════════
731    // EXTRACT_JSON_FROM_OUTPUT TESTS
732    // ══════════════════════════════════════════════════════════════
733
734    #[test]
735    fn extract_json_direct_parse() {
736        let input = r#"{"key": "value"}"#;
737        let result = extract_json_from_output(input).unwrap();
738        assert_eq!(result["key"], "value");
739    }
740
741    #[test]
742    fn extract_json_with_whitespace() {
743        let input = r#"
744            {"key": "value"}
745        "#;
746        let result = extract_json_from_output(input).unwrap();
747        assert_eq!(result["key"], "value");
748    }
749
750    #[test]
751    fn extract_json_from_markdown_json_block() {
752        let input = r#"Here's the JSON:
753
754```json
755{"name": "Thibaut", "score": 42}
756```
757
758Hope this helps!"#;
759        let result = extract_json_from_output(input).unwrap();
760        assert_eq!(result["name"], "Thibaut");
761        assert_eq!(result["score"], 42);
762    }
763
764    #[test]
765    fn extract_json_from_markdown_plain_block() {
766        let input = r#"The result:
767
768```
769{"items": [1, 2, 3]}
770```
771"#;
772        let result = extract_json_from_output(input).unwrap();
773        assert!(result["items"].is_array());
774    }
775
776    #[test]
777    fn extract_json_from_prose_with_braces() {
778        let input = r#"I'll generate the fortune for you:
779
780The cosmic reading reveals: {"sign": "scorpio", "lucky_number": 7, "message": "Great things await"}
781
782This is based on ancient wisdom."#;
783        let result = extract_json_from_output(input).unwrap();
784        assert_eq!(result["sign"], "scorpio");
785        assert_eq!(result["lucky_number"], 7);
786    }
787
788    #[test]
789    fn extract_json_array_from_markdown() {
790        let input = r#"```json
791[{"id": 1}, {"id": 2}, {"id": 3}]
792```"#;
793        let result = extract_json_from_output(input).unwrap();
794        assert!(result.is_array());
795        assert_eq!(result.as_array().unwrap().len(), 3);
796    }
797
798    #[test]
799    fn extract_json_nested_objects() {
800        let input = r#"Result: {"outer": {"inner": {"deep": "value"}}}"#;
801        let result = extract_json_from_output(input).unwrap();
802        assert_eq!(result["outer"]["inner"]["deep"], "value");
803    }
804
805    #[test]
806    fn extract_json_with_escaped_braces_in_strings() {
807        let input = r#"{"template": "Use {{variable}} syntax", "count": 1}"#;
808        let result = extract_json_from_output(input).unwrap();
809        assert_eq!(result["template"], "Use {{variable}} syntax");
810    }
811
812    #[test]
813    fn extract_json_unbalanced_braces_in_string_values() {
814        // Edge case: closing brace inside a JSON string value confuses depth tracking.
815        // Strategy 4b (greedy fallback) should handle this.
816        let input = r#"Here is the result: {"msg": "close brace } here", "ok": true}"#;
817        let result = extract_json_from_output(input).unwrap();
818        assert_eq!(result["msg"], "close brace } here");
819        assert_eq!(result["ok"], true);
820    }
821
822    #[test]
823    fn extract_json_no_json_found() {
824        let input = "This is just plain text without any JSON.";
825        let result = extract_json_from_output(input);
826        assert!(result.is_err());
827        assert!(result
828            .unwrap_err()
829            .contains("No JSON object or array found"));
830    }
831
832    #[tokio::test]
833    async fn make_task_result_handles_markdown_wrapped_json() {
834        use crate::ast::OutputPolicy;
835
836        let policy = OutputPolicy {
837            format: OutputFormat::Json,
838            schema: None,
839            from_example: None,
840            max_retries: None,
841            source_structured_spec: None,
842        };
843
844        // Simulate LLM output with markdown code block
845        let llm_output = r#"Here's your fortune:
846
847```json
848{
849  "sign": "scorpio",
850  "lucky_number": 7,
851  "message": "The stars align in your favor"
852}
853```
854
855Enjoy your reading!"#;
856
857        let result = make_task_result(
858            llm_output.to_string(),
859            Some(&policy),
860            Duration::from_millis(100),
861        )
862        .await;
863
864        assert!(result.is_success(), "Should parse JSON from markdown block");
865        assert_eq!(result.output["sign"], "scorpio");
866        assert_eq!(result.output["lucky_number"], 7);
867    }
868
869    #[tokio::test]
870    async fn make_task_result_empty_output_returns_null() {
871        use crate::ast::OutputPolicy;
872
873        let policy = OutputPolicy {
874            format: OutputFormat::Json,
875            schema: None,
876            from_example: None,
877            max_retries: None,
878            source_structured_spec: None,
879        };
880
881        // Empty output with JSON format returns null
882        let empty_output = "".to_string();
883        let result = make_task_result(empty_output, Some(&policy), std::time::Duration::ZERO).await;
884
885        assert!(result.is_success(), "Empty output should succeed with null");
886        assert!(result.output.is_null(), "Empty output should return null");
887    }
888
889    #[tokio::test]
890    async fn make_task_result_whitespace_output_returns_null() {
891        use crate::ast::OutputPolicy;
892
893        let policy = OutputPolicy {
894            format: OutputFormat::Json,
895            schema: None,
896            from_example: None,
897            max_retries: None,
898            source_structured_spec: None,
899        };
900
901        // Whitespace-only output also returns null
902        let whitespace_output = "   \n\t  ".to_string();
903        let result =
904            make_task_result(whitespace_output, Some(&policy), std::time::Duration::ZERO).await;
905
906        assert!(
907            result.is_success(),
908            "Whitespace-only output should succeed with null"
909        );
910        assert!(
911            result.output.is_null(),
912            "Whitespace-only output should return null"
913        );
914    }
915
916    // ══════════════════════════════════════════════════════════════
917    // INLINE SCHEMA VALIDATION TESTS
918    // ══════════════════════════════════════════════════════════════
919
920    #[tokio::test]
921    async fn validate_schema_ref_inline_success() {
922        let schema = serde_json::json!({
923            "type": "object",
924            "properties": {
925                "name": { "type": "string" }
926            },
927            "required": ["name"]
928        });
929        let value = serde_json::json!({"name": "test"});
930        let result = validate_schema_ref(&value, &SchemaRef::Inline(schema)).await;
931        assert!(result.is_ok());
932    }
933
934    #[tokio::test]
935    async fn validate_schema_ref_inline_failure() {
936        let schema = serde_json::json!({
937            "type": "object",
938            "properties": {
939                "name": { "type": "string" }
940            },
941            "required": ["name"]
942        });
943        let value = serde_json::json!({"other": "field"});
944        let result = validate_schema_ref(&value, &SchemaRef::Inline(schema)).await;
945        assert!(result.is_err());
946        let err = result.unwrap_err().to_string();
947        assert!(
948            err.contains("required") || err.contains("name"),
949            "Error should mention missing required field: {}",
950            err
951        );
952    }
953
954    #[test]
955    fn format_validation_errors_shows_details() {
956        let schema = serde_json::json!({
957            "type": "object",
958            "properties": {
959                "age": { "type": "integer", "minimum": 0 }
960            },
961            "required": ["age"]
962        });
963        let value = serde_json::json!({"age": -5});
964        let errors = format_validation_errors(&value, &schema);
965        assert!(errors.contains("-5"), "Should show the invalid value");
966    }
967}
nika_engine/runtime/output.rs

nika_engine/runtime/
output.rs
