motosan_agent_loop/

streaming_executor.rs

//! Streaming tool executor — starts executing tools as LLM streams tool_use blocks.
//!
//! Instead of waiting for the LLM to finish streaming before executing any tools,
//! [`StreamingToolExecutor`] spawns tool execution tasks the moment a complete
//! `tool_use` block arrives. This overlaps tool execution with ongoing LLM output,
//! reducing wall-clock latency.

use std::collections::HashMap;
use std::sync::Arc;

use motosan_agent_tool::{Tool, ToolContext, ToolResult};
use tokio::task::JoinHandle;

use crate::llm::ToolCallItem;

/// Executes tools eagerly as they arrive from a streaming LLM response.
///
/// Call [`submit`](Self::submit) each time a complete `tool_use` block is received
/// from the stream. After the stream ends, call [`collect`](Self::collect) to
/// await all pending tasks and retrieve results in submission order.
pub struct StreamingToolExecutor {
    pending: Vec<(ToolCallItem, JoinHandle<ToolResult>)>,
}

impl StreamingToolExecutor {
    /// Create a new executor with no pending tasks.
    pub fn new() -> Self {
        Self {
            pending: Vec::new(),
        }
    }

    /// Submit a tool for immediate execution.
    ///
    /// The tool is looked up in `tool_map` and spawned as a background task.
    /// If the tool is not found, the task will resolve with an error result.
    pub fn submit(
        &mut self,
        item: ToolCallItem,
        tool_map: &HashMap<String, Arc<dyn Tool>>,
        timeout: Option<std::time::Duration>,
        ctx: &ToolContext,
    ) {
        let tool = tool_map.get(&item.name).cloned();
        let name = item.name.clone();
        let args = item.args.clone();
        let ctx = ctx.clone();

        let handle = tokio::spawn(async move {
            let fut = async {
                if let Some(tool) = tool {
                    tool.call(args, &ctx).await
                } else {
                    ToolResult::error(format!("unknown tool: {name}"))
                }
            };
            if let Some(dur) = timeout {
                match tokio::time::timeout(dur, fut).await {
                    Ok(result) => result,
                    Err(_) => ToolResult::error(format!("tool '{name}' timed out after {dur:?}")),
                }
            } else {
                fut.await
            }
        });

        self.pending.push((item, handle));
    }

    /// Returns true if any tools have been submitted.
    pub fn has_pending(&self) -> bool {
        !self.pending.is_empty()
    }

    /// Await all submitted tasks and return `(items, results)` in submission order.
    ///
    /// Items and results are returned as separate vectors so callers can use
    /// the existing `execute_and_record_tool_calls` pipeline unchanged.
    pub async fn collect(self) -> (Vec<ToolCallItem>, Vec<ToolResult>) {
        let mut items = Vec::with_capacity(self.pending.len());
        let mut results = Vec::with_capacity(self.pending.len());

        for (item, handle) in self.pending {
            let result = handle
                .await
                .unwrap_or_else(|e| ToolResult::error(format!("tool task panicked: {e}")));
            items.push(item);
            results.push(result);
        }

        (items, results)
    }
}

impl Default for StreamingToolExecutor {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use motosan_agent_tool::ToolDef;
    use std::sync::atomic::{AtomicBool, Ordering};

    /// A tool that records when it starts executing.
    struct TimestampTool {
        name: String,
        started: Arc<AtomicBool>,
        result: String,
    }

    impl TimestampTool {
        fn new(name: &str, result: &str, started: Arc<AtomicBool>) -> Self {
            Self {
                name: name.to_string(),
                started,
                result: result.to_string(),
            }
        }
    }

    impl Tool for TimestampTool {
        fn def(&self) -> ToolDef {
            ToolDef {
                name: self.name.clone(),
                description: "test tool".into(),
                input_schema: serde_json::json!({
                    "type": "object",
                    "properties": {},
                    "required": []
                }),
            }
        }

