ironflow_core/operations/
agent.rs

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

ironflow_core/operations/
agent.rs
