From 94811c5a899adae225ebf9e0c3d18718e0030d3c Mon Sep 17 00:00:00 2001 From: milan Date: Tue, 7 Apr 2026 17:53:32 +0200 Subject: [PATCH] add permission support refactoring --- .gitignore | 3 +- mcp_server_collection/Cargo.lock | 364 ++++++++++++++++++++++++++++ mcp_server_collection/Cargo.toml | 3 + mcp_server_collection/src/config.rs | 10 +- mcp_server_collection/src/main.rs | 37 ++- own_mcp/src/mcp/chat.rs | 191 +++++++++++---- src/config.rs | 5 +- src/main.rs | 16 +- 8 files changed, 566 insertions(+), 63 deletions(-) diff --git a/.gitignore b/.gitignore index 9126183..aafa6c3 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ /target */target -assist.toml \ No newline at end of file +assist.toml +mcp_server_collection/server.toml \ No newline at end of file diff --git a/mcp_server_collection/Cargo.lock b/mcp_server_collection/Cargo.lock index 94798a5..4b18138 100644 --- a/mcp_server_collection/Cargo.lock +++ b/mcp_server_collection/Cargo.lock @@ -2,6 +2,15 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + [[package]] name = "android_system_properties" version = "0.1.5" @@ -206,6 +215,17 @@ dependencies = [ "syn", ] +[[package]] +name = "displaydoc" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97369cbbc041bc366949bc74d34658d6cda5621039731c6310521892a3a20ae0" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "dyn-clone" version = "1.0.20" @@ -476,6 +496,88 @@ dependencies = [ "cc", ] +[[package]] +name = "icu_collections" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2984d1cd16c883d7935b9e07e44071dca8d917fd52ecc02c04d5fa0b5a3f191c" +dependencies = [ + "displaydoc", + "potential_utf", + "utf8_iter", + "yoke", + "zerofrom", + "zerovec", +] + +[[package]] +name = "icu_locale_core" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92219b62b3e2b4d88ac5119f8904c10f8f61bf7e95b640d25ba3075e6cac2c29" +dependencies = [ + "displaydoc", + "litemap", + "tinystr", + "writeable", + "zerovec", +] + +[[package]] +name = "icu_normalizer" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c56e5ee99d6e3d33bd91c5d85458b6005a22140021cc324cea84dd0e72cff3b4" +dependencies = [ + "icu_collections", + "icu_normalizer_data", + "icu_properties", + "icu_provider", + "smallvec", + "zerovec", +] + +[[package]] +name = "icu_normalizer_data" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da3be0ae77ea334f4da67c12f149704f19f81d1adf7c51cf482943e84a2bad38" + +[[package]] +name = "icu_properties" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bee3b67d0ea5c2cca5003417989af8996f8604e34fb9ddf96208a033901e70de" +dependencies = [ + "icu_collections", + "icu_locale_core", + "icu_properties_data", + "icu_provider", + "zerotrie", + "zerovec", +] + +[[package]] +name = "icu_properties_data" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e2bbb201e0c04f7b4b3e14382af113e17ba4f63e2c9d2ee626b720cbce54a14" + +[[package]] +name = "icu_provider" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "139c4cf31c8b5f33d7e199446eff9c1e02decfc2f0eec2c8d71f65befa45b421" +dependencies = [ + "displaydoc", + "icu_locale_core", + "writeable", + "yoke", + "zerofrom", + "zerotrie", + "zerovec", +] + [[package]] name = "id-arena" version = "2.3.0" @@ -488,6 +590,27 @@ version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39" +[[package]] +name = "idna" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b0875f23caa03898994f6ddc501886a45c7d3d62d04d2d90788d47be1b1e4de" +dependencies = [ + "idna_adapter", + "smallvec", + "utf8_iter", +] + +[[package]] +name = "idna_adapter" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3acae9609540aa318d1bc588455225fb2085b9ed0c4f6bd0d9d5bcd86f1a0344" +dependencies = [ + "icu_normalizer", + "icu_properties", +] + [[package]] name = "indexmap" version = "2.13.1" @@ -516,6 +639,12 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + [[package]] name = "leb128fmt" version = "0.1.0" @@ -528,6 +657,12 @@ version = "0.2.184" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "48f5d2a454e16a5ea0f4ced81bd44e4cfc7bd3a507b61887c99fd3538b28e4af" +[[package]] +name = "litemap" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" + [[package]] name = "lock_api" version = "0.4.14" @@ -543,6 +678,15 @@ version = "0.4.29" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" +[[package]] +name = "matchers" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1525a2a28c7f4fa0fc98bb91ae755d1e2d1505079e05539e35bc876b5d65ae9" +dependencies = [ + "regex-automata", +] + [[package]] name = "matchit" version = "0.8.4" @@ -560,6 +704,8 @@ dependencies = [ "serde", "tokio", "tokio-util", + "tracing-subscriber", + "url", ] [[package]] @@ -585,6 +731,15 @@ dependencies = [ "windows-sys", ] +[[package]] +name = "nu-ansi-term" +version = "0.50.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" +dependencies = [ + "windows-sys", +] + [[package]] name = "num-traits" version = "0.2.19" @@ -650,6 +805,15 @@ version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" +[[package]] +name = "potential_utf" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0103b1cef7ec0cf76490e969665504990193874ea05c85ff9bab8b911d0a0564" +dependencies = [ + "zerovec", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -730,6 +894,23 @@ dependencies = [ "syn", ] +[[package]] +name = "regex-automata" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc897dd8d9e8bd1ed8cdad82b5966c3e0ecae09fb1907d58efaa013543185d0a" + [[package]] name = "rmcp" version = "1.3.0" @@ -910,6 +1091,15 @@ dependencies = [ "serde", ] +[[package]] +name = "sharded-slab" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" +dependencies = [ + "lazy_static", +] + [[package]] name = "shlex" version = "1.3.0" @@ -961,6 +1151,12 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "stable_deref_trait" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" + [[package]] name = "strsim" version = "0.11.1" @@ -984,6 +1180,17 @@ version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0bf256ce5efdfa370213c1dabab5935a12e49f2c58d15e9eac2870d3b4f27263" +[[package]] +name = "synstructure" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "thiserror" version = "2.0.18" @@ -1004,6 +1211,25 @@ dependencies = [ "syn", ] +[[package]] +name = "thread_local" +version = "1.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f60246a4944f24f6e018aa17cdeffb7818b76356965d03b07d6a9886e8962185" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "tinystr" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8323304221c2a851516f22236c5722a72eaa19749016521d6dff0824447d96d" +dependencies = [ + "displaydoc", + "zerovec", +] + [[package]] name = "tokio" version = "1.51.0" @@ -1153,6 +1379,36 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" dependencies = [ "once_cell", + "valuable", +] + +[[package]] +name = "tracing-log" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3" +dependencies = [ + "log", + "once_cell", + "tracing-core", +] + +[[package]] +name = "tracing-subscriber" +version = "0.3.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319" +dependencies = [ + "matchers", + "nu-ansi-term", + "once_cell", + "regex-automata", + "sharded-slab", + "smallvec", + "thread_local", + "tracing", + "tracing-core", + "tracing-log", ] [[package]] @@ -1167,6 +1423,25 @@ version = "0.2.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" +[[package]] +name = "url" +version = "2.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff67a8a4397373c3ef660812acab3268222035010ab8680ec4215f38ba3d0eed" +dependencies = [ + "form_urlencoded", + "idna", + "percent-encoding", + "serde", + "serde_derive", +] + +[[package]] +name = "utf8_iter" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" + [[package]] name = "uuid" version = "1.23.0" @@ -1178,6 +1453,12 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "valuable" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" + [[package]] name = "wasi" version = "0.11.1+wasi-snapshot-preview1" @@ -1443,6 +1724,89 @@ dependencies = [ "wasmparser", ] +[[package]] +name = "writeable" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4" + +[[package]] +name = "yoke" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "abe8c5fda708d9ca3df187cae8bfb9ceda00dd96231bed36e445a1a48e66f9ca" +dependencies = [ + "stable_deref_trait", + "yoke-derive", + "zerofrom", +] + +[[package]] +name = "yoke-derive" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[package]] +name = "zerofrom" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69faa1f2a1ea75661980b013019ed6687ed0e83d069bc1114e2cc74c6c04c4df" +dependencies = [ + "zerofrom-derive", +] + +[[package]] +name = "zerofrom-derive" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[package]] +name = "zerotrie" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f9152d31db0792fa83f70fb2f83148effb5c1f5b8c7686c3459e361d9bc20bf" +dependencies = [ + "displaydoc", + "yoke", + "zerofrom", +] + +[[package]] +name = "zerovec" +version = "0.11.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90f911cbc359ab6af17377d242225f4d75119aec87ea711a880987b18cd7b239" +dependencies = [ + "yoke", + "zerofrom", + "zerovec-derive", +] + +[[package]] +name = "zerovec-derive" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "zmij" version = "1.0.21" diff --git a/mcp_server_collection/Cargo.toml b/mcp_server_collection/Cargo.toml index 596f014..007bab6 100644 --- a/mcp_server_collection/Cargo.toml +++ b/mcp_server_collection/Cargo.toml @@ -11,6 +11,9 @@ serde = { version = "1.0.228", features = ["derive"] } tokio = { version = "1.51.0", features = ["full"] } own_assist_common = {path="../common"} tokio-util = "0.7.18" +tracing-subscriber = {version = "0.3.23", features = ["env-filter"]} +url = {version = "2.5.8", features = ["serde"]} [features] +default = ["datetime"] datetime = [] \ No newline at end of file diff --git a/mcp_server_collection/src/config.rs b/mcp_server_collection/src/config.rs index db8300c..aa24241 100644 --- a/mcp_server_collection/src/config.rs +++ b/mcp_server_collection/src/config.rs @@ -1,14 +1,16 @@ -use serde::Deserialize; +use serde::{Deserialize, Serialize}; use own_assist_common::config_from_file; use own_assist_common::config_loader::ConfigLoadingError; use crate::McpServiceType; -#[derive(Debug, Deserialize)] +#[derive(Debug, Deserialize, Serialize)] pub(crate) struct Config { + #[serde(rename = "bind-address")] + pub(crate) bind_address: Option, pub(crate) servers: Vec, } -#[derive(Debug, Deserialize)] +#[derive(Debug, Clone, Deserialize, Serialize)] pub(crate) struct ServerConfig { pub(crate) path: String, pub(crate) r#type: McpServiceType @@ -16,6 +18,6 @@ pub(crate) struct ServerConfig { impl Config { pub fn from_file() -> Result { - config_from_file("servers.toml") + config_from_file("server.toml") } } \ No newline at end of file diff --git a/mcp_server_collection/src/main.rs b/mcp_server_collection/src/main.rs index 8c72c40..fe053ae 100644 --- a/mcp_server_collection/src/main.rs +++ b/mcp_server_collection/src/main.rs @@ -2,13 +2,16 @@ pub mod datetime; mod config; +use axum::response::Json; use axum::Router; use rmcp::transport::{ StreamableHttpServerConfig, streamable_http_server::{session::local::LocalSessionManager, tower::StreamableHttpService}, }; -use serde::Deserialize; +use serde::{Deserialize, Serialize}; use tokio::main; +use tracing_subscriber::layer::SubscriberExt; +use tracing_subscriber::util::SubscriberInitExt; use crate::config::{Config, ServerConfig}; #[cfg(feature = "datetime")] use crate::datetime::DateTimeHandler; @@ -21,16 +24,16 @@ pub trait McpServerHandler: rmcp::ServerHandler { } } -#[derive(Debug, Deserialize)] +#[derive(Debug, Clone, Deserialize, Serialize)] #[non_exhaustive] pub enum McpServiceType { DateTime } -pub(crate) fn mcp_router(server_configs: Vec) -> Router { +pub(crate) fn mcp_router(server_configs: &Vec) -> Router { let mut router = Router::new(); - for config in &server_configs { + for config in server_configs { router = match config.r#type { McpServiceType::DateTime => { router.route_service(config.path.as_str(), DateTimeHandler::mcp_service()) @@ -41,16 +44,36 @@ pub(crate) fn mcp_router(server_configs: Vec) -> Router { router } +fn bind_address_format(url: url::Url) -> String { + format!("{}:{}", url.host_str().expect("No host specified"), url.port().unwrap_or(8000)) +} + #[main] async fn main() { + tracing_subscriber::registry() + .with( + tracing_subscriber::EnvFilter::try_from_default_env() + .unwrap_or_else(|_| "debug".to_string().into()), + ) + .with(tracing_subscriber::fmt::layer()) + .init(); + let config = Config::from_file().inspect_err(|e|{ - eprintln!("Error loading config: {}", e); // FIXME: use actual logger instead + eprintln!("Error loading config: {e}"); std::process::exit(1); }).unwrap(); - let router = mcp_router(config.servers); + println!("config: {config:#?}"); - let tcp_listener = tokio::net::TcpListener::bind("127.0.0.1:8000").await.unwrap(); // use toml config + let router = mcp_router(&config.servers).route("/", axum::routing::get(|| async { Json(config.servers) })); + + let bind_address = config.bind_address.unwrap_or(url::Url::parse("localhost:8000").unwrap()); + println!("binding address at {bind_address}"); + + let tcp_listener = tokio::net::TcpListener::bind(bind_address_format(bind_address)).await.inspect_err(|e|{ + eprintln!("Error bind tcp listener: {e:#?}"); + std::process::exit(1); + }).unwrap(); let ct = tokio_util::sync::CancellationToken::new(); diff --git a/own_mcp/src/mcp/chat.rs b/own_mcp/src/mcp/chat.rs index b731317..86eaa7d 100644 --- a/own_mcp/src/mcp/chat.rs +++ b/own_mcp/src/mcp/chat.rs @@ -5,9 +5,9 @@ use ollama_rs::generation::chat::ChatMessage; use ollama_rs::generation::chat::request::ChatMessageRequest; use ollama_rs::generation::tools::ToolInfo; use rmcp::ServiceError; -use rmcp::model::{CallToolRequestParams, CallToolResult, }; +use rmcp::model::{CallToolRequestParams, CallToolResult}; use std::collections::HashMap; - +use std::ops::Index; use thiserror::Error; #[derive(Debug, Copy, Clone, Default)] @@ -18,9 +18,15 @@ pub enum ToolPermission { Denied, } +#[derive(Debug, Copy, Clone)] +pub enum PermissionAnswer { + Granted, + Denied, +} + #[derive(Debug, Clone)] pub struct RestrictedTool { - permission: ToolPermission, // TODO: USE PERMISSION + permission: ToolPermission, tool_info: ToolInfo, } @@ -28,7 +34,7 @@ pub struct RestrictedTool { pub struct MCPServerData { name: String, client: MCPClient, - /// tools in the ollama format. mcp resources are also translated into tools + /// tools in the ollama format. mcp resources are also translated into tools (not yet) translated_tools: Vec, } @@ -43,6 +49,8 @@ impl MCPServerData { }) } + // Resource into tool translation commented out, since resource reading is not yet solved. `AgentChat` would instead try to call tool with the resource name. + /* match crate::mcp::translation::get_server_resource_tool_info(&client, &name).await { Ok(resources_as_tool_infos) => { for tool_info in resources_as_tool_infos { @@ -56,6 +64,7 @@ impl MCPServerData { log::warn!("Server `{name}` does not support resources"); } } + */ Ok(MCPServerData { name, @@ -85,6 +94,8 @@ pub enum ChatError { ArgumentParsingError, #[error("service could not be found")] ServiceNotFoundError(String), + #[error("tool could not be found")] + ToolNotFoundError(String), } impl AgentChat { @@ -92,7 +103,7 @@ impl AgentChat { ollama_client: Ollama, model: String, mcp_clients: HashMap, - system_prompt: String + system_prompt: String, ) -> Result { let history = vec![ChatMessage::system(system_prompt)]; @@ -111,22 +122,74 @@ impl AgentChat { }) } - pub fn get_all_tools(&self) -> Vec { - let mut all_tools: Vec = Vec::new(); - - for server in self.mcp_servers.values() { - for tool in server.translated_tools.iter() { - all_tools.push(tool.tool_info.clone()); - } - } - - all_tools + pub fn get_all_tools(&self) -> impl Iterator { + self.mcp_servers + .values() + .map(|server| server.translated_tools.iter()) + .flatten() } - pub async fn message(&mut self, user_message: String) -> Result { - let all_tools = self.get_all_tools(); + /// + /// + /// # Arguments + /// + /// * `name`: must be in the format "mcp_server_name::tool_or_resource_name". + /// + /// returns: + /// Ok(tuple) => tuple of the server name and the tool or resource name + /// Err(error) => ChatError::FunctionParseError(name) + fn parse_tool_name(name: String) -> Result<(String, String), ChatError> { + let (mcp_server_name, tool_name) = name + .split_once("::") + .ok_or(ChatError::FunctionParseError(name.clone()))?; - log::debug!("all tools: {:#?}", all_tools); + Ok((mcp_server_name.to_string(), tool_name.to_string())) + } + + pub fn get_mcp_server_by_name(&self, name: String) -> Option<&MCPServerData> { + self.mcp_servers.get(&name) + } + + /// + /// + /// # Arguments + /// + /// * `mpc_server_data`: + /// * `name`: must be in the format "mcp_server_name::tool_or_resource_name". + /// + /// returns: Option<&RestrictedTool> + pub fn get_tool<'a>( + &self, + mpc_server_data: &'a MCPServerData, + name: String, + ) -> Option<&'a RestrictedTool> { + mpc_server_data + .translated_tools + .iter() + .find(|tool| tool.tool_info.function.name == name) + } + + pub fn set_permission(&mut self, full_tool_name: String, permission: ToolPermission) -> Result<(), ChatError> { + let (mcp_server_name, tool_name) = Self::parse_tool_name(full_tool_name.clone())?; + + let mcp_server_data: &mut MCPServerData = self.mcp_servers.get_mut(&mcp_server_name.clone()).ok_or(ChatError::ServiceNotFoundError(mcp_server_name.clone()))?; + let tool_index = mcp_server_data.translated_tools.iter().position(|tool| tool.tool_info.function.name == full_tool_name).ok_or(ChatError::ToolNotFoundError(tool_name))?; + mcp_server_data.translated_tools[tool_index].permission = permission; + Ok(()) + } + + pub async fn message>( + &mut self, + user_message: String, + permission_request_callback: fn(mcp_server_name: String, tool_name: String) -> C, + ) -> Result + { + let all_tools: Vec = self + .get_all_tools() + .map(|restricted_tool| restricted_tool.tool_info.clone()) + .collect(); + + log::debug!("all tools: {all_tools:#?}"); let mut response = self .ollama_client @@ -145,13 +208,16 @@ impl AgentChat { for tool_call in &response.message.tool_calls { log::debug!("calling tool {}", tool_call.function.name); - let result = self.call_tool(tool_call).await?; - let contents = result - .content - .iter() - .filter_map(|content| content.as_text()) - .map(|text_content| text_content.text.clone()) - .collect::>(); + let result = self.call_tool(tool_call, permission_request_callback).await?; + + let contents = match result { + Some(result) => result.content + .iter() + .filter_map(|content| content.as_text()) + .map(|text_content| text_content.text.clone()) + .collect::>(), + None => vec![String::from("Tool Permission Denied by the user")], + }; log::debug!("contents: {:#?}", contents); @@ -163,7 +229,8 @@ impl AgentChat { .ollama_client .send_chat_messages_with_history( &mut self.message_history, - ChatMessageRequest::new(self.model.clone(), Vec::new()).tools(all_tools.clone()), + ChatMessageRequest::new(self.model.clone(), Vec::new()) + .tools(all_tools.clone()), ) .await?; log::debug!("received message: {:#?}", response.message); @@ -172,10 +239,25 @@ impl AgentChat { Ok(response.message) } - async fn call_tool( + /// + /// + /// # Arguments + /// + /// * `tool_call`: + /// * `permission_request_callback`: + /// + /// returns: Result, ChatError> + /// if permission was denied: Ok(None) + /// # Examples + /// + /// ``` + /// + /// ``` + async fn call_tool>( &self, tool_call: &ollama_rs::generation::tools::ToolCall, - ) -> Result { + permission_request_callback: fn(mcp_server_name: String, tool_name: String) -> C, + ) -> Result, ChatError> { let arguments_json_object = tool_call .function .arguments @@ -184,18 +266,44 @@ impl AgentChat { .inspect_err(|_| log::error!("arguments are not of type object"))?; log::trace!("arguments_json_object: {:?}", arguments_json_object); - let (mcp_name, function_name) = - tool_call - .function - .name - .split_once("::") - .ok_or(ChatError::FunctionParseError( - tool_call.function.name.clone(), - ))?; - let (mcp_name, function_name) = (mcp_name.to_string(), function_name.to_string()); + let full_unparsed_tool_name = &tool_call.function.name; // "mcp_server_name::tool_name" + + let (mcp_server_name, tool_name) = Self::parse_tool_name(full_unparsed_tool_name.clone())?; + + let mcp_server_data = self + .get_mcp_server_by_name(mcp_server_name.clone()) + .ok_or(ChatError::ServiceNotFoundError(mcp_server_name.clone()))?; + let restricted_tool = self + .get_tool(&mcp_server_data, full_unparsed_tool_name.clone()) + .ok_or(ChatError::ToolNotFoundError( + full_unparsed_tool_name.clone(), + ))?; + + match restricted_tool.permission { + ToolPermission::Allowed => { + log::trace!("allowed restricted tool: {:?}", restricted_tool.tool_info); + } + ToolPermission::Denied => { + log::info!("denied restricted tool: {:?}", restricted_tool); + return Ok(None); + } + ToolPermission::Ask => { + log::trace!("ask restricted tool: {:?}", restricted_tool); + let response = permission_request_callback(mcp_server_name, tool_name.clone()).await; + match response { + PermissionAnswer::Granted => { + log::trace!("granted restricted tool: {:?}", restricted_tool); + } + PermissionAnswer::Denied => { + log::info!("denied restricted tool: {:?}", restricted_tool); + return Ok(None); + } + } + } + } let request_params = - CallToolRequestParams::new(function_name).with_arguments(arguments_json_object.clone()); + CallToolRequestParams::new(tool_name).with_arguments(arguments_json_object.clone()); log::debug!( "calling tool {} with request_params: {:#?}", @@ -203,11 +311,6 @@ impl AgentChat { request_params ); - Ok(self - .mcp_servers - .get(&mcp_name) - .ok_or(ChatError::ServiceNotFoundError(mcp_name))? - .client.call_tool(request_params) - .await?) + Ok(Some(mcp_server_data.client.call_tool(request_params).await?)) } } diff --git a/src/config.rs b/src/config.rs index cdca07c..dac2d99 100644 --- a/src/config.rs +++ b/src/config.rs @@ -26,14 +26,13 @@ impl Config { } pub fn ollama_instance(&self) -> Ollama { - let ollama_default_url = Url::parse("http://127.0.0.1:11434").unwrap(); - let mut ollama_headers = HeaderMap::new(); if let Some(authorization_header) = &self.ollama.authorization { ollama_headers.append("Authorization", HeaderValue::from_str(&authorization_header).unwrap()); } + let ollama_default_url = Url::parse("http://127.0.0.1:11434").unwrap(); let mut ollama = Ollama::from_url(self.ollama.url.clone().unwrap_or(ollama_default_url)); ollama.set_headers(Some(ollama_headers)); @@ -64,7 +63,7 @@ impl Config { pub struct OllamaConfig { pub url: Option, pub model: OllamaModelConfig, - pub authorization: Option, // TODO: implement + pub authorization: Option, } #[derive(Debug, Deserialize)] diff --git a/src/main.rs b/src/main.rs index 1c02655..8367d7d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,9 +1,14 @@ mod config; mod model; +use own_mcp::mcp::chat::{PermissionAnswer, ToolPermission}; use crate::config::Config; use crate::model::create_model_from_config; +async fn permission_callback(mcp_server_name: String, tool_name:String) -> PermissionAnswer { + PermissionAnswer::Denied // TODO: implement +} + #[tokio::main] async fn main() { env_logger::init(); @@ -27,7 +32,8 @@ async fn main() { let system_prompt = "Du bist ein Assistent, der per Sprache bedient wird. Du erhälst die Transkription. \ Wichtiger als deine Antworten sind deine Aktionen.\ Nutze bitte die tools, falls du sie brauchst um Informationen zu bekommen (z.B: über das aktuelle Datum oder den aktuellen Wochentag). \ - Du bist in einem Agent Loop und kannst mehrere Tools hintereinander nutzen.".to_string(); + Du bist in einem Agent Loop und kannst mehrere Tools hintereinander nutzen. \ + Falls der Nutzer das Nutzen eines Tools ablehnt, sag ihm bescheid, dass du es brauchst.".to_string(); let mut agent_chat = own_mcp::AgentChat::new(ollama, model_name.clone(), mcp_clients, system_prompt).await.inspect_err( |e| { @@ -36,10 +42,12 @@ async fn main() { } ).unwrap(); - log::info!("all tools: {:#?}", agent_chat.get_all_tools()); + log::info!("all tools: {:#?}", agent_chat.get_all_tools().collect::>()); - let answer = agent_chat.message("Welche Tools kannst du benutzten?".to_string()).await; - let answer = agent_chat.message("Teste beide Server aus. Nutze bei fetch https://example.com/. Melde mir die Ergebnisse zurück.".to_string()).await; + agent_chat.set_permission("mcp-playground-echo-server::echo".to_string(), ToolPermission::Allowed).unwrap(); + + let answer = agent_chat.message("Welche Tools kannst du benutzten?".to_string(), permission_callback).await; + let answer = agent_chat.message("Teste beide Server aus. Nutze bei fetch https://example.com/. Melde mir die Ergebnisse zurück.".to_string(), permission_callback).await; dbg!(answer); }