adds subagent spawn context

Author: dingfeli

crates/agent/src/agent/mod.rs

@@ -127,6 +127,7 @@ use types::{
     AgentSnapshot,
     ConversationMetadata,
     ConversationState,
+    EmbeddedUserMessages,
 };
 use util::path::canonicalize_path_sys;
 use util::providers::{
@@ -222,6 +223,7 @@ pub struct Agent {
     id: AgentId,
     agent_config: AgentConfig,
 
+    embedded_user_msgs: EmbeddedUserMessages,
     conversation_state: ConversationState,
     conversation_metadata: ConversationMetadata,
     execution_state: ExecutionState,
@@ -303,6 +305,7 @@ impl Agent {
         Ok(Self {
             id: snapshot.id,
             agent_config,
+            embedded_user_msgs: Default::default(),
             conversation_state: snapshot.conversation_state,
             conversation_metadata: snapshot.conversation_metadata,
             execution_state: snapshot.execution_state,
@@ -323,6 +326,10 @@ impl Agent {
         })
     }
 
+    pub fn set_embedded_user_msg(&mut self, msg: &str) {
+        self.embedded_user_msgs.messages.push(msg.to_string());
+    }
+
     pub fn set_sys_provider(&mut self, provider: impl SystemProvider) {
         self.sys_provider = Arc::new(provider);
     }
@@ -981,6 +988,7 @@ impl Agent {
             self.make_tool_spec().await,
             &self.agent_config,
             self.agent_spawn_hooks.iter().map(|(_, c)| c),
+            &self.embedded_user_msgs,
             &self.sys_provider,
         )
         .await
@@ -1720,6 +1728,7 @@ async fn format_request<T, U, P>(
     mut tool_spec: Vec<ToolSpec>,
     agent_config: &AgentConfig,
     agent_spawn_hooks: T,
+    embedded_user_msgs: &EmbeddedUserMessages,
     provider: &P,
 ) -> SendRequestArgs
 where
@@ -1734,10 +1743,20 @@ where
         messages.push_front(msg);
     }
 
+    let system_prompt = match (agent_config.system_prompt(), embedded_user_msgs.messages.is_empty()) {
+        (Some(system_prompt), false) => {
+            let embedded_user_msgs_as_str = embedded_user_msgs.to_string();
+            Some(format!("{embedded_user_msgs_as_str}\n{system_prompt}"))
+        },
+        (Some(system_prompt), true) => Some(system_prompt.to_string()),
+        (None, false) => Some(embedded_user_msgs.to_string()),
+        (_, _) => None,
+    };
+
     SendRequestArgs::new(
         messages.into(),
         if tool_spec.is_empty() { None } else { Some(tool_spec) },
-        agent_config.system_prompt().map(String::from),
+        system_prompt,
     )
 }
 

crates/agent/src/agent/tools/summary.rs

@@ -19,7 +19,7 @@ use crate::protocol::AgentEvent;
 
 /// A tool for conveying information from subagent to its main agent
 #[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
-/// A tool for conveying information from subagent to its main agent
+#[serde(rename_all = "camelCase")]
 pub struct Summary {
     /// Description of the task that was assigned to the subagent
     pub task_description: String,
@@ -41,19 +41,41 @@ HOW TO USE:
 - Provide the result of the task performed
 "#;
 
+const SUMMARY_TOOL_SCHEMA: &str = r#"
+{
+    "type": "object",
+    "properties": {
+        "taskDescription": {
+            "type": "string",
+            "description": "Description of the task that was assigned to the subagent"
+        },
+        "contextSummary": {
+            "type": "string",
+            "description": "Relevant context and information gathered during task execution"
+        },
+        "taskResult": {
+            "type": "string",
+            "description": "The final result or outcome of the completed task"
+        }
+    },
+    "required": [
+        "taskDescription",
+        "taskResult"
+    ]
+}
+"#;
+
 impl BuiltInToolTrait for Summary {
     fn name() -> super::BuiltInToolName {
-        BuiltInToolName::Ls
+        BuiltInToolName::Summary
     }
 
     fn description() -> std::borrow::Cow<'static, str> {
         SUMMARY_TOOL_DESCRIPTION.into()
     }
 
     fn input_schema() -> std::borrow::Cow<'static, str> {
-        serde_json::to_string(&Self::tool_schema())
-            .expect("serializing schema should not fail")
-            .into()
+        SUMMARY_TOOL_SCHEMA.into()
     }
 }
 