        fn call(
            &self,
            _args: serde_json::Value,
            _ctx: &ToolContext,
        ) -> std::pin::Pin<Box<dyn std::future::Future<Output = ToolResult> + Send + '_>> {
            let started = self.started.clone();
            let result = self.result.clone();
            Box::pin(async move {
                started.store(true, Ordering::SeqCst);
                // Small delay to simulate work
                tokio::time::sleep(std::time::Duration::from_millis(10)).await;
                ToolResult::text(result)
            })
        }
    }

    #[tokio::test]
    async fn submit_and_collect_returns_results_in_order() {
        let started_a = Arc::new(AtomicBool::new(false));
        let started_b = Arc::new(AtomicBool::new(false));

        let tool_a: Arc<dyn Tool> =
            Arc::new(TimestampTool::new("tool_a", "result_a", started_a.clone()));
        let tool_b: Arc<dyn Tool> =
            Arc::new(TimestampTool::new("tool_b", "result_b", started_b.clone()));

        let mut tool_map: HashMap<String, Arc<dyn Tool>> = HashMap::new();
        tool_map.insert("tool_a".to_string(), tool_a);
        tool_map.insert("tool_b".to_string(), tool_b);

        let ctx = ToolContext::default();
        let mut executor = StreamingToolExecutor::new();

        let item_a = ToolCallItem {
            id: "call_1".to_string(),
            name: "tool_a".to_string(),
            args: serde_json::json!({}),
        };
        let item_b = ToolCallItem {
            id: "call_2".to_string(),
            name: "tool_b".to_string(),
            args: serde_json::json!({}),
        };

        executor.submit(item_a, &tool_map, None, &ctx);
        executor.submit(item_b, &tool_map, None, &ctx);

        assert!(executor.has_pending());

        let (items, results) = executor.collect().await;

        assert_eq!(items.len(), 2);
        assert_eq!(items[0].id, "call_1");
        assert_eq!(items[1].id, "call_2");
        assert_eq!(results.len(), 2);
        // Both tools should have started
        assert!(started_a.load(Ordering::SeqCst));
        assert!(started_b.load(Ordering::SeqCst));
    }

    #[tokio::test]
    async fn unknown_tool_returns_error_result() {
        let tool_map: HashMap<String, Arc<dyn Tool>> = HashMap::new();
        let ctx = ToolContext::default();
        let mut executor = StreamingToolExecutor::new();

        let item = ToolCallItem {
            id: "call_1".to_string(),
            name: "nonexistent".to_string(),
            args: serde_json::json!({}),
        };

        executor.submit(item, &tool_map, None, &ctx);
        let (items, results) = executor.collect().await;

        assert_eq!(items.len(), 1);
        assert_eq!(results.len(), 1);
        // The result should be an error about unknown tool
        let text = format!("{:?}", results[0]);
        assert!(text.contains("unknown tool"), "got: {text}");
    }

    #[tokio::test]
    async fn tools_start_executing_immediately_after_submit() {
        let started = Arc::new(AtomicBool::new(false));
        let tool: Arc<dyn Tool> =
            Arc::new(TimestampTool::new("slow_tool", "done", started.clone()));

        let mut tool_map: HashMap<String, Arc<dyn Tool>> = HashMap::new();
        tool_map.insert("slow_tool".to_string(), tool);

        let ctx = ToolContext::default();
        let mut executor = StreamingToolExecutor::new();

        let item = ToolCallItem {
            id: "call_1".to_string(),
            name: "slow_tool".to_string(),
            args: serde_json::json!({}),
        };

        executor.submit(item, &tool_map, None, &ctx);

        // Yield to let the spawned task start
        tokio::task::yield_now().await;
        tokio::time::sleep(std::time::Duration::from_millis(5)).await;

        // Tool should have started executing before we call collect()
        assert!(
            started.load(Ordering::SeqCst),
            "tool should start executing immediately after submit, before collect()"
        );

        let (_, results) = executor.collect().await;
        assert_eq!(results.len(), 1);
    }

    #[tokio::test]
    async fn empty_executor_collects_nothing() {
        let executor = StreamingToolExecutor::new();
        assert!(!executor.has_pending());

        let (items, results) = executor.collect().await;
        assert!(items.is_empty());
        assert!(results.is_empty());
    }
}