ironflow_core/operations/
agent.rs

1//! Agent operation - build and execute AI agent calls.
2//!
3//! The [`Agent`] builder lets you configure a single agent invocation (model,
4//! prompt, tools, budget, permissions, etc.) and execute it through any
5//! [`AgentProvider`]. The result is an [`AgentResult`] that provides typed
6//! access to the agent's response, session metadata, and usage statistics.
7//!
8//! # Examples
9//!
10//! ```no_run
11//! use ironflow_core::prelude::*;
12//!
13//! # async fn example() -> Result<(), OperationError> {
14//! let provider = ClaudeCodeProvider::new();
15//!
16//! let result = Agent::new()
17//!     .prompt("Summarize the README.md file")
18//!     .model(Model::SONNET)
19//!     .max_turns(3)
20//!     .run(&provider)
21//!     .await?;
22//!
23//! println!("{}", result.text());
24//! # Ok(())
25//! # }
26//! ```
27
28use std::any;
29use std::sync::Arc;
30
31use schemars::{JsonSchema, schema_for};
32use serde::de::DeserializeOwned;
33use serde::{Deserialize, Serialize};
34use serde_json::{Value, from_value, to_string};
35use tokio::time;
36use tracing::{info, warn};
37
38use crate::error::OperationError;
39#[cfg(feature = "prometheus")]
40use crate::metric_names;
41use crate::provider::{AgentConfig, AgentOutput, AgentProvider, DebugMessage, LogSink};
42use crate::retry::RetryPolicy;
43
44/// Provider-agnostic model identifiers.
45///
46/// Constants are provided for well-known Claude models, but any string
47/// is accepted - custom [`AgentProvider`] implementations interpret the
48/// model identifier however they wish.
49///
50/// # Examples
51///
52/// ```no_run
53/// use ironflow_core::prelude::*;
54///
55/// # async fn example() -> Result<(), OperationError> {
56/// let provider = ClaudeCodeProvider::new();
57///
58/// // Using a built-in constant
59/// let r = Agent::new()
60///     .prompt("hi")
61///     .model(Model::SONNET)
62///     .run(&provider)
63///     .await?;
64///
65/// // Using a custom model string
66/// let r = Agent::new()
67///     .prompt("hi")
68///     .model("mistral-large-latest")
69///     .run(&provider)
70///     .await?;
71/// # Ok(())
72/// # }
73/// ```
74pub struct Model;
75
76impl Model {
77    // ── Aliases (latest version, CLI resolves to current) ───────────
78
79    /// Claude Sonnet - balanced speed and capability (default).
80    pub const SONNET: &str = "sonnet";
81    /// Claude Opus - highest capability.
82    pub const OPUS: &str = "opus";
83    /// Claude Haiku - fastest and cheapest.
84    pub const HAIKU: &str = "haiku";
85
86    // ── Claude 4.5 ─────────────────────────────────────────────────
87
88    /// Claude Haiku 4.5.
89    pub const HAIKU_45: &str = "claude-haiku-4-5-20251001";
90
91    // ── Claude 4.6 - 200K context ──────────────────────────────────
92
93    /// Claude Sonnet 4.6.
94    pub const SONNET_46: &str = "claude-sonnet-4-6";
95    /// Claude Opus 4.6.
96    pub const OPUS_46: &str = "claude-opus-4-6";
97
98    // ── Claude 4.6 - 1M context ────────────────────────────────────
99
100    /// Claude Sonnet 4.6 with 1M token context window.
101    pub const SONNET_46_1M: &str = "claude-sonnet-4-6[1m]";
102    /// Claude Opus 4.6 with 1M token context window.
103    pub const OPUS_46_1M: &str = "claude-opus-4-6[1m]";
104
105    // ── Claude 4.7 - 1M context native ─────────────────────────────
106
107    /// Claude Opus 4.7 - previous flagship, 1M token context native.
108    pub const OPUS_47: &str = "claude-opus-4-7";
109    /// Claude Opus 4.7 with 1M token context window explicit.
110    pub const OPUS_47_1M: &str = "claude-opus-4-7[1m]";
111
112    // ── Claude 4.8 - 1M context native ─────────────────────────────
113
114    /// Claude Opus 4.8 - previous Opus flagship, 1M token context native.
115    pub const OPUS_48: &str = "claude-opus-4-8";
116    /// Claude Opus 4.8 with 1M token context window explicit.
117    pub const OPUS_48_1M: &str = "claude-opus-4-8[1m]";
118
119    // ── Claude 5 - 1M context native ───────────────────────────────
120
121    /// Claude Fable 5 - most capable widely released model, 1M token context native.
122    pub const FABLE_5: &str = "claude-fable-5";
123    /// Claude Mythos 5 - Fable 5 capabilities, limited availability (Project Glasswing).
124    pub const MYTHOS_5: &str = "claude-mythos-5";
125    /// Claude Opus 5 - current flagship for agentic coding, 1M token context native.
126    pub const OPUS_5: &str = "claude-opus-5";
127    /// Claude Opus 5 with 1M token context window explicit.
128    pub const OPUS_5_1M: &str = "claude-opus-5[1m]";
129    /// Claude Sonnet 5 - best speed/intelligence balance, 1M token context native.
130    pub const SONNET_5: &str = "claude-sonnet-5";
131    /// Claude Sonnet 5 with 1M token context window explicit.
132    pub const SONNET_5_1M: &str = "claude-sonnet-5[1m]";
133}
134
135/// Controls how the agent handles tool-use permission prompts.
136///
137/// These map to the `--permission-mode` and `--dangerously-skip-permissions`
138/// flags in the Claude CLI.
139#[derive(Debug, Default, Clone, Copy, Serialize)]
140pub enum PermissionMode {
141    /// Use the CLI default permission behavior.
142    #[default]
143    Default,
144    /// Automatically approve tool-use requests.
145    Auto,
146    /// Suppress all permission prompts (the agent proceeds without asking).
147    DontAsk,
148    /// Skip all permission checks entirely.
149    ///
150    /// **Warning**: the agent will have unrestricted filesystem and shell access.
151    BypassPermissions,
152}
153
154impl<'de> Deserialize<'de> for PermissionMode {
155    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
156    where
157        D: serde::Deserializer<'de>,
158    {
159        let s = String::deserialize(deserializer)?;
160        Ok(match s.to_lowercase().replace('_', "").as_str() {
161            "auto" => Self::Auto,
162            "dontask" => Self::DontAsk,
163            "bypass" | "bypasspermissions" => Self::BypassPermissions,
164            _ => Self::Default,
165        })
166    }
167}
168
169/// Builder for a single agent invocation.
170///
171/// Create with [`Agent::new`], chain configuration methods, then call
172/// [`run`](Agent::run) with an [`AgentProvider`] to execute.
173///
174/// # Examples
175///
176/// ```no_run
177/// use ironflow_core::prelude::*;
178///
179/// # async fn example() -> Result<(), OperationError> {
180/// let provider = ClaudeCodeProvider::new();
181///
182/// let result = Agent::new()
183///     .system_prompt("You are a Rust expert.")
184///     .prompt("Review this code for safety issues.")
185///     .model(Model::OPUS)
186///     .allowed_tools(&["Read", "Grep"])
187///     .max_turns(5)
188///     .max_budget_usd(0.50)
189///     .working_dir("/tmp/project")
190///     .permission_mode(PermissionMode::Auto)
191///     .run(&provider)
192///     .await?;
193///
194/// println!("Cost: ${:.4}", result.cost_usd().unwrap_or(0.0));
195/// # Ok(())
196/// # }
197/// ```
198#[must_use = "an Agent does nothing until .run() is awaited"]
199pub struct Agent {
200    config: AgentConfig,
201    dry_run: Option<bool>,
202    retry_policy: Option<RetryPolicy>,
203    log_sink: Option<Arc<dyn LogSink>>,
204}
205
206impl Agent {
207    /// Create a new agent builder with default settings.
208    ///
209    /// Defaults: [`Model::SONNET`], no system prompt, no tool restrictions,
210    /// no budget/turn limits, [`PermissionMode::Default`].
211    pub fn new() -> Self {
212        Self {
213            config: AgentConfig::new(""),
214            dry_run: None,
215            retry_policy: None,
216            log_sink: None,
217        }
218    }
219
220    /// Create an agent builder from an existing [`AgentConfig`].
221    ///
222    /// Useful when the config comes from a serialized workflow definition
223    /// rather than being built programmatically.
224    ///
225    /// # Examples
226    ///
227    /// ```no_run
228    /// use ironflow_core::prelude::*;
229    /// use ironflow_core::provider::AgentConfig;
230    ///
231    /// # async fn example() -> Result<(), OperationError> {
232    /// let provider = ClaudeCodeProvider::new();
233    /// let config = AgentConfig::new("Summarize the README");
234    /// let result = Agent::from_config(config).run(&provider).await?;
235    /// # Ok(())
236    /// # }
237    /// ```
238    pub fn from_config(config: impl Into<AgentConfig>) -> Self {
239        Self {
240            config: config.into(),
241            dry_run: None,
242            retry_policy: None,
243            log_sink: None,
244        }
245    }
246
247    /// Set the system prompt that defines the agent's persona or constraints.
248    pub fn system_prompt(mut self, prompt: &str) -> Self {
249        self.config.system_prompt = Some(prompt.to_string());
250        self
251    }
252
253    /// Set the user prompt - the main instruction sent to the agent.
254    pub fn prompt(mut self, prompt: &str) -> Self {
255        self.config.prompt = prompt.to_string();
256        self
257    }
258
259    /// Set the model to use for this invocation.
260    ///
261    /// Accepts any string-like value. Use [`Model`] constants for well-known
262    /// Claude models, or pass an arbitrary string for custom providers.
263    ///
264    /// Defaults to [`Model::SONNET`] if not called.
265    pub fn model(mut self, model: impl Into<String>) -> Self {
266        self.config.model = model.into();
267        self
268    }
269
270    /// Restrict which tools the agent may invoke.
271    ///
272    /// Pass an empty slice (or do not call this method) to allow the provider
273    /// default set of tools.
274    pub fn allowed_tools(mut self, tools: &[&str]) -> Self {
275        self.config.allowed_tools = tools.iter().map(|s| s.to_string()).collect();
276        self
277    }
278
279    /// Set the maximum number of agentic turns.
280    ///
281    /// # Panics
282    ///
283    /// Panics if `turns` is `0`.
284    pub fn max_turns(mut self, turns: u32) -> Self {
285        assert!(turns > 0, "max_turns must be greater than 0");
286        self.config.max_turns = Some(turns);
287        self
288    }
289
290    /// Set the maximum spend in USD for this invocation.
291    ///
292    /// # Panics
293    ///
294    /// Panics if `budget` is negative, NaN, or infinity.
295    pub fn max_budget_usd(mut self, budget: f64) -> Self {
296        assert!(
297            budget.is_finite() && budget > 0.0,
298            "budget must be a positive finite number, got {budget}"
299        );
300        self.config.max_budget_usd = Some(budget);
301        self
302    }
303
304    /// Set the working directory for the agent process.
305    pub fn working_dir(mut self, dir: &str) -> Self {
306        self.config.working_dir = Some(dir.to_string());
307        self
308    }
309
310    /// Set the path to an MCP (Model Context Protocol) server configuration file.
311    pub fn mcp_config(mut self, config: &str) -> Self {
312        self.config.mcp_config = Some(config.to_string());
313        self
314    }
315
316    /// Set the permission mode controlling tool-use approval behavior.
317    ///
318    /// See [`PermissionMode`] for details on each variant.
319    pub fn permission_mode(mut self, mode: PermissionMode) -> Self {
320        self.config.permission_mode = mode;
321        self
322    }
323
324    /// Request structured (typed) output from the agent.
325    ///
326    /// The type `T` must implement [`JsonSchema`]. The generated schema is sent
327    /// to the provider so the model returns JSON conforming to `T`, which can
328    /// then be deserialized with [`AgentResult::json`].
329    ///
330    /// # Examples
331    ///
332    /// ```no_run
333    /// use ironflow_core::prelude::*;
334    ///
335    /// #[derive(Deserialize, JsonSchema)]
336    /// struct Review {
337    ///     score: u8,
338    ///     summary: String,
339    /// }
340    ///
341    /// # async fn example() -> Result<(), OperationError> {
342    /// let provider = ClaudeCodeProvider::new();
343    /// let result = Agent::new()
344    ///     .prompt("Review the codebase")
345    ///     .output::<Review>()
346    ///     .run(&provider)
347    ///     .await?;
348    ///
349    /// let review: Review = result.json().expect("schema-validated output");
350    /// println!("Score: {}/10 - {}", review.score, review.summary);
351    /// # Ok(())
352    /// # }
353    /// ```
354    pub fn output<T: JsonSchema>(mut self) -> Self {
355        let schema = schema_for!(T);
356        self.config.json_schema = match to_string(&schema) {
357            Ok(s) => Some(s),
358            Err(e) => {
359                warn!(error = %e, type_name = any::type_name::<T>(), "failed to serialize JSON schema, structured output disabled");
360                None
361            }
362        };
363        self
364    }
365
366    /// Set structured output from a pre-serialized JSON Schema string.
367    ///
368    /// Use this when the schema comes from configuration or another source
369    /// rather than a Rust type. For type-safe schema generation, prefer
370    /// [`output`](Agent::output).
371    ///
372    /// **Important:** structured output requires `max_turns >= 2`. The Claude CLI
373    /// uses the first turn for reasoning and a second turn to produce the
374    /// schema-conforming JSON.
375    ///
376    /// # Examples
377    ///
378    /// ```no_run
379    /// use ironflow_core::prelude::*;
380    ///
381    /// # async fn example() -> Result<(), OperationError> {
382    /// let schema = r#"{"type":"object","properties":{"labels":{"type":"array","items":{"type":"string"}}}}"#;
383    /// let agent = Agent::new()
384    ///     .prompt("Classify this email")
385    ///     .output_schema_raw(schema);
386    /// # Ok(())
387    /// # }
388    /// ```
389    pub fn output_schema_raw(mut self, schema: &str) -> Self {
390        self.config.json_schema = Some(schema.to_string());
391        self
392    }
393
394    /// Retry the agent invocation up to `max_retries` times on transient failures.
395    ///
396    /// Uses default exponential backoff settings (200ms initial, 2x multiplier,
397    /// 30s cap). For custom backoff parameters, use [`retry_policy`](Agent::retry_policy).
398    ///
399    /// Only transient errors are retried: process failures and timeouts.
400    /// Deterministic errors (prompt too large, schema validation) are never retried.
401    ///
402    /// # Panics
403    ///
404    /// Panics if `max_retries` is `0`.
405    ///
406    /// # Examples
407    ///
408    /// ```no_run
409    /// use ironflow_core::prelude::*;
410    ///
411    /// # async fn example() -> Result<(), OperationError> {
412    /// let provider = ClaudeCodeProvider::new();
413    /// let result = Agent::new()
414    ///     .prompt("Summarize the codebase")
415    ///     .retry(2)
416    ///     .run(&provider)
417    ///     .await?;
418    /// # Ok(())
419    /// # }
420    /// ```
421    pub fn retry(mut self, max_retries: u32) -> Self {
422        self.retry_policy = Some(RetryPolicy::new(max_retries));
423        self
424    }
425
426    /// Set a custom [`RetryPolicy`] for this agent invocation.
427    ///
428    /// Allows full control over backoff duration, multiplier, and max delay.
429    /// See [`RetryPolicy`] for details.
430    ///
431    /// # Examples
432    ///
433    /// ```no_run
434    /// use std::time::Duration;
435    /// use ironflow_core::prelude::*;
436    /// use ironflow_core::retry::RetryPolicy;
437    ///
438    /// # async fn example() -> Result<(), OperationError> {
439    /// let provider = ClaudeCodeProvider::new();
440    /// let result = Agent::new()
441    ///     .prompt("Analyze the code")
442    ///     .retry_policy(
443    ///         RetryPolicy::new(3)
444    ///             .backoff(Duration::from_secs(1))
445    ///             .max_backoff(Duration::from_secs(60))
446    ///     )
447    ///     .run(&provider)
448    ///     .await?;
449    /// # Ok(())
450    /// # }
451    /// ```
452    pub fn retry_policy(mut self, policy: RetryPolicy) -> Self {
453        self.retry_policy = Some(policy);
454        self
455    }
456
457    /// Enable or disable dry-run mode for this specific operation.
458    ///
459    /// When dry-run is active, the agent call is logged but not executed.
460    /// A synthetic [`AgentResult`] is returned with a placeholder text,
461    /// zero cost, and zero tokens.
462    ///
463    /// If not set, falls back to the global dry-run setting
464    /// (see [`set_dry_run`](crate::dry_run::set_dry_run)).
465    pub fn dry_run(mut self, enabled: bool) -> Self {
466        self.dry_run = Some(enabled);
467        self
468    }
469
470    /// Attach a [`LogSink`] for real-time log streaming.
471    ///
472    /// When set, [`invoke_with_logs`](AgentProvider::invoke_with_logs) is called
473    /// instead of [`invoke`](AgentProvider::invoke), allowing providers that
474    /// support streaming to pipe output lines in real time.
475    ///
476    /// # Examples
477    ///
478    /// ```no_run
479    /// use std::sync::Arc;
480    /// use ironflow_core::prelude::*;
481    ///
482    /// # async fn example() -> Result<(), OperationError> {
483    /// # struct MySink;
484    /// # impl LogSink for MySink { fn log(&self, _: &str, _: &str) {} }
485    /// let provider = ClaudeCodeProvider::new();
486    /// let sink: Arc<dyn LogSink> = Arc::new(MySink);
487    ///
488    /// let result = Agent::new()
489    ///     .prompt("Analyze src/")
490    ///     .log_sink(sink)
491    ///     .run(&provider)
492    ///     .await?;
493    /// # Ok(())
494    /// # }
495    /// ```
496    pub fn log_sink(mut self, sink: Arc<dyn LogSink>) -> Self {
497        self.log_sink = Some(sink);
498        self
499    }
500
501    /// Enable verbose/debug mode to capture the full conversation trace.
502    ///
503    /// When enabled, the provider captures every assistant message and tool
504    /// call into [`AgentResult::debug_messages`]. Useful for understanding
505    /// why the agent returned an unexpected result.
506    ///
507    /// # Examples
508    ///
509    /// ```no_run
510    /// use ironflow_core::prelude::*;
511    ///
512    /// # async fn example() -> Result<(), OperationError> {
513    /// let provider = ClaudeCodeProvider::new();
514    ///
515    /// let result = Agent::new()
516    ///     .prompt("Analyze src/")
517    ///     .verbose()
518    ///     .max_budget_usd(0.10)
519    ///     .run(&provider)
520    ///     .await?;
521    ///
522    /// if let Some(messages) = result.debug_messages() {
523    ///     for msg in messages {
524    ///         println!("{msg}");
525    ///     }
526    /// }
527    /// # Ok(())
528    /// # }
529    /// ```
530    pub fn verbose(mut self) -> Self {
531        self.config.verbose = true;
532        self
533    }
534
535    /// Resume a previous agent conversation by session ID.
536    ///
537    /// Pass the session ID from a previous [`AgentResult::session_id()`] to
538    /// continue the multi-turn conversation.
539    ///
540    /// # Examples
541    ///
542    /// ```no_run
543    /// use ironflow_core::prelude::*;
544    ///
545    /// # async fn example() -> Result<(), OperationError> {
546    /// let provider = ClaudeCodeProvider::new();
547    ///
548    /// let first = Agent::new()
549    ///     .prompt("Analyze the src/ directory")
550    ///     .max_budget_usd(0.10)
551    ///     .run(&provider)
552    ///     .await?;
553    ///
554    /// let session = first.session_id().expect("provider returned session ID");
555    ///
556    /// let followup = Agent::new()
557    ///     .prompt("Now suggest improvements")
558    ///     .resume(session)
559    ///     .max_budget_usd(0.10)
560    ///     .run(&provider)
561    ///     .await?;
562    /// # Ok(())
563    /// # }
564    /// ```
565    ///
566    /// # Panics
567    ///
568    /// Panics if `session_id` is empty or contains characters other than
569    /// alphanumerics, hyphens, and underscores.
570    pub fn resume(mut self, session_id: &str) -> Self {
571        assert!(!session_id.is_empty(), "session_id must not be empty");
572        assert!(
573            session_id
574                .chars()
575                .all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_'),
576            "session_id must only contain alphanumeric characters, hyphens, or underscores, got: {session_id}"
577        );
578        self.config.resume_session_id = Some(session_id.to_string());
579        self
580    }
581
582    /// Execute the agent invocation using the given [`AgentProvider`].
583    ///
584    /// If a [`retry_policy`](Agent::retry_policy) is configured, transient
585    /// failures (process crashes, timeouts, schema validation) are retried
586    /// with exponential backoff. When structured output is requested
587    /// (`json_schema` is set) and no explicit retry policy is configured,
588    /// an automatic retry policy of 2 retries is applied to handle
589    /// non-deterministic `structured_output: null` responses from the CLI.
590    /// Deterministic errors (prompt too large) are returned immediately
591    /// without retry.
592    ///
593    /// # Errors
594    ///
595    /// Returns [`OperationError::Agent`] if the provider reports a failure
596    /// (process crash, timeout, or schema validation error).
597    ///
598    /// # Panics
599    ///
600    /// Panics if [`prompt`](Agent::prompt) was never called or the prompt is
601    /// empty (whitespace-only counts as empty).
602    #[tracing::instrument(name = "agent", skip_all, fields(model = %self.config.model, prompt_len = self.config.prompt.len()))]
603    pub async fn run(self, provider: &dyn AgentProvider) -> Result<AgentResult, OperationError> {
604        assert!(
605            !self.config.prompt.trim().is_empty(),
606            "prompt must not be empty - call .prompt(\"...\") before .run()"
607        );
608
609        if crate::dry_run::effective_dry_run(self.dry_run) {
610            info!(
611                prompt_len = self.config.prompt.len(),
612                "[dry-run] agent call skipped"
613            );
614            let mut output =
615                AgentOutput::new(Value::String("[dry-run] agent call skipped".to_string()));
616            output.cost_usd = Some(0.0);
617            output.input_tokens = Some(0);
618            output.output_tokens = Some(0);
619            return Ok(AgentResult { output });
620        }
621
622        let result = self.invoke_once(provider).await;
623
624        let default_schema_retry = RetryPolicy::new(2);
625        let policy = match &self.retry_policy {
626            Some(p) => p,
627            None if self.config.json_schema.is_some() => &default_schema_retry,
628            None => return result,
629        };
630
631        // Non-retryable errors are returned immediately.
632        if let Err(ref err) = result {
633            if !crate::retry::is_retryable(err) {
634                return result;
635            }
636        } else {
637            return result;
638        }
639
640        let mut last_result = result;
641
642        for attempt in 0..policy.max_retries {
643            let delay = policy.delay_for_attempt(attempt);
644            let retry_reason = if matches!(
645                &last_result,
646                Err(OperationError::Agent(
647                    crate::error::AgentError::SchemaValidation { .. }
648                ))
649            ) {
650                "structured_output was null (CLI non-determinism)"
651            } else {
652                "transient failure"
653            };
654            warn!(
655                attempt = attempt + 1,
656                max_retries = policy.max_retries,
657                delay_ms = delay.as_millis() as u64,
658                reason = retry_reason,
659                "retrying agent invocation"
660            );
661            time::sleep(delay).await;
662
663            last_result = self.invoke_once(provider).await;
664
665            match &last_result {
666                Ok(_) => return last_result,
667                Err(err) if !crate::retry::is_retryable(err) => return last_result,
668                _ => {}
669            }
670        }
671
672        last_result
673    }
674
675    /// Execute a single agent invocation attempt (no retry logic).
676    async fn invoke_once(
677        &self,
678        provider: &dyn AgentProvider,
679    ) -> Result<AgentResult, OperationError> {
680        #[cfg(feature = "prometheus")]
681        let model_label = self.config.model.to_string();
682
683        let invoke_result = match self.log_sink {
684            Some(ref sink) => provider.invoke_with_logs(&self.config, sink.clone()).await,
685            None => provider.invoke(&self.config).await,
686        };
687        let output = match invoke_result {
688            Ok(output) => output,
689            Err(e) => {
690                #[cfg(feature = "prometheus")]
691                {
692                    metrics::counter!(metric_names::AGENT_TOTAL, "model" => model_label.clone(), "status" => metric_names::STATUS_ERROR).increment(1);
693                }
694                return Err(OperationError::Agent(e));
695            }
696        };
697
698        info!(
699            duration_ms = output.duration_ms,
700            cost_usd = output.cost_usd,
701            input_tokens = output.input_tokens,
702            output_tokens = output.output_tokens,
703            model = output.model,
704            "agent completed"
705        );
706
707        #[cfg(feature = "prometheus")]
708        {
709            metrics::counter!(metric_names::AGENT_TOTAL, "model" => model_label.clone(), "status" => metric_names::STATUS_SUCCESS).increment(1);
710            metrics::histogram!(metric_names::AGENT_DURATION_SECONDS, "model" => model_label.clone())
711                .record(output.duration_ms as f64 / 1000.0);
712            if let Some(cost) = output.cost_usd {
713                metrics::gauge!(metric_names::AGENT_COST_USD_TOTAL, "model" => model_label.clone())
714                    .increment(cost);
715            }
716            if let Some(tokens) = output.input_tokens {
717                metrics::counter!(metric_names::AGENT_TOKENS_INPUT_TOTAL, "model" => model_label.clone()).increment(tokens);
718            }
719            if let Some(tokens) = output.output_tokens {
720                metrics::counter!(metric_names::AGENT_TOKENS_OUTPUT_TOTAL, "model" => model_label)
721                    .increment(tokens);
722            }
723        }
724
725        Ok(AgentResult { output })
726    }
727}
728
729impl Default for Agent {
730    fn default() -> Self {
731        Self::new()
732    }
733}
734
735/// The result of a successful agent invocation.
736///
737/// Wraps the raw [`AgentOutput`] and provides convenience accessors for the
738/// response text, typed JSON deserialization, session metadata, and usage stats.
739#[derive(Debug)]
740pub struct AgentResult {
741    output: AgentOutput,
742}
743
744impl AgentResult {
745    /// Return the agent's response as a plain text string.
746    ///
747    /// If the underlying value is not a JSON string (e.g. when structured output
748    /// was requested), returns an empty string and logs a warning.
749    pub fn text(&self) -> &str {
750        match self.output.value.as_str() {
751            Some(s) => s,
752            None => {
753                warn!(
754                    value_type = self.output.value.to_string(),
755                    "agent output is not a string, returning empty"
756                );
757                ""
758            }
759        }
760    }
761
762    /// Return the raw JSON [`Value`] of the agent's response.
763    pub fn value(&self) -> &Value {
764        &self.output.value
765    }
766
767    /// Deserialize the agent's response into the given type `T`.
768    ///
769    /// This clones the underlying JSON value. If you no longer need the
770    /// `AgentResult` afterwards, use [`into_json`](AgentResult::into_json)
771    /// instead to avoid the clone.
772    ///
773    /// # Errors
774    ///
775    /// Returns [`OperationError::Deserialize`] if the JSON value does not match `T`.
776    pub fn json<T: DeserializeOwned>(&self) -> Result<T, OperationError> {
777        from_value(self.output.value.clone()).map_err(OperationError::deserialize::<T>)
778    }
779
780    /// Consume the result and deserialize the response into `T` without cloning.
781    ///
782    /// # Errors
783    ///
784    /// Returns [`OperationError::Deserialize`] if the JSON value does not match `T`.
785    pub fn into_json<T: DeserializeOwned>(self) -> Result<T, OperationError> {
786        from_value(self.output.value).map_err(OperationError::deserialize::<T>)
787    }
788
789    /// Build an `AgentResult` from a raw [`AgentOutput`].
790    ///
791    /// This is available only in test builds to simplify test setup without
792    /// going through the full record/replay pipeline.
793    #[cfg(test)]
794    pub(crate) fn from_output(output: AgentOutput) -> Self {
795        Self { output }
796    }
797
798    /// Return the provider-assigned session ID, if available.
799    pub fn session_id(&self) -> Option<&str> {
800        self.output.session_id.as_deref()
801    }
802
803    /// Return the cost of this invocation in USD, if reported by the provider.
804    pub fn cost_usd(&self) -> Option<f64> {
805        self.output.cost_usd
806    }
807
808    /// Return the number of input tokens consumed, if reported.
809    pub fn input_tokens(&self) -> Option<u64> {
810        self.output.input_tokens
811    }
812
813    /// Return the number of output tokens generated, if reported.
814    pub fn output_tokens(&self) -> Option<u64> {
815        self.output.output_tokens
816    }
817
818    /// Return the wall-clock duration of the invocation in milliseconds.
819    pub fn duration_ms(&self) -> u64 {
820        self.output.duration_ms
821    }
822
823    /// Return the concrete model identifier used, if reported by the provider.
824    pub fn model(&self) -> Option<&str> {
825        self.output.model.as_deref()
826    }
827
828    /// Return the conversation trace captured during a verbose invocation.
829    ///
830    /// Returns `None` when [`Agent::verbose`] was not called. When present,
831    /// each [`DebugMessage`] contains the
832    /// assistant's text and tool calls for one conversation turn.
833    pub fn debug_messages(&self) -> Option<&[DebugMessage]> {
834        self.output.debug_messages.as_deref()
835    }
836}
837
838#[cfg(test)]
839mod tests {
840    use super::*;
841    use crate::error::AgentError;
842    use crate::provider::InvokeFuture;
843    use serde_json::json;
844
845    struct TestProvider {
846        output: AgentOutput,
847    }
848
849    impl AgentProvider for TestProvider {
850        fn invoke<'a>(&'a self, _config: &'a AgentConfig) -> InvokeFuture<'a> {
851            Box::pin(async move {
852                Ok(AgentOutput {
853                    value: self.output.value.clone(),
854                    session_id: self.output.session_id.clone(),
855                    cost_usd: self.output.cost_usd,
856                    input_tokens: self.output.input_tokens,
857                    output_tokens: self.output.output_tokens,
858                    model: self.output.model.clone(),
859                    duration_ms: self.output.duration_ms,
860                    debug_messages: None,
861                })
862            })
863        }
864    }
865
866    struct ConfigCapture {
867        output: AgentOutput,
868    }
869
870    impl AgentProvider for ConfigCapture {
871        fn invoke<'a>(&'a self, config: &'a AgentConfig) -> InvokeFuture<'a> {
872            let config_json = serde_json::to_value(config).unwrap();
873            Box::pin(async move {
874                Ok(AgentOutput {
875                    value: config_json,
876                    session_id: self.output.session_id.clone(),
877                    cost_usd: self.output.cost_usd,
878                    input_tokens: self.output.input_tokens,
879                    output_tokens: self.output.output_tokens,
880                    model: self.output.model.clone(),
881                    duration_ms: self.output.duration_ms,
882                    debug_messages: None,
883                })
884            })
885        }
886    }
887
888    fn default_output() -> AgentOutput {
889        AgentOutput {
890            value: json!("test output"),
891            session_id: Some("sess-123".to_string()),
892            cost_usd: Some(0.05),
893            input_tokens: Some(100),
894            output_tokens: Some(50),
895            model: Some("sonnet".to_string()),
896            duration_ms: 1500,
897            debug_messages: None,
898        }
899    }
900
901    // --- Model constants ---
902
903    #[test]
904    fn model_constants_have_expected_values() {
905        assert_eq!(Model::SONNET, "sonnet");
906        assert_eq!(Model::OPUS, "opus");
907        assert_eq!(Model::HAIKU, "haiku");
908        assert_eq!(Model::HAIKU_45, "claude-haiku-4-5-20251001");
909        assert_eq!(Model::SONNET_46, "claude-sonnet-4-6");
910        assert_eq!(Model::OPUS_46, "claude-opus-4-6");
911        assert_eq!(Model::SONNET_46_1M, "claude-sonnet-4-6[1m]");
912        assert_eq!(Model::OPUS_46_1M, "claude-opus-4-6[1m]");
913        assert_eq!(Model::OPUS_47, "claude-opus-4-7");
914        assert_eq!(Model::OPUS_47_1M, "claude-opus-4-7[1m]");
915        assert_eq!(Model::OPUS_48, "claude-opus-4-8");
916        assert_eq!(Model::OPUS_48_1M, "claude-opus-4-8[1m]");
917        assert_eq!(Model::FABLE_5, "claude-fable-5");
918        assert_eq!(Model::MYTHOS_5, "claude-mythos-5");
919        assert_eq!(Model::OPUS_5, "claude-opus-5");
920        assert_eq!(Model::OPUS_5_1M, "claude-opus-5[1m]");
921        assert_eq!(Model::SONNET_5, "claude-sonnet-5");
922        assert_eq!(Model::SONNET_5_1M, "claude-sonnet-5[1m]");
923    }
924
925    // --- Agent::new() defaults via ConfigCapture ---
926
927    #[tokio::test]
928    async fn agent_new_default_values() {
929        let provider = ConfigCapture {
930            output: default_output(),
931        };
932        let result = Agent::new().prompt("hi").run(&provider).await.unwrap();
933
934        let config = result.value();
935        assert_eq!(config["system_prompt"], json!(null));
936        assert_eq!(config["prompt"], json!("hi"));
937        assert_eq!(config["model"], json!("sonnet"));
938        assert_eq!(config["allowed_tools"], json!([]));
939        assert_eq!(config["max_turns"], json!(null));
940        assert_eq!(config["max_budget_usd"], json!(null));
941        assert_eq!(config["working_dir"], json!(null));
942        assert_eq!(config["mcp_config"], json!(null));
943        assert_eq!(config["permission_mode"], json!("Default"));
944        assert_eq!(config["json_schema"], json!(null));
945    }
946
947    #[tokio::test]
948    async fn agent_default_matches_new() {
949        let provider = ConfigCapture {
950            output: default_output(),
951        };
952        let result_new = Agent::new().prompt("x").run(&provider).await.unwrap();
953        let result_default = Agent::default().prompt("x").run(&provider).await.unwrap();
954
955        assert_eq!(result_new.value(), result_default.value());
956    }
957
958    // --- Builder methods ---
959
960    #[tokio::test]
961    async fn builder_methods_store_values_correctly() {
962        let provider = ConfigCapture {
963            output: default_output(),
964        };
965        let result = Agent::new()
966            .system_prompt("you are a bot")
967            .prompt("do something")
968            .model(Model::OPUS)
969            .allowed_tools(&["Read", "Write"])
970            .max_turns(5)
971            .max_budget_usd(1.5)
972            .working_dir("/tmp")
973            .mcp_config("{}")
974            .permission_mode(PermissionMode::Auto)
975            .run(&provider)
976            .await
977            .unwrap();
978
979        let config = result.value();
980        assert_eq!(config["system_prompt"], json!("you are a bot"));
981        assert_eq!(config["prompt"], json!("do something"));
982        assert_eq!(config["model"], json!("opus"));
983        assert_eq!(config["allowed_tools"], json!(["Read", "Write"]));
984        assert_eq!(config["max_turns"], json!(5));
985        assert_eq!(config["max_budget_usd"], json!(1.5));
986        assert_eq!(config["working_dir"], json!("/tmp"));
987        assert_eq!(config["mcp_config"], json!("{}"));
988        assert_eq!(config["permission_mode"], json!("Auto"));
989    }
990
991    // --- Panics ---
992
993    #[test]
994    #[should_panic(expected = "max_turns must be greater than 0")]
995    fn max_turns_zero_panics() {
996        let _ = Agent::new().max_turns(0);
997    }
998
999    #[test]
1000    #[should_panic(expected = "budget must be a positive finite number")]
1001    fn max_budget_negative_panics() {
1002        let _ = Agent::new().max_budget_usd(-1.0);
1003    }
1004
1005    #[test]
1006    #[should_panic(expected = "budget must be a positive finite number")]
1007    fn max_budget_nan_panics() {
1008        let _ = Agent::new().max_budget_usd(f64::NAN);
1009    }
1010
1011    #[test]
1012    #[should_panic(expected = "budget must be a positive finite number")]
1013    fn max_budget_infinity_panics() {
1014        let _ = Agent::new().max_budget_usd(f64::INFINITY);
1015    }
1016
1017    // --- AgentResult accessors ---
1018
1019    #[tokio::test]
1020    async fn agent_result_text_with_string_value() {
1021        let provider = TestProvider {
1022            output: AgentOutput {
1023                value: json!("hello world"),
1024                ..default_output()
1025            },
1026        };
1027        let result = Agent::new().prompt("test").run(&provider).await.unwrap();
1028        assert_eq!(result.text(), "hello world");
1029    }
1030
1031    #[tokio::test]
1032    async fn agent_result_text_with_non_string_value() {
1033        let provider = TestProvider {
1034            output: AgentOutput {
1035                value: json!(42),
1036                ..default_output()
1037            },
1038        };
1039        let result = Agent::new().prompt("test").run(&provider).await.unwrap();
1040        assert_eq!(result.text(), "");
1041    }
1042
1043    #[tokio::test]
1044    async fn agent_result_text_with_null_value() {
1045        let provider = TestProvider {
1046            output: AgentOutput {
1047                value: json!(null),
1048                ..default_output()
1049            },
1050        };
1051        let result = Agent::new().prompt("test").run(&provider).await.unwrap();
1052        assert_eq!(result.text(), "");
1053    }
1054
1055    #[tokio::test]
1056    async fn agent_result_json_successful_deserialize() {
1057        #[derive(Deserialize, PartialEq, Debug)]
1058        struct MyOutput {
1059            name: String,
1060            count: u32,
1061        }
1062        let provider = TestProvider {
1063            output: AgentOutput {
1064                value: json!({"name": "test", "count": 7}),
1065                ..default_output()
1066            },
1067        };
1068        let result = Agent::new().prompt("test").run(&provider).await.unwrap();
1069        let parsed: MyOutput = result.json().unwrap();
1070        assert_eq!(parsed.name, "test");
1071        assert_eq!(parsed.count, 7);
1072    }
1073
1074    #[tokio::test]
1075    async fn agent_result_json_failed_deserialize() {
1076        #[derive(Debug, Deserialize)]
1077        #[allow(dead_code)]
1078        struct MyOutput {
1079            name: String,
1080        }
1081        let provider = TestProvider {
1082            output: AgentOutput {
1083                value: json!(42),
1084                ..default_output()
1085            },
1086        };
1087        let result = Agent::new().prompt("test").run(&provider).await.unwrap();
1088        let err = result.json::<MyOutput>().unwrap_err();
1089        assert!(matches!(err, OperationError::Deserialize { .. }));
1090    }
1091
1092    #[tokio::test]
1093    async fn agent_result_accessors() {
1094        let provider = TestProvider {
1095            output: AgentOutput {
1096                value: json!("v"),
1097                session_id: Some("s-1".to_string()),
1098                cost_usd: Some(0.123),
1099                input_tokens: Some(999),
1100                output_tokens: Some(456),
1101                model: Some("opus".to_string()),
1102                duration_ms: 2000,
1103                debug_messages: None,
1104            },
1105        };
1106        let result = Agent::new().prompt("test").run(&provider).await.unwrap();
1107        assert_eq!(result.session_id(), Some("s-1"));
1108        assert_eq!(result.cost_usd(), Some(0.123));
1109        assert_eq!(result.input_tokens(), Some(999));
1110        assert_eq!(result.output_tokens(), Some(456));
1111        assert_eq!(result.duration_ms(), 2000);
1112        assert_eq!(result.model(), Some("opus"));
1113    }
1114
1115    // --- Session resume ---
1116
1117    #[tokio::test]
1118    async fn resume_passes_session_id_in_config() {
1119        let provider = ConfigCapture {
1120            output: default_output(),
1121        };
1122        let result = Agent::new()
1123            .prompt("followup")
1124            .resume("sess-abc")
1125            .run(&provider)
1126            .await
1127            .unwrap();
1128
1129        let config = result.value();
1130        assert_eq!(config["resume_session_id"], json!("sess-abc"));
1131    }
1132
1133    #[tokio::test]
1134    async fn no_resume_has_null_session_id() {
1135        let provider = ConfigCapture {
1136            output: default_output(),
1137        };
1138        let result = Agent::new()
1139            .prompt("first call")
1140            .run(&provider)
1141            .await
1142            .unwrap();
1143
1144        let config = result.value();
1145        assert_eq!(config["resume_session_id"], json!(null));
1146    }
1147
1148    #[test]
1149    #[should_panic(expected = "session_id must not be empty")]
1150    fn resume_empty_session_id_panics() {
1151        let _ = Agent::new().resume("");
1152    }
1153
1154    #[test]
1155    #[should_panic(expected = "session_id must only contain")]
1156    fn resume_invalid_chars_panics() {
1157        let _ = Agent::new().resume("sess;rm -rf /");
1158    }
1159
1160    #[test]
1161    fn resume_valid_formats_accepted() {
1162        let _ = Agent::new().resume("sess-abc123");
1163        let _ = Agent::new().resume("a1b2c3d4_session");
1164        let _ = Agent::new().resume("abc-DEF-123_456");
1165    }
1166
1167    #[tokio::test]
1168    #[should_panic(expected = "prompt must not be empty")]
1169    async fn run_without_prompt_panics() {
1170        let provider = TestProvider {
1171            output: default_output(),
1172        };
1173        let _ = Agent::new().run(&provider).await;
1174    }
1175
1176    #[tokio::test]
1177    #[should_panic(expected = "prompt must not be empty")]
1178    async fn run_with_whitespace_only_prompt_panics() {
1179        let provider = TestProvider {
1180            output: default_output(),
1181        };
1182        let _ = Agent::new().prompt("   ").run(&provider).await;
1183    }
1184
1185    // --- Model accepts arbitrary strings ---
1186
1187    #[tokio::test]
1188    async fn model_accepts_custom_string() {
1189        let provider = ConfigCapture {
1190            output: default_output(),
1191        };
1192        let result = Agent::new()
1193            .prompt("hi")
1194            .model("mistral-large-latest")
1195            .run(&provider)
1196            .await
1197            .unwrap();
1198        assert_eq!(result.value()["model"], json!("mistral-large-latest"));
1199    }
1200
1201    #[tokio::test]
1202    async fn verbose_sets_config_flag() {
1203        let provider = ConfigCapture {
1204            output: default_output(),
1205        };
1206        let result = Agent::new()
1207            .prompt("hi")
1208            .verbose()
1209            .run(&provider)
1210            .await
1211            .unwrap();
1212        assert_eq!(result.value()["verbose"], json!(true));
1213    }
1214
1215    #[tokio::test]
1216    async fn verbose_not_set_by_default() {
1217        let provider = ConfigCapture {
1218            output: default_output(),
1219        };
1220        let result = Agent::new().prompt("hi").run(&provider).await.unwrap();
1221        assert_eq!(result.value()["verbose"], json!(false));
1222    }
1223
1224    #[tokio::test]
1225    async fn debug_messages_none_without_verbose() {
1226        let provider = TestProvider {
1227            output: default_output(),
1228        };
1229        let result = Agent::new().prompt("test").run(&provider).await.unwrap();
1230        assert!(result.debug_messages().is_none());
1231    }
1232
1233    #[tokio::test]
1234    async fn model_accepts_owned_string() {
1235        let provider = ConfigCapture {
1236            output: default_output(),
1237        };
1238        let model_name = String::from("gpt-4o");
1239        let result = Agent::new()
1240            .prompt("hi")
1241            .model(model_name)
1242            .run(&provider)
1243            .await
1244            .unwrap();
1245        assert_eq!(result.value()["model"], json!("gpt-4o"));
1246    }
1247
1248    #[tokio::test]
1249    async fn into_json_success() {
1250        #[derive(Deserialize, PartialEq, Debug)]
1251        struct Out {
1252            name: String,
1253        }
1254        let provider = TestProvider {
1255            output: AgentOutput {
1256                value: json!({"name": "test"}),
1257                ..default_output()
1258            },
1259        };
1260        let result = Agent::new().prompt("test").run(&provider).await.unwrap();
1261        let parsed: Out = result.into_json().unwrap();
1262        assert_eq!(parsed.name, "test");
1263    }
1264
1265    #[tokio::test]
1266    async fn into_json_failure() {
1267        #[derive(Debug, Deserialize)]
1268        #[allow(dead_code)]
1269        struct Out {
1270            name: String,
1271        }
1272        let provider = TestProvider {
1273            output: AgentOutput {
1274                value: json!(42),
1275                ..default_output()
1276            },
1277        };
1278        let result = Agent::new().prompt("test").run(&provider).await.unwrap();
1279        let err = result.into_json::<Out>().unwrap_err();
1280        assert!(matches!(err, OperationError::Deserialize { .. }));
1281    }
1282
1283    #[test]
1284    fn from_output_creates_result() {
1285        let output = AgentOutput {
1286            value: json!("hello"),
1287            ..default_output()
1288        };
1289        let result = AgentResult::from_output(output);
1290        assert_eq!(result.text(), "hello");
1291        assert_eq!(result.cost_usd(), Some(0.05));
1292    }
1293
1294    #[test]
1295    #[should_panic(expected = "budget must be a positive finite number")]
1296    fn max_budget_zero_panics() {
1297        let _ = Agent::new().max_budget_usd(0.0);
1298    }
1299
1300    #[test]
1301    fn model_constant_equality() {
1302        assert_eq!(Model::SONNET, "sonnet");
1303        assert_ne!(Model::SONNET, Model::OPUS);
1304    }
1305
1306    #[test]
1307    fn permission_mode_serialize_deserialize_roundtrip() {
1308        for mode in [
1309            PermissionMode::Default,
1310            PermissionMode::Auto,
1311            PermissionMode::DontAsk,
1312            PermissionMode::BypassPermissions,
1313        ] {
1314            let json = to_string(&mode).unwrap();
1315            let back: PermissionMode = serde_json::from_str(&json).unwrap();
1316            assert_eq!(format!("{:?}", mode), format!("{:?}", back));
1317        }
1318    }
1319
1320    // --- Retry builder ---
1321
1322    #[test]
1323    fn retry_builder_stores_policy() {
1324        let agent = Agent::new().retry(3);
1325        assert!(agent.retry_policy.is_some());
1326        assert_eq!(agent.retry_policy.unwrap().max_retries(), 3);
1327    }
1328
1329    #[test]
1330    fn retry_policy_builder_stores_custom_policy() {
1331        use crate::retry::RetryPolicy;
1332        let policy = RetryPolicy::new(5).backoff(Duration::from_secs(1));
1333        let agent = Agent::new().retry_policy(policy);
1334        let p = agent.retry_policy.unwrap();
1335        assert_eq!(p.max_retries(), 5);
1336    }
1337
1338    #[test]
1339    fn no_retry_by_default() {
1340        let agent = Agent::new();
1341        assert!(agent.retry_policy.is_none());
1342    }
1343
1344    // --- Retry behavior ---
1345
1346    use std::sync::Arc;
1347    use std::sync::atomic::{AtomicU32, Ordering};
1348    use std::time::Duration;
1349
1350    struct FailNTimesProvider {
1351        fail_count: AtomicU32,
1352        failures_before_success: u32,
1353        output: AgentOutput,
1354    }
1355
1356    impl AgentProvider for FailNTimesProvider {
1357        fn invoke<'a>(&'a self, _config: &'a AgentConfig) -> InvokeFuture<'a> {
1358            Box::pin(async move {
1359                let current = self.fail_count.fetch_add(1, Ordering::SeqCst);
1360                if current < self.failures_before_success {
1361                    Err(AgentError::ProcessFailed {
1362                        exit_code: 1,
1363                        stderr: format!("transient failure #{}", current + 1),
1364                    })
1365                } else {
1366                    Ok(AgentOutput {
1367                        value: self.output.value.clone(),
1368                        session_id: self.output.session_id.clone(),
1369                        cost_usd: self.output.cost_usd,
1370                        input_tokens: self.output.input_tokens,
1371                        output_tokens: self.output.output_tokens,
1372                        model: self.output.model.clone(),
1373                        duration_ms: self.output.duration_ms,
1374                        debug_messages: None,
1375                    })
1376                }
1377            })
1378        }
1379    }
1380
1381    #[tokio::test]
1382    async fn retry_succeeds_after_transient_failures() {
1383        let provider = FailNTimesProvider {
1384            fail_count: AtomicU32::new(0),
1385            failures_before_success: 2,
1386            output: default_output(),
1387        };
1388        let result = Agent::new()
1389            .prompt("test")
1390            .retry_policy(crate::retry::RetryPolicy::new(3).backoff(Duration::from_millis(1)))
1391            .run(&provider)
1392            .await;
1393
1394        assert!(result.is_ok());
1395        assert_eq!(provider.fail_count.load(Ordering::SeqCst), 3); // 1 initial + 2 retries
1396    }
1397
1398    #[tokio::test]
1399    async fn retry_exhausted_returns_last_error() {
1400        let provider = FailNTimesProvider {
1401            fail_count: AtomicU32::new(0),
1402            failures_before_success: 10, // always fails
1403            output: default_output(),
1404        };
1405        let result = Agent::new()
1406            .prompt("test")
1407            .retry_policy(crate::retry::RetryPolicy::new(2).backoff(Duration::from_millis(1)))
1408            .run(&provider)
1409            .await;
1410
1411        assert!(result.is_err());
1412        // 1 initial + 2 retries = 3 total
1413        assert_eq!(provider.fail_count.load(Ordering::SeqCst), 3);
1414    }
1415
1416    #[tokio::test]
1417    async fn retry_does_not_retry_prompt_too_large() {
1418        let call_count = Arc::new(AtomicU32::new(0));
1419        let count = call_count.clone();
1420
1421        struct CountingNonRetryable {
1422            count: Arc<AtomicU32>,
1423        }
1424        impl AgentProvider for CountingNonRetryable {
1425            fn invoke<'a>(&'a self, _config: &'a AgentConfig) -> InvokeFuture<'a> {
1426                self.count.fetch_add(1, Ordering::SeqCst);
1427                Box::pin(async move {
1428                    Err(AgentError::PromptTooLarge {
1429                        chars: 1_000_000,
1430                        estimated_tokens: 250_000,
1431                        model_limit: 200_000,
1432                    })
1433                })
1434            }
1435        }
1436
1437        let provider = CountingNonRetryable { count };
1438        let result = Agent::new()
1439            .prompt("test")
1440            .retry_policy(crate::retry::RetryPolicy::new(3).backoff(Duration::from_millis(1)))
1441            .run(&provider)
1442            .await;
1443
1444        assert!(result.is_err());
1445        assert_eq!(call_count.load(Ordering::SeqCst), 1);
1446    }
1447
1448    #[tokio::test]
1449    async fn retry_retries_schema_validation_errors() {
1450        let call_count = Arc::new(AtomicU32::new(0));
1451        let count = call_count.clone();
1452
1453        struct SchemaFailProvider {
1454            count: Arc<AtomicU32>,
1455        }
1456        impl AgentProvider for SchemaFailProvider {
1457            fn invoke<'a>(&'a self, _config: &'a AgentConfig) -> InvokeFuture<'a> {
1458                self.count.fetch_add(1, Ordering::SeqCst);
1459                Box::pin(async move {
1460                    Err(AgentError::SchemaValidation {
1461                        expected: "object".to_string(),
1462                        got: "null".to_string(),
1463                        debug_messages: Vec::new(),
1464                        partial_usage: Box::default(),
1465                        raw_response: None,
1466                    })
1467                })
1468            }
1469        }
1470
1471        let provider = SchemaFailProvider { count };
1472        let result = Agent::new()
1473            .prompt("test")
1474            .retry_policy(crate::retry::RetryPolicy::new(2).backoff(Duration::from_millis(1)))
1475            .run(&provider)
1476            .await;
1477
1478        assert!(result.is_err());
1479        // 1 initial + 2 retries = 3 total
1480        assert_eq!(call_count.load(Ordering::SeqCst), 3);
1481    }
1482
1483    #[tokio::test]
1484    async fn schema_validation_succeeds_on_retry() {
1485        let call_count = Arc::new(AtomicU32::new(0));
1486        let count = call_count.clone();
1487
1488        struct SchemaFailThenSucceed {
1489            count: Arc<AtomicU32>,
1490            output: AgentOutput,
1491        }
1492        impl AgentProvider for SchemaFailThenSucceed {
1493            fn invoke<'a>(&'a self, _config: &'a AgentConfig) -> InvokeFuture<'a> {
1494                let current = self.count.fetch_add(1, Ordering::SeqCst);
1495                let output = self.output.clone();
1496                Box::pin(async move {
1497                    if current == 0 {
1498                        Err(AgentError::SchemaValidation {
1499                            expected: "structured_output field".to_string(),
1500                            got: "null".to_string(),
1501                            debug_messages: Vec::new(),
1502                            partial_usage: Box::default(),
1503                            raw_response: None,
1504                        })
1505                    } else {
1506                        Ok(output)
1507                    }
1508                })
1509            }
1510        }
1511
1512        let provider = SchemaFailThenSucceed {
1513            count,
1514            output: default_output(),
1515        };
1516        let result = Agent::new()
1517            .prompt("test")
1518            .retry_policy(crate::retry::RetryPolicy::new(1).backoff(Duration::from_millis(1)))
1519            .run(&provider)
1520            .await;
1521
1522        assert!(result.is_ok());
1523        assert_eq!(call_count.load(Ordering::SeqCst), 2);
1524    }
1525
1526    #[tokio::test]
1527    async fn auto_retry_applied_when_json_schema_set() {
1528        let call_count = Arc::new(AtomicU32::new(0));
1529        let count = call_count.clone();
1530
1531        struct AlwaysSchemaFail {
1532            count: Arc<AtomicU32>,
1533        }
1534        impl AgentProvider for AlwaysSchemaFail {
1535            fn invoke<'a>(&'a self, _config: &'a AgentConfig) -> InvokeFuture<'a> {
1536                self.count.fetch_add(1, Ordering::SeqCst);
1537                Box::pin(async move {
1538                    Err(AgentError::SchemaValidation {
1539                        expected: "object".to_string(),
1540                        got: "null".to_string(),
1541                        debug_messages: Vec::new(),
1542                        partial_usage: Box::default(),
1543                        raw_response: None,
1544                    })
1545                })
1546            }
1547        }
1548
1549        let provider = AlwaysSchemaFail { count };
1550        let result = Agent::new()
1551            .prompt("test")
1552            .output_schema_raw(r#"{"type":"object"}"#)
1553            .run(&provider)
1554            .await;
1555
1556        assert!(result.is_err());
1557        // auto-retry(2) : 1 initial + 2 retries = 3 total
1558        assert_eq!(call_count.load(Ordering::SeqCst), 3);
1559    }
1560
1561    #[tokio::test]
1562    async fn no_retry_without_policy() {
1563        let provider = FailNTimesProvider {
1564            fail_count: AtomicU32::new(0),
1565            failures_before_success: 1,
1566            output: default_output(),
1567        };
1568        let result = Agent::new().prompt("test").run(&provider).await;
1569
1570        assert!(result.is_err());
1571        assert_eq!(provider.fail_count.load(Ordering::SeqCst), 1);
1572    }
1573
1574    // ── log_sink tests ────────────────────────────────────────────
1575
1576    use crate::test_support::VecSink;
1577
1578    struct SinkCapture {
1579        output: AgentOutput,
1580        saw_logs: Arc<AtomicU32>,
1581    }
1582
1583    impl AgentProvider for SinkCapture {
1584        fn invoke<'a>(&'a self, _config: &'a AgentConfig) -> InvokeFuture<'a> {
1585            Box::pin(async {
1586                Ok(AgentOutput {
1587                    value: self.output.value.clone(),
1588                    session_id: self.output.session_id.clone(),
1589                    cost_usd: self.output.cost_usd,
1590                    input_tokens: self.output.input_tokens,
1591                    output_tokens: self.output.output_tokens,
1592                    model: self.output.model.clone(),
1593                    duration_ms: self.output.duration_ms,
1594                    debug_messages: None,
1595                })
1596            })
1597        }
1598
1599        fn invoke_with_logs<'a>(
1600            &'a self,
1601            config: &'a AgentConfig,
1602            log_sink: Arc<dyn LogSink>,
1603        ) -> InvokeFuture<'a> {
1604            self.saw_logs.fetch_add(1, Ordering::SeqCst);
1605            log_sink.log("stdout", "streaming line");
1606            self.invoke(config)
1607        }
1608    }
1609
1610    #[tokio::test]
1611    async fn log_sink_routes_to_invoke_with_logs() {
1612        let saw_logs = Arc::new(AtomicU32::new(0));
1613        let provider = SinkCapture {
1614            output: default_output(),
1615            saw_logs: saw_logs.clone(),
1616        };
1617        let sink: Arc<dyn LogSink> = VecSink::new();
1618
1619        let result = Agent::new()
1620            .prompt("test")
1621            .log_sink(sink)
1622            .run(&provider)
1623            .await;
1624
1625        assert!(result.is_ok());
1626        assert_eq!(saw_logs.load(Ordering::SeqCst), 1);
1627    }
1628
1629    #[tokio::test]
1630    async fn no_log_sink_routes_to_invoke() {
1631        let saw_logs = Arc::new(AtomicU32::new(0));
1632        let provider = SinkCapture {
1633            output: default_output(),
1634            saw_logs: saw_logs.clone(),
1635        };
1636
1637        let result = Agent::new().prompt("test").run(&provider).await;
1638
1639        assert!(result.is_ok());
1640        assert_eq!(saw_logs.load(Ordering::SeqCst), 0);
1641    }
1642
1643    #[tokio::test]
1644    async fn log_sink_receives_provider_lines() {
1645        let saw_logs = Arc::new(AtomicU32::new(0));
1646        let provider = SinkCapture {
1647            output: default_output(),
1648            saw_logs: saw_logs.clone(),
1649        };
1650        let sink = VecSink::new();
1651
1652        let _ = Agent::new()
1653            .prompt("test")
1654            .log_sink(sink.clone() as Arc<dyn LogSink>)
1655            .run(&provider)
1656            .await;
1657
1658        let lines = sink.0.lock().unwrap();
1659        assert_eq!(lines.len(), 1);
1660        assert_eq!(lines[0].0, "stdout");
1661        assert_eq!(lines[0].1, "streaming line");
1662    }
1663}
ironflow_core/operations/agent.rs

ironflow_core/operations/
agent.rs
