Owlibou/owlen
2
0
Fork 0
Code Issues 20 Pull Requests Actions Packages Projects Releases 5 Activity

feat(mcp): add MCP client abstraction and feature flag

Browse Source 
Introduce the foundation for the Multi-Client Provider (MCP) architecture.

This phase includes:
- A new `McpClient` trait to abstract tool execution.
- A `LocalMcpClient` that executes tools in-process for backward compatibility ("legacy mode").
- A placeholder `RemoteMcpClient` for future development.
- An `McpMode` enum in the configuration (`mcp.mode`) to toggle between `legacy` and `enabled` modes, defaulting to `legacy`.
- Refactoring of `SessionController` to use the `McpClient` abstraction, decoupling it from the tool registry.

This lays the groundwork for routing tool calls to a remote MCP server in subsequent phases.
This commit is contained in:
vikingowl
2025-10-06 20:03:01 +02:00
parent 235f84fa19
commit 67381b02db
5 changed files with 211 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files

213
crates/owlen-core/src/session.rs
View File

@@ -1,10 +1,12 @@
use crate::config::Config;
use crate::config::{Config, McpMode};
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::mcp::client::{McpClient, RemoteMcpClient};
use crate::mcp::{LocalMcpClient, McpToolCall};
use crate::model::ModelManager;
use crate::provider::{ChatStream, Provider};
use crate::storage::{SessionMeta, StorageManager};
@@ -102,6 +104,7 @@ pub struct SessionController {
consent_manager: Arc<Mutex<ConsentManager>>,
tool_registry: Arc<ToolRegistry>,
schema_validator: Arc<SchemaValidator>,
mcp_client: Arc<dyn McpClient>,
storage: Arc<StorageManager>,
vault: Option<Arc<Mutex<VaultHandle>>>,
master_key: Option<Arc<Vec<u8>>>,
@@ -109,6 +112,83 @@ pub struct SessionController {
enable_code_tools: bool, // Whether to enable code execution tools (code client only)
}
fn build_tools(
config: &Config,
enable_code_tools: bool,
consent_manager: Arc<Mutex<ConsentManager>>,
credential_manager: Option<Arc<CredentialManager>>,
vault: Option<Arc<Mutex<VaultHandle>>>,
) -> Result<(Arc<ToolRegistry>, Arc<SchemaValidator>)> {
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 config
.security
.allowed_tools
.iter()
.any(|tool| tool == "web_search")
&& config.tools.web_search.enabled
&& config.privacy.enable_remote_search
{
let tool = WebSearchTool::new(
consent_manager.clone(),
credential_manager.clone(),
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 config
.security
.allowed_tools
.iter()
.any(|tool| tool == "web_search") // Same permission as web_search
&& config.tools.web_search.enabled
&& config.privacy.enable_remote_search
{
let tool = WebSearchDetailedTool::new(
consent_manager.clone(),
credential_manager.clone(),
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 enable_code_tools
&& config
.security
.allowed_tools
.iter()
.any(|tool| tool == "code_exec")
&& config.tools.code_exec.enabled
{
let tool = CodeExecTool::new(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);
}
Ok((Arc::new(registry), Arc::new(validator)))
}
impl SessionController {
/// Create a new controller with the given provider and configuration
///
@@ -175,7 +255,23 @@ impl SessionController {
let model_manager = ModelManager::new(config.general.model_cache_ttl());
let mut controller = Self {
let (tool_registry, schema_validator) = build_tools(
&config,
enable_code_tools,
consent_manager.clone(),
credential_manager.clone(),
vault_handle.clone(),
)?;
let mcp_client: Arc<dyn McpClient> = match config.mcp.mode {
McpMode::Legacy => Arc::new(LocalMcpClient::new(
tool_registry.clone(),
schema_validator.clone(),
)),
McpMode::Enabled => Arc::new(RemoteMcpClient {}),
};
let controller = Self {
provider,
conversation,
model_manager,
@@ -183,8 +279,9 @@ impl SessionController {
formatter,
config,
consent_manager,
tool_registry: Arc::new(ToolRegistry::new()),
schema_validator: Arc::new(SchemaValidator::new()),
tool_registry,
schema_validator,
mcp_client,
storage,
vault: vault_handle,
master_key,
@@ -192,8 +289,6 @@ impl SessionController {
enable_code_tools,
};
controller.rebuild_tools()?;
Ok(controller)
}
@@ -416,78 +511,24 @@ impl SessionController {
}
fn rebuild_tools(&mut self) -> Result<()> {
let mut registry = ToolRegistry::new();
let mut validator = SchemaValidator::new();
let (registry, validator) = build_tools(
&self.config,
self.enable_code_tools,
self.consent_manager.clone(),
self.credential_manager.clone(),
self.vault.clone(),
)?;
self.tool_registry = registry;
self.schema_validator = validator;
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}");
}
}
self.mcp_client = match self.config.mcp.mode {
McpMode::Legacy => Arc::new(LocalMcpClient::new(
self.tool_registry.clone(),
self.schema_validator.clone(),
)),
McpMode::Enabled => Arc::new(RemoteMcpClient {}),
};
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(())
}
@@ -592,10 +633,13 @@ impl SessionController {
// 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 mcp_tool_call = McpToolCall {
name: tool_call.name.clone(),
arguments: tool_call.arguments.clone(),
};
let tool_result =
self.mcp_client.call_tool(mcp_tool_call).await;
let tool_response_content = match tool_result {
Ok(result) => serde_json::to_string_pretty(&result.output)
@@ -694,10 +738,11 @@ impl SessionController {
) -> 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 mcp_tool_call = McpToolCall {
name: tool_call.name.clone(),
arguments: tool_call.arguments.clone(),
};
let tool_result = self.mcp_client.call_tool(mcp_tool_call).await;
let tool_response_content = match tool_result {
Ok(result) => serde_json::to_string_pretty(&result.output)