ironflow_core/
provider.rs

1//! Provider trait and configuration types for agent invocations.
2//!
3//! The [`AgentProvider`] trait is the primary extension point in ironflow: implement it
4//! to plug in any AI backend (local model, HTTP API, mock, etc.) without changing
5//! your workflow code.
6//!
7//! Built-in implementations:
8//!
9//! * [`ClaudeCodeProvider`](crate::providers::claude::ClaudeCodeProvider) - local `claude` CLI.
10//! * `SshProvider` - remote via SSH (requires `transport-ssh` feature).
11//! * `DockerProvider` - Docker container (requires `transport-docker` feature).
12//! * `K8sEphemeralProvider` - ephemeral K8s pod (requires `transport-k8s` feature).
13//! * `K8sPersistentProvider` - persistent K8s pod (requires `transport-k8s` feature).
14//! * [`RecordReplayProvider`](crate::providers::record_replay::RecordReplayProvider) -
15//!   records and replays fixtures for deterministic testing.
16
17use std::fmt;
18use std::future::Future;
19use std::pin::Pin;
20
21use schemars::JsonSchema;
22use serde::{Deserialize, Serialize};
23use serde_json::Value;
24use tracing::warn;
25
26use crate::error::AgentError;
27use crate::operations::agent::{Model, PermissionMode};
28
29/// Boxed future returned by [`AgentProvider::invoke`].
30pub type InvokeFuture<'a> =
31    Pin<Box<dyn Future<Output = Result<AgentOutput, AgentError>> + Send + 'a>>;
32
33/// Serializable configuration passed to an [`AgentProvider`] for a single invocation.
34///
35/// Built by [`Agent::run`](crate::operations::agent::Agent::run) from the builder state.
36/// Provider implementations translate these fields into whatever format the underlying
37/// backend expects.
38#[derive(Debug, Clone, Serialize, Deserialize)]
39#[non_exhaustive]
40pub struct AgentConfig {
41    /// Optional system prompt that sets the agent's persona or constraints.
42    pub system_prompt: Option<String>,
43
44    /// The user prompt - the main instruction to the agent.
45    pub prompt: String,
46
47    /// Which model to use for this invocation.
48    ///
49    /// Accepts any string. Use [`Model`] constants for well-known Claude models
50    /// (e.g. `Model::SONNET`), or pass a custom identifier for other providers.
51    #[serde(default = "default_model")]
52    pub model: String,
53
54    /// Allowlist of tool names the agent may invoke (empty = provider default).
55    #[serde(default)]
56    pub allowed_tools: Vec<String>,
57
58    /// Maximum number of agentic turns before the provider should stop.
59    pub max_turns: Option<u32>,
60
61    /// Maximum spend in USD for this single invocation.
62    pub max_budget_usd: Option<f64>,
63
64    /// Working directory for the agent process.
65    pub working_dir: Option<String>,
66
67    /// Path to an MCP server configuration file.
68    pub mcp_config: Option<String>,
69
70    /// Permission mode controlling how the agent handles tool-use approvals.
71    #[serde(default)]
72    pub permission_mode: PermissionMode,
73
74    /// Optional JSON Schema string. When set, the provider should request
75    /// structured (typed) output from the model.
76    #[serde(alias = "output_schema")]
77    pub json_schema: Option<String>,
78
79    /// Optional session ID to resume a previous conversation.
80    ///
81    /// When set, the provider should continue the conversation from the
82    /// specified session rather than starting a new one.
83    pub resume_session_id: Option<String>,
84
85    /// Enable verbose/debug mode to capture the full conversation trace.
86    ///
87    /// When `true`, the provider uses streaming output (`stream-json`) to
88    /// record every assistant message and tool call. The resulting
89    /// [`AgentOutput::debug_messages`] field will contain the conversation
90    /// trace for inspection.
91    #[serde(default)]
92    pub verbose: bool,
93}
94
95fn default_model() -> String {
96    Model::SONNET.to_string()
97}
98
99/// Raw output returned by an [`AgentProvider`] after a successful invocation.
100///
101/// Carries the agent's response value together with usage and billing metadata.
102#[derive(Clone, Debug, Serialize, Deserialize)]
103#[non_exhaustive]
104pub struct AgentOutput {
105    /// The agent's response. A plain [`Value::String`] for text mode, or an
106    /// arbitrary JSON value when a JSON schema was requested.
107    pub value: Value,
108
109    /// Provider-assigned session identifier, useful for resuming conversations.
110    pub session_id: Option<String>,
111
112    /// Total cost in USD for this invocation, if reported by the provider.
113    pub cost_usd: Option<f64>,
114
115    /// Number of input tokens consumed, if reported.
116    pub input_tokens: Option<u64>,
117
118    /// Number of output tokens generated, if reported.
119    pub output_tokens: Option<u64>,
120
121    /// The concrete model identifier used (e.g. `"claude-sonnet-4-20250514"`).
122    pub model: Option<String>,
123
124    /// Wall-clock duration of the invocation in milliseconds.
125    pub duration_ms: u64,
126
127    /// Conversation trace captured when [`AgentConfig::verbose`] is `true`.
128    ///
129    /// Contains every assistant message and tool call made during the
130    /// invocation, in chronological order. `None` when verbose mode is off.
131    pub debug_messages: Option<Vec<DebugMessage>>,
132}
133
134/// A single assistant turn captured during a verbose invocation.
135///
136/// Each `DebugMessage` represents one assistant response, which may contain
137/// free-form text, tool calls, or both.
138///
139/// # Examples
140///
141/// ```no_run
142/// use ironflow_core::prelude::*;
143///
144/// # async fn example() -> Result<(), OperationError> {
145/// let provider = ClaudeCodeProvider::new();
146/// let result = Agent::new()
147///     .prompt("List files in src/")
148///     .verbose()
149///     .run(&provider)
150///     .await?;
151///
152/// if let Some(messages) = result.debug_messages() {
153///     for msg in messages {
154///         println!("{msg}");
155///     }
156/// }
157/// # Ok(())
158/// # }
159/// ```
160#[derive(Debug, Clone, Serialize, Deserialize)]
161#[non_exhaustive]
162pub struct DebugMessage {
163    /// Free-form text produced by the assistant in this turn, if any.
164    pub text: Option<String>,
165
166    /// Tool calls made by the assistant in this turn.
167    pub tool_calls: Vec<DebugToolCall>,
168
169    /// The model's stop reason for this turn (e.g. `"end_turn"`, `"tool_use"`).
170    pub stop_reason: Option<String>,
171}
172
173impl fmt::Display for DebugMessage {
174    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
175        if let Some(ref text) = self.text {
176            writeln!(f, "[assistant] {text}")?;
177        }
178        for tc in &self.tool_calls {
179            write!(f, "{tc}")?;
180        }
181        Ok(())
182    }
183}
184
185/// A single tool call captured during a verbose invocation.
186///
187/// Records the tool name and its input arguments as a raw JSON value.
188#[derive(Debug, Clone, Serialize, Deserialize)]
189#[non_exhaustive]
190pub struct DebugToolCall {
191    /// Name of the tool invoked (e.g. `"Read"`, `"Bash"`, `"Grep"`).
192    pub name: String,
193
194    /// Input arguments passed to the tool, as raw JSON.
195    pub input: Value,
196}
197
198impl fmt::Display for DebugToolCall {
199    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
200        writeln!(f, "  [tool_use] {} -> {}", self.name, self.input)
201    }
202}
203
204impl AgentConfig {
205    /// Create an `AgentConfig` with required fields and defaults for the rest.
206    pub fn new(prompt: &str) -> Self {
207        Self {
208            system_prompt: None,
209            prompt: prompt.to_string(),
210            model: Model::SONNET.to_string(),
211            allowed_tools: Vec::new(),
212            max_turns: None,
213            max_budget_usd: None,
214            working_dir: None,
215            mcp_config: None,
216            permission_mode: PermissionMode::Default,
217            json_schema: None,
218            resume_session_id: None,
219            verbose: false,
220        }
221    }
222
223    /// Set the system prompt.
224    pub fn system_prompt(mut self, prompt: &str) -> Self {
225        self.system_prompt = Some(prompt.to_string());
226        self
227    }
228
229    /// Set the model name.
230    pub fn model(mut self, model: &str) -> Self {
231        self.model = model.to_string();
232        self
233    }
234
235    /// Set the maximum budget in USD.
236    pub fn max_budget_usd(mut self, budget: f64) -> Self {
237        self.max_budget_usd = Some(budget);
238        self
239    }
240
241    /// Set the maximum number of turns.
242    pub fn max_turns(mut self, turns: u32) -> Self {
243        self.max_turns = Some(turns);
244        self
245    }
246
247    /// Add an allowed tool.
248    pub fn allow_tool(mut self, tool: &str) -> Self {
249        self.allowed_tools.push(tool.to_string());
250        self
251    }
252
253    /// Set the working directory.
254    pub fn working_dir(mut self, dir: &str) -> Self {
255        self.working_dir = Some(dir.to_string());
256        self
257    }
258
259    /// Set the permission mode.
260    pub fn permission_mode(mut self, mode: PermissionMode) -> Self {
261        self.permission_mode = mode;
262        self
263    }
264
265    /// Enable verbose/debug mode.
266    pub fn verbose(mut self, enabled: bool) -> Self {
267        self.verbose = enabled;
268        self
269    }
270
271    /// Set structured output from a Rust type implementing [`JsonSchema`].
272    ///
273    /// The schema is serialized once at build time. When set, the provider
274    /// will request typed output conforming to this schema.
275    ///
276    /// **Important:** structured output requires `max_turns >= 2`.
277    pub fn output<T: JsonSchema>(mut self) -> Self {
278        let schema = schemars::schema_for!(T);
279        self.json_schema = match serde_json::to_string(&schema) {
280            Ok(s) => Some(s),
281            Err(e) => {
282                warn!(
283                    error = %e,
284                    type_name = std::any::type_name::<T>(),
285                    "failed to serialize JSON schema, structured output disabled"
286                );
287                None
288            }
289        };
290        self
291    }
292
293    /// Set structured output from a pre-serialized JSON Schema string.
294    pub fn output_schema_raw(mut self, schema: &str) -> Self {
295        self.json_schema = Some(schema.to_string());
296        self
297    }
298
299    /// Set the MCP server configuration file path.
300    pub fn mcp_config(mut self, config: &str) -> Self {
301        self.mcp_config = Some(config.to_string());
302        self
303    }
304
305    /// Set a session ID to resume a previous conversation.
306    pub fn resume(mut self, session_id: &str) -> Self {
307        self.resume_session_id = Some(session_id.to_string());
308        self
309    }
310}
311
312impl AgentOutput {
313    /// Create an `AgentOutput` with the given value and sensible defaults.
314    pub fn new(value: Value) -> Self {
315        Self {
316            value,
317            session_id: None,
318            cost_usd: None,
319            input_tokens: None,
320            output_tokens: None,
321            model: None,
322            duration_ms: 0,
323            debug_messages: None,
324        }
325    }
326}
327
328/// Trait for AI agent backends.
329///
330/// Implement this trait to provide a custom AI backend for [`Agent`](crate::operations::agent::Agent).
331/// The only required method is [`invoke`](AgentProvider::invoke), which takes an
332/// [`AgentConfig`] and returns an [`AgentOutput`] (or an [`AgentError`]).
333///
334/// # Examples
335///
336/// ```no_run
337/// use ironflow_core::provider::{AgentConfig, AgentOutput, AgentProvider, InvokeFuture};
338///
339/// struct MyProvider;
340///
341/// impl AgentProvider for MyProvider {
342///     fn invoke<'a>(&'a self, config: &'a AgentConfig) -> InvokeFuture<'a> {
343///         Box::pin(async move {
344///             // Call your custom backend here...
345///             todo!()
346///         })
347///     }
348/// }
349/// ```
350pub trait AgentProvider: Send + Sync {
351    /// Execute a single agent invocation with the given configuration.
352    ///
353    /// # Errors
354    ///
355    /// Returns [`AgentError`] if the underlying backend process fails,
356    /// times out, or produces output that does not match the requested schema.
357    fn invoke<'a>(&'a self, config: &'a AgentConfig) -> InvokeFuture<'a>;
358}
359
360#[cfg(test)]
361mod tests {
362    use super::*;
363    use serde_json::json;
364
365    fn full_config() -> AgentConfig {
366        AgentConfig {
367            system_prompt: Some("you are helpful".to_string()),
368            prompt: "do stuff".to_string(),
369            model: Model::OPUS.to_string(),
370            allowed_tools: vec!["Read".to_string(), "Write".to_string()],
371            max_turns: Some(10),
372            max_budget_usd: Some(2.5),
373            working_dir: Some("/tmp".to_string()),
374            mcp_config: Some("{}".to_string()),
375            permission_mode: PermissionMode::Auto,
376            json_schema: Some(r#"{"type":"object"}"#.to_string()),
377            resume_session_id: None,
378            verbose: false,
379        }
380    }
381
382    #[test]
383    fn agent_config_serialize_deserialize_roundtrip() {
384        let config = full_config();
385        let json = serde_json::to_string(&config).unwrap();
386        let back: AgentConfig = serde_json::from_str(&json).unwrap();
387
388        assert_eq!(back.system_prompt, Some("you are helpful".to_string()));
389        assert_eq!(back.prompt, "do stuff");
390        assert_eq!(back.allowed_tools, vec!["Read", "Write"]);
391        assert_eq!(back.max_turns, Some(10));
392        assert_eq!(back.max_budget_usd, Some(2.5));
393        assert_eq!(back.working_dir, Some("/tmp".to_string()));
394        assert_eq!(back.mcp_config, Some("{}".to_string()));
395        assert_eq!(back.json_schema, Some(r#"{"type":"object"}"#.to_string()));
396    }
397
398    #[test]
399    fn agent_config_with_all_optional_fields_none() {
400        let config = AgentConfig {
401            system_prompt: None,
402            prompt: "hello".to_string(),
403            model: Model::HAIKU.to_string(),
404            allowed_tools: vec![],
405            max_turns: None,
406            max_budget_usd: None,
407            working_dir: None,
408            mcp_config: None,
409            permission_mode: PermissionMode::Default,
410            json_schema: None,
411            resume_session_id: None,
412            verbose: false,
413        };
414        let json = serde_json::to_string(&config).unwrap();
415        let back: AgentConfig = serde_json::from_str(&json).unwrap();
416
417        assert_eq!(back.system_prompt, None);
418        assert_eq!(back.prompt, "hello");
419        assert!(back.allowed_tools.is_empty());
420        assert_eq!(back.max_turns, None);
421        assert_eq!(back.max_budget_usd, None);
422        assert_eq!(back.working_dir, None);
423        assert_eq!(back.mcp_config, None);
424        assert_eq!(back.json_schema, None);
425    }
426
427    #[test]
428    fn agent_output_serialize_deserialize_roundtrip() {
429        let output = AgentOutput {
430            value: json!({"key": "value"}),
431            session_id: Some("sess-abc".to_string()),
432            cost_usd: Some(0.01),
433            input_tokens: Some(500),
434            output_tokens: Some(200),
435            model: Some("claude-sonnet".to_string()),
436            duration_ms: 3000,
437            debug_messages: None,
438        };
439        let json = serde_json::to_string(&output).unwrap();
440        let back: AgentOutput = serde_json::from_str(&json).unwrap();
441
442        assert_eq!(back.value, json!({"key": "value"}));
443        assert_eq!(back.session_id, Some("sess-abc".to_string()));
444        assert_eq!(back.cost_usd, Some(0.01));
445        assert_eq!(back.input_tokens, Some(500));
446        assert_eq!(back.output_tokens, Some(200));
447        assert_eq!(back.model, Some("claude-sonnet".to_string()));
448        assert_eq!(back.duration_ms, 3000);
449    }
450
451    #[test]
452    fn agent_config_new_has_correct_defaults() {
453        let config = AgentConfig::new("test prompt");
454        assert_eq!(config.prompt, "test prompt");
455        assert_eq!(config.system_prompt, None);
456        assert_eq!(config.model, Model::SONNET);
457        assert!(config.allowed_tools.is_empty());
458        assert_eq!(config.max_turns, None);
459        assert_eq!(config.max_budget_usd, None);
460        assert_eq!(config.working_dir, None);
461        assert_eq!(config.mcp_config, None);
462        assert!(matches!(config.permission_mode, PermissionMode::Default));
463        assert_eq!(config.json_schema, None);
464        assert_eq!(config.resume_session_id, None);
465        assert!(!config.verbose);
466    }
467
468    #[test]
469    fn agent_output_new_has_correct_defaults() {
470        let output = AgentOutput::new(json!("test"));
471        assert_eq!(output.value, json!("test"));
472        assert_eq!(output.session_id, None);
473        assert_eq!(output.cost_usd, None);
474        assert_eq!(output.input_tokens, None);
475        assert_eq!(output.output_tokens, None);
476        assert_eq!(output.model, None);
477        assert_eq!(output.duration_ms, 0);
478        assert!(output.debug_messages.is_none());
479    }
480
481    #[test]
482    fn agent_config_resume_session_roundtrip() {
483        let mut config = AgentConfig::new("test");
484        config.resume_session_id = Some("sess-xyz".to_string());
485        let json = serde_json::to_string(&config).unwrap();
486        let back: AgentConfig = serde_json::from_str(&json).unwrap();
487        assert_eq!(back.resume_session_id, Some("sess-xyz".to_string()));
488    }
489
490    #[test]
491    fn agent_output_debug_does_not_panic() {
492        let output = AgentOutput {
493            value: json!(null),
494            session_id: None,
495            cost_usd: None,
496            input_tokens: None,
497            output_tokens: None,
498            model: None,
499            duration_ms: 0,
500            debug_messages: None,
501        };
502        let debug_str = format!("{:?}", output);
503        assert!(!debug_str.is_empty());
504    }
505}
ironflow_core/provider.rs

ironflow_core/
provider.rs
