adk_agent/workflow/

loop_agent.rs

#[cfg(feature = "skills")]
use crate::skill_shim::load_skill_index;
use crate::skill_shim::{SelectionPolicy, SkillIndex};
use adk_core::{
    AfterAgentCallback, Agent, BeforeAgentCallback, CallbackContext, Content, Event, EventStream,
    InvocationContext, ReadonlyContext, Result, Session, State,
};
use async_stream::stream;
use async_trait::async_trait;
use std::collections::HashMap;
use std::sync::{Arc, RwLock};

/// Default maximum iterations for LoopAgent when none is specified.
/// Prevents infinite loops from consuming unbounded resources.
pub const DEFAULT_LOOP_MAX_ITERATIONS: u32 = 1000;

/// Loop agent executes sub-agents repeatedly for N iterations or until escalation
pub struct LoopAgent {
    name: String,
    description: String,
    sub_agents: Vec<Arc<dyn Agent>>,
    max_iterations: u32,
    skills_index: Option<Arc<SkillIndex>>,
    skill_policy: SelectionPolicy,
    max_skill_chars: usize,
    before_callbacks: Arc<Vec<BeforeAgentCallback>>,
    after_callbacks: Arc<Vec<AfterAgentCallback>>,
}

impl LoopAgent {
    /// Create a new loop agent with the given name and sub-agents.
    pub fn new(name: impl Into<String>, sub_agents: Vec<Arc<dyn Agent>>) -> Self {
        Self {
            name: name.into(),
            description: String::new(),
            sub_agents,
            max_iterations: DEFAULT_LOOP_MAX_ITERATIONS,
            skills_index: None,
            skill_policy: SelectionPolicy::default(),
            max_skill_chars: 2000,
            before_callbacks: Arc::new(Vec::new()),
            after_callbacks: Arc::new(Vec::new()),
        }
    }

    /// Set the agent description.
    pub fn with_description(mut self, desc: impl Into<String>) -> Self {
        self.description = desc.into();
        self
    }

    /// Set the maximum number of loop iterations.
    ///
    /// Each iteration runs every sub-agent once, in order; an escalation ends the loop early.
    /// `0` runs no iterations: the before- and after-agent callbacks still run and the agent
    /// completes immediately without invoking any sub-agent.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// use adk_agent::LoopAgent;
    ///
    /// let refine = LoopAgent::new("refine", vec![writer, critic]).with_max_iterations(3);
    /// ```
    pub fn with_max_iterations(mut self, max: u32) -> Self {
        self.max_iterations = max;
        self
    }

    /// Set a preloaded skills index for this agent.
    #[cfg(feature = "skills")]
    pub fn with_skills(mut self, index: SkillIndex) -> Self {
        self.skills_index = Some(Arc::new(index));
        self
    }

    /// Auto-load skills from `.skills/` in the current working directory.
    #[cfg(feature = "skills")]
    pub fn with_auto_skills(self) -> Result<Self> {
        self.with_skills_from_root(".")
    }

    /// Auto-load skills from `.skills/` under a custom root directory.
    #[cfg(feature = "skills")]
    pub fn with_skills_from_root(mut self, root: impl AsRef<std::path::Path>) -> Result<Self> {
        let index = load_skill_index(root).map_err(|e| adk_core::AdkError::agent(e.to_string()))?;
        self.skills_index = Some(Arc::new(index));
        Ok(self)
    }

    /// Customize skill selection behavior.
    #[cfg(feature = "skills")]
    pub fn with_skill_policy(mut self, policy: SelectionPolicy) -> Self {
        self.skill_policy = policy;
        self
    }

    /// Limit injected skill content length.
    #[cfg(feature = "skills")]
    pub fn with_skill_budget(mut self, max_chars: usize) -> Self {
        self.max_skill_chars = max_chars;
        self
    }

    /// Add a before-agent callback.
    pub fn before_callback(mut self, callback: BeforeAgentCallback) -> Self {
        if let Some(callbacks) = Arc::get_mut(&mut self.before_callbacks) {
            callbacks.push(callback);
        }
        self
    }

    /// Add an after-agent callback.
    pub fn after_callback(mut self, callback: AfterAgentCallback) -> Self {
        if let Some(callbacks) = Arc::get_mut(&mut self.after_callbacks) {
            callbacks.push(callback);
        }
        self
    }
}

