agents_runtime/agent/

runtime.rs

//! Deep Agent runtime implementation
//!
//! This module contains the core DeepAgent struct and its runtime behavior,
//! including message handling, tool execution, HITL support, and state management.

use super::config::DeepAgentConfig;
use crate::middleware::{
    AgentMiddleware, AnthropicPromptCachingMiddleware, BaseSystemPromptMiddleware,
    DeepAgentPromptMiddleware, FilesystemMiddleware, HumanInLoopMiddleware, MiddlewareContext,
    ModelRequest, PlanningMiddleware, SubAgentDescriptor, SubAgentMiddleware, SubAgentRegistration,
    SummarizationMiddleware,
};
use crate::planner::LlmBackedPlanner;
use agents_core::agent::{
    AgentDescriptor, AgentHandle, PlannerAction, PlannerContext, PlannerHandle,
};
use agents_core::hitl::{AgentInterrupt, HitlAction};
use agents_core::messaging::{AgentMessage, MessageContent, MessageMetadata, MessageRole};
use agents_core::persistence::{Checkpointer, ThreadId};
use agents_core::state::AgentStateSnapshot;
use agents_core::tools::{ToolBox, ToolContext, ToolResult};
use async_trait::async_trait;
use serde_json::Value;
use std::collections::{HashMap, HashSet};
use std::sync::{Arc, RwLock};

// Built-in tool names exposed by middlewares. The `task` tool for subagents is not gated.
const BUILTIN_TOOL_NAMES: &[&str] = &["write_todos", "ls", "read_file", "write_file", "edit_file"];

// (no streaming types in baseline)

/// Core Deep Agent runtime implementation
///
/// This struct contains all the runtime state and behavior for a Deep Agent,
/// including middleware management, tool execution, HITL support, and state persistence.
pub struct DeepAgent {
    descriptor: AgentDescriptor,
    instructions: String,
    planner: Arc<dyn PlannerHandle>,
    middlewares: Vec<Arc<dyn AgentMiddleware>>,
    base_tools: Vec<ToolBox>,
    state: Arc<RwLock<AgentStateSnapshot>>,
    history: Arc<RwLock<Vec<AgentMessage>>>,
    _summarization: Option<Arc<SummarizationMiddleware>>,
    _hitl: Option<Arc<HumanInLoopMiddleware>>,
    builtin_tools: Option<HashSet<String>>,
    checkpointer: Option<Arc<dyn Checkpointer>>,
}

impl DeepAgent {
    fn collect_tools(&self) -> HashMap<String, ToolBox> {
        let mut tools: HashMap<String, ToolBox> = HashMap::new();
        for tool in &self.base_tools {
            tools.insert(tool.schema().name.clone(), tool.clone());
        }
        for middleware in &self.middlewares {
            for tool in middleware.tools() {
                let tool_name = tool.schema().name.clone();
                if self.should_include(&tool_name) {
                    tools.insert(tool_name, tool);
                }
            }
        }
        tools
    }
    // no streaming path in baseline

    fn should_include(&self, name: &str) -> bool {
        let is_builtin = BUILTIN_TOOL_NAMES.contains(&name);
        if !is_builtin {
            return true;
        }
        match &self.builtin_tools {
            None => true,
            Some(selected) => selected.contains(name),
        }
    }

    fn append_history(&self, message: AgentMessage) {
        if let Ok(mut history) = self.history.write() {
            history.push(message);
        }
    }

    fn current_history(&self) -> Vec<AgentMessage> {
        self.history.read().map(|h| h.clone()).unwrap_or_default()
    }

    /// Save the current agent state to the configured checkpointer.
    pub async fn save_state(&self, thread_id: &ThreadId) -> anyhow::Result<()> {
        if let Some(ref checkpointer) = self.checkpointer {
            let state = self
                .state
                .read()
                .map_err(|_| anyhow::anyhow!("Failed to read agent state"))?
                .clone();
            checkpointer.save_state(thread_id, &state).await
        } else {
            tracing::warn!("Attempted to save state but no checkpointer is configured");
            Ok(())
        }
    }

