vasceannie/openagent
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: MCP tool integration in missions

Browse Source 
- Pass McpRegistry to mission runner and agent context
- Route tool calls to MCP servers when built-in tool not found
- Include MCP tool descriptions and schemas in system prompt
- Add has_tool() method to ToolRegistry for routing
This commit is contained in:
Thomas Marchand
2025-12-20 12:11:42 +00:00
parent 334711b2fe
commit c2c505a2f8
7 changed files with 122 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
src/agents/context.rs
View File

@@ -8,6 +8,7 @@ use uuid::Uuid;
use crate::budget::{ModelPricing, SharedBenchmarkRegistry, SharedModelResolver};
use crate::config::Config;
use crate::llm::LlmClient;
use crate::mcp::McpRegistry;
use crate::memory::MemorySystem;
use crate::tools::ToolRegistry;
use crate::tools::mission::MissionControl;
@@ -77,6 +78,9 @@ pub struct AgentContext {
/// Mission ID for tagging events (used in parallel mission execution).
pub mission_id: Option<Uuid>,
/// MCP registry for dynamic tool discovery and execution.
pub mcp: Option<Arc<McpRegistry>>,
}
impl AgentContext {
@@ -107,6 +111,7 @@ impl AgentContext {
tree_snapshot: None,
progress_snapshot: None,
mission_id: None,
mcp: None,
}
}
@@ -138,6 +143,7 @@ impl AgentContext {
tree_snapshot: None,
progress_snapshot: None,
mission_id: None,
mcp: None,
}
}
@@ -165,6 +171,7 @@ impl AgentContext {
tree_snapshot: self.tree_snapshot.clone(),
progress_snapshot: self.progress_snapshot.clone(),
mission_id: self.mission_id,
mcp: self.mcp.clone(),
}
}

58
src/agents/leaf/executor.rs
View File

