heartbit_core/agent/
workflow.rs

1//! Deterministic workflow agent primitives.
2//!
3//! These composable agents orchestrate sub-agents without LLM calls:
4//! - [`SequentialAgent`]: runs agents in order, piping output as input
5//! - [`ParallelAgent`]: runs agents concurrently via `tokio::JoinSet`
6//! - [`LoopAgent`]: repeats a single agent until a condition is met
7
8use std::sync::Arc;
9
10use serde::{Deserialize, Serialize};
11use tokio::task::JoinSet;
12
13use crate::error::Error;
14use crate::llm::LlmProvider;
15use crate::llm::types::TokenUsage;
16
17use super::AgentOutput;
18use super::AgentRunner;
19use super::dag::DagAgent;
20use super::debate::DebateAgent;
21use super::mixture::MixtureOfAgentsAgent;
22use super::voting::VotingAgent;
23
24/// Termination condition for [`LoopAgent`]. Returns `true` to stop the loop.
25type StopCondition = Box<dyn Fn(&str) -> bool + Send + Sync>;
26
27// ---------------------------------------------------------------------------
28// SequentialAgent
29// ---------------------------------------------------------------------------
30
31/// Runs a list of sub-agents in order. Each agent receives the previous
32/// agent's text output as its task input. Returns the final agent's output
33/// with accumulated `TokenUsage`.
34pub struct SequentialAgent<P: LlmProvider> {
35    agents: Vec<AgentRunner<P>>,
36}
37
38impl<P: LlmProvider> std::fmt::Debug for SequentialAgent<P> {
39    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
40        f.debug_struct("SequentialAgent")
41            .field("agent_count", &self.agents.len())
42            .finish()
43    }
44}
45
46/// Builder for [`SequentialAgent`].
47pub struct SequentialAgentBuilder<P: LlmProvider> {
48    agents: Vec<AgentRunner<P>>,
49}
50
51impl<P: LlmProvider> SequentialAgent<P> {
52    /// Create a new [`SequentialAgentBuilder`].
53    ///
54    /// Add agents with `.agent(...)` in execution order; each agent receives
55    /// the previous agent's text output as its task input.
56    ///
57    /// # Example
58    ///
59    /// ```rust,no_run
60    /// use std::sync::Arc;
61    /// use heartbit_core::{
62    ///     AgentRunner, AnthropicProvider, BoxedProvider, SequentialAgent,
63    /// };
64    ///
65    /// # async fn run() -> Result<(), heartbit_core::Error> {
66    /// let provider = Arc::new(BoxedProvider::new(AnthropicProvider::new(
67    ///     "sk-...",
68    ///     "claude-sonnet-4-20250514",
69    /// )));
70    /// let researcher = AgentRunner::builder(provider.clone())
71    ///     .system_prompt("Summarize the topic in 3 bullet points.")
72    ///     .build()?;
73    /// let writer = AgentRunner::builder(provider)
74    ///     .system_prompt("Rewrite as a single engaging paragraph.")
75    ///     .build()?;
76    ///
77    /// let pipeline = SequentialAgent::builder()
78    ///     .agent(researcher)
79    ///     .agent(writer)
80    ///     .build()?;
81    /// let output = pipeline.execute("History of Rust").await?;
82    /// println!("{}", output.result);
83    /// # Ok(()) }
84    /// ```
85    pub fn builder() -> SequentialAgentBuilder<P> {
86        SequentialAgentBuilder { agents: Vec::new() }
87    }
88
89    /// Execute the sequential pipeline, feeding each agent's output as the
90    /// next agent's input.
91    pub async fn execute(&self, task: &str) -> Result<AgentOutput, Error> {
92        let mut current_input = task.to_string();
93        let mut total_usage = TokenUsage::default();
94        let mut total_tool_calls = 0usize;
95        let mut total_cost: Option<f64> = None;
96        let mut last_output: Option<AgentOutput> = None;
97
98        for agent in &self.agents {
99            let result = agent
100                .execute(&current_input)
101                .await
102                .map_err(|e| e.accumulate_usage(total_usage))?;
103            result.accumulate_into(&mut total_usage, &mut total_tool_calls, &mut total_cost);
104            current_input = result.result.clone();
105            last_output = Some(result);
106        }
107
108        // Safety: builder guarantees at least one agent
109        let mut output = last_output.expect("at least one agent");
110        output.tokens_used = total_usage;
111        output.tool_calls_made = total_tool_calls;
112        output.estimated_cost_usd = total_cost;
113        Ok(output)
114    }
115}
116
117impl<P: LlmProvider> SequentialAgentBuilder<P> {
118    /// Add an agent to the sequential pipeline.
119    pub fn agent(mut self, agent: AgentRunner<P>) -> Self {
120        self.agents.push(agent);
121        self
122    }
123
124    /// Add multiple agents to the sequential pipeline.
125    pub fn agents(mut self, agents: Vec<AgentRunner<P>>) -> Self {
126        self.agents.extend(agents);
127        self
128    }
129
130    /// Build the [`SequentialAgent`]. Requires at least one agent.
131    pub fn build(self) -> Result<SequentialAgent<P>, Error> {
132        if self.agents.is_empty() {
133            return Err(Error::Config(
134                "SequentialAgent requires at least one agent".into(),
135            ));
136        }
137        Ok(SequentialAgent {
138            agents: self.agents,
139        })
140    }
141}
142
143// ---------------------------------------------------------------------------
144// ParallelAgent
145// ---------------------------------------------------------------------------
146
147/// Runs multiple sub-agents concurrently via `tokio::JoinSet`. All agents
148/// receive the same input task. Returns merged results with accumulated
149/// `TokenUsage`.
150pub struct ParallelAgent<P: LlmProvider + 'static> {
151    agents: Vec<Arc<AgentRunner<P>>>,
152}
153
154impl<P: LlmProvider + 'static> std::fmt::Debug for ParallelAgent<P> {
155    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
156        f.debug_struct("ParallelAgent")
157            .field("agent_count", &self.agents.len())
158            .finish()
159    }
160}
161
162/// Builder for [`ParallelAgent`].
163pub struct ParallelAgentBuilder<P: LlmProvider + 'static> {
164    agents: Vec<Arc<AgentRunner<P>>>,
165}
166
167impl<P: LlmProvider + 'static> ParallelAgent<P> {
168    /// Create a new [`ParallelAgentBuilder`].
169    pub fn builder() -> ParallelAgentBuilder<P> {
170        ParallelAgentBuilder { agents: Vec::new() }
171    }
172
173    /// Execute all agents concurrently. Fails fast on first error.
174    pub async fn execute(&self, task: &str) -> Result<AgentOutput, Error> {
175        let mut set = JoinSet::new();
176
177        for agent in &self.agents {
178            let agent = Arc::clone(agent);
179            let task = task.to_string();
180            set.spawn(async move {
181                let name = agent.name().to_string();
182                let result = agent.execute(&task).await;
183                (name, result)
184            });
185        }
186
187        let mut results: Vec<(String, AgentOutput)> = Vec::with_capacity(self.agents.len());
188        let mut total_usage = TokenUsage::default();
189        let mut total_tool_calls = 0usize;
190        let mut total_cost: Option<f64> = None;
191
192        while let Some(join_result) = set.join_next().await {
193            let (name, agent_result) = join_result.map_err(|e| {
194                // AP6: a panicked sibling must not discard the tokens already
195                // burned by completed siblings (the domain-error arm below
196                // accumulates; the panic arm previously did not).
197                Error::Agent(format!("parallel agent task panicked: {e}"))
198                    .accumulate_usage(total_usage)
199            })?;
200            let output = agent_result.map_err(|e| e.accumulate_usage(total_usage))?;
201            output.accumulate_into(&mut total_usage, &mut total_tool_calls, &mut total_cost);
202            results.push((name, output));
203        }
204
205        // Sort by agent name for deterministic output ordering
206        results.sort_by(|a, b| a.0.cmp(&b.0));
207
208        let merged_text = results
209            .iter()
210            .map(|(name, output)| format!("## {name}\n{}", output.result))
211            .collect::<Vec<_>>()
212            .join("\n\n");
213
214        Ok(AgentOutput {
215            result: merged_text,
216            tool_calls_made: total_tool_calls,
217            tokens_used: total_usage,
218            structured: None,
219            estimated_cost_usd: total_cost,
220            model_name: None,
221            // Composite agents don't track per-sub-agent tool calls in this aggregate output.
222            tool_call_results: Vec::new(),
223            goal_met: None,
224        })
225    }
226}
227
228impl<P: LlmProvider + 'static> ParallelAgentBuilder<P> {
229    /// Add an agent. Wraps it in `Arc` for concurrent sharing.
230    pub fn agent(mut self, agent: AgentRunner<P>) -> Self {
231        self.agents.push(Arc::new(agent));
232        self
233    }
234
235    /// Add multiple agents.
236    pub fn agents(mut self, agents: Vec<AgentRunner<P>>) -> Self {
237        self.agents.extend(agents.into_iter().map(Arc::new));
238        self
239    }
240
241    /// Build the [`ParallelAgent`]. Requires at least one agent.
242    pub fn build(self) -> Result<ParallelAgent<P>, Error> {
243        if self.agents.is_empty() {
244            return Err(Error::Config(
245                "ParallelAgent requires at least one agent".into(),
246            ));
247        }
248        Ok(ParallelAgent {
249            agents: self.agents,
250        })
251    }
252}
253
254// ---------------------------------------------------------------------------
255// LoopAgent
256// ---------------------------------------------------------------------------
257
258/// Runs a single agent in a loop. Stops when `should_stop` returns `true`
259/// on the output text, or when `max_iterations` is reached. Returns the
260/// final iteration's output with accumulated `TokenUsage`.
261pub struct LoopAgent<P: LlmProvider> {
262    agent: AgentRunner<P>,
263    max_iterations: usize,
264    should_stop: StopCondition,
265}
266
267impl<P: LlmProvider> std::fmt::Debug for LoopAgent<P> {
268    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
269        f.debug_struct("LoopAgent")
270            .field("max_iterations", &self.max_iterations)
271            .finish()
272    }
273}
274
275/// Builder for [`LoopAgent`].
276pub struct LoopAgentBuilder<P: LlmProvider> {
277    agent: Option<AgentRunner<P>>,
278    max_iterations: Option<usize>,
279    should_stop: Option<StopCondition>,
280}
281
282impl<P: LlmProvider> LoopAgent<P> {
283    /// Create a new [`LoopAgentBuilder`].
284    pub fn builder() -> LoopAgentBuilder<P> {
285        LoopAgentBuilder {
286            agent: None,
287            max_iterations: None,
288            should_stop: None,
289        }
290    }
291
292    /// Execute the loop, feeding each iteration's output as the next input.
293    pub async fn execute(&self, task: &str) -> Result<AgentOutput, Error> {
294        let mut current_input = task.to_string();
295        let mut total_usage = TokenUsage::default();
296        let mut total_tool_calls = 0usize;
297        let mut total_cost: Option<f64> = None;
298        let mut last_output: Option<AgentOutput> = None;
299
300        for _ in 0..self.max_iterations {
301            let result = self
302                .agent
303                .execute(&current_input)
304                .await
305                .map_err(|e| e.accumulate_usage(total_usage))?;
306            result.accumulate_into(&mut total_usage, &mut total_tool_calls, &mut total_cost);
307            current_input = result.result.clone();
308            let should_stop = (self.should_stop)(&result.result);
309            last_output = Some(result);
310            if should_stop {
311                break;
312            }
313        }
314
315        // Safety: max_iterations >= 1 guarantees at least one iteration
316        let mut output = last_output.expect("at least one iteration");
317        output.tokens_used = total_usage;
318        output.tool_calls_made = total_tool_calls;
319        output.estimated_cost_usd = total_cost;
320        Ok(output)
321    }
322}
323
324impl<P: LlmProvider> LoopAgentBuilder<P> {
325    /// Set the agent to loop.
326    pub fn agent(mut self, agent: AgentRunner<P>) -> Self {
327        self.agent = Some(agent);
328        self
329    }
330
331    /// Set the maximum number of iterations (must be >= 1).
332    pub fn max_iterations(mut self, n: usize) -> Self {
333        self.max_iterations = Some(n);
334        self
335    }
336
337    /// Set the termination condition. The closure receives the agent's output
338    /// text and returns `true` to stop the loop.
339    pub fn should_stop(mut self, f: impl Fn(&str) -> bool + Send + Sync + 'static) -> Self {
340        self.should_stop = Some(Box::new(f));
341        self
342    }
343
344    /// Build the [`LoopAgent`].
345    pub fn build(self) -> Result<LoopAgent<P>, Error> {
346        let agent = self
347            .agent
348            .ok_or_else(|| Error::Config("LoopAgent requires an agent".into()))?;
349        let max_iterations = self
350            .max_iterations
351            .ok_or_else(|| Error::Config("LoopAgent requires max_iterations".into()))?;
352        if max_iterations == 0 {
353            return Err(Error::Config(
354                "LoopAgent max_iterations must be at least 1".into(),
355            ));
356        }
357        let should_stop = self
358            .should_stop
359            .ok_or_else(|| Error::Config("LoopAgent requires a should_stop condition".into()))?;
360        Ok(LoopAgent {
361            agent,
362            max_iterations,
363            should_stop,
364        })
365    }
366}
367
368// ---------------------------------------------------------------------------
369// WorkflowType
370// ---------------------------------------------------------------------------
371
372/// Identifies which workflow pattern to use.
373#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
374#[serde(rename_all = "snake_case")]
375pub enum WorkflowType {
376    /// Run agents one after another, piping output to input.
377    Sequential,
378    /// Run agents concurrently, merging results.
379    Parallel,
380    /// Repeat an agent until a stop condition is met.
381    Loop,
382    /// Run agents according to a dependency graph.
383    Dag,
384    /// Run agents as debaters with an optional judge.
385    Debate,
386    /// Run agents as voters and aggregate the result.
387    Voting,
388    /// Run proposer agents and synthesize with a final agent.
389    Mixture,
390}
391
392// ---------------------------------------------------------------------------
393// WorkflowRouter
394// ---------------------------------------------------------------------------
395
396/// Routes execution to one of the workflow agent types.
397/// Allows config-driven workflow selection without hardcoding the type.
398pub enum WorkflowRouter<P: LlmProvider + 'static> {
399    /// A sequential pipeline workflow agent.
400    Sequential(Box<SequentialAgent<P>>),
401    /// A parallel (concurrent) workflow agent.
402    Parallel(Box<ParallelAgent<P>>),
403    /// A looping workflow agent.
404    Loop(Box<LoopAgent<P>>),
405    /// A DAG-based workflow agent.
406    Dag(Box<DagAgent<P>>),
407    /// A debate workflow agent.
408    Debate(Box<DebateAgent<P>>),
409    /// A voting workflow agent.
410    Voting(Box<VotingAgent<P>>),
411    /// A mixture-of-agents workflow agent.
412    Mixture(Box<MixtureOfAgentsAgent<P>>),
413}
414
415impl<P: LlmProvider + 'static> WorkflowRouter<P> {
416    /// Execute the contained workflow agent.
417    ///
418    /// Note: `Voting` returns only the winning voter's `AgentOutput`; the
419    /// `VoteResult` metadata (winner string, tally) is discarded. Use
420    /// `VotingAgent::execute()` directly when you need the full result.
421    pub async fn execute(&self, task: &str) -> Result<AgentOutput, Error> {
422        match self {
423            Self::Sequential(a) => a.execute(task).await,
424            Self::Parallel(a) => a.execute(task).await,
425            Self::Loop(a) => a.execute(task).await,
426            Self::Dag(a) => a.execute(task).await,
427            Self::Debate(a) => a.execute(task).await,
428            Self::Mixture(a) => a.execute(task).await,
429            Self::Voting(a) => a.execute(task).await.map(|vr| vr.output),
430        }
431    }
432
433    /// Returns which workflow type this router contains.
434    pub fn workflow_type(&self) -> WorkflowType {
435        match self {
436            Self::Sequential(_) => WorkflowType::Sequential,
437            Self::Parallel(_) => WorkflowType::Parallel,
438            Self::Loop(_) => WorkflowType::Loop,
439            Self::Dag(_) => WorkflowType::Dag,
440            Self::Debate(_) => WorkflowType::Debate,
441            Self::Voting(_) => WorkflowType::Voting,
442            Self::Mixture(_) => WorkflowType::Mixture,
443        }
444    }
445}
446
447impl<P: LlmProvider + 'static> std::fmt::Debug for WorkflowRouter<P> {
448    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
449        f.debug_tuple("WorkflowRouter")
450            .field(&self.workflow_type())
451            .finish()
452    }
453}
454
455impl<P: LlmProvider + 'static> From<SequentialAgent<P>> for WorkflowRouter<P> {
456    fn from(agent: SequentialAgent<P>) -> Self {
457        Self::Sequential(Box::new(agent))
458    }
459}
460
461impl<P: LlmProvider + 'static> From<ParallelAgent<P>> for WorkflowRouter<P> {
462    fn from(agent: ParallelAgent<P>) -> Self {
463        Self::Parallel(Box::new(agent))
464    }
465}
466
467impl<P: LlmProvider + 'static> From<LoopAgent<P>> for WorkflowRouter<P> {
468    fn from(agent: LoopAgent<P>) -> Self {
469        Self::Loop(Box::new(agent))
470    }
471}
472
473impl<P: LlmProvider + 'static> From<DagAgent<P>> for WorkflowRouter<P> {
474    fn from(agent: DagAgent<P>) -> Self {
475        Self::Dag(Box::new(agent))
476    }
477}
478
479impl<P: LlmProvider + 'static> From<DebateAgent<P>> for WorkflowRouter<P> {
480    fn from(agent: DebateAgent<P>) -> Self {
481        Self::Debate(Box::new(agent))
482    }
483}
484
485impl<P: LlmProvider + 'static> From<VotingAgent<P>> for WorkflowRouter<P> {
486    fn from(agent: VotingAgent<P>) -> Self {
487        Self::Voting(Box::new(agent))
488    }
489}
490
491impl<P: LlmProvider + 'static> From<MixtureOfAgentsAgent<P>> for WorkflowRouter<P> {
492    fn from(agent: MixtureOfAgentsAgent<P>) -> Self {
493        Self::Mixture(Box::new(agent))
494    }
495}
496
497// ===========================================================================
498// Tests
499// ===========================================================================
500
501#[cfg(test)]
502mod tests {
503    use super::*;
504    use crate::agent::test_helpers::{MockProvider, make_agent};
505
506    // -----------------------------------------------------------------------
507    // SequentialAgent builder tests
508    // -----------------------------------------------------------------------
509
510    #[test]
511    fn sequential_builder_rejects_empty_agents() {
512        let result = SequentialAgent::<MockProvider>::builder().build();
513        assert!(result.is_err());
514        assert!(
515            result
516                .unwrap_err()
517                .to_string()
518                .contains("at least one agent")
519        );
520    }
521
522    #[test]
523    fn sequential_builder_accepts_one_agent() {
524        let provider = Arc::new(MockProvider::new(vec![MockProvider::text_response(
525            "done", 10, 5,
526        )]));
527        let agent = make_agent(provider, "a");
528        let seq = SequentialAgent::builder().agent(agent).build();
529        assert!(seq.is_ok());
530    }
531
532    // -----------------------------------------------------------------------
533    // SequentialAgent execution tests
534    // -----------------------------------------------------------------------
535
536    #[tokio::test]
537    async fn sequential_single_agent() {
538        let provider = Arc::new(MockProvider::new(vec![MockProvider::text_response(
539            "hello world",
540            100,
541            50,
542        )]));
543        let agent = make_agent(provider, "step1");
544        let seq = SequentialAgent::builder().agent(agent).build().unwrap();
545
546        let output = seq.execute("start").await.unwrap();
547        assert_eq!(output.result, "hello world");
548        assert_eq!(output.tokens_used.input_tokens, 100);
549        assert_eq!(output.tokens_used.output_tokens, 50);
550    }
551
552    #[tokio::test]
553    async fn sequential_chains_output_as_input() {
554        // Agent A responds with "step-a-output", Agent B responds with "step-b-output".
555        // We verify the second agent runs (and its output is final).
556        let provider_a = Arc::new(MockProvider::new(vec![MockProvider::text_response(
557            "step-a-output",
558            100,
559            50,
560        )]));
561        let provider_b = Arc::new(MockProvider::new(vec![MockProvider::text_response(
562            "step-b-output",
563            200,
564            80,
565        )]));
566
567        let agent_a = make_agent(provider_a, "agent-a");
568        let agent_b = make_agent(provider_b, "agent-b");
569
570        let seq = SequentialAgent::builder()
571            .agent(agent_a)
572            .agent(agent_b)
573            .build()
574            .unwrap();
575
576        let output = seq.execute("initial task").await.unwrap();
577        assert_eq!(output.result, "step-b-output");
578        // Usage should be accumulated
579        assert_eq!(output.tokens_used.input_tokens, 300);
580        assert_eq!(output.tokens_used.output_tokens, 130);
581    }
582
583    #[tokio::test]
584    async fn sequential_three_agents_accumulates_usage() {
585        let p1 = Arc::new(MockProvider::new(vec![MockProvider::text_response(
586            "out1", 10, 5,
587        )]));
588        let p2 = Arc::new(MockProvider::new(vec![MockProvider::text_response(
589            "out2", 20, 10,
590        )]));
591        let p3 = Arc::new(MockProvider::new(vec![MockProvider::text_response(
592            "out3", 30, 15,
593        )]));
594
595        let seq = SequentialAgent::builder()
596            .agent(make_agent(p1, "a"))
597            .agent(make_agent(p2, "b"))
598            .agent(make_agent(p3, "c"))
599            .build()
600            .unwrap();
601
602        let output = seq.execute("go").await.unwrap();
603        assert_eq!(output.result, "out3");
604        assert_eq!(output.tokens_used.input_tokens, 60);
605        assert_eq!(output.tokens_used.output_tokens, 30);
606    }
607
608    #[tokio::test]
609    async fn sequential_error_carries_partial_usage() {
610        let p1 = Arc::new(MockProvider::new(vec![MockProvider::text_response(
611            "ok", 100, 50,
612        )]));
613        // Second provider has no responses -> will error
614        let p2 = Arc::new(MockProvider::new(vec![]));
615
616        let seq = SequentialAgent::builder()
617            .agent(make_agent(p1, "good"))
618            .agent(make_agent(p2, "bad"))
619            .build()
620            .unwrap();
621
622        let err = seq.execute("task").await.unwrap_err();
623        let partial = err.partial_usage();
624        // Should include the first agent's usage
625        assert!(partial.input_tokens >= 100);
626    }
627
628    // -----------------------------------------------------------------------
629    // ParallelAgent builder tests
630    // -----------------------------------------------------------------------
631
632    #[test]
633    fn parallel_builder_rejects_empty_agents() {
634        let result = ParallelAgent::<MockProvider>::builder().build();
635        assert!(result.is_err());
636        assert!(
637            result
638                .unwrap_err()
639                .to_string()
640                .contains("at least one agent")
641        );
642    }
643
644    #[test]
645    fn parallel_builder_accepts_one_agent() {
646        let provider = Arc::new(MockProvider::new(vec![MockProvider::text_response(
647            "ok", 10, 5,
648        )]));
649        let agent = make_agent(provider, "a");
650        let par = ParallelAgent::builder().agent(agent).build();
651        assert!(par.is_ok());
652    }
653
654    // -----------------------------------------------------------------------
655    // ParallelAgent execution tests
656    // -----------------------------------------------------------------------
657
658    #[tokio::test]
659    async fn parallel_single_agent() {
660        let provider = Arc::new(MockProvider::new(vec![MockProvider::text_response(
661            "result-a", 100, 50,
662        )]));
663        let agent = make_agent(provider, "agent-a");
664        let par = ParallelAgent::builder().agent(agent).build().unwrap();
665
666        let output = par.execute("task").await.unwrap();
667        assert!(output.result.contains("agent-a"));
668        assert!(output.result.contains("result-a"));
669        assert_eq!(output.tokens_used.input_tokens, 100);
670        assert_eq!(output.tokens_used.output_tokens, 50);
671    }
672
673    #[tokio::test]
674    async fn parallel_multiple_agents_accumulates_usage() {
675        let p1 = Arc::new(MockProvider::new(vec![MockProvider::text_response(
676            "out-a", 100, 50,
677        )]));
678        let p2 = Arc::new(MockProvider::new(vec![MockProvider::text_response(
679            "out-b", 200, 80,
680        )]));
681
682        let par = ParallelAgent::builder()
683            .agent(make_agent(p1, "alpha"))
684            .agent(make_agent(p2, "beta"))
685            .build()
686            .unwrap();
687
688        let output = par.execute("same task").await.unwrap();
689        // Both agent outputs should appear
690        assert!(output.result.contains("out-a"));
691        assert!(output.result.contains("out-b"));
692        // Both headers should appear
693        assert!(output.result.contains("## alpha"));
694        assert!(output.result.contains("## beta"));
695        // Usage accumulated
696        assert_eq!(output.tokens_used.input_tokens, 300);
697        assert_eq!(output.tokens_used.output_tokens, 130);
698    }
699
700    #[tokio::test]
701    async fn parallel_output_sorted_by_name() {
702        let p1 = Arc::new(MockProvider::new(vec![MockProvider::text_response(
703            "out-z", 10, 5,
704        )]));
705        let p2 = Arc::new(MockProvider::new(vec![MockProvider::text_response(
706            "out-a", 10, 5,
707        )]));
708
709        let par = ParallelAgent::builder()
710            .agent(make_agent(p1, "zebra"))
711            .agent(make_agent(p2, "alpha"))
712            .build()
713            .unwrap();
714
715        let output = par.execute("task").await.unwrap();
716        // "alpha" should come before "zebra" in the output
717        let alpha_pos = output.result.find("## alpha").unwrap();
718        let zebra_pos = output.result.find("## zebra").unwrap();
719        assert!(alpha_pos < zebra_pos);
720    }
721
722    #[tokio::test]
723    async fn parallel_error_fails_fast() {
724        let p1 = Arc::new(MockProvider::new(vec![MockProvider::text_response(
725            "ok", 100, 50,
726        )]));
727        // Second provider will error
728        let p2 = Arc::new(MockProvider::new(vec![]));
729
730        let par = ParallelAgent::builder()
731            .agent(make_agent(p1, "good"))
732            .agent(make_agent(p2, "bad"))
733            .build()
734            .unwrap();
735
736        let result = par.execute("task").await;
737        assert!(result.is_err());
738    }
739
740    // -----------------------------------------------------------------------
741    // LoopAgent builder tests
742    // -----------------------------------------------------------------------
743
744    #[test]
745    fn loop_builder_rejects_missing_agent() {
746        let result = LoopAgent::<MockProvider>::builder()
747            .max_iterations(3)
748            .should_stop(|_| true)
749            .build();
750        assert!(result.is_err());
751        assert!(
752            result
753                .unwrap_err()
754                .to_string()
755                .contains("requires an agent")
756        );
757    }
758
759    #[test]
760    fn loop_builder_rejects_missing_max_iterations() {
761        let provider = Arc::new(MockProvider::new(vec![]));
762        let agent = make_agent(provider, "a");
763        let result = LoopAgent::builder()
764            .agent(agent)
765            .should_stop(|_| true)
766            .build();
767        assert!(result.is_err());
768        assert!(
769            result
770                .unwrap_err()
771                .to_string()
772                .contains("requires max_iterations")
773        );
774    }
775
776    #[test]
777    fn loop_builder_rejects_zero_max_iterations() {
778        let provider = Arc::new(MockProvider::new(vec![]));
779        let agent = make_agent(provider, "a");
780        let result = LoopAgent::builder()
781            .agent(agent)
782            .max_iterations(0)
783            .should_stop(|_| true)
784            .build();
785        assert!(result.is_err());
786        assert!(result.unwrap_err().to_string().contains("at least 1"));
787    }
788
789    #[test]
790    fn loop_builder_rejects_missing_should_stop() {
791        let provider = Arc::new(MockProvider::new(vec![]));
792        let agent = make_agent(provider, "a");
793        let result = LoopAgent::builder().agent(agent).max_iterations(3).build();
794        assert!(result.is_err());
795        assert!(
796            result
797                .unwrap_err()
798                .to_string()
799                .contains("requires a should_stop")
800        );
801    }
802
803    #[test]
804    fn loop_builder_accepts_valid_config() {
805        let provider = Arc::new(MockProvider::new(vec![MockProvider::text_response(
806            "x", 1, 1,
807        )]));
808        let agent = make_agent(provider, "a");
809        let result = LoopAgent::builder()
810            .agent(agent)
811            .max_iterations(5)
812            .should_stop(|_| true)
813            .build();
814        assert!(result.is_ok());
815    }
816
817    // -----------------------------------------------------------------------
818    // LoopAgent execution tests
819    // -----------------------------------------------------------------------
820
821    #[tokio::test]
822    async fn loop_stops_on_condition() {
823        // Provide 3 responses: the second one contains "DONE"
824        let provider = Arc::new(MockProvider::new(vec![
825            MockProvider::text_response("working...", 10, 5),
826            MockProvider::text_response("DONE", 10, 5),
827            MockProvider::text_response("should not reach", 10, 5),
828        ]));
829        let agent = make_agent(provider, "worker");
830
831        let loop_agent = LoopAgent::builder()
832            .agent(agent)
833            .max_iterations(10)
834            .should_stop(|text| text.contains("DONE"))
835            .build()
836            .unwrap();
837
838        let output = loop_agent.execute("start").await.unwrap();
839        assert_eq!(output.result, "DONE");
840        // Only 2 iterations ran
841        assert_eq!(output.tokens_used.input_tokens, 20);
842        assert_eq!(output.tokens_used.output_tokens, 10);
843    }
844
845    #[tokio::test]
846    async fn loop_stops_at_max_iterations() {
847        let provider = Arc::new(MockProvider::new(vec![
848            MockProvider::text_response("iter1", 10, 5),
849            MockProvider::text_response("iter2", 10, 5),
850            MockProvider::text_response("iter3", 10, 5),
851        ]));
852        let agent = make_agent(provider, "worker");
853
854        let loop_agent = LoopAgent::builder()
855            .agent(agent)
856            .max_iterations(3)
857            .should_stop(|_| false) // never stop
858            .build()
859            .unwrap();
860
861        let output = loop_agent.execute("start").await.unwrap();
862        assert_eq!(output.result, "iter3");
863        assert_eq!(output.tokens_used.input_tokens, 30);
864        assert_eq!(output.tokens_used.output_tokens, 15);
865    }
866
867    #[tokio::test]
868    async fn loop_single_iteration() {
869        let provider = Arc::new(MockProvider::new(vec![MockProvider::text_response(
870            "once", 50, 25,
871        )]));
872        let agent = make_agent(provider, "worker");
873
874        let loop_agent = LoopAgent::builder()
875            .agent(agent)
876            .max_iterations(1)
877            .should_stop(|_| false)
878            .build()
879            .unwrap();
880
881        let output = loop_agent.execute("go").await.unwrap();
882        assert_eq!(output.result, "once");
883        assert_eq!(output.tokens_used.input_tokens, 50);
884    }
885
886    #[tokio::test]
887    async fn loop_error_carries_partial_usage() {
888        // First response succeeds, second errors
889        let provider = Arc::new(MockProvider::new(vec![MockProvider::text_response(
890            "ok", 100, 50,
891        )]));
892        let agent = make_agent(provider, "worker");
893
894        let loop_agent = LoopAgent::builder()
895            .agent(agent)
896            .max_iterations(5)
897            .should_stop(|_| false) // never stop, will error on 2nd iteration
898            .build()
899            .unwrap();
900
901        let err = loop_agent.execute("go").await.unwrap_err();
902        let partial = err.partial_usage();
903        assert!(partial.input_tokens >= 100);
904    }
905
906    // -----------------------------------------------------------------------
907    // SequentialAgent builder .agents() method
908    // -----------------------------------------------------------------------
909
910    #[test]
911    fn sequential_builder_agents_method() {
912        let p1 = Arc::new(MockProvider::new(vec![MockProvider::text_response(
913            "a", 1, 1,
914        )]));
915        let p2 = Arc::new(MockProvider::new(vec![MockProvider::text_response(
916            "b", 1, 1,
917        )]));
918        let agents = vec![make_agent(p1, "x"), make_agent(p2, "y")];
919        let seq = SequentialAgent::builder().agents(agents).build();
920        assert!(seq.is_ok());
921    }
922
923    // -----------------------------------------------------------------------
924    // ParallelAgent builder .agents() method
925    // -----------------------------------------------------------------------
926
927    #[test]
928    fn parallel_builder_agents_method() {
929        let p1 = Arc::new(MockProvider::new(vec![MockProvider::text_response(
930            "a", 1, 1,
931        )]));
932        let p2 = Arc::new(MockProvider::new(vec![MockProvider::text_response(
933            "b", 1, 1,
934        )]));
935        let agents = vec![make_agent(p1, "x"), make_agent(p2, "y")];
936        let par = ParallelAgent::builder().agents(agents).build();
937        assert!(par.is_ok());
938    }
939
940    // -----------------------------------------------------------------------
941    // AgentRunner::name() getter test
942    // -----------------------------------------------------------------------
943
944    #[test]
945    fn agent_runner_name_getter() {
946        let provider = Arc::new(MockProvider::new(vec![]));
947        let agent = make_agent(provider, "test-agent");
948        assert_eq!(agent.name(), "test-agent");
949    }
950
951    // -----------------------------------------------------------------------
952    // WorkflowType tests
953    // -----------------------------------------------------------------------
954
955    #[test]
956    fn workflow_type_serde_roundtrip() {
957        for wt in [
958            WorkflowType::Sequential,
959            WorkflowType::Parallel,
960            WorkflowType::Loop,
961            WorkflowType::Dag,
962            WorkflowType::Debate,
963            WorkflowType::Voting,
964            WorkflowType::Mixture,
965        ] {
966            let json = serde_json::to_string(&wt).unwrap();
967            let back: WorkflowType = serde_json::from_str(&json).unwrap();
968            assert_eq!(wt, back);
969        }
970    }
971
972    #[test]
973    fn workflow_type_snake_case() {
974        assert_eq!(
975            serde_json::to_string(&WorkflowType::Sequential).unwrap(),
976            "\"sequential\""
977        );
978        assert_eq!(
979            serde_json::to_string(&WorkflowType::Parallel).unwrap(),
980            "\"parallel\""
981        );
982        assert_eq!(
983            serde_json::to_string(&WorkflowType::Loop).unwrap(),
984            "\"loop\""
985        );
986        assert_eq!(
987            serde_json::to_string(&WorkflowType::Dag).unwrap(),
988            "\"dag\""
989        );
990        assert_eq!(
991            serde_json::to_string(&WorkflowType::Debate).unwrap(),
992            "\"debate\""
993        );
994        assert_eq!(
995            serde_json::to_string(&WorkflowType::Voting).unwrap(),
996            "\"voting\""
997        );
998        assert_eq!(
999            serde_json::to_string(&WorkflowType::Mixture).unwrap(),
1000            "\"mixture\""
1001        );
1002    }
1003
1004    // -----------------------------------------------------------------------
1005    // WorkflowRouter tests
1006    // -----------------------------------------------------------------------
1007
1008    #[tokio::test]
1009    async fn router_sequential() {
1010        let provider = Arc::new(MockProvider::new(vec![MockProvider::text_response(
1011            "seq-out", 10, 5,
1012        )]));
1013        let seq = SequentialAgent::builder()
1014            .agent(make_agent(provider, "s"))
1015            .build()
1016            .unwrap();
1017        let router = WorkflowRouter::Sequential(Box::new(seq));
1018        assert_eq!(router.workflow_type(), WorkflowType::Sequential);
1019        let output = router.execute("task").await.unwrap();
1020        assert_eq!(output.result, "seq-out");
1021    }
1022
1023    #[tokio::test]
1024    async fn router_parallel() {
1025        let provider = Arc::new(MockProvider::new(vec![MockProvider::text_response(
1026            "par-out", 10, 5,
1027        )]));
1028        let par = ParallelAgent::builder()
1029            .agent(make_agent(provider, "p"))
1030            .build()
1031            .unwrap();
1032        let router = WorkflowRouter::Parallel(Box::new(par));
1033        assert_eq!(router.workflow_type(), WorkflowType::Parallel);
1034        let output = router.execute("task").await.unwrap();
1035        assert!(output.result.contains("par-out"));
1036    }
1037
1038    #[tokio::test]
1039    async fn router_loop() {
1040        let provider = Arc::new(MockProvider::new(vec![MockProvider::text_response(
1041            "loop-out", 10, 5,
1042        )]));
1043        let lp = LoopAgent::builder()
1044            .agent(make_agent(provider, "l"))
1045            .max_iterations(1)
1046            .should_stop(|_| true)
1047            .build()
1048            .unwrap();
1049        let router = WorkflowRouter::Loop(Box::new(lp));
1050        assert_eq!(router.workflow_type(), WorkflowType::Loop);
1051        let output = router.execute("task").await.unwrap();
1052        assert_eq!(output.result, "loop-out");
1053    }
1054
1055    #[tokio::test]
1056    async fn router_dag() {
1057        use crate::agent::dag::DagAgent;
1058
1059        let provider = Arc::new(MockProvider::new(vec![MockProvider::text_response(
1060            "dag-out", 10, 5,
1061        )]));
1062        let dag = DagAgent::builder()
1063            .node("A", make_agent(provider, "A"))
1064            .build()
1065            .unwrap();
1066        let router = WorkflowRouter::Dag(Box::new(dag));
1067        assert_eq!(router.workflow_type(), WorkflowType::Dag);
1068        let output = router.execute("task").await.unwrap();
1069        assert_eq!(output.result, "dag-out");
1070    }
1071
1072    #[test]
1073    fn router_from_sequential() {
1074        let provider = Arc::new(MockProvider::new(vec![MockProvider::text_response(
1075            "x", 1, 1,
1076        )]));
1077        let seq = SequentialAgent::builder()
1078            .agent(make_agent(provider, "s"))
1079            .build()
1080            .unwrap();
1081        let router: WorkflowRouter<MockProvider> = seq.into();
1082        assert_eq!(router.workflow_type(), WorkflowType::Sequential);
1083    }
1084
1085    #[test]
1086    fn router_from_dag() {
1087        use crate::agent::dag::DagAgent;
1088
1089        let provider = Arc::new(MockProvider::new(vec![MockProvider::text_response(
1090            "x", 1, 1,
1091        )]));
1092        let dag = DagAgent::builder()
1093            .node("A", make_agent(provider, "A"))
1094            .build()
1095            .unwrap();
1096        let router: WorkflowRouter<MockProvider> = dag.into();
1097        assert_eq!(router.workflow_type(), WorkflowType::Dag);
1098    }
1099
1100    #[test]
1101    fn router_debug() {
1102        let provider = Arc::new(MockProvider::new(vec![MockProvider::text_response(
1103            "x", 1, 1,
1104        )]));
1105        let seq = SequentialAgent::builder()
1106            .agent(make_agent(provider, "s"))
1107            .build()
1108            .unwrap();
1109        let router = WorkflowRouter::Sequential(Box::new(seq));
1110        let debug = format!("{router:?}");
1111        assert!(debug.contains("WorkflowRouter"));
1112        assert!(debug.contains("Sequential"));
1113    }
1114}
heartbit_core/agent/workflow.rs

heartbit_core/agent/
workflow.rs