@@ -75,6 +97,12 @@ impl Summary {
 mod tests {
     use super::*;
 
+    #[test]
+    fn test_summary_tool_schema() {
+        let schema = Summary::input_schema();
+        println!("{:#?}", schema);
+    }
+
     #[tokio::test]
     async fn test_summary_tool_execute() {
         let (tx, mut rx) = broadcast::channel(10);

crates/agent/src/agent/types.rs

@@ -143,6 +143,52 @@ impl Default for AgentSettings {
     }
 }
 
+pub struct DecoratedEmbeddedUserMessagesIterator<'a> {
+    pub inner_iter: std::slice::Iter<'a, String>,
+}
+
+impl<'a> Iterator for DecoratedEmbeddedUserMessagesIterator<'a> {
+    type Item = String;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.inner_iter.next().map(|message| {
+            format!(
+                "{}{message}\n{}",
+                EmbeddedUserMessages::START_HEADER,
+                EmbeddedUserMessages::END_HEADER
+            )
+        })
+    }
+}
+
+#[derive(Default, Debug, Clone, Serialize, Deserialize)]
+pub struct EmbeddedUserMessages {
+    pub messages: Vec<String>,
+}
+
+impl EmbeddedUserMessages {
+    const END_HEADER: &str = "--- EMBEDDED USER MESSAGE SEGMENT END ---\n\n";
+    const START_HEADER: &str = "--- EMBEDDED USER MESSAGE SEGMENT BEGIN ---\n";
+
+    pub fn to_string(&self) -> String {
+        self.into_iter().fold(String::new(), |mut acc, msg| {
+            acc.push_str(&msg);
+            acc
+        })
+    }
+}
+
+impl<'a> IntoIterator for &'a EmbeddedUserMessages {
+    type IntoIter = DecoratedEmbeddedUserMessagesIterator<'a>;
+    type Item = String;
+
+    fn into_iter(self) -> Self::IntoIter {
+        DecoratedEmbeddedUserMessagesIterator {
+            inner_iter: self.messages.iter(),
+        }
+    }
+}
+
 /// State associated with a history of messages.
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct ConversationState {

crates/chat-cli/src/agent/mod.rs

@@ -37,6 +37,12 @@ use crate::os::Os;
 
 mod rts;
 
+// TODO: use the one supplied by science (this one has been modified for testing)
+const SUBAGENT_EMBEDDED_USER_MSG: &str = r#"
+You are a subagent executing a task delegated to you by the main agent.
+After what is asked of you has concluded, call the summary tool to convey your findings to the main agent.
+"#;
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 struct JsonOutput {
     /// Whether or not the user turn completed successfully
@@ -102,9 +108,14 @@ impl<'a> SubAgent<'a> {
         };
 
         let mcp_manager_handle = McpManager::default().spawn();
-        let agent = agent::Agent::new(snapshot, model, mcp_manager_handle).await?.spawn();
+        let mut agent = agent::Agent::new(snapshot, model, mcp_manager_handle).await?;
+        if let Some(msg) = self.embedded_user_msg {
+            agent.set_embedded_user_msg(msg);
+        }
 
-        self.main_loop(agent, output).await
+        let agent_handle = agent.spawn();
+
+        self.main_loop(agent_handle, output).await
     }
 
     async fn main_loop(&self, mut agent: AgentHandle, output: &mut impl Write) -> Result<QueryResult> {
@@ -165,6 +176,9 @@ impl<'a> SubAgent<'a> {
                 AgentEvent::Mcp(evt) => {
                     info!(?evt, "received mcp agent event");
                 },
+                AgentEvent::SubagentSummary(summary) => {
+                    println!("Summary: {:#?}", summary);
+                },
                 _ => {},
             }
         }
@@ -203,7 +217,7 @@ async fn test_sub_agent_routine() {
     let sub_agent = SubAgent {
         query: "What notion docs do I have",
         agent_name: Some("test_test"),
-        embedded_user_msg: None,
+        embedded_user_msg: Some(SUBAGENT_EMBEDDED_USER_MSG),
         dangerously_trust_all_tools: true,
     };
 

crates/chat-cli/src/agent/rts/mod.rs

@@ -201,7 +201,7 @@ impl RtsModel {
         &self,
         mut messages: Vec<Message>,
         tool_specs: Option<Vec<ToolSpec>>,
-        _system_prompt: Option<String>,
+        system_prompt: Option<String>,
     ) -> Result<ConversationState, String> {
         debug!(?messages, ?tool_specs, "creating conversation state");
         let tools = tool_specs.map(|v| {
@@ -235,7 +235,22 @@ impl RtsModel {
             None => return Err("Empty conversation".to_string()),
         };
 
-        let history = messages
+        let mut history = Vec::<rts::ChatMessage>::new();
+        if let Some(system_prompt) = system_prompt {
+            history.push(rts::ChatMessage::UserInputMessage(UserInputMessage {
+                content: system_prompt,
+                user_input_message_context: None,
+                user_intent: None,
+                images: None,
+                model_id: None,
+            }));
+            history.push(rts::ChatMessage::AssistantResponseMessage(rts::AssistantResponseMessage { 
+                message_id: None, 
+                content: "I will fully incorporate this information when generating my responses, and explicitly acknowledge relevant parts of the summary when answering questions.".to_string(),
+                tool_uses: None }));
+        }
+
+        history.append(&mut messages
             .into_iter()
             .map(|m| match m.role {
                 Role::User => {
@@ -269,7 +284,7 @@ impl RtsModel {
                     rts::ChatMessage::AssistantResponseMessage(msg)
                 },
             })
-            .collect();
+            .collect());
 
         Ok(ConversationState {
             conversation_id: Some(self.conversation_id.to_string()),