    /// Load agent state from the configured checkpointer.
    pub async fn load_state(&self, thread_id: &ThreadId) -> anyhow::Result<bool> {
        if let Some(ref checkpointer) = self.checkpointer {
            if let Some(saved_state) = checkpointer.load_state(thread_id).await? {
                *self
                    .state
                    .write()
                    .map_err(|_| anyhow::anyhow!("Failed to write agent state"))? = saved_state;
                tracing::info!(thread_id = %thread_id, "Loaded agent state from checkpointer");
                Ok(true)
            } else {
                tracing::debug!(thread_id = %thread_id, "No saved state found for thread");
                Ok(false)
            }
        } else {
            tracing::warn!("Attempted to load state but no checkpointer is configured");
            Ok(false)
        }
    }

    /// Delete saved state for the specified thread.
    pub async fn delete_thread(&self, thread_id: &ThreadId) -> anyhow::Result<()> {
        if let Some(ref checkpointer) = self.checkpointer {
            checkpointer.delete_thread(thread_id).await
        } else {
            tracing::warn!("Attempted to delete thread state but no checkpointer is configured");
            Ok(())
        }
    }

    /// List all threads with saved state.
    pub async fn list_threads(&self) -> anyhow::Result<Vec<ThreadId>> {
        if let Some(ref checkpointer) = self.checkpointer {
            checkpointer.list_threads().await
        } else {
            Ok(Vec::new())
        }
    }

    async fn execute_tool(
        &self,
        tool: ToolBox,
        _tool_name: String,
        payload: Value,
    ) -> anyhow::Result<AgentMessage> {
        let state_snapshot = self.state.read().unwrap().clone();
        let ctx = ToolContext::with_mutable_state(Arc::new(state_snapshot), self.state.clone());

        let result = tool.execute(payload, ctx).await?;
        Ok(self.apply_tool_result(result))
    }

    fn apply_tool_result(&self, result: ToolResult) -> AgentMessage {
        match result {
            ToolResult::Message(message) => {
                // Tool results are not added to conversation history
                // Only the final LLM response after tool execution is added
                message
            }
            ToolResult::WithStateUpdate {
                message,
                state_diff,
            } => {
                if let Ok(mut state) = self.state.write() {
                    let command = agents_core::command::Command::with_state(state_diff);
                    command.apply_to(&mut state);
                }
                // Tool results are not added to conversation history
                // Only the final LLM response after tool execution is added
                message
            }
        }
    }

    /// Get the current pending interrupt, if any.
    pub fn current_interrupt(&self) -> Option<AgentInterrupt> {
        self.state
            .read()
            .ok()
            .and_then(|guard| guard.pending_interrupts.first().cloned())
    }

