implements api to invoke new session with agent crate

Author: dingfeli

crates/chat-cli/Cargo.toml

@@ -8,6 +8,10 @@ version.workspace = true
 license.workspace = true
 default-run = "chat_cli"
 
+[[bin]]
+name = "temp_bin"
+path = "src/bin/temp_bin.rs"
+
 [lints]
 workspace = true
 

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

@@ -1 +1,214 @@
+use std::io::Write;
+use std::sync::Arc;
+
+use agent::AgentHandle;
+use agent::agent_config::load_agents;
+use agent::agent_loop::protocol::LoopEndReason;
+use agent::mcp::McpManager;
+use agent::protocol::{
+    AgentEvent,
+    AgentStopReason,
+    ApprovalResult,
+    ContentChunk,
+    SendApprovalResultArgs,
+    SendPromptArgs,
+};
+use agent::types::AgentSnapshot;
+use eyre::{
+    Result,
+    bail,
+};
+use rts::{
+    RtsModel,
+    RtsModelState,
+};
+use serde::{
+    Deserialize,
+    Serialize,
+};
+use tracing::{
+    debug,
+    error,
+    info,
+    warn,
+};
+
+use crate::os::Os;
+
 mod rts;
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+struct JsonOutput {
+    /// Whether or not the user turn completed successfully
+    is_error: bool,
+    /// Text from the final message, if available
+    result: Option<String>,
+    /// The number of requests sent to the model
+    number_of_requests: u32,
+    /// The number of tool use / tool result pairs in the turn
+    ///
+    /// This could be less than the number of requests in the case of retries
+    number_of_cycles: u32,
+    /// Duration of the turn, in milliseconds
+    duration_ms: u32,
+}
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct QueryResult {
+    pub context_summary: String,
+    pub task_summary: String,
+}
+
+#[derive(Debug)]
+pub struct SubAgent<'a> {
+    pub query: &'a str,
+    pub agent_name: Option<&'a str>,
+    pub embedded_user_msg: Option<&'a str>,
+    pub dangerously_trust_all_tools: bool,
+}
+
+impl<'a> SubAgent<'a> {
+    pub async fn query(self, os: &mut Os, output: &mut impl Write) -> Result<QueryResult> {
+        let mut snapshot = AgentSnapshot::default();
+
+        let model = {
+            let rts_state: RtsModelState = snapshot
+                .model_state
+                .as_ref()
+                .and_then(|s| {
+                    serde_json::from_value(s.clone())
+                        .map_err(|err| error!(?err, ?s, "failed to deserialize RTS state"))
+                        .ok()
+                })
+                .unwrap_or({
+                    let state = RtsModelState::new();
+                    info!(?state.conversation_id, "generated new conversation id");
+                    state
+                });
+            Arc::new(RtsModel::new(
+                os.client.clone(),
+                rts_state.conversation_id,
+                rts_state.model_id,
+            ))
+        };
+
+        if let Some(name) = self.agent_name {
+            let (configs, _) = load_agents().await?;
+            if let Some(cfg) = configs.into_iter().find(|c| c.name() == name) {
+                snapshot.agent_config = cfg.config().clone();
+            } else {
+                bail!("unable to find agent with name: {}", name);
+            }
+        };
+
+        let mcp_manager_handle = McpManager::default().spawn();
+        let agent = agent::Agent::new(snapshot, model, mcp_manager_handle).await?.spawn();
+
+        self.main_loop(agent, output).await
+    }
+
+    async fn main_loop(&self, mut agent: AgentHandle, output: &mut impl Write) -> Result<QueryResult> {
+        // First, wait for agent initialization
+        while let Ok(evt) = agent.recv().await {
+            if matches!(evt, AgentEvent::Mcp(_)) {
+                info!(?evt, "received mcp agent event");
+                // TODO: Send it through conduit
+            }
+            if matches!(evt, AgentEvent::Initialized) {
+                break;
+            }
+        }
+
+        agent
+            .send_prompt(SendPromptArgs {
+                content: vec![ContentChunk::Text(self.query.to_string())],
+                should_continue_turn: None,
+            })
+            .await?;
+
+        // Holds the final result of the user turn.
+        #[allow(unused_assignments)]
+        let mut user_turn_metadata = None;
+
+        loop {
+            let Ok(evt) = agent.recv().await else {
+                bail!("channel closed");
+            };
+            debug!(?evt, "received new agent event");
+
+            // Check for exit conditions
+            match &evt {
+                AgentEvent::Update(evt) => {
+                    info!(?evt, "received update event");
+                    println!("received update event {:?}", evt);
+                },
+                AgentEvent::EndTurn(metadata) => {
+                    user_turn_metadata = Some(metadata.clone());
+                    break;
+                },
+                AgentEvent::Stop(AgentStopReason::Error(agent_error)) => {
+                    bail!("agent encountered an error: {:?}", agent_error)
+                },
+                AgentEvent::ApprovalRequest { id, tool_use, .. } => {
+                    if !self.dangerously_trust_all_tools {
+                        bail!("Tool approval is required: {:?}", tool_use);
+                    } else {
+                        warn!(?tool_use, "trust all is enabled, ignoring approval request");
+                        agent
+                            .send_tool_use_approval_result(SendApprovalResultArgs {
+                                id: id.clone(),
+                                result: ApprovalResult::Approve,
+                            })
+                            .await?;
+                    }
+                },
+                AgentEvent::Mcp(evt) => {
+                    info!(?evt, "received mcp agent event");
+                },
+                _ => {},
+            }
+        }
+
+        let md = user_turn_metadata.expect("user turn metadata should exist");
+        let is_error = md.end_reason != LoopEndReason::UserTurnEnd || md.result.as_ref().is_none_or(|v| v.is_err());
+        let result = md.result.and_then(|r| r.ok().map(|m| m.text()));
+
+        let output = JsonOutput {
+            result,
+            is_error,
+            number_of_requests: md.total_request_count,
+            number_of_cycles: md.number_of_cycles,
+            duration_ms: md.turn_duration.map(|d| d.as_millis() as u32).unwrap_or_default(),
+        };
+
+        info!(?output, "sub agent routine completed");
+
+        Ok(QueryResult {
+            context_summary: Default::default(),
+            task_summary: Default::default(),
+        })
+    }
+}
+
+pub fn temp_func() {
+    let rt = tokio::runtime::Builder::new_multi_thread()
+        .enable_all()
+        .build()
+        .expect("failed to build runtime");
+
+    rt.block_on(test_sub_agent_routine());
+}
+
+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,
+        dangerously_trust_all_tools: true,
+    };
+
+    let mut os = Os::new().await.expect("failed to spawn os");
+    let mut output = Vec::<u8>::new();
+
+    _ = sub_agent.query(&mut os, &mut output).await;
+}

crates/chat-cli/src/bin/temp_bin.rs

@@ -0,0 +1,5 @@
+use chat_cli::agent::temp_func;
+
+fn main() {
+    temp_func();
+}

crates/chat-cli/src/lib.rs

@@ -2,6 +2,7 @@
 //! This lib.rs is only here for testing purposes.
 //! `test_mcp_server/test_server.rs` is declared as a separate binary and would need a way to
 //! reference types defined inside of this crate, hence the export.
+pub mod agent;
 pub mod api_client;
 pub mod auth;
 pub mod aws_common;