use super::McpHandler;
use crate::daemon::hook::AgentScope;
use parking_lot::RwLock;
use runtime::Hook;
use schemars::JsonSchema;
use serde::Deserialize;
use std::{collections::BTreeMap, sync::Arc};
use wcore::{ToolDispatch, ToolFuture, agent::AsTool};
#[derive(Deserialize, JsonSchema)]
pub struct Mcp {
pub name: String,
#[serde(default)]
pub args: Option<String>,
}
pub struct McpHook {
mcp: Arc<McpHandler>,
scopes: Arc<RwLock<BTreeMap<String, AgentScope>>>,
prompt: String,
}
impl McpHook {
pub fn new(
mcp: Arc<McpHandler>,
scopes: Arc<RwLock<BTreeMap<String, AgentScope>>>,
prompt: String,
) -> Self {
Self {
mcp,
scopes,
prompt,
}
}
}
impl Hook for McpHook {
fn schema(&self) -> Vec<wcore::model::Tool> {
vec![Mcp::as_tool()]
}
fn scoped_tools(&self, config: &wcore::AgentConfig) -> (Vec<String>, Option<String>) {
if config.mcps.is_empty() {
return (vec![], None);
}
let tools = self
.schema()
.iter()
.map(|t| t.function.name.clone())
.collect();
let names: Vec<&str> = config.mcps.iter().map(|s| s.as_str()).collect();
let line = format!("mcp servers: {}", names.join(", "));
(tools, Some(line))
}
fn system_prompt(&self) -> Option<String> {
Some(self.prompt.clone())
}
fn dispatch<'a>(&'a self, name: &'a str, call: ToolDispatch) -> Option<ToolFuture<'a>> {
if name != "mcp" {
return None;
}
Some(Box::pin(async move {
let allowed_mcps: Vec<String> = self
.scopes
.read()
.get(&call.agent)
.filter(|s| !s.mcps.is_empty())
.map(|s| s.mcps.clone())
.unwrap_or_default();
super::dispatch::dispatch_mcp(&self.mcp, &call.args, &allowed_mcps).await
}))
}
}