    /// Resume execution after human approval of an interrupt.
    pub async fn resume_with_approval(&self, action: HitlAction) -> anyhow::Result<AgentMessage> {
        // Get the first pending interrupt
        let interrupt = {
            let state_guard = self
                .state
                .read()
                .map_err(|_| anyhow::anyhow!("Failed to acquire read lock on state"))?;
            state_guard
                .pending_interrupts
                .first()
                .cloned()
                .ok_or_else(|| anyhow::anyhow!("No pending interrupts"))?
        };

        let result_message = match action {
            HitlAction::Accept => {
                // Execute with original args
                let AgentInterrupt::HumanInLoop(hitl) = interrupt;
                tracing::info!(
                    tool_name = %hitl.tool_name,
                    call_id = %hitl.call_id,
                    "✅ HITL: Tool approved, executing with original arguments"
                );

                let tools = self.collect_tools();
                let tool = tools
                    .get(&hitl.tool_name)
                    .cloned()
                    .ok_or_else(|| anyhow::anyhow!("Tool '{}' not found", hitl.tool_name))?;

                self.execute_tool(tool, hitl.tool_name, hitl.tool_args)
                    .await?
            }

            HitlAction::Edit {
                tool_name,
                tool_args,
            } => {
                // Execute with modified args
                tracing::info!(
                    tool_name = %tool_name,
                    "✏️ HITL: Tool edited, executing with modified arguments"
                );

                let tools = self.collect_tools();
                let tool = tools
                    .get(&tool_name)
                    .cloned()
                    .ok_or_else(|| anyhow::anyhow!("Tool '{}' not found", tool_name))?;

                self.execute_tool(tool, tool_name, tool_args).await?
            }

            HitlAction::Reject { reason } => {
                // Don't execute - return rejection message
                tracing::info!("❌ HITL: Tool rejected");

                let text = reason
                    .unwrap_or_else(|| "Tool execution rejected by human reviewer.".to_string());

                let message = AgentMessage {
                    role: MessageRole::Tool,
                    content: MessageContent::Text(text),
                    metadata: None,
                };

                self.append_history(message.clone());
                message
            }

            HitlAction::Respond { message } => {
                // Don't execute - return custom message
                tracing::info!("💬 HITL: Custom response provided");

                self.append_history(message.clone());
                message
            }
        };

        // Clear the interrupt from state
        {
            let mut state_guard = self
                .state
                .write()
                .map_err(|_| anyhow::anyhow!("Failed to acquire write lock on state"))?;
            state_guard.clear_interrupts();
        }

        // Persist cleared state
        if let Some(checkpointer) = &self.checkpointer {
            let state_clone = self
                .state
                .read()
                .map_err(|_| anyhow::anyhow!("Failed to acquire read lock on state"))?
                .clone();
            checkpointer
                .save_state(&ThreadId::default(), &state_clone)
                .await?;
        }

        Ok(result_message)
    }

    /// Handle message from string input - converts string to AgentMessage internally
    pub async fn handle_message(
        &self,
        input: impl AsRef<str>,
        state: Arc<AgentStateSnapshot>,
    ) -> anyhow::Result<AgentMessage> {
        self.handle_message_with_metadata(input, None, state).await
    }

    /// Handle message from string input with metadata - converts string to AgentMessage internally
    pub async fn handle_message_with_metadata(
        &self,
        input: impl AsRef<str>,
        metadata: Option<MessageMetadata>,
        state: Arc<AgentStateSnapshot>,
    ) -> anyhow::Result<AgentMessage> {
        let agent_message = AgentMessage {
            role: MessageRole::User,
            content: MessageContent::Text(input.as_ref().to_string()),
            metadata,
        };
        self.handle_message_internal(agent_message, state).await
    }

