From 159f2d3330a2c03db921e1e708ff355e85078948 Mon Sep 17 00:00:00 2001 From: vikingowl Date: Fri, 26 Dec 2025 19:24:03 +0100 Subject: [PATCH] feat(core): Implement Permission Interceptor and async request/response --- crates/app/cli/src/agent_manager.rs | 87 ++++++++++++++--- crates/app/cli/src/engine.rs | 143 +++++++++++++--------------- crates/app/cli/src/messages.rs | 2 + crates/app/cli/src/state.rs | 11 ++- 4 files changed, 153 insertions(+), 90 deletions(-) diff --git a/crates/app/cli/src/agent_manager.rs b/crates/app/cli/src/agent_manager.rs index 6f6bcdf..ef8c8d2 100644 --- a/crates/app/cli/src/agent_manager.rs +++ b/crates/app/cli/src/agent_manager.rs @@ -1,20 +1,27 @@ use std::sync::Arc; -use tokio::sync::Mutex; -use crate::state::AppState; +use tokio::sync::{mpsc, Mutex}; +use crate::state::{AppState, AppMode}; use llm_core::{LlmProvider, ChatMessage, ChatOptions}; use color_eyre::eyre::Result; -use futures::StreamExt; +use crate::messages::{Message, AgentResponse}; /// Manages the lifecycle and state of the agent pub struct AgentManager { client: Arc, state: Arc>, + tx_ui: Option>, } impl AgentManager { /// Create a new AgentManager pub fn new(client: Arc, state: Arc>) -> Self { - Self { client, state } + Self { client, state, tx_ui: None } + } + + /// Set the UI message sender + pub fn with_ui_sender(mut self, tx: mpsc::Sender) -> Self { + self.tx_ui = Some(tx); + self } /// Get a reference to the LLM client @@ -27,11 +34,48 @@ impl AgentManager { &self.state } + /// Check if a tool execution is permitted + pub async fn check_permission(&self, tool: &str, context: Option<&str>) -> Result { + let mode = { + let guard = self.state.lock().await; + guard.mode + }; + + match mode { + AppMode::AcceptAll => Ok(true), + AppMode::Normal | AppMode::Plan => { + // Request permission from UI + if let Some(tx) = &self.tx_ui { + let _ = tx.send(Message::AgentResponse(AgentResponse::PermissionRequest { + tool: tool.to_string(), + context: context.map(|s| s.to_string()), + })).await; + + // Wait for result + let notify = { + let guard = self.state.lock().await; + guard.permission_notify.clone() + }; + + notify.notified().await; + + let result = { + let mut guard = self.state.lock().await; + guard.last_permission_result.take().unwrap_or(false) + }; + + Ok(result) + } else { + // No UI sender, default to deny for safety + Ok(false) + } + } + } + } + /// Spawn a sub-agent for a specific task pub async fn spawn_sub_agent(&self, _task: &str) -> Result { // Basic placeholder implementation for Sub-Agent support - // In a real implementation, this would create a new AgentManager instance - // with a focused system prompt and context. Ok("Sub-agent task completed (mock)".to_string()) } @@ -43,16 +87,12 @@ impl AgentManager { guard.add_message(ChatMessage::user(input.to_string())); } - // 2. Prepare context (Sliding Window logic would go here) + // 2. Prepare context let messages = { let guard = self.state.lock().await; guard.messages.clone() }; - // 3. Call LLM (Non-streaming for now to simplify step logic, or use stream internally) - // For the Reasoning Loop, we need to handle potential tool calls. - // This initial implementation just chats. - let options = ChatOptions::default(); let response = self.client.chat(&messages, &options, None).await?; @@ -70,7 +110,7 @@ impl AgentManager { #[cfg(test)] mod tests { use super::*; - use llm_core::{LlmProvider, ChatMessage, ChatOptions, Tool, ChunkStream, ChatResponse}; + use llm_core::{Tool, ChunkStream, ChatResponse}; use async_trait::async_trait; struct MockProvider; @@ -119,4 +159,25 @@ mod tests { let result = manager.spawn_sub_agent("Do something").await.unwrap(); assert_eq!(result, "Sub-agent task completed (mock)"); } -} \ No newline at end of file + + #[tokio::test] + async fn test_permission_request() { + let client = Arc::new(MockProvider); + let state = Arc::new(Mutex::new(AppState::new())); + let (tx, mut rx) = mpsc::channel(1); + + let manager = AgentManager::new(client, state.clone()).with_ui_sender(tx); + + let state_clone = state.clone(); + tokio::spawn(async move { + // Simulate UI receiving request and granting permission + if let Some(Message::AgentResponse(AgentResponse::PermissionRequest { .. })) = rx.recv().await { + let mut guard = state_clone.lock().await; + guard.set_permission_result(true); + } + }); + + let allowed = manager.check_permission("bash", Some("ls")).await.unwrap(); + assert!(allowed); + } +} diff --git a/crates/app/cli/src/engine.rs b/crates/app/cli/src/engine.rs index dcb970d..f7aca7e 100644 --- a/crates/app/cli/src/engine.rs +++ b/crates/app/cli/src/engine.rs @@ -4,6 +4,7 @@ use tokio::sync::{mpsc, Mutex}; use std::sync::Arc; use llm_core::{LlmProvider, ChatMessage, ChatOptions}; use futures::StreamExt; +use crate::agent_manager::AgentManager; /// The main background task that handles logic, API calls, and state updates. pub async fn run_engine_loop( @@ -12,52 +13,64 @@ pub async fn run_engine_loop( client: Arc, state: Arc>, ) { + let agent_manager = Arc::new(AgentManager::new(client, state.clone()).with_ui_sender(tx_ui.clone())); + while let Some(msg) = rx.recv().await { match msg { Message::UserAction(UserAction::Input(text)) => { - // Update history with user message - let messages = { - let mut guard = state.lock().await; - guard.add_message(ChatMessage::user(text.clone())); - guard.messages.clone() - }; + let tx_ui_clone = tx_ui.clone(); + let agent_manager_clone = agent_manager.clone(); + + // Spawn a task for the agent interaction so the engine loop can + // continue receiving messages (like PermissionResult). + tokio::spawn(async move { + let messages = { + let mut guard = agent_manager_clone.state().lock().await; + guard.add_message(ChatMessage::user(text.clone())); + guard.messages.clone() + }; - // Use default options for now - let options = ChatOptions::default(); + let options = ChatOptions::default(); - match client.chat_stream(&messages, &options, None).await { - Ok(mut stream) => { - let mut full_response = String::new(); + match agent_manager_clone.client().chat_stream(&messages, &options, None).await { + Ok(mut stream) => { + let mut full_response = String::new(); - while let Some(result) = stream.next().await { - match result { - Ok(chunk) => { - if let Some(content) = chunk.content { - full_response.push_str(&content); - if let Err(e) = tx_ui.send(Message::AgentResponse(AgentResponse::Token(content))).await { - eprintln!("Failed to send token to UI: {}", e); - break; + while let Some(result) = stream.next().await { + match result { + Ok(chunk) => { + if let Some(content) = chunk.content { + full_response.push_str(&content); + let _ = tx_ui_clone.send(Message::AgentResponse(AgentResponse::Token(content))).await; } } - } - Err(e) => { - let _ = tx_ui.send(Message::AgentResponse(AgentResponse::Error(e.to_string()))).await; + Err(e) => { + let _ = tx_ui_clone.send(Message::AgentResponse(AgentResponse::Error(e.to_string()))).await; + } } } - } - - // Add assistant response to history - { - let mut guard = state.lock().await; - guard.add_message(ChatMessage::assistant(full_response)); - } + + { + let mut guard = agent_manager_clone.state().lock().await; + guard.add_message(ChatMessage::assistant(full_response)); + } - let _ = tx_ui.send(Message::AgentResponse(AgentResponse::Complete)).await; + let _ = tx_ui_clone.send(Message::AgentResponse(AgentResponse::Complete)).await; + } + Err(e) => { + let _ = tx_ui_clone.send(Message::AgentResponse(AgentResponse::Error(e.to_string()))).await; + } } - Err(e) => { - let _ = tx_ui.send(Message::AgentResponse(AgentResponse::Error(e.to_string()))).await; - } - } + }); + } + Message::UserAction(UserAction::PermissionResult(res)) => { + let mut guard = state.lock().await; + guard.set_permission_result(res); + } + Message::UserAction(UserAction::Exit) => { + let mut guard = state.lock().await; + guard.running = false; + break; } _ => {} } @@ -67,8 +80,7 @@ pub async fn run_engine_loop( #[cfg(test)] mod tests { use super::*; - use crate::messages::{Message, UserAction, AgentResponse}; - use llm_core::{LlmProvider, LlmError, ChatMessage, ChatOptions, Tool, ChunkStream, StreamChunk}; + use llm_core::{LlmError, Tool, ChunkStream, StreamChunk}; use async_trait::async_trait; use futures::stream; @@ -110,29 +122,21 @@ mod tests { tx_in.send(Message::UserAction(UserAction::Input("Hi".to_string()))).await.unwrap(); // Verify streaming responses - if let Some(Message::AgentResponse(AgentResponse::Token(s))) = rx_out.recv().await { - assert_eq!(s, "Hello"); - } else { - panic!("Expected Token(Hello)"); - } - - if let Some(Message::AgentResponse(AgentResponse::Token(s))) = rx_out.recv().await { - assert_eq!(s, " World"); - } else { - panic!("Expected Token( World)"); - } - - if let Some(Message::AgentResponse(AgentResponse::Complete)) = rx_out.recv().await { - // OK - } else { - panic!("Expected Complete"); + let mut tokens = Vec::new(); + while let Some(msg) = rx_out.recv().await { + match msg { + Message::AgentResponse(AgentResponse::Token(s)) => tokens.push(s), + Message::AgentResponse(AgentResponse::Complete) => break, + _ => {} + } } + assert_eq!(tokens, vec!["Hello", " World"]); } - + #[tokio::test] - async fn test_engine_state_update() { - let (tx_in, rx_in) = mpsc::channel(1); - let (tx_out, mut rx_out) = mpsc::channel(10); + async fn test_engine_permission_result() { + let (tx_in, rx_in) = mpsc::channel(1); + let (tx_out, _rx_out) = mpsc::channel(10); let client = Arc::new(MockProvider); let state = Arc::new(Mutex::new(AppState::new())); @@ -143,26 +147,13 @@ mod tests { run_engine_loop(rx_in, tx_out, client, state_clone).await; }); - // Send a message - tx_in.send(Message::UserAction(UserAction::Input("Hi".to_string()))).await.unwrap(); - - // Wait for completion - while let Some(msg) = rx_out.recv().await { - if let Message::AgentResponse(AgentResponse::Complete) = msg { - break; - } - } - - // Verify state + // Send a PermissionResult + tx_in.send(Message::UserAction(UserAction::PermissionResult(true))).await.unwrap(); + + // Give it a moment to process + tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; + let guard = state.lock().await; - assert_eq!(guard.messages.len(), 2); // User + Assistant - match &guard.messages[0].role { - llm_core::Role::User => {}, - _ => panic!("First message should be User"), - } - match &guard.messages[1].role { - llm_core::Role::Assistant => {}, - _ => panic!("Second message should be Assistant"), - } + assert_eq!(guard.last_permission_result, Some(true)); } } diff --git a/crates/app/cli/src/messages.rs b/crates/app/cli/src/messages.rs index 370c538..718cee8 100644 --- a/crates/app/cli/src/messages.rs +++ b/crates/app/cli/src/messages.rs @@ -11,6 +11,7 @@ pub enum Message { pub enum UserAction { Input(String), Command(String), + PermissionResult(bool), Exit, } @@ -18,6 +19,7 @@ pub enum UserAction { pub enum AgentResponse { Token(String), ToolCall { name: String, args: String }, + PermissionRequest { tool: String, context: Option }, Complete, Error(String), } diff --git a/crates/app/cli/src/state.rs b/crates/app/cli/src/state.rs index 4367ec0..678a45b 100644 --- a/crates/app/cli/src/state.rs +++ b/crates/app/cli/src/state.rs @@ -1,5 +1,5 @@ use std::sync::Arc; -use tokio::sync::Mutex; +use tokio::sync::{Mutex, Notify}; use llm_core::ChatMessage; use serde::{Deserialize, Serialize}; @@ -22,6 +22,8 @@ pub struct AppState { pub messages: Vec, pub running: bool, pub mode: AppMode, + pub last_permission_result: Option, + pub permission_notify: Arc, } impl AppState { @@ -30,12 +32,19 @@ impl AppState { messages: Vec::new(), running: true, mode: AppMode::Normal, + last_permission_result: None, + permission_notify: Arc::new(Notify::new()), } } pub fn add_message(&mut self, message: ChatMessage) { self.messages.push(message); } + + pub fn set_permission_result(&mut self, result: bool) { + self.last_permission_result = Some(result); + self.permission_notify.notify_one(); + } } #[cfg(test)]