struct HistoryTrackingSession {
    parent_ctx: Arc<dyn InvocationContext>,
    history: Arc<RwLock<Vec<Content>>>,
    /// Event id of the streamed response that produced the last history entry.
    streamed_id: RwLock<Option<String>>,
    state: StateTrackingState,
}

struct StateTrackingState {
    values: RwLock<HashMap<String, serde_json::Value>>,
}

impl StateTrackingState {
    fn new(parent_ctx: &Arc<dyn InvocationContext>) -> Self {
        Self { values: RwLock::new(parent_ctx.session().state().all()) }
    }

    fn apply_delta(&self, delta: &HashMap<String, serde_json::Value>) {
        if delta.is_empty() {
            return;
        }

        let mut values = self.values.write().unwrap_or_else(|e| e.into_inner());
        for (key, value) in delta {
            values.insert(key.clone(), value.clone());
        }
    }
}

impl State for StateTrackingState {
    fn get(&self, key: &str) -> Option<serde_json::Value> {
        self.values.read().unwrap_or_else(|e| e.into_inner()).get(key).cloned()
    }

    fn set(&mut self, key: String, value: serde_json::Value) {
        if let Err(msg) = adk_core::validate_state_key(&key) {
            tracing::warn!(key = %key, "rejecting invalid state key: {msg}");
            return;
        }
        self.values.write().unwrap_or_else(|e| e.into_inner()).insert(key, value);
    }

    fn all(&self) -> HashMap<String, serde_json::Value> {
        self.values.read().unwrap_or_else(|e| e.into_inner()).clone()
    }
}

impl HistoryTrackingSession {
    fn new(parent_ctx: Arc<dyn InvocationContext>) -> Self {
        Self {
            history: Arc::new(RwLock::new(parent_ctx.session().conversation_history())),
            streamed_id: RwLock::new(None),
            state: StateTrackingState::new(&parent_ctx),
            parent_ctx,
        }
    }

    fn apply_event(&self, event: &Event) {
        if let Some(content) = &event.llm_response.content {
            // Consolidate streaming chunks: chunks of one response merge into the
            // entry its first chunk started instead of creating new entries.
            // This prevents N streaming chunks from becoming N separate Content
            // entries that bloat context for subsequent agents.
            let mut history = self.history.write().unwrap_or_else(|e| e.into_inner());
            let mut streamed_id = self.streamed_id.write().unwrap_or_else(|e| e.into_inner());
            let continues_stream = streamed_id.as_deref() == Some(event.id.as_str());

            if event.llm_response.partial {
                // Partial chunk — merge into the entry this response started, so a
                // later complete snapshot replaces only this response.
                if continues_stream
                    && let Some(last) = history.last_mut()
                    && last.role == content.role
                {
                    for part in &content.parts {
                        if let adk_core::Part::Text { text } = part {
                            // Append text to the last Text part
                            if let Some(adk_core::Part::Text { text: existing }) =
                                last.parts.last_mut()
                            {
                                existing.push_str(text);
                            } else {
                                last.parts.push(part.clone());
                            }
                        } else {
                            last.parts.push(part.clone());
                        }
                    }
                    return;
                }
                // No matching last entry — start a new one
                history.push(content.clone());
                *streamed_id = Some(event.id.clone());
            } else {
                let content_complete = event
                    .llm_response
                    .provider_metadata
                    .as_ref()
                    .and_then(|metadata| metadata.get("content_complete"))
                    .and_then(serde_json::Value::as_bool)
                    == Some(true);
                let replaced_stream = continues_stream && content_complete;
                *streamed_id = None;
                if replaced_stream && let Some(last) = history.last_mut() {
                    // A complete snapshot already contains every streamed delta.
                    *last = content.clone();
                } else if let Some(last) = history.last_mut() {
                    // Final event (partial=false) — append as-is.
                    // For non-streaming mode this carries the full content.
                    // Without a complete snapshot the accumulated text is already in
                    // the last history entry from partial merges above, so the final
                    // chunk (which may carry the last fragment or be empty) is
                    // merged if same role, or appended if different.
                    if last.role == content.role && !content.parts.is_empty() {
                        // Merge any remaining text from the final chunk
                        for part in &content.parts {
                            if let adk_core::Part::Text { text } = part {
                                if let Some(adk_core::Part::Text { text: existing }) =
                                    last.parts.last_mut()
                                {
                                    existing.push_str(text);
                                } else {
                                    last.parts.push(part.clone());
                                }
                            } else {
                                last.parts.push(part.clone());
                            }
                        }
                    } else if !content.parts.is_empty() {
                        history.push(content.clone());
                    }
                } else {
                    history.push(content.clone());
                }
            }
        }
        self.state.apply_delta(&event.actions.state_delta);
    }
}