    /// Internal method that contains the actual message handling logic
    async fn handle_message_internal(
        &self,
        input: AgentMessage,
        _state: Arc<AgentStateSnapshot>,
    ) -> anyhow::Result<AgentMessage> {
        self.append_history(input.clone());

        // Build request with current history
        let mut request = ModelRequest::new(&self.instructions, self.current_history());
        let tools = self.collect_tools();
        for middleware in &self.middlewares {
            let mut ctx = MiddlewareContext::with_request(&mut request, self.state.clone());
            middleware.modify_model_request(&mut ctx).await?;
        }

        let tool_schemas: Vec<_> = tools.values().map(|t| t.schema()).collect();
        let context = PlannerContext {
            history: request.messages.clone(),
            system_prompt: request.system_prompt.clone(),
            tools: tool_schemas,
        };
        let state_snapshot = Arc::new(self.state.read().map(|s| s.clone()).unwrap_or_default());

        // Ask LLM what to do
        let decision = self.planner.plan(context, state_snapshot).await?;

        match decision.next_action {
            PlannerAction::Respond { message } => {
                self.append_history(message.clone());
                Ok(message)
            }
            PlannerAction::CallTool { tool_name, payload } => {
                if let Some(tool) = tools.get(&tool_name).cloned() {
                    // Check all middleware for interrupts before executing tool
                    let call_id = format!("call_{}", uuid::Uuid::new_v4());
                    for middleware in &self.middlewares {
                        if let Some(interrupt) = middleware
                            .before_tool_execution(&tool_name, &payload, &call_id)
                            .await?
                        {
                            // Save interrupt to state
                            {
                                let mut state_guard = self.state.write().map_err(|_| {
                                    anyhow::anyhow!("Failed to acquire write lock on state")
                                })?;
                                state_guard.add_interrupt(interrupt.clone());
                            }

                            // Persist state with checkpointer
                            if let Some(checkpointer) = &self.checkpointer {
                                let state_clone = self
                                    .state
                                    .read()
                                    .map_err(|_| {
                                        anyhow::anyhow!("Failed to acquire read lock on state")
                                    })?
                                    .clone();
                                checkpointer
                                    .save_state(&ThreadId::default(), &state_clone)
                                    .await?;
                            }

                            // Return interrupt message - execution pauses here
                            let interrupt_message = AgentMessage {
                                role: MessageRole::System,
                                content: MessageContent::Text(format!(
                                    "⏸️ Execution paused: Tool '{}' requires human approval",
                                    tool_name
                                )),
                                metadata: None,
                            };
                            self.append_history(interrupt_message.clone());
                            return Ok(interrupt_message);
                        }
                    }

                    // No interrupt - execute tool
                    let start_time = std::time::Instant::now();
                    tracing::warn!(
                        "⚙️ EXECUTING TOOL: {} with payload: {}",
                        tool_name,
                        serde_json::to_string(&payload)
                            .unwrap_or_else(|_| "invalid json".to_string())
                    );

                    let result = self
                        .execute_tool(tool.clone(), tool_name.clone(), payload.clone())
                        .await;

                    let duration = start_time.elapsed();
                    match result {
                        Ok(tool_result_message) => {
                            let content_preview = match &tool_result_message.content {
                                MessageContent::Text(t) => {
                                    if t.len() > 100 {
                                        format!("{}... ({} chars)", &t[..100], t.len())
                                    } else {
                                        t.clone()
                                    }
                                }
                                MessageContent::Json(v) => {
                                    format!("JSON: {} bytes", v.to_string().len())
                                }
                            };
                            tracing::warn!(
                                "✅ TOOL COMPLETED: {} in {:?} - Result: {}",
                                tool_name,
                                duration,
                                content_preview
                            );

                            // Tool executed successfully - now respond naturally
                            // Create a natural response incorporating the tool result
                            let natural_response = match &tool_result_message.content {
                                MessageContent::Text(text) => {
                                    if text.is_empty() {
                                        format!(
                                            "I've executed the {} tool successfully.",
                                            tool_name
                                        )
                                    } else {
                                        // Include the tool result in the response
                                        text.clone()
                                    }
                                }
                                MessageContent::Json(json) => {
                                    format!("Tool result: {}", json)
                                }
                            };

                            let response = AgentMessage {
                                role: MessageRole::Agent,
                                content: MessageContent::Text(natural_response),
                                metadata: None,
                            };
                            self.append_history(response.clone());
                            Ok(response)
                        }
                        Err(e) => {
                            tracing::error!(
                                "❌ TOOL FAILED: {} in {:?} - Error: {}",
                                tool_name,
                                duration,
                                e
                            );

                            // Tool failed - respond with error message
                            let error_response = AgentMessage {
                                role: MessageRole::Agent,
                                content: MessageContent::Text(format!(
                                    "I encountered an error while executing {}: {}",
                                    tool_name, e
                                )),
                                metadata: None,
                            };
                            self.append_history(error_response.clone());
                            Ok(error_response)
                        }
                    }
                } else {
                    // Tool not found
                    tracing::warn!("⚠️ Tool '{}' not found", tool_name);
                    let error_response = AgentMessage {
                        role: MessageRole::Agent,
                        content: MessageContent::Text(format!(
                            "I don't have access to the '{}' tool.",
                            tool_name
                        )),
                        metadata: None,
                    };
                    self.append_history(error_response.clone());
                    Ok(error_response)
                }
            }
            PlannerAction::Terminate => {
                tracing::debug!("🛑 Agent terminated");
                let message = AgentMessage {
                    role: MessageRole::Agent,
                    content: MessageContent::Text("Task completed.".into()),
                    metadata: None,
                };
                self.append_history(message.clone());
                Ok(message)
            }
        }
    }
}

#[async_trait]
impl AgentHandle for DeepAgent {
    async fn describe(&self) -> AgentDescriptor {
        self.descriptor.clone()
    }

