use super::*;
impl ServerHandler for ProxyHandler {
fn get_info(&self) -> ServerInfo {
self.discovery.info()
}
async fn initialize(
&self,
request: InitializeRequestParams,
context: RequestContext<RoleServer>,
) -> Result<InitializeResult, ErrorData> {
context.peer.set_peer_info(request.clone());
if let Some(session) = &self.per_session {
session.notify_slot.set(context.peer.clone());
self.ensure_connected().await?;
info!(
"[initialize] Per-session backend ready - MCP ID: {}",
self.mcp_id
);
}
let mut info = self.get_info();
info.protocol_version =
negotiate_protocol_version(&request.protocol_version, info.protocol_version);
Ok(info)
}
#[tracing::instrument(skip(self, request, context), fields(
mcp_id = %self.mcp_id,
request = ?request,
))]
async fn list_tools(
&self,
request: Option<PaginatedRequestParams>,
context: RequestContext<RoleServer>,
) -> Result<ListToolsResult, ErrorData> {
if context.ct.is_cancelled() {
return Err(ErrorData::internal_error(
"Request cancelled".to_string(),
None,
));
}
let is_first_page = request
.as_ref()
.and_then(|params| params.cursor.as_ref())
.is_none();
match self.capabilities().tools {
Some(_) => {
let result = match self
.forward_backend(
&context,
|peer| async move { peer.list_tools(request).await },
)
.await
{
Ok(result) => result,
Err(error) if context.ct.is_cancelled() => return Err(error),
Err(error) => {
if let Some(cached) = self.discovery.tools() {
warn!(%error, "real tools/list failed; returning cached tools snapshot");
return Ok(self.filter_tools(cached));
}
return Err(error);
}
};
if is_first_page && result.next_cursor.is_none() {
self.update_cached_tools(result.clone());
}
let result = self.filter_tools(result);
info!(
"[list_tools] Tool list results - MCP ID: {}, number of tools: {}{}",
self.mcp_id,
result.tools.len(),
if self.tool_filter.is_enabled() {
" (filtered)"
} else {
""
}
);
debug!(
"Proxying list_tools response with {} tools",
result.tools.len()
);
Ok(result)
}
None => {
warn!("Server doesn't support tools capability");
Ok(ListToolsResult::default())
}
}
}
#[tracing::instrument(skip(self, request, context), fields(
mcp_id = %self.mcp_id,
tool_name = %request.name,
tool_arguments = ?request.arguments,
))]
async fn call_tool(
&self,
mut request: CallToolRequestParams,
context: RequestContext<RoleServer>,
) -> Result<CallToolResponse, ErrorData> {
let request_id = REQUEST_COUNTER.fetch_add(1, Ordering::Relaxed);
let start = Instant::now();
let tool_name = request.name.clone();
info!(
"[call_tool:{}] Start - Tool: {}, MCP ID: {}",
request_id, tool_name, self.mcp_id
);
if !self.tool_filter.is_allowed(&tool_name) {
info!(
"[call_tool:{}] Tool is filtered - MCP ID: {}, Tool: {}",
request_id, self.mcp_id, tool_name
);
return Ok(CallToolResult::error(vec![ContentBlock::text(format!(
"Tool '{tool_name}' is not allowed by filter configuration"
))])
.into());
}
if self.capabilities().tools.is_none() {
error!(
"[call_tool:{}] The server does not support tools capability - MCP ID: {}",
request_id, self.mcp_id
);
return Ok(CallToolResult::error(vec![ContentBlock::text(
"Server doesn't support tools capability",
)])
.into());
}
merge_context_meta_into_params(&mut request, &context.meta);
let _progress_guard = self.maybe_progress_guard(&context);
info!(
"[call_tool:{}] Send request to backend... - Tool: {}, Elapsed time: {}ms",
request_id,
tool_name,
start.elapsed().as_millis()
);
let call_result = self
.forward_backend_with_heartbeat(
&context,
request_id,
tool_name.as_ref(),
|peer| async move { peer.call_tool_once(request).await },
)
.await;
let result =
self.finish_call_tool(request_id, tool_name.as_ref(), start.elapsed(), call_result);
info!(
"[call_tool:{}] Completed - Tool: {}, total time taken: {}ms",
request_id,
tool_name,
start.elapsed().as_millis()
);
Ok(result)
}
async fn list_resources(
&self,
request: Option<PaginatedRequestParams>,
context: RequestContext<RoleServer>,
) -> Result<rmcp::model::ListResourcesResult, ErrorData> {
match self.capabilities().resources {
Some(_) => {
let result = self
.forward_backend(&context, |peer| async move {
peer.list_resources(request).await
})
.await?;
info!(
"[list_resources] Resource list results - MCP ID: {}, resource quantity: {}",
self.mcp_id,
result.resources.len()
);
debug!("Proxying list_resources response");
Ok(result)
}
None => {
warn!("Server doesn't support resources capability");
Ok(rmcp::model::ListResourcesResult::default())
}
}
}
async fn read_resource(
&self,
mut request: rmcp::model::ReadResourceRequestParams,
context: RequestContext<RoleServer>,
) -> Result<ReadResourceResponse, ErrorData> {
match self.capabilities().resources {
Some(_) => {
merge_context_meta_into_params(&mut request, &context.meta);
let _progress_guard = self.maybe_progress_guard(&context);
let uri = request.uri.clone();
let result = self
.forward_backend(&context, |peer| async move {
peer.read_resource_once(request).await
})
.await?;
info!(
"[read_resource] Resource read result - MCP ID: {}, URI: {}",
self.mcp_id, uri
);
debug!("Proxying read_resource response for {}", uri);
Ok(result)
}
None => {
error!("Server doesn't support resources capability");
Ok(rmcp::model::ReadResourceResult::new(vec![]).into())
}
}
}
async fn list_resource_templates(
&self,
request: Option<PaginatedRequestParams>,
context: RequestContext<RoleServer>,
) -> Result<rmcp::model::ListResourceTemplatesResult, ErrorData> {
match self.capabilities().resources {
Some(_) => {
let result = self
.forward_backend(&context, |peer| async move {
peer.list_resource_templates(request).await
})
.await?;
debug!("Proxying list_resource_templates response");
Ok(result)
}
None => {
warn!("Server doesn't support resources capability");
Ok(rmcp::model::ListResourceTemplatesResult::default())
}
}
}
async fn list_prompts(
&self,
request: Option<PaginatedRequestParams>,
context: RequestContext<RoleServer>,
) -> Result<rmcp::model::ListPromptsResult, ErrorData> {
match self.capabilities().prompts {
Some(_) => {
let result =
self.forward_backend(&context, |peer| async move {
peer.list_prompts(request).await
})
.await?;
debug!("Proxying list_prompts response");
Ok(result)
}
None => {
warn!("Server doesn't support prompts capability");
Ok(rmcp::model::ListPromptsResult::default())
}
}
}
async fn get_prompt(
&self,
mut request: rmcp::model::GetPromptRequestParams,
context: RequestContext<RoleServer>,
) -> Result<GetPromptResponse, ErrorData> {
match self.capabilities().prompts {
Some(_) => {
merge_context_meta_into_params(&mut request, &context.meta);
let _progress_guard = self.maybe_progress_guard(&context);
let result = self
.forward_backend(&context, |peer| async move {
peer.get_prompt_once(request).await
})
.await?;
debug!("Proxying get_prompt response");
Ok(result)
}
None => {
warn!("Server doesn't support prompts capability");
let messages = Vec::new();
Ok(rmcp::model::GetPromptResult::new(messages).into())
}
}
}
async fn complete(
&self,
request: rmcp::model::CompleteRequestParams,
context: RequestContext<RoleServer>,
) -> Result<rmcp::model::CompleteResult, ErrorData> {
let result = self
.forward_backend(&context, |peer| async move { peer.complete(request).await })
.await?;
debug!("Proxying complete response");
Ok(result)
}
async fn subscribe(
&self,
request: SubscribeRequestParams,
context: RequestContext<RoleServer>,
) -> Result<(), ErrorData> {
match self
.capabilities()
.resources
.as_ref()
.and_then(|r| r.subscribe)
{
Some(true) => {
self.forward_backend(
&context,
|peer| async move { peer.subscribe(request).await },
)
.await
}
_ => Err(ErrorData::method_not_found::<SubscribeRequestMethod>()),
}
}
async fn unsubscribe(
&self,
request: UnsubscribeRequestParams,
context: RequestContext<RoleServer>,
) -> Result<(), ErrorData> {
match self
.capabilities()
.resources
.as_ref()
.and_then(|r| r.subscribe)
{
Some(true) => {
self.forward_backend(
&context,
|peer| async move { peer.unsubscribe(request).await },
)
.await
}
_ => Err(ErrorData::method_not_found::<UnsubscribeRequestMethod>()),
}
}
#[allow(deprecated)]
async fn set_level(
&self,
request: rmcp::model::SetLevelRequestParams,
context: RequestContext<RoleServer>,
) -> Result<(), ErrorData> {
if self.capabilities().logging.is_none() {
return Err(ErrorData::method_not_found::<SetLevelRequestMethod>());
}
self.forward_backend(
&context,
|peer| async move { peer.set_level(request).await },
)
.await
}
async fn enqueue_task(
&self,
mut request: CallToolRequestParams,
context: RequestContext<RoleServer>,
) -> Result<CreateTaskResult, ErrorData> {
let tasks = self.capabilities().tasks.clone();
if tasks.as_ref().is_none_or(|t| !t.supports_tools_call()) {
return Err(ErrorData::internal_error(
"Backend does not support task-based tools/call".to_string(),
None,
));
}
merge_context_meta_into_params(&mut request, &context.meta);
let _progress_guard = self.maybe_progress_guard(&context);
let result = self
.forward_backend(&context, |peer| async move {
let result = peer
.send_request(ClientRequest::CallToolRequest(CallToolRequest::new(
request,
)))
.await?;
match result {
ServerResult::CreateTaskResult(r) => Ok(r),
_ => Err(ServiceError::UnexpectedResponse),
}
})
.await?;
if self.isolation == BackendIsolation::Shared {
self.upstream_peers
.register_task(result.task.task_id.clone(), context.peer.clone());
}
Ok(result)
}
async fn list_tasks(
&self,
request: Option<PaginatedRequestParams>,
context: RequestContext<RoleServer>,
) -> Result<ListTasksResult, ErrorData> {
let tasks = self.capabilities().tasks.clone();
if tasks.as_ref().is_none_or(|t| !t.supports_list()) {
return Err(ErrorData::method_not_found::<rmcp::model::ListTasksMethod>());
}
self.forward_backend(&context, |peer| async move {
let result = peer
.send_request(ClientRequest::ListTasksRequest(ListTasksRequest {
method: Default::default(),
params: request,
extensions: Default::default(),
}))
.await?;
match result {
ServerResult::ListTasksResult(r) => Ok(r),
_ => Err(ServiceError::UnexpectedResponse),
}
})
.await
}
async fn get_task_info(
&self,
request: GetTaskParams,
context: RequestContext<RoleServer>,
) -> Result<GetTaskResult, ErrorData> {
if self.capabilities().tasks.is_none() {
return Err(ErrorData::method_not_found::<rmcp::model::GetTaskMethod>());
}
self.forward_backend(&context, |peer| async move {
let result = peer
.send_request(ClientRequest::GetTaskRequest(GetTaskRequest::new(request)))
.await?;
match result {
ServerResult::GetTaskResult(r) => Ok(r),
_ => Err(ServiceError::UnexpectedResponse),
}
})
.await
}
async fn get_task_result(
&self,
request: GetTaskPayloadParams,
context: RequestContext<RoleServer>,
) -> Result<rmcp::model::GetTaskPayloadResult, ErrorData> {
if self.capabilities().tasks.is_none() {
return Err(ErrorData::method_not_found::<
rmcp::model::GetTaskPayloadMethod,
>());
}
self.forward_backend(&context, |peer| async move {
let result = peer
.send_request(ClientRequest::GetTaskPayloadRequest(
GetTaskPayloadRequest::new(request),
))
.await?;
match result {
ServerResult::GetTaskPayloadResult(r) => Ok(r),
_ => Err(ServiceError::UnexpectedResponse),
}
})
.await
}
async fn cancel_task(
&self,
request: CancelTaskParams,
context: RequestContext<RoleServer>,
) -> Result<rmcp::model::CancelTaskResult, ErrorData> {
let tasks = self.capabilities().tasks.clone();
if tasks.as_ref().is_none_or(|t| !t.supports_cancel()) {
return Err(ErrorData::method_not_found::<rmcp::model::CancelTaskMethod>());
}
let task_id = request.task_id.clone();
let result = self
.forward_backend(&context, |peer| async move {
let result = peer
.send_request(ClientRequest::CancelTaskRequest(CancelTaskRequest::new(
request,
)))
.await?;
match result {
ServerResult::CancelTaskResult(r) => Ok(r),
_ => Err(ServiceError::UnexpectedResponse),
}
})
.await?;
self.upstream_peers.unregister_task(&task_id);
Ok(result)
}
async fn on_initialized(&self, context: NotificationContext<RoleServer>) {
if let Some(session) = &self.per_session {
session.notify_slot.set(context.peer.clone());
if let Err(err) = self.ensure_connected().await {
error!(
"[on_initialized] Per-session backend ensure failed - MCP ID: {}, err: {:?}",
self.mcp_id, err
);
}
return;
}
let id = self.upstream_peers.register(context.peer.clone());
info!(
"[on_initialized] Registered upstream peer {} - MCP ID: {}, peers: {}",
id,
self.mcp_id,
self.upstream_peers.len()
);
}
async fn on_progress(
&self,
notification: rmcp::model::ProgressNotificationParam,
_context: NotificationContext<RoleServer>,
) {
let peer = {
let inner_guard = self.peer.load();
let inner = match inner_guard.as_ref() {
Some(inner) => inner,
None => {
error!(
"Backend connection is not available, cannot forward progress notification"
);
return;
}
};
if inner.peer.is_transport_closed() {
error!("Backend transport is closed, cannot forward progress notification");
return;
}
inner.peer.clone()
};
match peer.notify_progress(notification).await {
Ok(_) => {
debug!("Proxying progress notification");
}
Err(err) => {
error!("Error notifying progress: {:?}", err);
}
}
}
async fn on_cancelled(
&self,
notification: rmcp::model::CancelledNotificationParam,
_context: NotificationContext<RoleServer>,
) {
let peer = {
let inner_guard = self.peer.load();
let inner = match inner_guard.as_ref() {
Some(inner) => inner,
None => {
error!(
"Backend connection is not available, cannot forward cancelled notification"
);
return;
}
};
if inner.peer.is_transport_closed() {
error!("Backend transport is closed, cannot forward cancelled notification");
return;
}
inner.peer.clone()
};
match peer.notify_cancelled(notification).await {
Ok(_) => {
debug!("Proxying cancelled notification");
}
Err(err) => {
error!("Error notifying cancelled: {:?}", err);
}
}
}
}