Skip to main content

ironflow_core/providers/claude/
common.rs

1//! Shared utilities for all Claude Code transport providers.
2//!
3//! This module contains command-line argument building, JSON response parsing,
4//! and structured-output extraction logic shared across local, SSH, Docker,
5//! and Kubernetes transports.
6
7use std::env;
8#[cfg(any(
9    feature = "transport-ssh",
10    feature = "transport-docker",
11    feature = "transport-k8s"
12))]
13use std::sync::Arc;
14use std::time::Duration;
15
16use serde::Deserialize;
17use serde_json::{Map, Value};
18#[cfg(any(
19    feature = "transport-ssh",
20    feature = "transport-docker",
21    feature = "transport-k8s"
22))]
23use tracing::debug;
24use tracing::{trace, warn};
25
26use crate::error::{AgentError, PartialUsage};
27use crate::operations::agent::PermissionMode;
28#[cfg(any(
29    feature = "transport-ssh",
30    feature = "transport-docker",
31    feature = "transport-k8s"
32))]
33use crate::provider::LogSink;
34use crate::provider::{AgentConfig, AgentOutput, DebugMessage, DebugToolCall, DebugToolResult};
35use crate::schema_transform::transform_schema;
36use crate::utils::estimate_tokens;
37
38/// Default timeout for a single Claude CLI invocation (5 minutes).
39pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(300);
40
41/// Byte threshold above which the prompt is piped via stdin instead of
42/// passed as a `-p` CLI argument. This avoids the OS `ARG_MAX` limit
43/// (typically 256 KB on macOS, 2 MB on Linux) when reviewing large diffs.
44pub const PROMPT_STDIN_THRESHOLD: usize = 100_000;
45
46/// Maximum byte length for raw response data included in error diagnostics.
47const RAW_RESPONSE_MAX_LEN: usize = 4000;
48
49/// Approximate byte limit for stdout fallback in error diagnostics.
50const RAW_RESPONSE_FALLBACK_MAX_LEN: usize = 2000;
51
52/// Maximum byte length for error detail in `ProcessFailed` errors.
53///
54/// 50 KB is enough for diagnostic context while staying well under typical
55/// DB column and API payload limits.
56const ERROR_DETAIL_MAX_LEN: usize = 50_000;
57
58/// Truncate a string to at most `max_len` bytes on a char boundary.
59fn truncate_to(s: &str, max_len: usize) -> String {
60    if s.len() <= max_len {
61        return s.to_string();
62    }
63    let end = s.floor_char_boundary(max_len);
64    format!("{}...(truncated)", &s[..end])
65}
66
67/// Return the context window size for a known Claude model identifier.
68///
69/// Returns `200_000` for all standard models, `1_000_000` for `[1m]` variants,
70/// and `200_000` as a safe default for unrecognised identifiers.
71pub fn context_window_for_model(model: &str) -> usize {
72    if model.ends_with("[1m]") {
73        1_000_000
74    } else {
75        200_000
76    }
77}
78
79/// Validate that the prompt fits within the Claude model's context window.
80///
81/// # Errors
82///
83/// Returns [`AgentError::PromptTooLarge`] if the estimated token count exceeds
84/// the model's context window.
85pub fn validate_prompt_size(config: &AgentConfig) -> Result<(), AgentError> {
86    let total_chars = config.prompt.len() + config.system_prompt.as_ref().map_or(0, |s| s.len());
87    let estimated_tokens = estimate_tokens(total_chars);
88    let model_limit = context_window_for_model(&config.model);
89    if estimated_tokens > model_limit {
90        return Err(AgentError::PromptTooLarge {
91            chars: total_chars,
92            estimated_tokens,
93            model_limit,
94        });
95    }
96    Ok(())
97}
98
99/// Parsed JSON output from the `claude` CLI.
100#[derive(Deserialize)]
101pub struct ClaudeJsonOutput {
102    /// Conversation session identifier for resuming multi-turn calls.
103    pub session_id: Option<String>,
104    /// Response subtype (e.g. `"success"`, `"error_max_budget_usd"`).
105    pub subtype: Option<String>,
106    /// The model's text response, if any.
107    pub result: Option<Value>,
108    /// Typed JSON output when a JSON schema was requested.
109    pub structured_output: Option<Value>,
110    /// Token usage breakdown.
111    pub usage: Option<ClaudeUsage>,
112    /// Total cost in USD for this invocation.
113    pub total_cost_usd: Option<f64>,
114    /// Wall-clock duration in milliseconds.
115    pub duration_ms: Option<u64>,
116    /// Per-model token usage keyed by model identifier.
117    #[serde(rename = "modelUsage")]
118    pub model_usage: Option<Map<String, Value>>,
119}
120
121/// Token usage statistics from the `claude` CLI.
122#[derive(Deserialize)]
123pub struct ClaudeUsage {
124    /// Direct input tokens consumed.
125    pub input_tokens: Option<u64>,
126    /// Output tokens generated.
127    pub output_tokens: Option<u64>,
128    /// Tokens used to populate the prompt cache.
129    pub cache_creation_input_tokens: Option<u64>,
130    /// Tokens served from the prompt cache.
131    pub cache_read_input_tokens: Option<u64>,
132}
133
134impl ClaudeUsage {
135    /// Total input tokens including cache creation and read tokens.
136    pub fn total_input_tokens(&self) -> u64 {
137        self.input_tokens.unwrap_or(0)
138            + self.cache_creation_input_tokens.unwrap_or(0)
139            + self.cache_read_input_tokens.unwrap_or(0)
140    }
141
142    /// Total output tokens.
143    pub fn total_output_tokens(&self) -> u64 {
144        self.output_tokens.unwrap_or(0)
145    }
146}
147
148/// Environment variable names that must be removed before spawning the
149/// `claude` CLI to prevent sub-agent mode interference.
150///
151/// When ironflow runs inside Claude Code (or cmux), the child process inherits
152/// variables like `CLAUDE_CODE_ENTRYPOINT`, `CLAUDE_CODE_SUBAGENT_MODEL`, etc.
153/// that force degraded/sub-agent behaviour, wrong models, or altered context
154/// handling. We strip all `CLAUDE*` vars plus `IRONFLOW_ALLOW_BYPASS`.
155///
156/// # Examples
157///
158/// ```no_run
159/// # fn example() {
160/// let vars = ironflow_core::providers::claude::common::env_vars_to_remove();
161/// assert!(vars.contains(&"IRONFLOW_ALLOW_BYPASS".to_string()));
162/// # }
163/// ```
164pub fn env_vars_to_remove() -> Vec<String> {
165    collect_vars_to_remove(env::vars().map(|(k, _)| k))
166}
167
168/// Filter environment variable names, keeping `CLAUDE*` prefixed ones
169/// and always including `IRONFLOW_ALLOW_BYPASS`.
170fn collect_vars_to_remove(keys: impl Iterator<Item = String>) -> Vec<String> {
171    let mut vars: Vec<String> = keys.filter(|key| key.starts_with("CLAUDE")).collect();
172    vars.push("IRONFLOW_ALLOW_BYPASS".to_string());
173    vars
174}
175
176/// Names of `CLAUDE*` env vars to unset in a remote shell command.
177///
178/// Returns a space-separated list suitable for `unset VAR1 VAR2 ...`.
179pub fn env_unset_shell_prefix() -> String {
180    let vars = env_vars_to_remove();
181    if vars.is_empty() {
182        return String::new();
183    }
184    format!("unset {} 2>/dev/null; ", vars.join(" "))
185}
186
187/// Return a clone of `config` with `verbose` forced to `true` when streaming
188/// is active but verbose is off. Returns `None` when no override is needed.
189#[cfg(any(
190    feature = "transport-ssh",
191    feature = "transport-docker",
192    feature = "transport-k8s"
193))]
194pub(super) fn force_verbose_for_streaming(
195    config: &AgentConfig,
196    streaming: bool,
197) -> Option<AgentConfig> {
198    if streaming && !config.verbose {
199        debug!("forcing verbose=true for log streaming (stream-json output required)");
200        Some(config.clone().verbose(true))
201    } else {
202        None
203    }
204}
205
206/// Forward raw output data to a [`LogSink`], splitting by line.
207#[cfg(any(
208    feature = "transport-ssh",
209    feature = "transport-docker",
210    feature = "transport-k8s"
211))]
212pub(super) fn stream_lines(data: &[u8], stream: &str, sink: Option<&Arc<dyn LogSink>>) {
213    if let Some(sink) = sink {
214        let text = String::from_utf8_lossy(data);
215        for line in text.lines() {
216            sink.log(stream, line);
217        }
218    }
219}
220
221/// Push a CLI flag and its value onto the argument list.
222pub fn push_flag(args: &mut Vec<String>, flag: &str, value: &str) {
223    args.push(flag.to_string());
224    args.push(value.to_string());
225}
226
227/// Push a CLI flag and its value onto the argument list, only if the value is `Some`.
228pub fn push_opt(args: &mut Vec<String>, flag: &str, value: &Option<impl ToString>) {
229    if let Some(v) = value {
230        push_flag(args, flag, &v.to_string());
231    }
232}
233
234/// Build the CLI argument list from an [`AgentConfig`].
235///
236/// Returns the list of arguments to pass after the `claude` binary name.
237///
238/// # Errors
239///
240/// Returns [`AgentError::ProcessFailed`] if `BypassPermissions` is requested
241/// without the `IRONFLOW_ALLOW_BYPASS=1` environment variable.
242pub fn build_args(config: &AgentConfig) -> Result<Vec<String>, AgentError> {
243    let output_format = if config.verbose {
244        "stream-json"
245    } else {
246        "json"
247    };
248
249    let mut args: Vec<String> = vec![
250        "-p".to_string(),
251        config.prompt.clone(),
252        "--output-format".to_string(),
253        output_format.to_string(),
254    ];
255
256    // Claude CLI requires --verbose when using --output-format=stream-json with -p
257    if config.verbose {
258        args.push("--verbose".to_string());
259    }
260
261    push_opt(&mut args, "--system-prompt", &config.system_prompt);
262    push_flag(&mut args, "--model", &config.model);
263    if !config.allowed_tools.is_empty() {
264        push_flag(&mut args, "--allowedTools", &config.allowed_tools.join(","));
265    }
266    if !config.disallowed_tools.is_empty() {
267        push_flag(
268            &mut args,
269            "--disallowedTools",
270            &config.disallowed_tools.join(","),
271        );
272    }
273    push_opt(&mut args, "--max-turns", &config.max_turns);
274    push_opt(&mut args, "--max-budget-usd", &config.max_budget_usd);
275    push_opt(&mut args, "--mcp-config", &config.mcp_config);
276    if config.strict_mcp_config {
277        args.push("--strict-mcp-config".to_string());
278    }
279    if config.bare {
280        args.push("--bare".to_string());
281    }
282
283    match config.permission_mode {
284        PermissionMode::Default => {}
285        PermissionMode::Auto => push_flag(&mut args, "--permission-mode", "auto"),
286        PermissionMode::DontAsk => push_flag(&mut args, "--permission-mode", "dontAsk"),
287        PermissionMode::BypassPermissions => {
288            if env::var("IRONFLOW_ALLOW_BYPASS").as_deref() != Ok("1") {
289                return Err(AgentError::ProcessFailed {
290                    exit_code: -1,
291                    stderr:
292                        "BypassPermissions requires IRONFLOW_ALLOW_BYPASS=1 environment variable"
293                            .to_string(),
294                });
295            }
296            warn!(
297                "using BypassPermissions: agent will have unrestricted filesystem and shell access"
298            );
299            args.push("--dangerously-skip-permissions".to_string());
300        }
301    }
302
303    let transformed_schema = config.json_schema.as_ref().map(|s| transform_schema(s));
304    push_opt(&mut args, "--json-schema", &transformed_schema);
305
306    if let Some(ref session_id) = config.resume_session_id {
307        args.push("--resume".to_string());
308        args.push(session_id.clone());
309    }
310
311    Ok(args)
312}
313
314/// CLI arguments and an optional prompt payload for stdin piping.
315///
316/// When the prompt exceeds [`PROMPT_STDIN_THRESHOLD`], it is excluded from
317/// the argument list and placed in [`stdin_prompt`](Self::stdin_prompt)
318/// instead. Each transport provider is responsible for piping it to the
319/// `claude` process's stdin.
320#[derive(Debug)]
321pub struct BuiltCommand {
322    /// Arguments to pass after the `claude` binary name.
323    pub args: Vec<String>,
324    /// Prompt text that must be written to stdin. `None` when the prompt
325    /// is small enough to fit in the CLI arguments.
326    pub stdin_prompt: Option<String>,
327}
328
329/// Build CLI arguments, splitting the prompt to stdin when it is too large
330/// for the OS argument limit.
331///
332/// Providers should use this instead of [`build_args`] to handle large
333/// prompts gracefully.
334///
335/// # Errors
336///
337/// Returns [`AgentError::ProcessFailed`] if `BypassPermissions` is requested
338/// without the `IRONFLOW_ALLOW_BYPASS=1` environment variable.
339pub fn build_command(config: &AgentConfig) -> Result<BuiltCommand, AgentError> {
340    let prompt_via_stdin = config.prompt.len() > PROMPT_STDIN_THRESHOLD;
341
342    let output_format = if config.verbose {
343        "stream-json"
344    } else {
345        "json"
346    };
347
348    let mut args: Vec<String> = Vec::with_capacity(16);
349    if !prompt_via_stdin {
350        args.push("-p".to_string());
351        args.push(config.prompt.clone());
352    }
353    args.push("--output-format".to_string());
354    args.push(output_format.to_string());
355
356    if config.verbose {
357        args.push("--verbose".to_string());
358    }
359
360    push_opt(&mut args, "--system-prompt", &config.system_prompt);
361    push_flag(&mut args, "--model", &config.model);
362    if !config.allowed_tools.is_empty() {
363        push_flag(&mut args, "--allowedTools", &config.allowed_tools.join(","));
364    }
365    if !config.disallowed_tools.is_empty() {
366        push_flag(
367            &mut args,
368            "--disallowedTools",
369            &config.disallowed_tools.join(","),
370        );
371    }
372    push_opt(&mut args, "--max-turns", &config.max_turns);
373    push_opt(&mut args, "--max-budget-usd", &config.max_budget_usd);
374    push_opt(&mut args, "--mcp-config", &config.mcp_config);
375    if config.strict_mcp_config {
376        args.push("--strict-mcp-config".to_string());
377    }
378    if config.bare {
379        args.push("--bare".to_string());
380    }
381
382    match config.permission_mode {
383        PermissionMode::Default => {}
384        PermissionMode::Auto => push_flag(&mut args, "--permission-mode", "auto"),
385        PermissionMode::DontAsk => push_flag(&mut args, "--permission-mode", "dontAsk"),
386        PermissionMode::BypassPermissions => {
387            if env::var("IRONFLOW_ALLOW_BYPASS").as_deref() != Ok("1") {
388                return Err(AgentError::ProcessFailed {
389                    exit_code: -1,
390                    stderr:
391                        "BypassPermissions requires IRONFLOW_ALLOW_BYPASS=1 environment variable"
392                            .to_string(),
393                });
394            }
395            warn!(
396                "using BypassPermissions: agent will have unrestricted filesystem and shell access"
397            );
398            args.push("--dangerously-skip-permissions".to_string());
399        }
400    }
401
402    let transformed_schema = config.json_schema.as_ref().map(|s| transform_schema(s));
403    push_opt(&mut args, "--json-schema", &transformed_schema);
404
405    if let Some(ref session_id) = config.resume_session_id {
406        args.push("--resume".to_string());
407        args.push(session_id.clone());
408    }
409
410    Ok(BuiltCommand {
411        args,
412        stdin_prompt: if prompt_via_stdin {
413            Some(config.prompt.clone())
414        } else {
415            None
416        },
417    })
418}
419
420/// Build a single shell command string from the `claude` binary path and arguments.
421///
422/// Each argument is escaped with single quotes for safe remote execution via `sh -c`.
423pub fn build_shell_command(claude_path: &str, args: &[String]) -> String {
424    let mut parts = vec![shell_escape(claude_path)];
425    for arg in args {
426        parts.push(shell_escape(arg));
427    }
428    parts.join(" ")
429}
430
431/// Escape a string for safe inclusion in a single-quoted shell argument.
432///
433/// Wraps the value in single quotes, escaping any embedded single quotes
434/// using the `'\''` idiom.
435fn shell_escape(s: &str) -> String {
436    format!("'{}'", s.replace('\'', "'\\''"))
437}
438
439/// Extract a structured JSON value from a parsed Claude CLI response.
440///
441/// Prefers `structured_output`; falls back to parsing `result` as JSON
442/// (direct parse, code-fence extraction, or brace extraction).
443///
444/// # Why the fallbacks exist
445///
446/// Claude CLI has several known bugs around structured output that make
447/// the `structured_output` field unreliable:
448///
449/// - When tools are used alongside `--json-schema`, `structured_output`
450///   is always `null` (the result lands in `result` as markdown text).
451///   See <https://github.com/anthropics/claude-code/issues/18536>.
452///   This case is blocked at compile time by the typestate, but defensive
453///   fallbacks remain for forward compatibility.
454/// - The CLI does not validate output against the schema; it may return
455///   malformed or non-conforming JSON.
456///   See <https://github.com/anthropics/claude-code/issues/9058>.
457/// - Wrapper objects with a single array field may be flattened to a bare
458///   array non-deterministically.
459///   See <https://github.com/anthropics/claude-agent-sdk-python/issues/502>
460///   and <https://github.com/anthropics/claude-agent-sdk-python/issues/374>.
461///
462/// Because of these issues, we try multiple extraction strategies in order:
463/// 1. `structured_output` field (when non-null)
464/// 2. Direct JSON parse of `result`
465/// 3. JSON code fence extraction from `result`
466/// 4. First `{...}` brace extraction from `result`
467pub fn extract_structured_value(parsed: &ClaudeJsonOutput) -> Option<Value> {
468    let from_structured = parsed.structured_output.as_ref().filter(|v| !v.is_null());
469    if let Some(v) = from_structured {
470        return Some(v.clone());
471    }
472
473    let text = parsed.result.as_ref()?.as_str()?;
474
475    if let Ok(v) = serde_json::from_str(text) {
476        return Some(v);
477    }
478
479    if let Some(start) = text.find("```json") {
480        let json_start = start + "```json".len();
481        if let Some(end) = text[json_start..].find("```") {
482            let json_str = text[json_start..json_start + end].trim();
483            if let Ok(v) = serde_json::from_str(json_str) {
484                return Some(v);
485            }
486        }
487    }
488
489    let start = text.find('{')?;
490    let end = text.rfind('}')?;
491    serde_json::from_str(&text[start..=end]).ok()
492}
493
494/// Extract the raw response text from a parsed CLI response for error diagnostics.
495///
496/// Prefers the `result` field (stringified and truncated); falls back to
497/// raw stdout when `result` is null.
498fn extract_raw_response_text(parsed: &ClaudeJsonOutput, stdout: &str) -> Option<String> {
499    if let Some(ref result) = parsed.result
500        && !result.is_null()
501    {
502        let text = match result.as_str() {
503            Some(s) => s.to_string(),
504            None => result.to_string(),
505        };
506        return Some(truncate_to(&text, RAW_RESPONSE_MAX_LEN));
507    }
508    if !stdout.is_empty() {
509        return Some(truncate_to(stdout, RAW_RESPONSE_FALLBACK_MAX_LEN));
510    }
511    None
512}
513
514/// Parse raw stdout from the `claude` CLI into an [`AgentOutput`].
515///
516/// # Errors
517///
518/// Returns [`AgentError::SchemaValidation`] if the JSON cannot be parsed or
519/// if structured output was requested but not present in the response.
520pub fn parse_response(
521    stdout: &str,
522    config: &AgentConfig,
523    fallback_duration_ms: u64,
524) -> Result<AgentOutput, AgentError> {
525    let parsed: ClaudeJsonOutput =
526        serde_json::from_str(stdout).map_err(|e| AgentError::SchemaValidation {
527            expected: "ClaudeJsonOutput".to_string(),
528            got: format!("parse error: {e}"),
529            debug_messages: Vec::new(),
530            partial_usage: Box::default(),
531            raw_response: Some(truncate_to(stdout, RAW_RESPONSE_MAX_LEN)),
532        })?;
533
534    let value = if config.json_schema.is_some() {
535        extract_structured_value(&parsed).ok_or_else(|| {
536            warn!(
537                subtype = ?parsed.subtype,
538                result_is_null = parsed.result.as_ref().is_none_or(|v| v.is_null()),
539                structured_output_is_null = parsed.structured_output.as_ref().is_none_or(|v| v.is_null()),
540                has_tools = !config.allowed_tools.is_empty(),
541                "structured_output extraction failed, dumping response fields for diagnosis"
542            );
543            if let Some(ref result) = parsed.result {
544                let preview = result.to_string();
545                let truncated = &preview[..preview.len().min(2000)];
546                warn!(result_preview = truncated, "result field content (truncated to 2000 chars)");
547            }
548            let usage = Box::new(PartialUsage {
549                cost_usd: parsed.total_cost_usd,
550                duration_ms: parsed.duration_ms,
551                input_tokens: parsed.usage.as_ref().map(|u| u.total_input_tokens()),
552                output_tokens: parsed.usage.as_ref().map(|u| u.total_output_tokens()),
553            });
554
555            // The budget running out is not a schema problem: retrying spends
556            // more money and cannot succeed. Report it as its own error so the
557            // retry layers can refuse to replay it.
558            if parsed.subtype.as_deref() == Some("error_max_budget_usd") {
559                return AgentError::BudgetExceeded {
560                    spent_usd: parsed.total_cost_usd.unwrap_or(0.0),
561                    limit_usd: config.max_budget_usd.unwrap_or(0.0),
562                    debug_messages: Vec::new(),
563                    partial_usage: usage,
564                };
565            }
566
567            let hint = match parsed.subtype.as_deref() {
568                Some("error_max_turns") => {
569                    " (max turns reached before structured output was generated - use max_turns >= 2 with structured output)"
570                }
571                Some(sub) => {
572                    warn!(subtype = sub, "claude returned no structured_output");
573                    ""
574                }
575                None => "",
576            };
577            let raw_response = extract_raw_response_text(&parsed, stdout);
578
579            AgentError::SchemaValidation {
580                expected: "structured_output field".to_string(),
581                got: format!("null{hint}"),
582                debug_messages: Vec::new(),
583                partial_usage: usage,
584                raw_response,
585            }
586        })?
587    } else {
588        parsed
589            .result
590            .filter(|v| !v.is_null())
591            .unwrap_or_else(|| Value::String(String::new()))
592    };
593
594    let model_name = parsed
595        .model_usage
596        .as_ref()
597        .and_then(|m| m.keys().next().cloned());
598
599    Ok(AgentOutput {
600        value,
601        session_id: parsed.session_id,
602        cost_usd: parsed.total_cost_usd,
603        input_tokens: parsed.usage.as_ref().map(|u| u.total_input_tokens()),
604        output_tokens: parsed.usage.as_ref().map(|u| u.total_output_tokens()),
605        model: model_name,
606        duration_ms: parsed.duration_ms.unwrap_or(fallback_duration_ms),
607        debug_messages: None,
608    })
609}
610
611/// Parse `stream-json` output from the `claude` CLI into an [`AgentOutput`]
612/// with conversation trace in [`AgentOutput::debug_messages`].
613///
614/// The `stream-json` format emits one JSON object per line. Lines with
615/// `"type":"assistant"` carry conversation content (text and tool calls).
616/// The final `"type":"result"` line carries the same payload as the `json`
617/// format and is used to populate the standard output fields.
618///
619/// # Errors
620///
621/// Returns [`AgentError::SchemaValidation`] if the result line is missing
622/// or cannot be parsed.
623pub fn parse_stream_response(
624    stdout: &str,
625    config: &AgentConfig,
626    fallback_duration_ms: u64,
627) -> Result<AgentOutput, AgentError> {
628    let mut debug_messages: Vec<DebugMessage> = Vec::new();
629    let mut result_line: Option<&str> = None;
630    // True when the next `assistant` line should merge into the last pushed
631    // `DebugMessage`. A Claude CLI turn can span several `assistant` lines
632    // (one per content_block: thinking, tool_use, text). We aggregate them
633    // until a `user` (tool_result) line or a terminal `stop_reason` arrives.
634    let mut assistant_turn_open = false;
635
636    for line in stdout.lines() {
637        let trimmed = line.trim();
638        if trimmed.is_empty() {
639            continue;
640        }
641
642        let parsed: Value = match serde_json::from_str(trimmed) {
643            Ok(v) => v,
644            Err(_) => continue,
645        };
646
647        let line_type = parsed
648            .get("type")
649            .and_then(|t| t.as_str())
650            .unwrap_or("<missing>");
651        let content_kinds: Vec<&str> = parsed
652            .get("message")
653            .and_then(|m| m.get("content"))
654            .and_then(|c| c.as_array())
655            .map(|blocks| {
656                blocks
657                    .iter()
658                    .filter_map(|b| b.get("type").and_then(|t| t.as_str()))
659                    .collect()
660            })
661            .unwrap_or_default();
662        trace!(
663            target: "ironflow_core::stream",
664            line_type,
665            content_kinds = ?content_kinds,
666            raw_len = trimmed.len(),
667            "stream-json line"
668        );
669
670        match parsed.get("type").and_then(|t| t.as_str()) {
671            Some("assistant") => {
672                let message = parsed.get("message");
673                let content = message
674                    .and_then(|m| m.get("content"))
675                    .and_then(|c| c.as_array());
676                let stop_reason = message
677                    .and_then(|m| m.get("stop_reason"))
678                    .and_then(|s| s.as_str())
679                    .map(|s| s.to_string());
680
681                let usage = message.and_then(|m| m.get("usage"));
682                let input_tokens = usage
683                    .and_then(|u| u.get("input_tokens"))
684                    .and_then(|v| v.as_u64());
685                let output_tokens = usage
686                    .and_then(|u| u.get("output_tokens"))
687                    .and_then(|v| v.as_u64());
688
689                let mut text_parts: Vec<String> = Vec::new();
690                let mut thinking_parts: Vec<String> = Vec::new();
691                let mut tool_calls: Vec<DebugToolCall> = Vec::new();
692                let mut thinking_redacted = false;
693
694                if let Some(blocks) = content {
695                    for block in blocks {
696                        match block.get("type").and_then(|t| t.as_str()) {
697                            Some("text") => {
698                                if let Some(t) = block.get("text").and_then(|t| t.as_str()) {
699                                    text_parts.push(t.to_string());
700                                }
701                            }
702                            Some("thinking") => {
703                                // Try the canonical field first, fall back to
704                                // `text` which some Claude CLI versions use.
705                                let text_value = block
706                                    .get("thinking")
707                                    .and_then(|t| t.as_str())
708                                    .or_else(|| block.get("text").and_then(|t| t.as_str()));
709                                if let Some(t) = text_value
710                                    && !t.is_empty()
711                                {
712                                    thinking_parts.push(t.to_string());
713                                } else {
714                                    // Opus 4.7 adaptive thinking and
715                                    // `display: "omitted"` yield signature-only
716                                    // thinking blocks. Flag them so the UI can
717                                    // still surface that reasoning happened.
718                                    thinking_redacted = true;
719                                }
720                            }
721                            Some("tool_use") => {
722                                let id = block
723                                    .get("id")
724                                    .and_then(|n| n.as_str())
725                                    .map(|s| s.to_string());
726                                let name = block
727                                    .get("name")
728                                    .and_then(|n| n.as_str())
729                                    .unwrap_or("unknown")
730                                    .to_string();
731                                let input = block.get("input").cloned().unwrap_or(Value::Null);
732                                tool_calls.push(DebugToolCall { id, name, input });
733                            }
734                            _ => {}
735                        }
736                    }
737                }
738
739                let text = if text_parts.is_empty() {
740                    None
741                } else {
742                    Some(text_parts.join("\n"))
743                };
744                let thinking = if thinking_parts.is_empty() {
745                    None
746                } else {
747                    Some(thinking_parts.join("\n"))
748                };
749
750                let stop_is_terminal = stop_reason.is_some();
751
752                if assistant_turn_open && let Some(last) = debug_messages.last_mut() {
753                    // Merge into the turn still being built.
754                    if let Some(t) = text {
755                        last.text = Some(match last.text.take() {
756                            Some(existing) if !existing.is_empty() => format!("{existing}\n{t}"),
757                            _ => t,
758                        });
759                    }
760                    if let Some(t) = thinking {
761                        last.thinking = Some(match last.thinking.take() {
762                            Some(existing) if !existing.is_empty() => format!("{existing}\n{t}"),
763                            _ => t,
764                        });
765                    }
766                    last.thinking_redacted = last.thinking_redacted || thinking_redacted;
767                    last.tool_calls.extend(tool_calls);
768                    if stop_is_terminal {
769                        last.stop_reason = stop_reason;
770                    }
771                    if let Some(v) = input_tokens {
772                        last.input_tokens = Some(last.input_tokens.unwrap_or(0) + v);
773                    }
774                    if let Some(v) = output_tokens {
775                        last.output_tokens = Some(last.output_tokens.unwrap_or(0) + v);
776                    }
777                } else {
778                    debug_messages.push(DebugMessage {
779                        text,
780                        thinking,
781                        thinking_redacted,
782                        tool_calls,
783                        tool_results: Vec::new(),
784                        stop_reason,
785                        input_tokens,
786                        output_tokens,
787                    });
788                }
789
790                // Keep aggregating unless the CLI signalled the end of the turn.
791                assistant_turn_open = !stop_is_terminal;
792            }
793            Some("user") => {
794                // A tool_result always closes the current assistant turn.
795                assistant_turn_open = false;
796                // Tool results come as user messages whose content is an array
797                // of `tool_result` blocks. Attach them to the most recent
798                // assistant turn that emitted matching tool_use entries so
799                // the timeline stays compact.
800                let content = parsed
801                    .get("message")
802                    .and_then(|m| m.get("content"))
803                    .and_then(|c| c.as_array());
804
805                if let Some(blocks) = content {
806                    for block in blocks {
807                        if block.get("type").and_then(|t| t.as_str()) == Some("tool_result") {
808                            let tool_use_id = block
809                                .get("tool_use_id")
810                                .and_then(|v| v.as_str())
811                                .map(|s| s.to_string());
812                            let content_value =
813                                block.get("content").cloned().unwrap_or(Value::Null);
814                            let is_error = block
815                                .get("is_error")
816                                .and_then(|v| v.as_bool())
817                                .unwrap_or(false);
818
819                            let result = DebugToolResult {
820                                tool_use_id: tool_use_id.clone(),
821                                content: content_value,
822                                is_error,
823                            };
824
825                            // Attach to the turn whose tool_calls include this id.
826                            let target = tool_use_id.as_deref().and_then(|id| {
827                                debug_messages.iter_mut().rev().find(|m| {
828                                    m.tool_calls.iter().any(|tc| tc.id.as_deref() == Some(id))
829                                })
830                            });
831
832                            if let Some(msg) = target {
833                                msg.tool_results.push(result);
834                            } else if let Some(last) = debug_messages.last_mut() {
835                                last.tool_results.push(result);
836                            }
837                        }
838                    }
839                }
840            }
841            Some("result") => {
842                result_line = Some(trimmed);
843            }
844            _ => {}
845        }
846    }
847
848    let result_str = match result_line {
849        Some(line) => line,
850        None => {
851            return Err(AgentError::SchemaValidation {
852                expected: "stream-json result line".to_string(),
853                got: "no result line found in stream output".to_string(),
854                debug_messages,
855                partial_usage: Box::default(),
856                raw_response: if stdout.is_empty() {
857                    None
858                } else {
859                    Some(truncate_to(stdout, RAW_RESPONSE_FALLBACK_MAX_LEN))
860                },
861            });
862        }
863    };
864
865    match parse_response(result_str, config, fallback_duration_ms) {
866        Ok(mut output) => {
867            output.debug_messages = Some(debug_messages);
868            Ok(output)
869        }
870        Err(AgentError::SchemaValidation {
871            expected,
872            got,
873            partial_usage,
874            raw_response,
875            ..
876        }) => Err(AgentError::SchemaValidation {
877            expected,
878            got,
879            debug_messages,
880            partial_usage,
881            raw_response,
882        }),
883        Err(AgentError::BudgetExceeded {
884            spent_usd,
885            limit_usd,
886            partial_usage,
887            ..
888        }) => Err(AgentError::BudgetExceeded {
889            spent_usd,
890            limit_usd,
891            debug_messages,
892            partial_usage,
893        }),
894        Err(other) => Err(other),
895    }
896}
897
898/// Parse CLI output, dispatching to the correct parser based on verbose mode.
899///
900/// When [`AgentConfig::verbose`] is `true`, uses [`parse_stream_response`] to
901/// extract the full conversation trace. Otherwise uses [`parse_response`] for
902/// the standard single-JSON output.
903///
904/// # Errors
905///
906/// Returns [`AgentError`] if parsing fails (see individual parsers).
907pub fn parse_output(
908    stdout: &str,
909    config: &AgentConfig,
910    fallback_duration_ms: u64,
911) -> Result<AgentOutput, AgentError> {
912    if config.verbose {
913        parse_stream_response(stdout, config, fallback_duration_ms)
914    } else {
915        parse_response(stdout, config, fallback_duration_ms)
916    }
917}
918
919/// Check whether a [`SchemaValidation`](AgentError::SchemaValidation) error
920/// contains real usage data from the CLI (cost or duration present).
921fn has_usage_data(err: &AgentError) -> bool {
922    if let AgentError::SchemaValidation { partial_usage, .. } = err {
923        return partial_usage.cost_usd.is_some() || partial_usage.duration_ms.is_some();
924    }
925    false
926}
927
928/// Handle a non-zero exit code from the Claude CLI.
929///
930/// The CLI exits with code 1 for budget/turn limit errors but still writes
931/// valid JSON (with usage data) to stdout or stderr. This function tries to
932/// parse that output so cost, duration, and tokens are preserved in the error.
933///
934/// Returns `Ok` if the JSON is a valid successful response (rare but possible),
935/// `Err(BudgetExceeded)` when the CLI reported `error_max_budget_usd`,
936/// `Err(SchemaValidation)` with partial usage when structured output was
937/// requested but missing, or `Err(ProcessFailed)` as a fallback.
938pub fn handle_nonzero_exit(
939    exit_code: i32,
940    stdout: &str,
941    stderr: &str,
942    config: &AgentConfig,
943    duration_ms: u64,
944    log_prefix: &str,
945) -> Result<AgentOutput, AgentError> {
946    let json_source = if stdout.is_empty() { stderr } else { stdout };
947
948    if !json_source.is_empty() {
949        match parse_output(json_source, config, duration_ms) {
950            ok @ Ok(_) => return ok,
951            // Always carries the usage the CLI reported before stopping.
952            Err(err @ AgentError::BudgetExceeded { .. }) => return Err(err),
953            Err(err @ AgentError::SchemaValidation { .. }) => {
954                if has_usage_data(&err) {
955                    return Err(err);
956                }
957                // No usage data means the JSON wasn't a real CLI response
958                // (e.g. a parse error). Fall through to ProcessFailed.
959            }
960            Err(_) => {} // not parseable, fall through
961        }
962    }
963
964    let error_detail = if stdout.is_empty() {
965        if stderr.is_empty() {
966            "(no output captured)".to_string()
967        } else {
968            truncate_to(stderr, ERROR_DETAIL_MAX_LEN)
969        }
970    } else {
971        truncate_to(stdout, ERROR_DETAIL_MAX_LEN)
972    };
973
974    tracing::error!(
975        exit_code,
976        error_detail_len = error_detail.len(),
977        "{log_prefix} claude process failed"
978    );
979
980    Err(AgentError::ProcessFailed {
981        exit_code,
982        stderr: error_detail,
983    })
984}
985
986#[cfg(test)]
987mod tests {
988    use super::*;
989    use serde_json::json;
990
991    #[test]
992    fn deserialize_full_claude_json_output() {
993        let raw = json!({
994            "session_id": "sess-abc123",
995            "subtype": "success",
996            "result": "Hello, world!",
997            "structured_output": null,
998            "usage": {
999                "input_tokens": 100,
1000                "output_tokens": 50,
1001                "cache_creation_input_tokens": 20,
1002                "cache_read_input_tokens": 30
1003            },
1004            "total_cost_usd": 0.042,
1005            "duration_ms": 1500,
1006            "modelUsage": {
1007                "claude-sonnet-4-20250514": {
1008                    "inputTokens": 100,
1009                    "outputTokens": 50
1010                }
1011            }
1012        });
1013
1014        let parsed: ClaudeJsonOutput = serde_json::from_value(raw).unwrap();
1015        assert_eq!(parsed.session_id, Some("sess-abc123".to_string()));
1016        assert_eq!(parsed.subtype, Some("success".to_string()));
1017        assert_eq!(
1018            parsed.result,
1019            Some(Value::String("Hello, world!".to_string()))
1020        );
1021        assert!(parsed.structured_output.is_none());
1022        assert_eq!(parsed.total_cost_usd, Some(0.042));
1023        assert_eq!(parsed.duration_ms, Some(1500));
1024
1025        let usage = parsed.usage.unwrap();
1026        assert_eq!(usage.total_input_tokens(), 150); // 100 + 20 + 30
1027        assert_eq!(usage.total_output_tokens(), 50);
1028
1029        let model_usage = parsed.model_usage.unwrap();
1030        assert!(model_usage.contains_key("claude-sonnet-4-20250514"));
1031    }
1032
1033    #[test]
1034    fn deserialize_minimal_claude_json_output() {
1035        let raw = json!({});
1036
1037        let parsed: ClaudeJsonOutput = serde_json::from_value(raw).unwrap();
1038        assert!(parsed.session_id.is_none());
1039        assert!(parsed.subtype.is_none());
1040        assert!(parsed.result.is_none());
1041        assert!(parsed.structured_output.is_none());
1042        assert!(parsed.usage.is_none());
1043        assert!(parsed.total_cost_usd.is_none());
1044        assert!(parsed.duration_ms.is_none());
1045        assert!(parsed.model_usage.is_none());
1046    }
1047
1048    #[test]
1049    fn deserialize_structured_output_response() {
1050        let raw = json!({
1051            "session_id": "sess-xyz",
1052            "subtype": "success",
1053            "result": null,
1054            "structured_output": {"score": 9, "summary": "good"},
1055            "usage": {
1056                "input_tokens": 200,
1057                "output_tokens": 80,
1058                "cache_creation_input_tokens": 0,
1059                "cache_read_input_tokens": 0
1060            },
1061            "total_cost_usd": 0.08,
1062            "duration_ms": 3000
1063        });
1064
1065        let parsed: ClaudeJsonOutput = serde_json::from_value(raw).unwrap();
1066        let structured = parsed.structured_output.unwrap();
1067        assert_eq!(structured["score"], 9);
1068        assert_eq!(structured["summary"], "good");
1069    }
1070
1071    #[test]
1072    fn deserialize_budget_exceeded_response() {
1073        let raw = json!({
1074            "subtype": "error_max_budget_usd",
1075            "result": null,
1076            "structured_output": null,
1077            "total_cost_usd": 0.10,
1078            "duration_ms": 5000
1079        });
1080
1081        let parsed: ClaudeJsonOutput = serde_json::from_value(raw).unwrap();
1082        assert_eq!(parsed.subtype, Some("error_max_budget_usd".to_string()));
1083        assert!(parsed.result.is_none());
1084        assert!(parsed.structured_output.is_none());
1085    }
1086
1087    #[test]
1088    fn claude_usage_with_all_none_tokens() {
1089        let usage = ClaudeUsage {
1090            input_tokens: None,
1091            output_tokens: None,
1092            cache_creation_input_tokens: None,
1093            cache_read_input_tokens: None,
1094        };
1095        assert_eq!(usage.total_input_tokens(), 0);
1096        assert_eq!(usage.total_output_tokens(), 0);
1097    }
1098
1099    #[test]
1100    fn claude_usage_sums_cache_tokens() {
1101        let usage = ClaudeUsage {
1102            input_tokens: Some(50),
1103            output_tokens: Some(25),
1104            cache_creation_input_tokens: Some(10),
1105            cache_read_input_tokens: Some(15),
1106        };
1107        assert_eq!(usage.total_input_tokens(), 75); // 50 + 10 + 15
1108        assert_eq!(usage.total_output_tokens(), 25);
1109    }
1110
1111    #[test]
1112    fn extract_structured_prefers_structured_output() {
1113        let parsed: ClaudeJsonOutput = serde_json::from_value(json!({
1114            "result": "{\"other\": 1}",
1115            "structured_output": {"score": 9},
1116        }))
1117        .unwrap();
1118        let v = extract_structured_value(&parsed).unwrap();
1119        assert_eq!(v["score"], 9);
1120    }
1121
1122    #[test]
1123    fn extract_structured_from_result_direct_parse() {
1124        let parsed: ClaudeJsonOutput = serde_json::from_value(json!({
1125            "result": "{\"score\": 9}",
1126            "structured_output": null,
1127        }))
1128        .unwrap();
1129        let v = extract_structured_value(&parsed).unwrap();
1130        assert_eq!(v["score"], 9);
1131    }
1132
1133    #[test]
1134    fn extract_structured_from_code_fence() {
1135        let parsed: ClaudeJsonOutput = serde_json::from_value(json!({
1136            "result": "Here is the result:\n```json\n{\"score\": 9}\n```\nDone.",
1137            "structured_output": null,
1138        }))
1139        .unwrap();
1140        let v = extract_structured_value(&parsed).unwrap();
1141        assert_eq!(v["score"], 9);
1142    }
1143
1144    #[test]
1145    fn extract_structured_from_brace_extraction() {
1146        let parsed: ClaudeJsonOutput = serde_json::from_value(json!({
1147            "result": "The answer is {\"score\": 9} as expected.",
1148            "structured_output": null,
1149        }))
1150        .unwrap();
1151        let v = extract_structured_value(&parsed).unwrap();
1152        assert_eq!(v["score"], 9);
1153    }
1154
1155    #[test]
1156    fn extract_structured_returns_none_when_both_null() {
1157        let parsed: ClaudeJsonOutput = serde_json::from_value(json!({
1158            "result": null,
1159            "structured_output": null,
1160        }))
1161        .unwrap();
1162        assert!(extract_structured_value(&parsed).is_none());
1163    }
1164
1165    #[test]
1166    fn extract_structured_returns_none_for_non_json_text() {
1167        let parsed: ClaudeJsonOutput = serde_json::from_value(json!({
1168            "result": "just plain text with no json",
1169            "structured_output": null,
1170        }))
1171        .unwrap();
1172        assert!(extract_structured_value(&parsed).is_none());
1173    }
1174
1175    #[test]
1176    fn model_name_extracted_from_model_usage() {
1177        let raw = json!({
1178            "result": "ok",
1179            "modelUsage": {
1180                "claude-opus-4-20250514": {"inputTokens": 100}
1181            }
1182        });
1183        let parsed: ClaudeJsonOutput = serde_json::from_value(raw).unwrap();
1184        let name = parsed
1185            .model_usage
1186            .as_ref()
1187            .and_then(|m| m.keys().next().cloned());
1188        assert_eq!(name, Some("claude-opus-4-20250514".to_string()));
1189    }
1190
1191    #[test]
1192    fn build_args_basic_prompt() {
1193        let config = AgentConfig::new("hello world");
1194        let args = build_args(&config).unwrap();
1195        assert_eq!(args[0], "-p");
1196        assert_eq!(args[1], "hello world");
1197        assert_eq!(args[2], "--output-format");
1198        assert_eq!(args[3], "json");
1199    }
1200
1201    #[test]
1202    fn env_vars_to_remove_always_includes_ironflow_allow_bypass() {
1203        let vars = env_vars_to_remove();
1204        assert!(
1205            vars.contains(&"IRONFLOW_ALLOW_BYPASS".to_string()),
1206            "IRONFLOW_ALLOW_BYPASS must always be removed"
1207        );
1208    }
1209
1210    #[test]
1211    fn collect_vars_to_remove_captures_claude_prefixed_vars() {
1212        let keys = vec![
1213            "CLAUDE_CODE_ENTRYPOINT",
1214            "CLAUDE_CODE_SUBAGENT_MODEL",
1215            "CLAUDE_AUTOCOMPACT_PCT_OVERRIDE",
1216            "CLAUDECODE",
1217            "PATH",
1218            "HOME",
1219        ];
1220        let vars = collect_vars_to_remove(keys.into_iter().map(String::from));
1221
1222        assert!(vars.contains(&"CLAUDE_CODE_ENTRYPOINT".to_string()));
1223        assert!(vars.contains(&"CLAUDE_CODE_SUBAGENT_MODEL".to_string()));
1224        assert!(vars.contains(&"CLAUDE_AUTOCOMPACT_PCT_OVERRIDE".to_string()));
1225        assert!(vars.contains(&"CLAUDECODE".to_string()));
1226        assert!(vars.contains(&"IRONFLOW_ALLOW_BYPASS".to_string()));
1227    }
1228
1229    #[test]
1230    fn collect_vars_to_remove_excludes_unrelated_vars() {
1231        let keys = vec!["PATH", "HOME", "RUST_LOG"];
1232        let vars = collect_vars_to_remove(keys.into_iter().map(String::from));
1233
1234        assert!(!vars.contains(&"PATH".to_string()));
1235        assert!(!vars.contains(&"HOME".to_string()));
1236        // IRONFLOW_ALLOW_BYPASS is always present
1237        assert_eq!(vars.len(), 1);
1238    }
1239
1240    #[test]
1241    fn env_unset_shell_prefix_format() {
1242        // env_unset_shell_prefix always includes IRONFLOW_ALLOW_BYPASS at minimum
1243        let prefix = env_unset_shell_prefix();
1244        assert!(prefix.starts_with("unset "));
1245        assert!(prefix.ends_with("2>/dev/null; "));
1246        assert!(prefix.contains("IRONFLOW_ALLOW_BYPASS"));
1247    }
1248
1249    #[test]
1250    fn build_args_bypass_without_env_fails() {
1251        let mut config = AgentConfig::new("test");
1252        config.permission_mode = PermissionMode::BypassPermissions;
1253        // SAFETY: This test runs single-threaded and only removes a test-specific
1254        // env var that no other test reads concurrently.
1255        unsafe { std::env::remove_var("IRONFLOW_ALLOW_BYPASS") };
1256        let result = build_args(&config);
1257        assert!(result.is_err());
1258    }
1259
1260    #[test]
1261    fn build_shell_command_escapes_quotes() {
1262        let args = vec!["-p".to_string(), "it's a test".to_string()];
1263        let cmd = build_shell_command("claude", &args);
1264        assert_eq!(cmd, "'claude' '-p' 'it'\\''s a test'");
1265    }
1266
1267    #[test]
1268    fn shell_escape_basic() {
1269        assert_eq!(shell_escape("hello"), "'hello'");
1270    }
1271
1272    #[test]
1273    fn shell_escape_with_single_quotes() {
1274        assert_eq!(shell_escape("it's"), "'it'\\''s'");
1275    }
1276
1277    #[test]
1278    fn parse_response_text_mode() {
1279        let stdout = r#"{"session_id":"s1","result":"Hello","usage":{"input_tokens":10,"output_tokens":5},"total_cost_usd":0.01,"duration_ms":100}"#;
1280        let config = AgentConfig::new("test");
1281        let output = parse_response(stdout, &config, 200).unwrap();
1282        assert_eq!(output.value, Value::String("Hello".to_string()));
1283        assert_eq!(output.session_id, Some("s1".to_string()));
1284        assert_eq!(output.duration_ms, 100);
1285    }
1286
1287    #[test]
1288    fn parse_response_uses_fallback_duration() {
1289        let stdout = r#"{"result":"ok"}"#;
1290        let config = AgentConfig::new("test");
1291        let output = parse_response(stdout, &config, 999).unwrap();
1292        assert_eq!(output.duration_ms, 999);
1293    }
1294
1295    #[test]
1296    fn parse_response_invalid_json() {
1297        let config = AgentConfig::new("test");
1298        let result = parse_response("not json", &config, 0);
1299        assert!(result.is_err());
1300    }
1301
1302    #[test]
1303    fn build_args_verbose_uses_stream_json_and_verbose_flag() {
1304        let mut config = AgentConfig::new("hello");
1305        config.verbose = true;
1306        let args = build_args(&config).unwrap();
1307        assert_eq!(args[2], "--output-format");
1308        assert_eq!(args[3], "stream-json");
1309        assert!(
1310            args.contains(&"--verbose".to_string()),
1311            "stream-json with -p requires --verbose flag, got: {args:?}"
1312        );
1313    }
1314
1315    #[test]
1316    fn build_args_non_verbose_uses_json() {
1317        let config = AgentConfig::new("hello");
1318        let args = build_args(&config).unwrap();
1319        assert_eq!(args[3], "json");
1320        assert!(
1321            !args.contains(&"--verbose".to_string()),
1322            "--verbose should not be present when verbose is false"
1323        );
1324    }
1325
1326    #[test]
1327    fn build_args_strict_mcp_config_flag_absent_by_default() {
1328        let config = AgentConfig::new("hello");
1329        let args = build_args(&config).unwrap();
1330        assert!(
1331            !args.contains(&"--strict-mcp-config".to_string()),
1332            "--strict-mcp-config must not appear unless opted-in, got: {args:?}"
1333        );
1334    }
1335
1336    #[test]
1337    fn build_args_strict_mcp_config_flag_pushed_when_enabled() {
1338        let config = AgentConfig::new("hello").strict_mcp_config(true);
1339        let args = build_args(&config).unwrap();
1340        assert!(
1341            args.contains(&"--strict-mcp-config".to_string()),
1342            "--strict-mcp-config must be pushed when strict_mcp_config is true, got: {args:?}"
1343        );
1344    }
1345
1346    #[test]
1347    fn build_args_disallowed_tools_flag_absent_when_empty() {
1348        let config = AgentConfig::new("hello");
1349        let args = build_args(&config).unwrap();
1350        assert!(
1351            !args.contains(&"--disallowedTools".to_string()),
1352            "--disallowedTools must not appear when list is empty, got: {args:?}"
1353        );
1354    }
1355
1356    #[test]
1357    fn build_args_disallowed_tools_flag_joined_with_commas() {
1358        let config = AgentConfig::new("hello").disallowed_tools(["Write", "Edit", "Bash"]);
1359        let args = build_args(&config).unwrap();
1360
1361        let pos = args
1362            .iter()
1363            .position(|a| a == "--disallowedTools")
1364            .expect("--disallowedTools missing");
1365        assert_eq!(args[pos + 1], "Write,Edit,Bash");
1366    }
1367
1368    #[test]
1369    fn build_args_disallowed_tools_combined_with_allowed_tools() {
1370        let config: AgentConfig = AgentConfig::new("hello")
1371            .allow_tool("Read")
1372            .allow_tool("Grep")
1373            .into();
1374        let config = config.disallowed_tools(["Write", "Edit"]);
1375        let args = build_args(&config).unwrap();
1376
1377        let allowed_pos = args
1378            .iter()
1379            .position(|a| a == "--allowedTools")
1380            .expect("--allowedTools missing");
1381        assert_eq!(args[allowed_pos + 1], "Read,Grep");
1382
1383        let disallowed_pos = args
1384            .iter()
1385            .position(|a| a == "--disallowedTools")
1386            .expect("--disallowedTools missing");
1387        assert_eq!(args[disallowed_pos + 1], "Write,Edit");
1388    }
1389
1390    #[test]
1391    fn build_args_bare_flag_absent_by_default() {
1392        let config = AgentConfig::new("hello");
1393        let args = build_args(&config).unwrap();
1394        assert!(
1395            !args.contains(&"--bare".to_string()),
1396            "--bare must not appear unless opted-in, got: {args:?}"
1397        );
1398    }
1399
1400    #[test]
1401    fn build_args_bare_flag_pushed_when_enabled() {
1402        let config = AgentConfig::new("hello").bare(true);
1403        let args = build_args(&config).unwrap();
1404        assert!(
1405            args.contains(&"--bare".to_string()),
1406            "--bare must be pushed when bare is true, got: {args:?}"
1407        );
1408    }
1409
1410    #[test]
1411    fn build_args_strict_mcp_config_with_mcp_config_includes_both() {
1412        let config = AgentConfig::new("hello")
1413            .mcp_config(r#"{"mcpServers":{}}"#)
1414            .strict_mcp_config(true);
1415        let args = build_args(&config).unwrap();
1416
1417        let mcp_pos = args
1418            .iter()
1419            .position(|a| a == "--mcp-config")
1420            .expect("--mcp-config missing");
1421        assert_eq!(args[mcp_pos + 1], r#"{"mcpServers":{}}"#);
1422        assert!(
1423            args.contains(&"--strict-mcp-config".to_string()),
1424            "--strict-mcp-config missing when both flags requested"
1425        );
1426    }
1427
1428    #[test]
1429    fn parse_stream_response_extracts_messages_and_result() {
1430        let stream = [
1431            r#"{"type":"assistant","message":{"role":"assistant","content":[{"type":"text","text":"Let me read that file."},{"type":"tool_use","id":"tu_1","name":"Read","input":{"file_path":"/tmp/test.rs"}}],"stop_reason":"tool_use"}}"#,
1432            r#"{"type":"assistant","message":{"role":"assistant","content":[{"type":"text","text":"Done."}],"stop_reason":"end_turn"}}"#,
1433            r#"{"type":"result","session_id":"s1","result":"Done.","usage":{"input_tokens":100,"output_tokens":50},"total_cost_usd":0.02,"duration_ms":500}"#,
1434        ]
1435        .join("\n");
1436
1437        let config = AgentConfig::new("test");
1438        let output = parse_stream_response(&stream, &config, 999).unwrap();
1439
1440        assert_eq!(output.value, Value::String("Done.".to_string()));
1441        assert_eq!(output.session_id, Some("s1".to_string()));
1442        assert_eq!(output.duration_ms, 500);
1443
1444        let messages = output.debug_messages.unwrap();
1445        assert_eq!(messages.len(), 2);
1446
1447        assert_eq!(messages[0].text.as_deref(), Some("Let me read that file."));
1448        assert_eq!(messages[0].tool_calls.len(), 1);
1449        assert_eq!(messages[0].tool_calls[0].name, "Read");
1450        assert_eq!(messages[0].stop_reason.as_deref(), Some("tool_use"));
1451
1452        assert_eq!(messages[1].text.as_deref(), Some("Done."));
1453        assert!(messages[1].tool_calls.is_empty());
1454        assert_eq!(messages[1].stop_reason.as_deref(), Some("end_turn"));
1455    }
1456
1457    #[test]
1458    fn parse_stream_response_no_result_line_errors() {
1459        let stream = r#"{"type":"assistant","message":{"role":"assistant","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn"}}"#;
1460        let config = AgentConfig::new("test");
1461        let result = parse_stream_response(stream, &config, 0);
1462        assert!(result.is_err());
1463    }
1464
1465    #[test]
1466    fn parse_stream_response_empty_stream_errors() {
1467        let config = AgentConfig::new("test");
1468        let result = parse_stream_response("", &config, 0);
1469        assert!(result.is_err());
1470    }
1471
1472    #[test]
1473    fn parse_stream_response_skips_invalid_lines() {
1474        let stream = [
1475            "not json",
1476            "",
1477            r#"{"type":"result","result":"ok","duration_ms":100}"#,
1478        ]
1479        .join("\n");
1480
1481        let config = AgentConfig::new("test");
1482        let output = parse_stream_response(&stream, &config, 999).unwrap();
1483        assert_eq!(output.value, Value::String("ok".to_string()));
1484        let messages = output.debug_messages.unwrap();
1485        assert!(messages.is_empty());
1486    }
1487
1488    #[test]
1489    fn parse_stream_response_multiple_tool_calls_in_one_turn() {
1490        let stream = [
1491            r#"{"type":"assistant","message":{"role":"assistant","content":[{"type":"tool_use","id":"t1","name":"Grep","input":{"pattern":"foo"}},{"type":"tool_use","id":"t2","name":"Read","input":{"file_path":"/tmp/bar"}}],"stop_reason":"tool_use"}}"#,
1492            r#"{"type":"result","result":"done","duration_ms":200}"#,
1493        ]
1494        .join("\n");
1495
1496        let config = AgentConfig::new("test");
1497        let output = parse_stream_response(&stream, &config, 0).unwrap();
1498        let messages = output.debug_messages.unwrap();
1499        assert_eq!(messages.len(), 1);
1500        assert_eq!(messages[0].tool_calls.len(), 2);
1501        assert_eq!(messages[0].tool_calls[0].name, "Grep");
1502        assert_eq!(messages[0].tool_calls[1].name, "Read");
1503        assert!(messages[0].text.is_none());
1504    }
1505
1506    #[test]
1507    fn debug_message_display_format() {
1508        let msg = DebugMessage {
1509            text: Some("Analyzing...".to_string()),
1510            thinking: None,
1511            thinking_redacted: false,
1512            tool_calls: vec![DebugToolCall {
1513                id: Some("tu_1".to_string()),
1514                name: "Read".to_string(),
1515                input: json!({"file_path": "/tmp/test.rs"}),
1516            }],
1517            tool_results: Vec::new(),
1518            stop_reason: Some("tool_use".to_string()),
1519            input_tokens: None,
1520            output_tokens: None,
1521        };
1522        let display = format!("{msg}");
1523        assert!(display.contains("[assistant] Analyzing..."));
1524        assert!(display.contains("[tool_use] Read"));
1525    }
1526
1527    #[test]
1528    fn parse_stream_response_flags_redacted_thinking() {
1529        let stream = [
1530            r#"{"type":"assistant","message":{"content":[{"type":"thinking","thinking":"","signature":"sig_abc"}]}}"#,
1531            r#"{"type":"assistant","message":{"content":[{"type":"tool_use","id":"tu_1","name":"Bash","input":{"command":"ls"}}],"stop_reason":"tool_use"}}"#,
1532            r#"{"type":"result","result":"","duration_ms":100}"#,
1533        ]
1534        .join("\n");
1535        let config = AgentConfig::new("test");
1536        let output = parse_stream_response(&stream, &config, 0).unwrap();
1537        let messages = output.debug_messages.unwrap();
1538        assert_eq!(messages.len(), 1);
1539        assert!(messages[0].thinking_redacted);
1540        assert!(messages[0].thinking.is_none());
1541        assert_eq!(messages[0].tool_calls.len(), 1);
1542    }
1543
1544    #[test]
1545    fn parse_stream_response_extracts_thinking_block() {
1546        let stream = [
1547            r#"{"type":"assistant","message":{"role":"assistant","content":[{"type":"thinking","thinking":"Let me reason about this step by step."},{"type":"text","text":"Answer: 42"}],"stop_reason":"end_turn","usage":{"input_tokens":120,"output_tokens":30}}}"#,
1548            r#"{"type":"result","result":"Answer: 42","duration_ms":250}"#,
1549        ]
1550        .join("\n");
1551        let config = AgentConfig::new("test");
1552        let output = parse_stream_response(&stream, &config, 0).unwrap();
1553        let messages = output.debug_messages.unwrap();
1554        assert_eq!(messages.len(), 1);
1555        assert_eq!(
1556            messages[0].thinking.as_deref(),
1557            Some("Let me reason about this step by step.")
1558        );
1559        assert_eq!(messages[0].text.as_deref(), Some("Answer: 42"));
1560        assert_eq!(messages[0].input_tokens, Some(120));
1561        assert_eq!(messages[0].output_tokens, Some(30));
1562    }
1563
1564    #[test]
1565    fn parse_stream_response_attaches_tool_results_to_matching_turn() {
1566        let stream = [
1567            r#"{"type":"assistant","message":{"content":[{"type":"tool_use","id":"tu_1","name":"Read","input":{"file_path":"/tmp/a"}}],"stop_reason":"tool_use"}}"#,
1568            r#"{"type":"user","message":{"content":[{"type":"tool_result","tool_use_id":"tu_1","content":"file contents here","is_error":false}]}}"#,
1569            r#"{"type":"assistant","message":{"content":[{"type":"text","text":"Done."}],"stop_reason":"end_turn"}}"#,
1570            r#"{"type":"result","result":"Done.","duration_ms":400}"#,
1571        ]
1572        .join("\n");
1573        let config = AgentConfig::new("test");
1574        let output = parse_stream_response(&stream, &config, 0).unwrap();
1575        let messages = output.debug_messages.unwrap();
1576        assert_eq!(messages.len(), 2);
1577        assert_eq!(messages[0].tool_calls.len(), 1);
1578        assert_eq!(messages[0].tool_calls[0].id.as_deref(), Some("tu_1"));
1579        assert_eq!(messages[0].tool_results.len(), 1);
1580        assert_eq!(
1581            messages[0].tool_results[0].tool_use_id.as_deref(),
1582            Some("tu_1")
1583        );
1584        assert!(!messages[0].tool_results[0].is_error);
1585        assert!(messages[1].tool_results.is_empty());
1586    }
1587
1588    #[test]
1589    fn parse_stream_response_merges_consecutive_assistant_content_blocks() {
1590        // The Claude CLI emits one `assistant` line per content block:
1591        // thinking first (no stop_reason), then tool_use (stop_reason=tool_use).
1592        // Both belong to the same logical turn and must be collapsed into one
1593        // DebugMessage.
1594        let stream = [
1595            r#"{"type":"assistant","message":{"content":[{"type":"thinking","thinking":"I should list files first."}],"usage":{"input_tokens":6,"output_tokens":0}}}"#,
1596            r#"{"type":"assistant","message":{"content":[{"type":"tool_use","id":"tu_1","name":"Bash","input":{"command":"ls"}}],"stop_reason":"tool_use","usage":{"input_tokens":1,"output_tokens":65}}}"#,
1597            r#"{"type":"user","message":{"content":[{"type":"tool_result","tool_use_id":"tu_1","content":"README.md","is_error":false}]}}"#,
1598            r#"{"type":"assistant","message":{"content":[{"type":"text","text":"Done."}],"stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":1}}}"#,
1599            r#"{"type":"result","result":"Done.","duration_ms":500}"#,
1600        ]
1601        .join("\n");
1602        let config = AgentConfig::new("test");
1603        let output = parse_stream_response(&stream, &config, 0).unwrap();
1604        let messages = output.debug_messages.unwrap();
1605
1606        assert_eq!(
1607            messages.len(),
1608            2,
1609            "expected 2 logical turns, got {messages:?}"
1610        );
1611
1612        // Turn 1: thinking + tool_use merged, tool_result attached.
1613        assert_eq!(
1614            messages[0].thinking.as_deref(),
1615            Some("I should list files first.")
1616        );
1617        assert_eq!(messages[0].tool_calls.len(), 1);
1618        assert_eq!(messages[0].tool_calls[0].id.as_deref(), Some("tu_1"));
1619        assert_eq!(messages[0].tool_results.len(), 1);
1620        assert_eq!(messages[0].stop_reason.as_deref(), Some("tool_use"));
1621        assert_eq!(messages[0].input_tokens, Some(7));
1622        assert_eq!(messages[0].output_tokens, Some(65));
1623
1624        // Turn 2: the final assistant text.
1625        assert_eq!(messages[1].text.as_deref(), Some("Done."));
1626        assert_eq!(messages[1].stop_reason.as_deref(), Some("end_turn"));
1627    }
1628
1629    #[test]
1630    fn parse_stream_response_marks_tool_result_error() {
1631        let stream = [
1632            r#"{"type":"assistant","message":{"content":[{"type":"tool_use","id":"tu_x","name":"Bash","input":{"command":"boom"}}],"stop_reason":"tool_use"}}"#,
1633            r#"{"type":"user","message":{"content":[{"type":"tool_result","tool_use_id":"tu_x","content":"command failed","is_error":true}]}}"#,
1634            r#"{"type":"result","result":"error","duration_ms":100}"#,
1635        ]
1636        .join("\n");
1637        let config = AgentConfig::new("test");
1638        let output = parse_stream_response(&stream, &config, 0).unwrap();
1639        let messages = output.debug_messages.unwrap();
1640        assert_eq!(messages[0].tool_results.len(), 1);
1641        assert!(messages[0].tool_results[0].is_error);
1642    }
1643
1644    #[test]
1645    fn build_args_includes_both_tools_and_json_schema() {
1646        use std::marker::PhantomData;
1647
1648        use crate::operations::agent::PermissionMode;
1649
1650        // Use direct field access to bypass typestate (testing CLI arg construction,
1651        // not the builder API -- this combination triggers a Claude CLI bug).
1652        let config = AgentConfig {
1653            prompt: "test prompt".to_string(),
1654            model: "sonnet".to_string(),
1655            json_schema: Some(
1656                r#"{"type":"object","properties":{"items":{"type":"array"}}}"#.to_string(),
1657            ),
1658            allowed_tools: vec!["WebSearch".to_string(), "WebFetch".to_string()],
1659            disallowed_tools: vec![],
1660            max_turns: Some(5),
1661            permission_mode: PermissionMode::Default,
1662            system_prompt: None,
1663            max_budget_usd: None,
1664            working_dir: None,
1665            mcp_config: None,
1666            strict_mcp_config: false,
1667            bare: false,
1668            resume_session_id: None,
1669            verbose: false,
1670            pod_labels: std::collections::BTreeMap::new(),
1671            inputs: Vec::new(),
1672            allow_failure: false,
1673            retry: None,
1674            _marker: PhantomData,
1675        };
1676
1677        let args = build_args(&config).unwrap();
1678
1679        assert!(args.contains(&"--allowedTools".to_string()));
1680        assert!(args.contains(&"WebSearch,WebFetch".to_string()));
1681        assert!(args.contains(&"--json-schema".to_string()));
1682
1683        let schema_pos = args
1684            .iter()
1685            .position(|a| a == "--json-schema")
1686            .expect("--json-schema missing");
1687        let schema_value: serde_json::Value =
1688            serde_json::from_str(&args[schema_pos + 1]).expect("schema is valid JSON");
1689        assert_eq!(schema_value["type"], "object");
1690        assert_eq!(schema_value["additionalProperties"], false);
1691
1692        assert!(args.contains(&"--output-format".to_string()));
1693        assert!(args.contains(&"json".to_string()));
1694    }
1695
1696    #[test]
1697    fn stream_response_preserves_debug_messages_on_schema_validation_error() {
1698        let assistant_line = r#"{"type":"assistant","message":{"content":[{"type":"text","text":"Searching..."},{"type":"tool_use","name":"WebSearch","input":{"query":"AI news"}}],"stop_reason":"tool_use"}}"#;
1699        let result_line = r#"{"type":"result","session_id":"s1","subtype":"success","result":"text response","usage":{"input_tokens":100,"output_tokens":50},"total_cost_usd":0.01,"duration_ms":500}"#;
1700
1701        let stdout = format!("{assistant_line}\n{result_line}");
1702
1703        let config: AgentConfig = AgentConfig::new("test")
1704            .output_schema_raw(r#"{"type":"object"}"#)
1705            .into();
1706        let config = config.verbose(true);
1707
1708        let err = parse_stream_response(&stdout, &config, 500).unwrap_err();
1709
1710        match err {
1711            AgentError::SchemaValidation { debug_messages, .. } => {
1712                assert_eq!(debug_messages.len(), 1);
1713                assert_eq!(debug_messages[0].text.as_deref(), Some("Searching..."));
1714                assert_eq!(debug_messages[0].tool_calls.len(), 1);
1715                assert_eq!(debug_messages[0].tool_calls[0].name, "WebSearch");
1716            }
1717            other => panic!("expected SchemaValidation, got {other:?}"),
1718        }
1719    }
1720
1721    #[test]
1722    fn parse_response_schema_validation_error_has_empty_debug_messages() {
1723        let stdout = r#"{"session_id":"s1","subtype":"success","result":"plain text","usage":{"input_tokens":10,"output_tokens":5},"total_cost_usd":0.01,"duration_ms":100}"#;
1724
1725        let config: AgentConfig = AgentConfig::new("test")
1726            .output_schema_raw(r#"{"type":"object"}"#)
1727            .into();
1728
1729        let err = parse_response(stdout, &config, 100).unwrap_err();
1730
1731        match err {
1732            AgentError::SchemaValidation { debug_messages, .. } => {
1733                assert!(debug_messages.is_empty());
1734            }
1735            other => panic!("expected SchemaValidation, got {other:?}"),
1736        }
1737    }
1738
1739    #[test]
1740    fn parse_response_schema_validation_preserves_usage_data() {
1741        let stdout = r#"{"session_id":"s1","subtype":"error_max_turns","result":null,"structured_output":null,"usage":{"input_tokens":500,"output_tokens":200,"cache_creation_input_tokens":50,"cache_read_input_tokens":30},"total_cost_usd":0.30,"duration_ms":4500}"#;
1742
1743        let config: AgentConfig = AgentConfig::new("test")
1744            .output_schema_raw(r#"{"type":"object"}"#)
1745            .into();
1746
1747        let err = parse_response(stdout, &config, 0).unwrap_err();
1748
1749        match err {
1750            AgentError::SchemaValidation { partial_usage, .. } => {
1751                assert_eq!(partial_usage.cost_usd, Some(0.30));
1752                assert_eq!(partial_usage.duration_ms, Some(4500));
1753                assert_eq!(partial_usage.input_tokens, Some(580)); // 500 + 50 + 30
1754                assert_eq!(partial_usage.output_tokens, Some(200));
1755            }
1756            other => panic!("expected SchemaValidation, got {other:?}"),
1757        }
1758    }
1759
1760    #[test]
1761    fn parse_response_schema_validation_no_usage_when_parse_fails() {
1762        let config: AgentConfig = AgentConfig::new("test")
1763            .output_schema_raw(r#"{"type":"object"}"#)
1764            .into();
1765
1766        let err = parse_response("not json at all", &config, 0).unwrap_err();
1767
1768        match err {
1769            AgentError::SchemaValidation { partial_usage, .. } => {
1770                assert!(partial_usage.cost_usd.is_none());
1771                assert!(partial_usage.duration_ms.is_none());
1772                assert!(partial_usage.input_tokens.is_none());
1773                assert!(partial_usage.output_tokens.is_none());
1774            }
1775            other => panic!("expected SchemaValidation, got {other:?}"),
1776        }
1777    }
1778
1779    #[test]
1780    fn stream_response_schema_validation_preserves_usage_and_debug() {
1781        let assistant_line = r#"{"type":"assistant","message":{"content":[{"type":"text","text":"Working..."}],"stop_reason":"end_turn"}}"#;
1782        let result_line = r#"{"type":"result","session_id":"s1","subtype":"error_max_turns","result":null,"structured_output":null,"usage":{"input_tokens":300,"output_tokens":100},"total_cost_usd":0.15,"duration_ms":3000}"#;
1783
1784        let stdout = format!("{assistant_line}\n{result_line}");
1785
1786        let config: AgentConfig = AgentConfig::new("test")
1787            .output_schema_raw(r#"{"type":"object"}"#)
1788            .into();
1789        let config = config.verbose(true);
1790
1791        let err = parse_stream_response(&stdout, &config, 0).unwrap_err();
1792
1793        match err {
1794            AgentError::SchemaValidation {
1795                debug_messages,
1796                partial_usage,
1797                ..
1798            } => {
1799                assert_eq!(debug_messages.len(), 1);
1800                assert_eq!(debug_messages[0].text.as_deref(), Some("Working..."));
1801                assert_eq!(partial_usage.cost_usd, Some(0.15));
1802                assert_eq!(partial_usage.duration_ms, Some(3000));
1803                assert_eq!(partial_usage.input_tokens, Some(300));
1804                assert_eq!(partial_usage.output_tokens, Some(100));
1805            }
1806            other => panic!("expected SchemaValidation, got {other:?}"),
1807        }
1808    }
1809
1810    #[test]
1811    fn handle_nonzero_exit_parses_schema_error_with_usage() {
1812        let stdout = r#"{"session_id":"s1","subtype":"error_max_turns","result":null,"structured_output":null,"usage":{"input_tokens":100,"output_tokens":50},"total_cost_usd":0.10,"duration_ms":2000}"#;
1813        let config: AgentConfig = AgentConfig::new("test")
1814            .output_schema_raw(r#"{"type":"object"}"#)
1815            .into();
1816
1817        let result = handle_nonzero_exit(1, stdout, "", &config, 2000, "test");
1818
1819        match result {
1820            Err(AgentError::SchemaValidation { partial_usage, .. }) => {
1821                assert_eq!(partial_usage.cost_usd, Some(0.10));
1822                assert_eq!(partial_usage.duration_ms, Some(2000));
1823            }
1824            other => panic!("expected Err(SchemaValidation), got {other:?}"),
1825        }
1826    }
1827
1828    #[test]
1829    fn parse_response_budget_exceeded_is_its_own_error() {
1830        let stdout = r#"{"session_id":"s1","subtype":"error_max_budget_usd","result":null,"structured_output":null,"usage":{"input_tokens":500,"output_tokens":200,"cache_creation_input_tokens":50,"cache_read_input_tokens":30},"total_cost_usd":0.30,"duration_ms":4500}"#;
1831
1832        let config: AgentConfig = AgentConfig::new("test")
1833            .output_schema_raw(r#"{"type":"object"}"#)
1834            .max_budget_usd(0.25)
1835            .into();
1836
1837        let err = parse_response(stdout, &config, 0).unwrap_err();
1838
1839        match err {
1840            AgentError::BudgetExceeded {
1841                spent_usd,
1842                limit_usd,
1843                partial_usage,
1844                debug_messages,
1845            } => {
1846                assert!((spent_usd - 0.30).abs() < f64::EPSILON);
1847                assert!((limit_usd - 0.25).abs() < f64::EPSILON);
1848                assert_eq!(partial_usage.cost_usd, Some(0.30));
1849                assert_eq!(partial_usage.duration_ms, Some(4500));
1850                assert_eq!(partial_usage.input_tokens, Some(580)); // 500 + 50 + 30
1851                assert_eq!(partial_usage.output_tokens, Some(200));
1852                assert!(debug_messages.is_empty());
1853            }
1854            other => panic!("expected BudgetExceeded, got {other:?}"),
1855        }
1856    }
1857
1858    #[test]
1859    fn stream_response_budget_exceeded_keeps_debug_messages() {
1860        let assistant_line = r#"{"type":"assistant","message":{"content":[{"type":"text","text":"Working..."}],"stop_reason":"end_turn"}}"#;
1861        let result_line = r#"{"type":"result","session_id":"s1","subtype":"error_max_budget_usd","result":null,"structured_output":null,"usage":{"input_tokens":300,"output_tokens":100},"total_cost_usd":0.15,"duration_ms":3000}"#;
1862
1863        let stdout = format!("{assistant_line}\n{result_line}");
1864
1865        let config: AgentConfig = AgentConfig::new("test")
1866            .output_schema_raw(r#"{"type":"object"}"#)
1867            .max_budget_usd(0.10)
1868            .into();
1869        let config = config.verbose(true);
1870
1871        let err = parse_stream_response(&stdout, &config, 0).unwrap_err();
1872
1873        match err {
1874            AgentError::BudgetExceeded {
1875                debug_messages,
1876                partial_usage,
1877                ..
1878            } => {
1879                assert_eq!(debug_messages.len(), 1);
1880                assert_eq!(debug_messages[0].text.as_deref(), Some("Working..."));
1881                assert_eq!(partial_usage.cost_usd, Some(0.15));
1882            }
1883            other => panic!("expected BudgetExceeded, got {other:?}"),
1884        }
1885    }
1886
1887    #[test]
1888    fn handle_nonzero_exit_propagates_budget_exceeded() {
1889        let stdout = r#"{"session_id":"s1","subtype":"error_max_budget_usd","result":null,"structured_output":null,"usage":{"input_tokens":100,"output_tokens":50},"total_cost_usd":0.10,"duration_ms":2000}"#;
1890        let config: AgentConfig = AgentConfig::new("test")
1891            .output_schema_raw(r#"{"type":"object"}"#)
1892            .max_budget_usd(0.05)
1893            .into();
1894
1895        let result = handle_nonzero_exit(1, stdout, "", &config, 2000, "test");
1896
1897        match result {
1898            Err(AgentError::BudgetExceeded { partial_usage, .. }) => {
1899                assert_eq!(partial_usage.cost_usd, Some(0.10));
1900                assert_eq!(partial_usage.duration_ms, Some(2000));
1901            }
1902            other => panic!("expected Err(BudgetExceeded), got {other:?}"),
1903        }
1904    }
1905
1906    #[test]
1907    fn handle_nonzero_exit_returns_ok_for_valid_text_response() {
1908        let stdout = r#"{"result":"Hello!","usage":{"input_tokens":10,"output_tokens":5},"total_cost_usd":0.01,"duration_ms":100}"#;
1909        let config = AgentConfig::new("test");
1910
1911        let result = handle_nonzero_exit(1, stdout, "", &config, 100, "test");
1912        assert!(result.is_ok());
1913    }
1914
1915    #[test]
1916    fn handle_nonzero_exit_falls_back_to_process_failed() {
1917        let config = AgentConfig::new("test");
1918        let result = handle_nonzero_exit(1, "", "some random error", &config, 0, "test");
1919
1920        match result {
1921            Err(AgentError::ProcessFailed { exit_code, stderr }) => {
1922                assert_eq!(exit_code, 1);
1923                assert_eq!(stderr, "some random error");
1924            }
1925            other => panic!("expected Err(ProcessFailed), got {other:?}"),
1926        }
1927    }
1928
1929    #[test]
1930    fn handle_nonzero_exit_prefers_stdout_over_stderr() {
1931        let stdout =
1932            r#"{"result":"ok","usage":{"input_tokens":10,"output_tokens":5},"duration_ms":100}"#;
1933        let stderr = "some error text";
1934        let config = AgentConfig::new("test");
1935
1936        let result = handle_nonzero_exit(1, stdout, stderr, &config, 100, "test");
1937        assert!(result.is_ok());
1938    }
1939
1940    #[test]
1941    fn handle_nonzero_exit_empty_both_returns_no_output() {
1942        let config = AgentConfig::new("test");
1943        let result = handle_nonzero_exit(1, "", "", &config, 0, "test");
1944
1945        match result {
1946            Err(AgentError::ProcessFailed { stderr, .. }) => {
1947                assert_eq!(stderr, "(no output captured)");
1948            }
1949            other => panic!("expected Err(ProcessFailed), got {other:?}"),
1950        }
1951    }
1952
1953    #[test]
1954    fn handle_nonzero_exit_truncates_large_stdout_in_error_detail() {
1955        let large_stdout = "x".repeat(ERROR_DETAIL_MAX_LEN + 10_000);
1956        let config = AgentConfig::new("test");
1957        let result = handle_nonzero_exit(1, &large_stdout, "", &config, 0, "test");
1958
1959        match result {
1960            Err(AgentError::ProcessFailed { stderr, .. }) => {
1961                assert!(
1962                    stderr.len() <= ERROR_DETAIL_MAX_LEN + 20,
1963                    "error_detail should be truncated, got {} bytes",
1964                    stderr.len()
1965                );
1966                assert!(stderr.ends_with("...(truncated)"));
1967            }
1968            other => panic!("expected Err(ProcessFailed), got {other:?}"),
1969        }
1970    }
1971
1972    #[test]
1973    fn handle_nonzero_exit_truncates_large_stderr_in_error_detail() {
1974        let large_stderr = "e".repeat(ERROR_DETAIL_MAX_LEN + 5_000);
1975        let config = AgentConfig::new("test");
1976        let result = handle_nonzero_exit(1, "", &large_stderr, &config, 0, "test");
1977
1978        match result {
1979            Err(AgentError::ProcessFailed { stderr, .. }) => {
1980                assert!(
1981                    stderr.len() <= ERROR_DETAIL_MAX_LEN + 20,
1982                    "error_detail should be truncated, got {} bytes",
1983                    stderr.len()
1984                );
1985                assert!(stderr.ends_with("...(truncated)"));
1986            }
1987            other => panic!("expected Err(ProcessFailed), got {other:?}"),
1988        }
1989    }
1990
1991    #[test]
1992    fn truncate_to_no_truncation_when_short() {
1993        assert_eq!(truncate_to("hello", 10), "hello");
1994    }
1995
1996    #[test]
1997    fn truncate_to_adds_marker_when_long() {
1998        let result = truncate_to("abcdefghij", 5);
1999        assert!(result.starts_with("abcde"));
2000        assert!(result.ends_with("...(truncated)"));
2001    }
2002
2003    #[test]
2004    fn truncate_to_handles_multibyte_chars() {
2005        let text = "cafe\u{0301}"; // e + combining accent = 5 bytes
2006        let result = truncate_to(text, 4);
2007        assert!(result.ends_with("...(truncated)"));
2008    }
2009
2010    #[test]
2011    fn schema_validation_includes_raw_response_from_result_field() {
2012        let stdout = r#"{"session_id":"s1","subtype":"success","result":"The model answered with plain text instead of JSON","usage":{"input_tokens":10,"output_tokens":5},"total_cost_usd":0.01,"duration_ms":100}"#;
2013
2014        let config: AgentConfig = AgentConfig::new("test")
2015            .output_schema_raw(r#"{"type":"object"}"#)
2016            .into();
2017
2018        let err = parse_response(stdout, &config, 100).unwrap_err();
2019
2020        match err {
2021            AgentError::SchemaValidation { raw_response, .. } => {
2022                let raw = raw_response.expect("raw_response should be Some when result is present");
2023                assert!(
2024                    raw.contains("The model answered with plain text instead of JSON"),
2025                    "raw_response should contain the result text, got: {raw}"
2026                );
2027            }
2028            other => panic!("expected SchemaValidation, got {other:?}"),
2029        }
2030    }
2031
2032    #[test]
2033    fn schema_validation_raw_response_from_stdout_when_result_null() {
2034        let stdout = r#"{"session_id":"s1","subtype":"error_max_turns","result":null,"structured_output":null,"usage":{"input_tokens":500,"output_tokens":200},"total_cost_usd":0.30,"duration_ms":4500}"#;
2035
2036        let config: AgentConfig = AgentConfig::new("test")
2037            .output_schema_raw(r#"{"type":"object"}"#)
2038            .into();
2039
2040        let err = parse_response(stdout, &config, 0).unwrap_err();
2041
2042        match err {
2043            AgentError::SchemaValidation { raw_response, .. } => {
2044                let raw = raw_response
2045                    .expect("raw_response should fall back to stdout when result is null");
2046                assert!(
2047                    raw.contains("error_max_turns"),
2048                    "raw_response should contain stdout content, got: {raw}"
2049                );
2050            }
2051            other => panic!("expected SchemaValidation, got {other:?}"),
2052        }
2053    }
2054
2055    #[test]
2056    fn schema_validation_raw_response_none_when_both_null_and_empty_stdout() {
2057        let stdout = r#"{"result":null,"structured_output":null}"#;
2058
2059        let config: AgentConfig = AgentConfig::new("test")
2060            .output_schema_raw(r#"{"type":"object"}"#)
2061            .into();
2062
2063        let err = parse_response(stdout, &config, 0).unwrap_err();
2064
2065        match err {
2066            AgentError::SchemaValidation { raw_response, .. } => {
2067                // stdout is non-empty (it's the JSON itself), so raw_response falls back to it
2068                assert!(raw_response.is_some());
2069            }
2070            other => panic!("expected SchemaValidation, got {other:?}"),
2071        }
2072    }
2073
2074    #[test]
2075    fn schema_validation_raw_response_truncated_to_max() {
2076        let long_text = "x".repeat(5000);
2077        let stdout = format!(
2078            r#"{{"result":"{long_text}","structured_output":null,"usage":{{"input_tokens":10,"output_tokens":5}},"total_cost_usd":0.01,"duration_ms":100}}"#,
2079        );
2080
2081        let config: AgentConfig = AgentConfig::new("test")
2082            .output_schema_raw(r#"{{"type":"object"}}"#)
2083            .into();
2084
2085        let err = parse_response(&stdout, &config, 0).unwrap_err();
2086
2087        match err {
2088            AgentError::SchemaValidation { raw_response, .. } => {
2089                let raw = raw_response.expect("raw_response should be present");
2090                assert!(
2091                    raw.len() <= RAW_RESPONSE_MAX_LEN + 20,
2092                    "raw_response should be truncated, got len={}",
2093                    raw.len()
2094                );
2095                assert!(raw.contains("...(truncated)"));
2096            }
2097            other => panic!("expected SchemaValidation, got {other:?}"),
2098        }
2099    }
2100
2101    #[test]
2102    fn stream_schema_validation_preserves_raw_response() {
2103        let assistant_line = r#"{"type":"assistant","message":{"content":[{"type":"text","text":"Searching..."}],"stop_reason":"end_turn"}}"#;
2104        let result_line = r#"{"type":"result","session_id":"s1","subtype":"success","result":"text response","usage":{"input_tokens":100,"output_tokens":50},"total_cost_usd":0.01,"duration_ms":500}"#;
2105
2106        let stdout = format!("{assistant_line}\n{result_line}");
2107
2108        let config: AgentConfig = AgentConfig::new("test")
2109            .output_schema_raw(r#"{"type":"object"}"#)
2110            .into();
2111        let config = config.verbose(true);
2112
2113        let err = parse_stream_response(&stdout, &config, 500).unwrap_err();
2114
2115        match err {
2116            AgentError::SchemaValidation {
2117                raw_response,
2118                debug_messages,
2119                ..
2120            } => {
2121                assert!(
2122                    raw_response.is_some(),
2123                    "raw_response should be propagated through stream parser"
2124                );
2125                assert!(
2126                    raw_response.unwrap().contains("text response"),
2127                    "raw_response should contain the result text"
2128                );
2129                assert_eq!(debug_messages.len(), 1);
2130            }
2131            other => panic!("expected SchemaValidation, got {other:?}"),
2132        }
2133    }
2134
2135    #[test]
2136    fn parse_response_json_parse_error_includes_raw_stdout() {
2137        let config: AgentConfig = AgentConfig::new("test")
2138            .output_schema_raw(r#"{"type":"object"}"#)
2139            .into();
2140
2141        let err = parse_response("this is not JSON at all", &config, 0).unwrap_err();
2142
2143        match err {
2144            AgentError::SchemaValidation { raw_response, .. } => {
2145                let raw = raw_response
2146                    .expect("raw_response should contain the raw stdout on parse failure");
2147                assert!(
2148                    raw.contains("this is not JSON at all"),
2149                    "raw_response should contain the unparseable input, got: {raw}"
2150                );
2151            }
2152            other => panic!("expected SchemaValidation, got {other:?}"),
2153        }
2154    }
2155
2156    #[test]
2157    fn extract_raw_response_text_prefers_result_string() {
2158        let parsed: ClaudeJsonOutput = serde_json::from_value(json!({
2159            "result": "The agent said hello",
2160            "structured_output": null,
2161        }))
2162        .unwrap();
2163        let raw = extract_raw_response_text(&parsed, "full stdout...");
2164        assert_eq!(raw, Some("The agent said hello".to_string()));
2165    }
2166
2167    #[test]
2168    fn extract_raw_response_text_falls_back_to_stdout() {
2169        let parsed: ClaudeJsonOutput = serde_json::from_value(json!({
2170            "result": null,
2171            "structured_output": null,
2172        }))
2173        .unwrap();
2174        let raw = extract_raw_response_text(&parsed, "raw stdout content");
2175        assert_eq!(raw, Some("raw stdout content".to_string()));
2176    }
2177
2178    #[test]
2179    fn extract_raw_response_text_none_when_all_empty() {
2180        let parsed: ClaudeJsonOutput = serde_json::from_value(json!({
2181            "result": null,
2182            "structured_output": null,
2183        }))
2184        .unwrap();
2185        let raw = extract_raw_response_text(&parsed, "");
2186        assert!(raw.is_none());
2187    }
2188
2189    #[test]
2190    fn extract_raw_response_text_stringifies_non_string_result() {
2191        let parsed: ClaudeJsonOutput = serde_json::from_value(json!({
2192            "result": {"partial": "data", "count": 42},
2193            "structured_output": null,
2194        }))
2195        .unwrap();
2196        let raw = extract_raw_response_text(&parsed, "");
2197        let raw = raw.expect("should extract stringified JSON object");
2198        assert!(raw.contains("partial"));
2199        assert!(raw.contains("42"));
2200    }
2201
2202    #[test]
2203    fn build_command_small_prompt_uses_arg() {
2204        let config = AgentConfig::new("hello world");
2205        let built = build_command(&config).unwrap();
2206        assert!(built.stdin_prompt.is_none());
2207        assert_eq!(built.args[0], "-p");
2208        assert_eq!(built.args[1], "hello world");
2209    }
2210
2211    #[test]
2212    fn build_command_large_prompt_uses_stdin() {
2213        let big = "x".repeat(PROMPT_STDIN_THRESHOLD + 1);
2214        let config = AgentConfig::new(&big);
2215        let built = build_command(&config).unwrap();
2216        assert!(built.stdin_prompt.is_some());
2217        assert_eq!(built.stdin_prompt.as_ref().unwrap().len(), big.len());
2218        assert!(!built.args.contains(&"-p".to_string()));
2219    }
2220
2221    #[test]
2222    fn build_command_exact_threshold_uses_arg() {
2223        let exact = "y".repeat(PROMPT_STDIN_THRESHOLD);
2224        let config = AgentConfig::new(&exact);
2225        let built = build_command(&config).unwrap();
2226        assert!(built.stdin_prompt.is_none());
2227        assert_eq!(built.args[0], "-p");
2228    }
2229
2230    #[test]
2231    fn build_command_preserves_other_flags() {
2232        let big = "z".repeat(PROMPT_STDIN_THRESHOLD + 1);
2233        let config = AgentConfig::new(&big)
2234            .model("claude-sonnet-4-20250514")
2235            .max_turns(5);
2236        let built = build_command(&config).unwrap();
2237        assert!(built.stdin_prompt.is_some());
2238        assert!(built.args.contains(&"--model".to_string()));
2239        assert!(built.args.contains(&"claude-sonnet-4-20250514".to_string()));
2240        assert!(built.args.contains(&"--max-turns".to_string()));
2241        assert!(built.args.contains(&"5".to_string()));
2242    }
2243}