    async fn handle_message(
        &self,
        input: AgentMessage,
        _state: Arc<AgentStateSnapshot>,
    ) -> anyhow::Result<AgentMessage> {
        self.handle_message_internal(input, _state).await
    }

    async fn handle_message_stream(
        &self,
        input: AgentMessage,
        _state: Arc<AgentStateSnapshot>,
    ) -> anyhow::Result<agents_core::agent::AgentStream> {
        use crate::planner::LlmBackedPlanner;
        use agents_core::llm::{LlmRequest, StreamChunk};

        // Add input to history
        self.append_history(input.clone());

        // Build the request similar to handle_message_internal
        let mut request = ModelRequest::new(&self.instructions, self.current_history());
        let tools = self.collect_tools();

        // Apply middleware modifications
        for middleware in &self.middlewares {
            let mut ctx = MiddlewareContext::with_request(&mut request, self.state.clone());
            middleware.modify_model_request(&mut ctx).await?;
        }

        // Convert ModelRequest to LlmRequest and add tools
        let tool_schemas: Vec<_> = tools.values().map(|t| t.schema()).collect();
        let llm_request = LlmRequest {
            system_prompt: request.system_prompt.clone(),
            messages: request.messages.clone(),
            tools: tool_schemas,
        };

        // Try to get the underlying LLM model for streaming
        let planner_any = self.planner.as_any();

        if let Some(llm_planner) = planner_any.downcast_ref::<LlmBackedPlanner>() {
            // We have an LlmBackedPlanner, use its model for streaming
            let model = llm_planner.model().clone();
            let stream = model.generate_stream(llm_request).await?;
            Ok(stream)
        } else {
            // Fallback to non-streaming
            let response = self.handle_message_internal(input, _state).await?;
            Ok(Box::pin(futures::stream::once(async move {
                Ok(StreamChunk::Done { message: response })
            })))
        }
    }

    async fn current_interrupt(&self) -> anyhow::Result<Option<AgentInterrupt>> {
        let state_guard = self
            .state
            .read()
            .map_err(|_| anyhow::anyhow!("Failed to acquire read lock on state"))?;
        Ok(state_guard.pending_interrupts.first().cloned())
    }

    async fn resume_with_approval(
        &self,
        action: agents_core::hitl::HitlAction,
    ) -> anyhow::Result<AgentMessage> {
        self.resume_with_approval(action).await
    }
}