impl Session for HistoryTrackingSession {
    fn id(&self) -> &str {
        self.parent_ctx.session().id()
    }

    fn app_name(&self) -> &str {
        self.parent_ctx.session().app_name()
    }

    fn user_id(&self) -> &str {
        self.parent_ctx.session().user_id()
    }

    fn state(&self) -> &dyn State {
        &self.state
    }

    fn conversation_history(&self) -> Vec<Content> {
        self.history.read().unwrap_or_else(|e| e.into_inner()).clone()
    }

    fn conversation_history_for_agent(&self, _agent_name: &str) -> Vec<Content> {
        self.conversation_history()
    }

    fn append_to_history(&self, content: Content) {
        self.history.write().unwrap_or_else(|e| e.into_inner()).push(content);
    }
}

/// Builds a [`HistoryTrackingContext`] for the wrapper conformance tests.
#[cfg(test)]
pub(crate) fn history_tracking_context_for_test(
    parent: Arc<dyn InvocationContext>,
) -> Arc<dyn InvocationContext> {
    Arc::new(HistoryTrackingContext::new(parent))
}

struct HistoryTrackingContext {
    parent_ctx: Arc<dyn InvocationContext>,
    session: HistoryTrackingSession,
}

impl HistoryTrackingContext {
    fn new(parent_ctx: Arc<dyn InvocationContext>) -> Self {
        let session = HistoryTrackingSession::new(parent_ctx.clone());
        Self { parent_ctx, session }
    }

    fn apply_event(&self, event: &Event) {
        self.session.apply_event(event);
    }
}

#[async_trait]
impl adk_core::ReadonlyContext for HistoryTrackingContext {
    fn invocation_id(&self) -> &str {
        self.parent_ctx.invocation_id()
    }

    fn agent_name(&self) -> &str {
        self.parent_ctx.agent_name()
    }

    fn user_id(&self) -> &str {
        self.parent_ctx.user_id()
    }

    fn app_name(&self) -> &str {
        self.parent_ctx.app_name()
    }

    fn session_id(&self) -> &str {
        self.parent_ctx.session_id()
    }

    fn branch(&self) -> &str {
        self.parent_ctx.branch()
    }

    fn user_content(&self) -> &Content {
        self.parent_ctx.user_content()
    }
}

#[async_trait]
impl CallbackContext for HistoryTrackingContext {
    fn artifacts(&self) -> Option<Arc<dyn adk_core::Artifacts>> {
        self.parent_ctx.artifacts()
    }

    fn shared_state(&self) -> Option<Arc<adk_core::SharedState>> {
        self.parent_ctx.shared_state()
    }
}

#[async_trait]
impl InvocationContext for HistoryTrackingContext {
    fn agent(&self) -> Arc<dyn Agent> {
        self.parent_ctx.agent()
    }

    fn memory(&self) -> Option<Arc<dyn adk_core::Memory>> {
        self.parent_ctx.memory()
    }

    fn session(&self) -> &dyn Session {
        &self.session
    }

    fn run_config(&self) -> &adk_core::RunConfig {
        self.parent_ctx.run_config()
    }

    fn end_invocation(&self) {
        self.parent_ctx.end_invocation();
    }

    fn ended(&self) -> bool {
        self.parent_ctx.ended()
    }

    fn is_cancelled(&self) -> bool {
        self.parent_ctx.is_cancelled()
    }

    fn user_scopes(&self) -> Vec<String> {
        self.parent_ctx.user_scopes()
    }

    fn request_metadata(&self) -> HashMap<String, serde_json::Value> {
        self.parent_ctx.request_metadata()
    }

    fn authoritative_transfer_targets(&self) -> bool {
        self.parent_ctx.authoritative_transfer_targets()
    }
    fn delegation_depth(&self) -> u32 {
        self.parent_ctx.delegation_depth()
    }
    fn max_delegation_depth(&self) -> Option<u32> {
        self.parent_ctx.max_delegation_depth()
    }