@@ -192,17 +192,24 @@ impl TaskExecutor {
&self,
working_dir: &str,
tools: &ToolRegistry,
mcp_tool_descriptions: &str,
reusable_tools: &str,
session_metadata: &str,
memory_context: &str,
) -> String {
let tool_descriptions = tools
let mut tool_descriptions = tools
.list_tools()
.iter()
.map(|t| format!("- **{}**: {}", t.name, t.description))
.collect::<Vec<_>>()
.join("\n");
// Append MCP tool descriptions if any
if !mcp_tool_descriptions.is_empty() {
tool_descriptions.push_str("\n\n### MCP Tools (External Integrations)\n");
tool_descriptions.push_str(mcp_tool_descriptions);
}
// Check if we're in mission mode (isolated working directory)
let is_mission_mode = working_dir.contains("mission-");
@@ -476,17 +483,33 @@ Use `search_memory` when you encounter a problem you might have solved before or
}
/// Execute a single tool call.
///
/// Routes to built-in tools first, then falls back to MCP tools.
async fn execute_tool_call(
&self,
tool_call: &ToolCall,
ctx: &AgentContext,
) -> anyhow::Result<String> {
let tool_name = &tool_call.function.name;
let args: serde_json::Value = serde_json::from_str(&tool_call.function.arguments)
.unwrap_or(serde_json::Value::Null);
ctx.tools
.execute(&tool_call.function.name, args, &ctx.working_dir)
.await
// Try built-in tools first
if ctx.tools.has_tool(tool_name) {
return ctx.tools
.execute(tool_name, args, &ctx.working_dir)
.await;
}
// Try MCP tools
if let Some(mcp) = &ctx.mcp {
if let Some(mcp_tool) = mcp.find_tool(tool_name).await {
tracing::debug!("Routing tool call '{}' to MCP server", tool_name);
return mcp.call_tool(mcp_tool.mcp_id, tool_name, args).await;
}
}
anyhow::bail!("Unknown tool: {}", tool_name)
}
/// Run the agent loop for a task.
@@ -536,10 +559,32 @@ Use `search_memory` when you encounter a problem you might have solved before or
// Get tool result truncation limit from config
let max_tool_result_chars = ctx.config.context.max_tool_result_chars;
// Get MCP tool descriptions and schemas
let (mcp_tool_descriptions, mcp_tool_schemas) = if let Some(mcp) = &ctx.mcp {
let mcp_tools = mcp.list_tools().await;
let enabled_tools: Vec<_> = mcp_tools.iter().filter(|t| t.enabled).collect();
if !enabled_tools.is_empty() {
tracing::info!("Discovered {} MCP tools", enabled_tools.len());
}
let descriptions = enabled_tools
.iter()
.map(|t| format!("- **{}**: {}", t.name, t.description))
.collect::<Vec<_>>()
.join("\n");
let schemas = mcp.get_tool_schemas().await;
(descriptions, schemas)
} else {
(String::new(), Vec::new())
};
// Build initial messages with all context
let system_prompt = self.build_system_prompt(
&ctx.working_dir_str(),
&ctx.tools,
&mcp_tool_descriptions,
&reusable_tools,
&session_metadata,
&memory_context,
@@ -549,8 +594,9 @@ Use `search_memory` when you encounter a problem you might have solved before or
ChatMessage::new(Role::User, task.description().to_string()),
];
// Get tool schemas
let tool_schemas = ctx.tools.get_tool_schemas();
// Get tool schemas (built-in + MCP)
let mut tool_schemas = ctx.tools.get_tool_schemas();
tool_schemas.extend(mcp_tool_schemas);
// Agent loop
for iteration in 0..ctx.max_iterations {

11
src/api/control.rs
View File

@@ -27,6 +27,7 @@ use crate::agents::{AgentContext, AgentRef};
use crate::budget::{Budget, ModelPricing};
use crate::config::Config;
use crate::llm::OpenRouterClient;
use crate::mcp::McpRegistry;
use crate::memory::{ContextBuilder, MemorySystem, MissionMessage};
use crate::task::VerificationCriteria;
use crate::tools::ToolRegistry;
@@ -949,6 +950,7 @@ pub fn spawn_control_session(
memory: Option<MemorySystem>,
benchmarks: crate::budget::SharedBenchmarkRegistry,
resolver: crate::budget::SharedModelResolver,
mcp: Arc<McpRegistry>,
) -> ControlState {
let (cmd_tx, cmd_rx) = mpsc::channel::<ControlCommand>(256);
let (events_tx, _events_rx) = broadcast::channel::<AgentEvent>(1024);
@@ -987,6 +989,7 @@ pub fn spawn_control_session(
memory.clone(),
benchmarks,
resolver,
mcp,
cmd_rx,
mission_cmd_rx,
mission_cmd_tx,
@@ -1067,6 +1070,7 @@ async fn control_actor_loop(
memory: Option<MemorySystem>,
benchmarks: crate::budget::SharedBenchmarkRegistry,
resolver: crate::budget::SharedModelResolver,
mcp: Arc<McpRegistry>,
mut cmd_rx: mpsc::Receiver<ControlCommand>,
mut mission_cmd_rx: mpsc::Receiver<crate::tools::mission::MissionControlCommand>,
mission_cmd_tx: mpsc::Sender<crate::tools::mission::MissionControlCommand>,
@@ -1257,6 +1261,7 @@ async fn control_actor_loop(
let mem = memory.clone();
let bench = Arc::clone(&benchmarks);
let res = Arc::clone(&resolver);
let mcp_ref = Arc::clone(&mcp);
let events = events_tx.clone();
let tools_hub = Arc::clone(&tool_hub);
let status_ref = Arc::clone(&status);
@@ -1277,6 +1282,7 @@ async fn control_actor_loop(
mem,
bench,
res,
mcp_ref,
pricing,
events,
tools_hub,
@@ -1410,6 +1416,7 @@ async fn control_actor_loop(
memory.clone(),
Arc::clone(&benchmarks),
Arc::clone(&resolver),
Arc::clone(&mcp),
Arc::clone(&pricing),
events_tx.clone(),
Arc::clone(&tool_hub),
@@ -1582,6 +1589,7 @@ async fn control_actor_loop(
let mem = memory.clone();
let bench = Arc::clone(&benchmarks);
let res = Arc::clone(&resolver);
let mcp_ref = Arc::clone(&mcp);
let events = events_tx.clone();
let tools_hub = Arc::clone(&tool_hub);
let status_ref = Arc::clone(&status);
@@ -1602,6 +1610,7 @@ async fn control_actor_loop(
mem,
bench,
res,
mcp_ref,
pricing,
events,
tools_hub,
@@ -1680,6 +1689,7 @@ async fn run_single_control_turn(
memory: Option<MemorySystem>,
benchmarks: crate::budget::SharedBenchmarkRegistry,
resolver: crate::budget::SharedModelResolver,
mcp: Arc<McpRegistry>,
pricing: Arc<ModelPricing>,
events_tx: broadcast::Sender<AgentEvent>,
tool_hub: Arc<FrontendToolHub>,
@@ -1746,6 +1756,7 @@ async fn run_single_control_turn(
ctx.resolver = Some(resolver);
ctx.tree_snapshot = Some(tree_snapshot);
ctx.progress_snapshot = Some(progress_snapshot);
ctx.mcp = Some(mcp);
let result = root_agent.execute(&mut task, &ctx).await;
result

5
src/api/mission_runner.rs
View File

@@ -23,6 +23,7 @@ use crate::agents::{AgentContext, AgentRef, AgentResult};
use crate::budget::{Budget, ModelPricing, SharedBenchmarkRegistry, SharedModelResolver};
use crate::config::Config;
use crate::llm::OpenRouterClient;
use crate::mcp::McpRegistry;
use crate::memory::{ContextBuilder, MemorySystem};
use crate::task::{VerificationCriteria, DeliverableSet, extract_deliverables};
use crate::tools::ToolRegistry;
@@ -224,6 +225,7 @@ impl MissionRunner {
memory: Option<MemorySystem>,
benchmarks: SharedBenchmarkRegistry,
resolver: SharedModelResolver,
mcp: Arc<McpRegistry>,
pricing: Arc<ModelPricing>,
events_tx: broadcast::Sender<AgentEvent>,
tool_hub: Arc<FrontendToolHub>,
@@ -275,6 +277,7 @@ impl MissionRunner {
memory,
benchmarks,
resolver,
mcp,
pricing,
events_tx,
tool_hub,
@@ -360,6 +363,7 @@ async fn run_mission_turn(
memory: Option<MemorySystem>,
benchmarks: SharedBenchmarkRegistry,
resolver: SharedModelResolver,
mcp: Arc<McpRegistry>,
pricing: Arc<ModelPricing>,
events_tx: broadcast::Sender<AgentEvent>,
tool_hub: Arc<FrontendToolHub>,
@@ -477,6 +481,7 @@ async fn run_mission_turn(
ctx.tree_snapshot = Some(tree_snapshot);
ctx.progress_snapshot = Some(progress_snapshot);
ctx.mission_id = Some(mission_id);
ctx.mcp = Some(mcp);
root_agent.execute(&mut task, &ctx).await
}

3
src/api/routes.rs
View File

@@ -84,7 +84,7 @@ pub async fn serve(config: Config) -> anyhow::Result<()> {
// Spawn the single global control session actor.
let control_state =
control::spawn_control_session(config.clone(), Arc::clone(&root_agent), memory.clone(), Arc::clone(&benchmarks), Arc::clone(&resolver));
control::spawn_control_session(config.clone(), Arc::clone(&root_agent), memory.clone(), Arc::clone(&benchmarks), Arc::clone(&resolver), Arc::clone(&mcp));
let state = Arc::new(AppState {
config: config.clone(),
@@ -365,6 +365,7 @@ async fn run_agent_task(
);
ctx.benchmarks = Some(Arc::clone(&state.benchmarks));
ctx.resolver = Some(Arc::clone(&state.resolver));
ctx.mcp = Some(Arc::clone(&state.mcp));
// Create a run in memory if available
let memory_run_id = if let Some(ref mem) = state.memory {

40
src/mcp/registry.rs
View File

@@ -749,4 +749,44 @@ impl McpRegistry {
pub async fn is_tool_enabled(&self, name: &str) -> bool {
!self.disabled_tools.read().await.contains(name)
}
/// Find a tool by name and return its MCP ID if found.
pub async fn find_tool(&self, name: &str) -> Option<McpTool> {
let states = self.states.read().await;
let disabled = self.disabled_tools.read().await;
for state in states.values() {
if state.config.enabled && state.status == McpStatus::Connected {
for descriptor in &state.config.tool_descriptors {
if descriptor.name == name && !disabled.contains(&descriptor.name) {
return Some(McpTool {
name: descriptor.name.clone(),
description: descriptor.description.clone(),
parameters_schema: descriptor.input_schema.clone(),
mcp_id: state.config.id,
enabled: true,
});
}
}
}
}
None
}
/// Get tool schemas in LLM-compatible format for all connected MCP tools.
pub async fn get_tool_schemas(&self) -> Vec<crate::llm::ToolDefinition> {
self.list_tools()
.await
.into_iter()
.filter(|t| t.enabled)
.map(|t| crate::llm::ToolDefinition {
tool_type: "function".to_string(),
function: crate::llm::FunctionDefinition {
name: t.name,
description: t.description,
parameters: t.parameters_schema,
},
})
.collect()
}
}

5
src/tools/mod.rs
View File

@@ -265,6 +265,11 @@ impl ToolRegistry {
.collect()
}
/// Check if a tool exists by name.
pub fn has_tool(&self, name: &str) -> bool {
self.tools.contains_key(name)
}
/// Get tool schemas in LLM-compatible format.
pub fn get_tool_schemas(&self) -> Vec<ToolDefinition> {
self.tools