/// Create a deep agent from configuration - matches Python middleware assembly exactly
///
/// This function assembles the middleware stack in the same order as the Python SDK:
/// planning → filesystem → subagents → summarization → prompt caching → optional HITL
pub fn create_deep_agent_from_config(config: DeepAgentConfig) -> DeepAgent {
    let state = Arc::new(RwLock::new(AgentStateSnapshot::default()));
    let history = Arc::new(RwLock::new(Vec::<AgentMessage>::new()));

    let planning = Arc::new(PlanningMiddleware::new(state.clone()));
    let filesystem = Arc::new(FilesystemMiddleware::new(state.clone()));

    // Build sub-agents from configurations
    let mut registrations: Vec<SubAgentRegistration> = Vec::new();

    // Build custom sub-agents from configs
    for subagent_config in &config.subagent_configs {
        // Determine the planner for this sub-agent
        let sub_planner = if let Some(ref model) = subagent_config.model {
            // Sub-agent has its own model - wrap it in a planner
            Arc::new(LlmBackedPlanner::new(model.clone())) as Arc<dyn PlannerHandle>
        } else {
            // Inherit parent's planner
            config.planner.clone()
        };

        // Create a DeepAgentConfig for this sub-agent
        let mut sub_cfg = DeepAgentConfig::new(subagent_config.instructions.clone(), sub_planner);

        // Configure tools
        if let Some(ref tools) = subagent_config.tools {
            for tool in tools {
                sub_cfg = sub_cfg.with_tool(tool.clone());
            }
        }

        // Configure built-in tools
        if let Some(ref builtin) = subagent_config.builtin_tools {
            sub_cfg = sub_cfg.with_builtin_tools(builtin.iter().cloned());
        }

        // Sub-agents should not have their own sub-agents
        sub_cfg = sub_cfg.with_auto_general_purpose(false);

        // Configure prompt caching
        sub_cfg = sub_cfg.with_prompt_caching(subagent_config.enable_prompt_caching);

        // Build the sub-agent recursively
        let sub_agent = create_deep_agent_from_config(sub_cfg);

        // Register the sub-agent
        registrations.push(SubAgentRegistration {
            descriptor: SubAgentDescriptor {
                name: subagent_config.name.clone(),
                description: subagent_config.description.clone(),
            },
            agent: Arc::new(sub_agent),
        });
    }

    // Optionally inject a general-purpose subagent
    if config.auto_general_purpose {
        let has_gp = registrations
            .iter()
            .any(|r| r.descriptor.name == "general-purpose");
        if !has_gp {
            // Create a subagent with inherited planner/tools and same instructions
            let mut sub_cfg =
                DeepAgentConfig::new(config.instructions.clone(), config.planner.clone())
                    .with_auto_general_purpose(false)
                    .with_prompt_caching(config.enable_prompt_caching);
            if let Some(ref selected) = config.builtin_tools {
                sub_cfg = sub_cfg.with_builtin_tools(selected.iter().cloned());
            }
            if let Some(ref sum) = config.summarization {
                sub_cfg = sub_cfg.with_summarization(sum.clone());
            }
            for t in &config.tools {
                sub_cfg = sub_cfg.with_tool(t.clone());
            }

            let gp = create_deep_agent_from_config(sub_cfg);
            registrations.push(SubAgentRegistration {
                descriptor: SubAgentDescriptor {
                    name: "general-purpose".into(),
                    description: "Default reasoning agent".into(),
                },
                agent: Arc::new(gp),
            });
        }
    }

    let subagent = Arc::new(SubAgentMiddleware::new(registrations));
    let base_prompt = Arc::new(BaseSystemPromptMiddleware);
    let deep_agent_prompt = Arc::new(DeepAgentPromptMiddleware::new(config.instructions.clone()));
    let summarization = config.summarization.as_ref().map(|cfg| {
        Arc::new(SummarizationMiddleware::new(
            cfg.messages_to_keep,
            cfg.summary_note.clone(),
        ))
    });
    let hitl = if config.tool_interrupts.is_empty() {
        None
    } else {
        // Validate that checkpointer is configured when HITL is enabled
        if config.checkpointer.is_none() {
            tracing::error!(
                "⚠️ HITL middleware requires a checkpointer to persist interrupt state. \
                 HITL will be disabled. Please configure a checkpointer to enable HITL."
            );
            None
        } else {
            tracing::info!("🔒 HITL enabled for {} tools", config.tool_interrupts.len());
            Some(Arc::new(HumanInLoopMiddleware::new(
                config.tool_interrupts.clone(),
            )))
        }
    };

    // Assemble middleware stack with Deep Agent prompt for automatic tool usage
    // Order: base → deep agent prompt → planning → filesystem → subagents → summarization → caching → HITL
    let mut middlewares: Vec<Arc<dyn AgentMiddleware>> = vec![
        base_prompt,
        deep_agent_prompt,
        planning,
        filesystem,
        subagent,
    ];
    if let Some(ref summary) = summarization {
        middlewares.push(summary.clone());
    }
    if config.enable_prompt_caching {
        middlewares.push(Arc::new(AnthropicPromptCachingMiddleware::with_defaults()));
    }
    if let Some(ref hitl_mw) = hitl {
        middlewares.push(hitl_mw.clone());
    }

    DeepAgent {
        descriptor: AgentDescriptor {
            name: "deep-agent".into(),
            version: "0.0.1".into(),
            description: Some("Rust deep agent".into()),
        },
        instructions: config.instructions,
        planner: config.planner,
        middlewares,
        base_tools: config.tools,
        state,
        history,
        _summarization: summarization,
        _hitl: hitl,
        builtin_tools: config.builtin_tools,
        checkpointer: config.checkpointer,
    }
}