Repository navigation
Expand file tree
/
Copy pathbuilder.rs
More file actions
155 lines (134 loc) · 4.91 KB
/
Copy pathbuilder.rs
File metadata and controls
155 lines (134 loc) · 4.91 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
//! MCP Server builder for easy server construction.
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU64};
use anyhow::Result;
use cortex_mcp_types::{CallToolResult, Implementation, LogLevel, ServerCapabilities, Tool};
use serde_json::Value;
use tokio::sync::RwLock;
use crate::handlers::{FnToolHandler, NoOpToolHandler, ToolHandler};
use crate::providers::{PromptProvider, ResourceProvider};
use crate::server::{McpServer, ServerState};
/// Builder for creating MCP servers.
pub struct McpServerBuilder {
name: String,
version: String,
capabilities: ServerCapabilities,
tools: Vec<Arc<dyn ToolHandler>>,
resource_provider: Option<Arc<dyn ResourceProvider>>,
prompt_provider: Option<Arc<dyn PromptProvider>>,
instructions: Option<String>,
}
impl McpServerBuilder {
/// Create a new server builder.
pub fn new(name: impl Into<String>, version: impl Into<String>) -> Self {
Self {
name: name.into(),
version: version.into(),
capabilities: ServerCapabilities::default(),
tools: Vec::new(),
resource_provider: None,
prompt_provider: None,
instructions: None,
}
}
/// Set server capabilities.
pub fn capabilities(mut self, capabilities: ServerCapabilities) -> Self {
self.capabilities = capabilities;
self
}
/// Enable tools capability.
pub fn with_tools_capability(mut self) -> Self {
self.capabilities = self.capabilities.with_tools();
self
}
/// Enable resources capability.
pub fn with_resources_capability(mut self) -> Self {
self.capabilities = self.capabilities.with_resources();
self
}
/// Enable prompts capability.
pub fn with_prompts_capability(mut self) -> Self {
self.capabilities = self.capabilities.with_prompts();
self
}
/// Enable logging capability.
pub fn with_logging_capability(mut self) -> Self {
self.capabilities = self.capabilities.with_logging();
self
}
/// Add a tool definition (creates a no-op handler).
pub fn tool(mut self, tool: Tool) -> Self {
let handler = Arc::new(NoOpToolHandler { tool });
self.tools.push(handler);
self.capabilities = self.capabilities.with_tools();
self
}
/// Add a tool handler.
pub fn tool_handler(mut self, handler: Arc<dyn ToolHandler>) -> Self {
self.tools.push(handler);
self.capabilities = self.capabilities.with_tools();
self
}
/// Add a tool with a synchronous handler function.
pub fn tool_fn<F>(mut self, tool: Tool, handler: F) -> Self
where
F: Fn(Value) -> Result<CallToolResult> + Send + Sync + 'static,
{
let handler = Arc::new(FnToolHandler::new(tool, handler));
self.tools.push(handler);
self.capabilities = self.capabilities.with_tools();
self
}
/// Set the resource provider.
pub fn resource_provider(mut self, provider: Arc<dyn ResourceProvider>) -> Self {
self.resource_provider = Some(provider);
self.capabilities = self.capabilities.with_resources();
self
}
/// Set the prompt provider.
pub fn prompt_provider(mut self, provider: Arc<dyn PromptProvider>) -> Self {
self.prompt_provider = Some(provider);
self.capabilities = self.capabilities.with_prompts();
self
}
/// Set instructions for clients.
pub fn instructions(mut self, instructions: impl Into<String>) -> Self {
self.instructions = Some(instructions.into());
self
}
/// Build the server.
pub fn build(self) -> Result<Arc<McpServer>> {
let info = Implementation::new(&self.name, &self.version);
let server = McpServer {
info,
capabilities: self.capabilities,
tools: RwLock::new(HashMap::new()),
resource_provider: RwLock::new(self.resource_provider),
prompt_provider: RwLock::new(self.prompt_provider),
log_level: RwLock::new(LogLevel::Info),
state: RwLock::new(ServerState::Uninitialized),
running: AtomicBool::new(false),
request_id: AtomicU64::new(1),
pending_requests: RwLock::new(HashMap::new()),
client_info: RwLock::new(None),
protocol_version: RwLock::new(None),
instructions: self.instructions,
};
let server = Arc::new(server);
// Register tools
let tools = self.tools;
let server_clone = server.clone();
tokio::spawn(async move {
for handler in tools {
server_clone.register_tool(handler).await;
}
});
Ok(server)
}
/// Build and run the server with stdio transport.
pub async fn build_and_run_stdio(self) -> Result<()> {
let server = self.build()?;
server.run_stdio().await
}
}