    fn orchestration_root_invocation_id(&self) -> &str {
        self.parent_ctx.orchestration_root_invocation_id()
    }

    fn orchestration_edge_id(&self) -> Option<&str> {
        self.parent_ctx.orchestration_edge_id()
    }

    fn requires_tool_confirmation(&self, tool_name: &str) -> bool {
        self.parent_ctx.requires_tool_confirmation(tool_name)
    }

    async fn get_secret(&self, name: &str) -> adk_core::Result<Option<String>> {
        self.parent_ctx.get_secret(name).await
    }

    async fn get_secret_for(
        &self,
        request: &adk_core::SecretRequest,
    ) -> adk_core::Result<Option<String>> {
        self.parent_ctx.get_secret_for(request).await
    }
}

#[async_trait]
impl Agent for LoopAgent {
    fn name(&self) -> &str {
        &self.name
    }

    fn description(&self) -> &str {
        &self.description
    }

    fn sub_agents(&self) -> &[Arc<dyn Agent>] {
        &self.sub_agents
    }

    fn supports_agent_transfer(&self) -> bool {
        // Deterministic workflow agent: on cross-turn resumption the runner
        // must restart from this root so every sub-agent runs again, rather
        // than resuming a single sub-agent that responded last.
        false
    }

    async fn run(&self, ctx: Arc<dyn InvocationContext>) -> Result<EventStream> {
        let sub_agents = self.sub_agents.clone();
        let max_iterations = self.max_iterations;
        let before_callbacks = self.before_callbacks.clone();
        let after_callbacks = self.after_callbacks.clone();
        let agent_name = self.name.clone();
        let run_ctx = super::skill_context::with_skill_injected_context(
            ctx,
            self.skills_index.as_ref(),
            &self.skill_policy,
            self.max_skill_chars,
        );
        let run_ctx = Arc::new(HistoryTrackingContext::new(run_ctx));

        let s = stream! {
            use futures::StreamExt;

            // ===== BEFORE AGENT CALLBACKS =====
            for callback in before_callbacks.as_ref() {
                match callback(run_ctx.clone() as Arc<dyn CallbackContext>).await {
                    Ok(Some(content)) => {
                        let mut early_event = Event::new(run_ctx.invocation_id());
                        early_event.author = agent_name.clone();
                        early_event.llm_response.content = Some(content);
                        yield Ok(early_event);

                        for after_cb in after_callbacks.as_ref() {
                            match after_cb(run_ctx.clone() as Arc<dyn CallbackContext>).await {
                                Ok(Some(after_content)) => {
                                    let mut after_event = Event::new(run_ctx.invocation_id());
                                    after_event.author = agent_name.clone();
                                    after_event.llm_response.content = Some(after_content);
                                    yield Ok(after_event);
                                    return;
                                }
                                Ok(None) => continue,
                                Err(e) => { yield Err(e); return; }
                            }
                        }
                        return;
                    }
                    Ok(None) => continue,
                    Err(e) => { yield Err(e); return; }
                }
            }

            // A range cannot underflow, so `max_iterations == 0` runs the body zero times
            // instead of wrapping to `u32::MAX`.
            for _ in 0..max_iterations {
                let mut should_exit = false;

                for agent in &sub_agents {
                    let mut stream = agent.run(run_ctx.clone() as Arc<dyn InvocationContext>).await?;

                    while let Some(result) = stream.next().await {
                        match result {
                            Ok(event) => {
                                run_ctx.apply_event(&event);
                                if event.actions.escalate {
                                    should_exit = true;
                                }
                                yield Ok(event);
                            }
                            Err(e) => {
                                yield Err(e);
                                return;
                            }
                        }
                    }

                    if should_exit {
                        break;
                    }
                }

                if should_exit {
                    break;
                }
            }

            // ===== AFTER AGENT CALLBACKS =====
            for callback in after_callbacks.as_ref() {
                match callback(run_ctx.clone() as Arc<dyn CallbackContext>).await {
                    Ok(Some(content)) => {
                        let mut after_event = Event::new(run_ctx.invocation_id());
                        after_event.author = agent_name.clone();
                        after_event.llm_response.content = Some(content);
                        yield Ok(after_event);
                        break;
                    }
                    Ok(None) => continue,
                    Err(e) => { yield Err(e); return; }
                }
            }
        };

        Ok(Box::pin(s))
    }
}