Integrate core functionality for tools, MCP, and enhanced session management

Adds consent management for tool execution, input validation, sandboxed process execution, and MCP server integration. Updates session management to support tool use, conversation persistence, and streaming responses.

Major additions:
- Database migrations for conversations and secure storage
- Encryption and credential management infrastructure
- Extensible tool system with code execution and web search
- Consent management and validation systems
- Sandboxed process execution
- MCP server integration

Infrastructure changes:
- Module registration and workspace dependencies
- ToolCall type and tool-related Message methods
- Privacy, security, and tool configuration structures
- Database-backed conversation persistence
- Tool call tracking in conversations

Provider and UI updates:
- Ollama provider updates for tool support and new Role types
- TUI chat and code app updates for async initialization
- CLI updates for new SessionController API
- Configuration documentation updates
- CHANGELOG updates

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
2025-10-06 18:36:42 +02:00
parent 9c777c8429
commit 235f84fa19
24 changed files with 4734 additions and 1549 deletions

View File

@@ -1,12 +1,26 @@
use crate::config::Config;
use crate::consent::ConsentManager;
use crate::conversation::ConversationManager;
use crate::credentials::CredentialManager;
use crate::encryption::{self, VaultHandle};
use crate::formatting::MessageFormatter;
use crate::input::InputBuffer;
use crate::model::ModelManager;
use crate::provider::{ChatStream, Provider};
use crate::types::{ChatParameters, ChatRequest, ChatResponse, Conversation, ModelInfo};
use crate::Result;
use std::sync::Arc;
use crate::storage::{SessionMeta, StorageManager};
use crate::tools::{
code_exec::CodeExecTool, registry::ToolRegistry, web_search::WebSearchTool,
web_search_detailed::WebSearchDetailedTool, Tool,
};
use crate::types::{
ChatParameters, ChatRequest, ChatResponse, Conversation, Message, ModelInfo, ToolCall,
};
use crate::validation::{get_builtin_schemas, SchemaValidator};
use crate::{Error, Result};
use log::warn;
use std::env;
use std::path::PathBuf;
use std::sync::{Arc, Mutex};
use uuid::Uuid;
/// Outcome of submitting a chat request
@@ -31,6 +45,7 @@ pub enum SessionOutcome {
/// use owlen_core::config::Config;
/// use owlen_core::provider::{Provider, ChatStream};
/// use owlen_core::session::{SessionController, SessionOutcome};
/// use owlen_core::storage::StorageManager;
/// use owlen_core::types::{ChatRequest, ChatResponse, ChatParameters, Message, ModelInfo};
/// use owlen_core::Result;
///
@@ -55,7 +70,9 @@ pub enum SessionOutcome {
/// async fn main() {
/// let provider = Arc::new(MockProvider);
/// let config = Config::default();
/// let mut session_controller = SessionController::new(provider, config);
/// let storage = Arc::new(StorageManager::new().await.unwrap());
/// let enable_code_tools = false; // Set to true for code client
/// let mut session_controller = SessionController::new(provider, config, storage, enable_code_tools).unwrap();
///
/// // Send a message
/// let outcome = session_controller.send_message(
@@ -82,17 +99,69 @@ pub struct SessionController {
input_buffer: InputBuffer,
formatter: MessageFormatter,
config: Config,
consent_manager: Arc<Mutex<ConsentManager>>,
tool_registry: Arc<ToolRegistry>,
schema_validator: Arc<SchemaValidator>,
storage: Arc<StorageManager>,
vault: Option<Arc<Mutex<VaultHandle>>>,
master_key: Option<Arc<Vec<u8>>>,
credential_manager: Option<Arc<CredentialManager>>,
enable_code_tools: bool, // Whether to enable code execution tools (code client only)
}
impl SessionController {
/// Create a new controller with the given provider and configuration
pub fn new(provider: Arc<dyn Provider>, config: Config) -> Self {
///
/// # Arguments
/// * `provider` - The LLM provider to use
/// * `config` - Application configuration
/// * `storage` - Storage manager for persistence
/// * `enable_code_tools` - Whether to enable code execution tools (should only be true for code client)
pub fn new(
provider: Arc<dyn Provider>,
config: Config,
storage: Arc<StorageManager>,
enable_code_tools: bool,
) -> Result<Self> {
let model = config
.general
.default_model
.clone()
.unwrap_or_else(|| "ollama/default".to_string());
let mut vault_handle: Option<Arc<Mutex<VaultHandle>>> = None;
let mut master_key: Option<Arc<Vec<u8>>> = None;
let mut credential_manager: Option<Arc<CredentialManager>> = None;
if config.privacy.encrypt_local_data {
let base_dir = storage
.database_path()
.parent()
.map(|p| p.to_path_buf())
.or_else(dirs::data_local_dir)
.unwrap_or_else(|| PathBuf::from("."));
let secure_path = base_dir.join("encrypted_data.json");
let handle = match env::var("OWLEN_MASTER_PASSWORD") {
Ok(password) if !password.is_empty() => {
encryption::unlock_with_password(secure_path, &password)?
}
_ => encryption::unlock_interactive(secure_path)?,
};
let master = Arc::new(handle.data.master_key.clone());
master_key = Some(master.clone());
vault_handle = Some(Arc::new(Mutex::new(handle)));
credential_manager = Some(Arc::new(CredentialManager::new(storage.clone(), master)));
}
// Load consent manager from vault if available, otherwise create new
let consent_manager = if let Some(ref vault) = vault_handle {
Arc::new(Mutex::new(ConsentManager::from_vault(vault)))
} else {
Arc::new(Mutex::new(ConsentManager::new()))
};
let conversation =
ConversationManager::with_history_capacity(model, config.storage.max_saved_sessions);
let formatter =
@@ -106,14 +175,26 @@ impl SessionController {
let model_manager = ModelManager::new(config.general.model_cache_ttl());
Self {
let mut controller = Self {
provider,
conversation,
model_manager,
input_buffer,
formatter,
config,
}
consent_manager,
tool_registry: Arc::new(ToolRegistry::new()),
schema_validator: Arc::new(SchemaValidator::new()),
storage,
vault: vault_handle,
master_key,
credential_manager,
enable_code_tools,
};
controller.rebuild_tools()?;
Ok(controller)
}
/// Access the active conversation
@@ -156,6 +237,260 @@ impl SessionController {
&mut self.config
}
/// Grant consent programmatically for a tool (for TUI consent dialog)
pub fn grant_consent(&self, tool_name: &str, data_types: Vec<String>, endpoints: Vec<String>) {
let mut consent = self
.consent_manager
.lock()
.expect("Consent manager mutex poisoned");
consent.grant_consent(tool_name, data_types, endpoints);
// Persist to vault if available
if let Some(vault) = &self.vault {
if let Err(e) = consent.persist_to_vault(vault) {
eprintln!("Warning: Failed to persist consent to vault: {}", e);
}
}
}
/// Check if consent is needed for tool calls (non-blocking check)
/// Returns a list of (tool_name, data_types, endpoints) tuples for tools that need consent
pub fn check_tools_consent_needed(
&self,
tool_calls: &[ToolCall],
) -> Vec<(String, Vec<String>, Vec<String>)> {
let consent = self
.consent_manager
.lock()
.expect("Consent manager mutex poisoned");
let mut needs_consent = Vec::new();
let mut seen_tools = std::collections::HashSet::new();
for tool_call in tool_calls {
// Skip if we already checked this tool
if seen_tools.contains(&tool_call.name) {
continue;
}
seen_tools.insert(tool_call.name.clone());
// Get tool metadata (data types and endpoints) based on tool name
let (data_types, endpoints) = match tool_call.name.as_str() {
"web_search" | "web_search_detailed" => (
vec!["search query".to_string()],
vec!["duckduckgo.com".to_string()],
),
"code_exec" => (
vec!["code to execute".to_string()],
vec!["local sandbox".to_string()],
),
_ => (vec![], vec![]),
};
if let Some((tool_name, dt, ep)) =
consent.check_if_consent_needed(&tool_call.name, data_types, endpoints)
{
needs_consent.push((tool_name, dt, ep));
}
}
needs_consent
}
/// Persist the active conversation to storage
pub async fn save_active_session(
&self,
name: Option<String>,
description: Option<String>,
) -> Result<Uuid> {
self.conversation
.save_active_with_description(&self.storage, name, description)
.await
}
/// Persist the active conversation without description override
pub async fn save_active_session_simple(&self, name: Option<String>) -> Result<Uuid> {
self.conversation.save_active(&self.storage, name).await
}
/// Load a saved conversation by ID and make it active
pub async fn load_saved_session(&mut self, id: Uuid) -> Result<()> {
self.conversation.load_saved(&self.storage, id).await
}
/// Retrieve session metadata from storage
pub async fn list_saved_sessions(&self) -> Result<Vec<SessionMeta>> {
ConversationManager::list_saved_sessions(&self.storage).await
}
pub async fn delete_session(&self, id: Uuid) -> Result<()> {
self.storage.delete_session(id).await
}
pub async fn clear_secure_data(&self) -> Result<()> {
self.storage.clear_secure_items().await?;
if let Some(vault) = &self.vault {
let mut guard = vault.lock().expect("Vault mutex poisoned");
guard.data.settings.clear();
guard.persist()?;
}
// Also clear consent records
{
let mut consent = self
.consent_manager
.lock()
.expect("Consent manager mutex poisoned");
consent.clear_all_consent();
}
self.persist_consent()?;
Ok(())
}
/// Persist current consent state to vault (if encryption is enabled)
pub fn persist_consent(&self) -> Result<()> {
if let Some(vault) = &self.vault {
let consent = self
.consent_manager
.lock()
.expect("Consent manager mutex poisoned");
consent.persist_to_vault(vault)?;
}
Ok(())
}
pub async fn set_tool_enabled(&mut self, tool: &str, enabled: bool) -> Result<()> {
match tool {
"web_search" => {
self.config.tools.web_search.enabled = enabled;
self.config.privacy.enable_remote_search = enabled;
}
"code_exec" => {
self.config.tools.code_exec.enabled = enabled;
}
other => {
return Err(Error::InvalidInput(format!("Unknown tool: {other}")));
}
}
self.rebuild_tools()?;
Ok(())
}
/// Access the consent manager shared across tools
pub fn consent_manager(&self) -> Arc<Mutex<ConsentManager>> {
self.consent_manager.clone()
}
/// Access the tool registry for executing registered tools
pub fn tool_registry(&self) -> Arc<ToolRegistry> {
Arc::clone(&self.tool_registry)
}
/// Access the schema validator used for tool input validation
pub fn schema_validator(&self) -> Arc<SchemaValidator> {
Arc::clone(&self.schema_validator)
}
/// Construct an MCP server facade for the active tool registry
pub fn mcp_server(&self) -> crate::mcp::McpServer {
crate::mcp::McpServer::new(self.tool_registry(), self.schema_validator())
}
/// Access the underlying storage manager
pub fn storage(&self) -> Arc<StorageManager> {
Arc::clone(&self.storage)
}
/// Retrieve the active master key if encryption is enabled
pub fn master_key(&self) -> Option<Arc<Vec<u8>>> {
self.master_key.as_ref().map(Arc::clone)
}
/// Access the vault handle for managing secure settings
pub fn vault(&self) -> Option<Arc<Mutex<VaultHandle>>> {
self.vault.as_ref().map(Arc::clone)
}
/// Access the credential manager if available
pub fn credential_manager(&self) -> Option<Arc<CredentialManager>> {
self.credential_manager.as_ref().map(Arc::clone)
}
fn rebuild_tools(&mut self) -> Result<()> {
let mut registry = ToolRegistry::new();
let mut validator = SchemaValidator::new();
for (name, schema) in get_builtin_schemas() {
if let Err(err) = validator.register_schema(&name, schema) {
warn!("Failed to register built-in schema {name}: {err}");
}
}
if self
.config
.security
.allowed_tools
.iter()
.any(|tool| tool == "web_search")
&& self.config.tools.web_search.enabled
&& self.config.privacy.enable_remote_search
{
let tool = WebSearchTool::new(
self.consent_manager.clone(),
self.credential_manager.clone(),
self.vault.clone(),
);
let schema = tool.schema();
if let Err(err) = validator.register_schema(tool.name(), schema) {
warn!("Failed to register schema for {}: {err}", tool.name());
}
registry.register(tool);
}
// Register web_search_detailed tool (provides snippets)
if self
.config
.security
.allowed_tools
.iter()
.any(|tool| tool == "web_search") // Same permission as web_search
&& self.config.tools.web_search.enabled
&& self.config.privacy.enable_remote_search
{
let tool = WebSearchDetailedTool::new(
self.consent_manager.clone(),
self.credential_manager.clone(),
self.vault.clone(),
);
let schema = tool.schema();
if let Err(err) = validator.register_schema(tool.name(), schema) {
warn!("Failed to register schema for {}: {err}", tool.name());
}
registry.register(tool);
}
// Code execution tool - only available in code client
if self.enable_code_tools
&& self
.config
.security
.allowed_tools
.iter()
.any(|tool| tool == "code_exec")
&& self.config.tools.code_exec.enabled
{
let tool = CodeExecTool::new(self.config.tools.code_exec.allowed_languages.clone());
let schema = tool.schema();
if let Err(err) = validator.register_schema(tool.name(), schema) {
warn!("Failed to register schema for {}: {err}", tool.name());
}
registry.register(tool);
}
self.tool_registry = Arc::new(registry);
self.schema_validator = Arc::new(validator);
Ok(())
}
/// Currently selected model identifier
pub fn selected_model(&self) -> &str {
&self.conversation.active().model
@@ -187,6 +522,13 @@ impl SessionController {
}
}
/// Replace the active provider at runtime and invalidate cached model listings
pub async fn switch_provider(&mut self, provider: Arc<dyn Provider>) -> Result<()> {
self.provider = provider;
self.model_manager.invalidate().await;
Ok(())
}
/// Submit a user message; optionally stream the response
pub async fn send_message(
&mut self,
@@ -210,38 +552,104 @@ impl SessionController {
let streaming = parameters.stream || self.config.general.enable_streaming;
parameters.stream = streaming;
let request = ChatRequest {
model: self.conversation.active().model.clone(),
messages: self.conversation.active().messages.clone(),
parameters,
// Get available tools
let tools = if !self.tool_registry.all().is_empty() {
Some(
self.tool_registry
.all()
.into_iter()
.map(|tool| crate::mcp::McpToolDescriptor {
name: tool.name().to_string(),
description: tool.description().to_string(),
input_schema: tool.schema(),
requires_network: tool.requires_network(),
requires_filesystem: tool.requires_filesystem(),
})
.collect(),
)
} else {
None
};
if streaming {
match self.provider.chat_stream(request).await {
Ok(stream) => {
let response_id = self.conversation.start_streaming_response();
Ok(SessionOutcome::Streaming {
response_id,
stream,
})
}
Err(err) => {
self.conversation
.push_assistant_message(format!("Error starting stream: {}", err));
Err(err)
let mut request = ChatRequest {
model: self.conversation.active().model.clone(),
messages: self.conversation.active().messages.clone(),
parameters: parameters.clone(),
tools: tools.clone(),
};
// Tool execution loop (non-streaming only for now)
if !streaming {
const MAX_TOOL_ITERATIONS: usize = 5;
for _iteration in 0..MAX_TOOL_ITERATIONS {
match self.provider.chat(request.clone()).await {
Ok(response) => {
// Check if the response has tool calls
if response.message.has_tool_calls() {
// Add assistant's tool call message to conversation
self.conversation.push_message(response.message.clone());
// Execute each tool call
if let Some(tool_calls) = &response.message.tool_calls {
for tool_call in tool_calls {
let tool_result = self
.tool_registry
.execute(&tool_call.name, tool_call.arguments.clone())
.await;
let tool_response_content = match tool_result {
Ok(result) => serde_json::to_string_pretty(&result.output)
.unwrap_or_else(|_| {
"Tool execution succeeded".to_string()
}),
Err(e) => format!("Tool execution failed: {}", e),
};
// Add tool response to conversation
let tool_msg =
Message::tool(tool_call.id.clone(), tool_response_content);
self.conversation.push_message(tool_msg);
}
}
// Update request with new messages for next iteration
request.messages = self.conversation.active().messages.clone();
continue;
} else {
// No more tool calls, return final response
self.conversation.push_message(response.message.clone());
return Ok(SessionOutcome::Complete(response));
}
}
Err(err) => {
self.conversation
.push_assistant_message(format!("Error: {}", err));
return Err(err);
}
}
}
} else {
match self.provider.chat(request).await {
Ok(response) => {
self.conversation.push_message(response.message.clone());
Ok(SessionOutcome::Complete(response))
}
Err(err) => {
self.conversation
.push_assistant_message(format!("Error: {}", err));
Err(err)
}
// Max iterations reached
self.conversation
.push_assistant_message("Maximum tool execution iterations reached".to_string());
return Err(crate::Error::Provider(anyhow::anyhow!(
"Maximum tool execution iterations reached"
)));
}
// Streaming mode with tool support
match self.provider.chat_stream(request).await {
Ok(stream) => {
let response_id = self.conversation.start_streaming_response();
Ok(SessionOutcome::Streaming {
response_id,
stream,
})
}
Err(err) => {
self.conversation
.push_assistant_message(format!("Error starting stream: {}", err));
Err(err)
}
}
}
@@ -254,10 +662,64 @@ impl SessionController {
/// Apply streaming chunk to the conversation
pub fn apply_stream_chunk(&mut self, message_id: Uuid, chunk: &ChatResponse) -> Result<()> {
// Check if this chunk contains tool calls
if chunk.message.has_tool_calls() {
// This is a tool call chunk - store the tool calls on the message
self.conversation.set_tool_calls_on_message(
message_id,
chunk.message.tool_calls.clone().unwrap_or_default(),
)?;
}
self.conversation
.append_stream_chunk(message_id, &chunk.message.content, chunk.is_final)
}
/// Check if a streaming message has complete tool calls that need execution
pub fn check_streaming_tool_calls(&self, message_id: Uuid) -> Option<Vec<ToolCall>> {
self.conversation
.active()
.messages
.iter()
.find(|m| m.id == message_id)
.and_then(|m| m.tool_calls.clone())
.filter(|calls| !calls.is_empty())
}
/// Execute tools for a streaming response and continue conversation
pub async fn execute_streaming_tools(
&mut self,
_message_id: Uuid,
tool_calls: Vec<ToolCall>,
) -> Result<SessionOutcome> {
// Execute each tool call
for tool_call in &tool_calls {
let tool_result = self
.tool_registry
.execute(&tool_call.name, tool_call.arguments.clone())
.await;
let tool_response_content = match tool_result {
Ok(result) => serde_json::to_string_pretty(&result.output)
.unwrap_or_else(|_| "Tool execution succeeded".to_string()),
Err(e) => format!("Tool execution failed: {}", e),
};
// Add tool response to conversation
let tool_msg = Message::tool(tool_call.id.clone(), tool_response_content);
self.conversation.push_message(tool_msg);
}
// Continue the conversation with tool results
let parameters = ChatParameters {
stream: self.config.general.enable_streaming,
..Default::default()
};
self.send_request_with_current_conversation(parameters)
.await
}
/// Access conversation history
pub fn history(&self) -> Vec<Conversation> {
self.conversation.history().cloned().collect()
@@ -335,6 +797,7 @@ impl SessionController {
stream: false,
extra: std::collections::HashMap::new(),
},
tools: None,
};
// Get the summary from the provider