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

ironflow_core/operations/
agent.rs
