feat: implement multi-agent handoff pipeline engine

src/lib.rs

@@ -3,11 +3,13 @@
 pub mod budget;
 pub mod client;
 pub mod persona;
+pub mod pipeline;
 pub mod registry;
 pub mod router;
 
 pub use budget::{BudgetExceeded, BudgetStatus, BudgetTracker, CreditTier};
 pub use client::{ClientError, ModelClient, OpenRouterClient};
 pub use persona::{DispatchResult, ModelAttempt, ModelCandidate, Persona, PersonaRegistry, Tier};
+pub use pipeline::{PipelineConfig, PipelineEvent, PipelineExecutor, PipelineStep};
 pub use registry::{ModelInfo, ModelPricing, ModelRegistry};
 pub use router::{KeywordRouter, LlmRouter, RouteDecision, RouteStage, Router};

src/pipeline.rs

@@ -0,0 +1,377 @@
+use crate::client::{ClientError, ModelClient};
+use crate::persona::{DispatchResult, ModelAttempt, PersonaRegistry, Tier};
+use crate::registry::ModelRegistry;
+use serde::{Deserialize, Serialize};
+use tokio::sync::mpsc;
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct PipelineStep {
+    pub persona: String,
+    #[serde(default)]
+    pub pass_previous_output: bool,
+    #[serde(default)]
+    pub custom_prompt_template: Option<String>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct PipelineConfig {
+    pub name: String,
+    pub steps: Vec<PipelineStep>,
+}
+
+#[derive(Debug, Clone)]
+pub enum PipelineEvent {
+    StepStarted {
+        step_index: usize,
+        total_steps: usize,
+        persona: String,
+        model: String,
+        tier: Tier,
+    },
+    Chunk {
+        step_index: usize,
+        delta: String,
+    },
+    StepCompleted {
+        step_index: usize,
+        persona: String,
+        full_reply: String,
+        fell_back: bool,
+        attempts: Vec<ModelAttempt>,
+    },
+}
+
+pub struct PipelineExecutor;
+
+impl PipelineExecutor {
+    /// Executes a single persona query with candidate fallback walk
+    pub async fn dispatch_persona(
+        persona_name: &str,
+        persona_registry: &PersonaRegistry,
+        model_registry: &ModelRegistry,
+        client: &dyn ModelClient,
+        prompt: &str,
+        allow_stealth: bool,
+        require_paid: bool,
+        chunk_sender: Option<mpsc::Sender<String>>,
+    ) -> Result<DispatchResult, ClientError> {
+        let persona = persona_registry
+            .get(persona_name)
+            .ok_or_else(|| ClientError::ModelUnavailable(format!("Persona '{persona_name}' not found")))?;
+
+        // Filter candidates against live registry and options
+        let eligible = model_registry
+            .filter_eligible(&persona.candidates, allow_stealth, require_paid)
+            .await;
+
+        let candidates_to_try = if !eligible.is_empty() {
+            eligible
+        } else {
+            // If live registry filtering yielded nothing (e.g. offline/strict), fall back to raw persona candidates
+            persona.candidates.iter().filter(|c| allow_stealth || !c.is_stealth).collect()
+        };
+
+        if candidates_to_try.is_empty() {
+            return Err(ClientError::ModelUnavailable(format!(
+                "No eligible candidates available for persona '{persona_name}'"
+            )));
+        }
+
+        let mut attempts = Vec::new();
+        let total_candidates = candidates_to_try.len();
+
+        for (idx, candidate) in candidates_to_try.iter().enumerate() {
+            let model_id = &candidate.id;
+            tracing::info!(
+                persona = %persona_name,
+                model = %model_id,
+                tier = ?candidate.tier,
+                attempt = idx + 1,
+                "Attempting candidate dispatch"
+            );
+
+            // Channel for streaming chunks if sender provided
+            let (tx, mut rx) = mpsc::channel::<String>(100);
+            let forwarder = {
+                let external_sender = chunk_sender.clone();
+                tokio::spawn(async move {
+                    while let Some(chunk) = rx.recv().await {
+                        if let Some(ref ext) = external_sender {
+                            let _ = ext.send(chunk).await;
+                        }
+                    }
+                })
+            };
+
+            let completion_result = client
+                .complete_stream(model_id, &persona.system_prompt, prompt, tx)
+                .await;
+
+            let _ = forwarder.await;
+
+            match completion_result {
+                Ok(reply) => {
+                    let fell_back = idx > 0;
+                    if fell_back {
+                        tracing::warn!(
+                            persona = %persona_name,
+                            model = %model_id,
+                            previous_failures = idx,
+                            "Successful fallback to candidate"
+                        );
+                    }
+                    return Ok(DispatchResult {
+                        persona: persona_name.to_string(),
+                        model_used: model_id.clone(),
+                        tier_used: candidate.tier,
+                        reply,
+                        fell_back,
+                        attempts,
+                    });
+                }
+                Err(e) => {
+                    let err_msg = e.to_string();
+                    tracing::warn!(
+                        persona = %persona_name,
+                        model = %model_id,
+                        error = %err_msg,
+                        "Candidate failed"
+                    );
+
+                    attempts.push(ModelAttempt {
+                        model_id: model_id.clone(),
+                        error: Some(err_msg.clone()),
+                    });
+
+                    // Fast-fail on auth or billing errors
+                    match e {
+                        ClientError::Unauthorized | ClientError::PaymentRequired => return Err(e),
+                        _ => {
+                            // Continue to next candidate
+                            if idx + 1 == total_candidates {
+                                return Err(ClientError::ModelUnavailable(format!(
+                                    "All {} candidates for persona '{persona_name}' failed. Last error: {err_msg}",
+                                    total_candidates
+                                )));
+                            }
+                        }
+                    }
+                }
+            }
+        }
+
+        Err(ClientError::ModelUnavailable(format!(
+            "All candidates for persona '{persona_name}' failed."
+        )))
+    }
+
+    /// Executes an entire pipeline sequentially, passing outputs between steps
+    pub async fn execute(
+        pipeline: &PipelineConfig,
+        persona_registry: &PersonaRegistry,
+        model_registry: &ModelRegistry,
+        client: &dyn ModelClient,
+        initial_prompt: &str,
+        allow_stealth: bool,
+        require_paid: bool,
+        event_sender: Option<mpsc::Sender<PipelineEvent>>,
+    ) -> Result<Vec<DispatchResult>, ClientError> {
+        let mut results = Vec::new();
+        let mut previous_output = String::new();
+        let total_steps = pipeline.steps.len();
+
+        for (step_idx, step) in pipeline.steps.iter().enumerate() {
+            let persona = persona_registry
+                .get(&step.persona)
+                .ok_or_else(|| ClientError::ModelUnavailable(format!("Persona '{}' not found", step.persona)))?;
+
+            // Prepare prompt for this step
+            let step_prompt = if step.pass_previous_output && !previous_output.is_empty() {
+                if let Some(ref template) = step.custom_prompt_template {
+                    template
+                        .replace("{previous_output}", &previous_output)
+                        .replace("{prompt}", initial_prompt)
+                } else {
+                    format!(
+                        "Context from previous step:\n```\n{}\n```\n\nTask:\n{}",
+                        previous_output, initial_prompt
+                    )
+                }
+            } else {
+                initial_prompt.to_string()
+            };
+
+            // Setup chunk forwarding to PipelineEvent::Chunk
+            let (chunk_tx, mut chunk_rx) = mpsc::channel::<String>(100);
+            let event_tx_clone = event_sender.clone();
+            let forwarder = tokio::spawn(async move {
+                while let Some(chunk) = chunk_rx.recv().await {
+                    if let Some(ref tx) = event_tx_clone {
+                        let _ = tx
+                            .send(PipelineEvent::Chunk {
+                                step_index: step_idx,
+                                delta: chunk,
+                            })
+                            .await;
+                    }
+                }
+            });
+
+            // Notify step start
+            if let Some(ref tx) = event_sender {
+                // Peek top candidate to inform UI
+                let top_candidate = persona
+                    .candidates
+                    .first()
+                    .map(|c| (c.id.clone(), c.tier))
+                    .unwrap_or_else(|| ("unknown".into(), Tier::Free));
+
+                let _ = tx
+                    .send(PipelineEvent::StepStarted {
+                        step_index: step_idx,
+                        total_steps,
+                        persona: step.persona.clone(),
+                        model: top_candidate.0,
+                        tier: top_candidate.1,
+                    })
+                    .await;
+            }
+
+            let dispatch_result = Self::dispatch_persona(
+                &step.persona,
+                persona_registry,
+                model_registry,
+                client,
+                &step_prompt,
+                allow_stealth,
+                require_paid,
+                Some(chunk_tx),
+            )
+            .await?;
+
+            let _ = forwarder.await;
+
+            previous_output = dispatch_result.reply.clone();
+
+            if let Some(ref tx) = event_sender {
+                let _ = tx
+                    .send(PipelineEvent::StepCompleted {
+                        step_index: step_idx,
+                        persona: step.persona.clone(),
+                        full_reply: dispatch_result.reply.clone(),
+                        fell_back: dispatch_result.fell_back,
+                        attempts: dispatch_result.attempts.clone(),
+                    })
+                    .await;
+            }
+
+            results.push(dispatch_result);
+        }
+
+        Ok(results)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::persona::{ModelCandidate, Persona};
+    use async_trait::async_trait;
+    use std::sync::atomic::{AtomicUsize, Ordering};
+    use std::sync::Arc;
+
+    struct MockPipelineClient {
+        calls: Arc<AtomicUsize>,
+    }
+
+    #[async_trait]
+    impl ModelClient for MockPipelineClient {
+        async fn complete(&self, model: &str, _system: &str, user: &str) -> Result<String, ClientError> {
+            self.calls.fetch_add(1, Ordering::SeqCst);
+            Ok(format!("Output from {} for: {}", model, user))
+        }
+
+        async fn complete_stream(
+            &self,
+            model: &str,
+            system: &str,
+            user: &str,
+            chunk_sender: mpsc::Sender<String>,
+        ) -> Result<String, ClientError> {
+            let res = self.complete(model, system, user).await?;
+            let _ = chunk_sender.send(res.clone()).await;
+            Ok(res)
+        }
+    }
+
+    #[tokio::test]
+    async fn test_pipeline_sequential_handoff() {
+        let coder = Persona {
+            name: "coder".into(),
+            system_prompt: "You are coder".into(),
+            trigger_words: vec![],
+            candidates: vec![ModelCandidate {
+                id: "coder-model".into(),
+                tier: Tier::Free,
+                is_stealth: false,
+            }],
+        };
+
+        let reviewer = Persona {
+            name: "reviewer".into(),
+            system_prompt: "You are reviewer".into(),
+            trigger_words: vec![],
+            candidates: vec![ModelCandidate {
+                id: "reviewer-model".into(),
+                tier: Tier::Free,
+                is_stealth: false,
+            }],
+        };
+
+        let persona_reg = PersonaRegistry::new(vec![coder, reviewer]);
+        let model_reg = ModelRegistry::new();
+
+        let pipeline = PipelineConfig {
+            name: "build-and-review".into(),
+            steps: vec![
+                PipelineStep {
+                    persona: "coder".into(),
+                    pass_previous_output: false,
+                    custom_prompt_template: None,
+                },
+                PipelineStep {
+                    persona: "reviewer".into(),
+                    pass_previous_output: true,
+                    custom_prompt_template: None,
+                },
+            ],
+        };
+
+        let calls = Arc::new(AtomicUsize::new(0));
+        let client = MockPipelineClient {
+            calls: calls.clone(),
+        };
+
+        let results = PipelineExecutor::execute(
+            &pipeline,
+            &persona_reg,
+            &model_reg,
+            &client,
+            "Build website",
+            false,
+            false,
+            None,
+        )
+        .await
+        .expect("pipeline execution failed");
+
+        assert_eq!(results.len(), 2);
+        assert_eq!(results[0].persona, "coder");
+        assert_eq!(results[0].model_used, "coder-model");
+        assert_eq!(results[1].persona, "reviewer");
+        assert_eq!(results[1].model_used, "reviewer-model");
+        // Reviewer step should contain output from coder step
+        assert!(results[1].reply.contains("Output from coder-model"));
+        assert_eq!(calls.load(Ordering::SeqCst), 2);
+    }
+}