1 files changed, 25 insertions, 4 deletions
diff --git a/src/function.rs b/src/function.rs
index 5907b71..778330a 100644
--- a/src/function.rs
+++ b/src/function.rs
@@ -1,5 +1,5 @@
 use crate::{
-    config::{Agent, Config, GlobalConfig},
+    config::{eval_subquery, Agent, Config, GlobalConfig, SubqueryContext},
     utils::*,
 };
 
@@ -18,7 +18,15 @@ const PATH_SEP: &str = ";";
 #[cfg(not(windows))]
 const PATH_SEP: &str = ":";
 
-pub fn eval_tool_calls(config: &GlobalConfig, mut calls: Vec<ToolCall>) -> Result<Vec<ToolResult>> {
+pub struct EvalToolCallsContext {
+    pub subquery_ctx: Option<SubqueryContext>,
+}
+
+pub async fn eval_tool_calls(
+    config: &GlobalConfig,
+    mut calls: Vec<ToolCall>,
+    ctx: &mut EvalToolCallsContext,
+) -> Result<Vec<ToolResult>> {
     let mut output = vec![];
     if calls.is_empty() {
         return Ok(output);
@@ -29,7 +37,7 @@ pub fn eval_tool_calls(config: &GlobalConfig, mut calls: Vec<ToolCall>) -> Resul
     }
     let mut is_all_null = true;
     for call in calls {
-        let mut result = call.eval(config)?;
+        let mut result = call.eval(config, ctx).await?;
         if result.is_null() {
             result = json!("DONE");
         } else {
@@ -170,7 +178,14 @@ impl ToolCall {
         }
     }
 
-    pub fn eval(&self, config: &GlobalConfig) -> Result<Value> {
+    #[async_recursion::async_recursion]
+    pub async fn eval(&self, config: &GlobalConfig, ctx: &EvalToolCallsContext) -> Result<Value> {
+        if self.name == SUBQUERY_TOOL_NAME {
+            return match &ctx.subquery_ctx {
+                Some(sub_ctx) => eval_subquery(config, sub_ctx, self.arguments.clone()).await,
+                None => bail!("Invalid subquery call: missing parent query context"),
+            };
+        }
         let (call_name, cmd_name, mut cmd_args, envs) = match &config.read().agent {
             Some(agent) => self.extract_call_config_from_agent(config, agent)?,
             None => self.extract_call_config_from_config(config)?,
@@ -245,6 +260,8 @@ impl ToolCall {
     }
 }
 
+const SUBQUERY_TOOL_NAME: &str = "subquery";
+
 pub fn run_llm_function(
     cmd_name: String,
     cmd_args: Vec<String>,
@@ -274,6 +291,10 @@ pub fn run_llm_function(
     #[cfg(windows)]
     let cmd_name = polyfill_cmd_name(&cmd_name, &bin_dirs);
     if *IS_STDOUT_TERMINAL {
+        let prompt = match current_query_scope() {
+            Some(scope) => format!("{scope} {prompt}"),
+            None => prompt,
+        };
         println!("{}", dimmed_text(&prompt));
     }
     let exit_code = run_command(&cmd_name, &cmd_args, Some(envs))