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

ironflow_core/operations/
agent.rs
