mcp_streamable_proxy/proxy_handler/
server_handler.rs

1use super::*;
2
3impl ServerHandler for ProxyHandler {
4    fn get_info(&self) -> ServerInfo {
5        self.discovery.info()
6    }
7
8    /// Connect per-session backends before returning capabilities.
9    async fn initialize(
10        &self,
11        request: InitializeRequestParams,
12        context: RequestContext<RoleServer>,
13    ) -> Result<InitializeResult, ErrorData> {
14        context.peer.set_peer_info(request.clone());
15
16        if let Some(session) = &self.per_session {
17            // Bind notify slot before connect so backend notifications during handshake work.
18            session.notify_slot.set(context.peer.clone());
19            self.ensure_connected().await?;
20            info!(
21                "[initialize] Per-session backend ready - MCP ID: {}",
22                self.mcp_id
23            );
24        }
25
26        let mut info = self.get_info();
27        info.protocol_version =
28            negotiate_protocol_version(&request.protocol_version, info.protocol_version);
29        Ok(info)
30    }
31
32    #[tracing::instrument(skip(self, request, context), fields(
33        mcp_id = %self.mcp_id,
34        request = ?request,
35    ))]
36    async fn list_tools(
37        &self,
38        request: Option<PaginatedRequestParams>,
39        context: RequestContext<RoleServer>,
40    ) -> Result<ListToolsResult, ErrorData> {
41        if context.ct.is_cancelled() {
42            return Err(ErrorData::internal_error(
43                "Request cancelled".to_string(),
44                None,
45            ));
46        }
47        let is_first_page = request
48            .as_ref()
49            .and_then(|params| params.cursor.as_ref())
50            .is_none();
51
52        match self.capabilities().tools {
53            Some(_) => {
54                let result = match self
55                    .forward_backend(
56                        &context,
57                        |peer| async move { peer.list_tools(request).await },
58                    )
59                    .await
60                {
61                    Ok(result) => result,
62                    Err(error) if context.ct.is_cancelled() => return Err(error),
63                    Err(error) => {
64                        if let Some(cached) = self.discovery.tools() {
65                            warn!(%error, "real tools/list failed; returning cached tools snapshot");
66                            return Ok(self.filter_tools(cached));
67                        }
68                        return Err(error);
69                    }
70                };
71                if is_first_page && result.next_cursor.is_none() {
72                    self.update_cached_tools(result.clone());
73                }
74                let result = self.filter_tools(result);
75
76                info!(
77                    "[list_tools] Tool list results - MCP ID: {}, number of tools: {}{}",
78                    self.mcp_id,
79                    result.tools.len(),
80                    if self.tool_filter.is_enabled() {
81                        " (filtered)"
82                    } else {
83                        ""
84                    }
85                );
86
87                debug!(
88                    "Proxying list_tools response with {} tools",
89                    result.tools.len()
90                );
91                Ok(result)
92            }
93            None => {
94                warn!("Server doesn't support tools capability");
95                Ok(ListToolsResult::default())
96            }
97        }
98    }
99
100    #[tracing::instrument(skip(self, request, context), fields(
101        mcp_id = %self.mcp_id,
102        tool_name = %request.name,
103        tool_arguments = ?request.arguments,
104    ))]
105    async fn call_tool(
106        &self,
107        mut request: CallToolRequestParams,
108        context: RequestContext<RoleServer>,
109    ) -> Result<CallToolResponse, ErrorData> {
110        let request_id = REQUEST_COUNTER.fetch_add(1, Ordering::Relaxed);
111        let start = Instant::now();
112        let tool_name = request.name.clone();
113
114        info!(
115            "[call_tool:{}] Start - Tool: {}, MCP ID: {}",
116            request_id, tool_name, self.mcp_id
117        );
118
119        if !self.tool_filter.is_allowed(&tool_name) {
120            info!(
121                "[call_tool:{}] Tool is filtered - MCP ID: {}, Tool: {}",
122                request_id, self.mcp_id, tool_name
123            );
124            return Ok(CallToolResult::error(vec![ContentBlock::text(format!(
125                "Tool '{tool_name}' is not allowed by filter configuration"
126            ))])
127            .into());
128        }
129
130        if self.capabilities().tools.is_none() {
131            error!(
132                "[call_tool:{}] The server does not support tools capability - MCP ID: {}",
133                request_id, self.mcp_id
134            );
135            return Ok(CallToolResult::error(vec![ContentBlock::text(
136                "Server doesn't support tools capability",
137            )])
138            .into());
139        }
140
141        merge_context_meta_into_params(&mut request, &context.meta);
142        let _progress_guard = self.maybe_progress_guard(&context);
143
144        info!(
145            "[call_tool:{}] Send request to backend... - Tool: {}, Elapsed time: {}ms",
146            request_id,
147            tool_name,
148            start.elapsed().as_millis()
149        );
150
151        let call_result = self
152            .forward_backend_with_heartbeat(
153                &context,
154                request_id,
155                tool_name.as_ref(),
156                |peer| async move { peer.call_tool_once(request).await },
157            )
158            .await;
159
160        let result =
161            self.finish_call_tool(request_id, tool_name.as_ref(), start.elapsed(), call_result);
162
163        info!(
164            "[call_tool:{}] Completed - Tool: {}, total time taken: {}ms",
165            request_id,
166            tool_name,
167            start.elapsed().as_millis()
168        );
169        Ok(result)
170    }
171
172    async fn list_resources(
173        &self,
174        request: Option<PaginatedRequestParams>,
175        context: RequestContext<RoleServer>,
176    ) -> Result<rmcp::model::ListResourcesResult, ErrorData> {
177        match self.capabilities().resources {
178            Some(_) => {
179                let result = self
180                    .forward_backend(&context, |peer| async move {
181                        peer.list_resources(request).await
182                    })
183                    .await?;
184                info!(
185                    "[list_resources] Resource list results - MCP ID: {}, resource quantity: {}",
186                    self.mcp_id,
187                    result.resources.len()
188                );
189                debug!("Proxying list_resources response");
190                Ok(result)
191            }
192            None => {
193                warn!("Server doesn't support resources capability");
194                Ok(rmcp::model::ListResourcesResult::default())
195            }
196        }
197    }
198
199    async fn read_resource(
200        &self,
201        mut request: rmcp::model::ReadResourceRequestParams,
202        context: RequestContext<RoleServer>,
203    ) -> Result<ReadResourceResponse, ErrorData> {
204        match self.capabilities().resources {
205            Some(_) => {
206                merge_context_meta_into_params(&mut request, &context.meta);
207                let _progress_guard = self.maybe_progress_guard(&context);
208                let uri = request.uri.clone();
209                let result = self
210                    .forward_backend(&context, |peer| async move {
211                        peer.read_resource_once(request).await
212                    })
213                    .await?;
214                info!(
215                    "[read_resource] Resource read result - MCP ID: {}, URI: {}",
216                    self.mcp_id, uri
217                );
218                debug!("Proxying read_resource response for {}", uri);
219                Ok(result)
220            }
221            None => {
222                error!("Server doesn't support resources capability");
223                Ok(rmcp::model::ReadResourceResult::new(vec![]).into())
224            }
225        }
226    }
227
228    async fn list_resource_templates(
229        &self,
230        request: Option<PaginatedRequestParams>,
231        context: RequestContext<RoleServer>,
232    ) -> Result<rmcp::model::ListResourceTemplatesResult, ErrorData> {
233        match self.capabilities().resources {
234            Some(_) => {
235                let result = self
236                    .forward_backend(&context, |peer| async move {
237                        peer.list_resource_templates(request).await
238                    })
239                    .await?;
240                debug!("Proxying list_resource_templates response");
241                Ok(result)
242            }
243            None => {
244                warn!("Server doesn't support resources capability");
245                Ok(rmcp::model::ListResourceTemplatesResult::default())
246            }
247        }
248    }
249
250    async fn list_prompts(
251        &self,
252        request: Option<PaginatedRequestParams>,
253        context: RequestContext<RoleServer>,
254    ) -> Result<rmcp::model::ListPromptsResult, ErrorData> {
255        match self.capabilities().prompts {
256            Some(_) => {
257                let result =
258                    self.forward_backend(&context, |peer| async move {
259                        peer.list_prompts(request).await
260                    })
261                    .await?;
262                debug!("Proxying list_prompts response");
263                Ok(result)
264            }
265            None => {
266                warn!("Server doesn't support prompts capability");
267                Ok(rmcp::model::ListPromptsResult::default())
268            }
269        }
270    }
271
272    async fn get_prompt(
273        &self,
274        mut request: rmcp::model::GetPromptRequestParams,
275        context: RequestContext<RoleServer>,
276    ) -> Result<GetPromptResponse, ErrorData> {
277        match self.capabilities().prompts {
278            Some(_) => {
279                merge_context_meta_into_params(&mut request, &context.meta);
280                let _progress_guard = self.maybe_progress_guard(&context);
281                let result = self
282                    .forward_backend(&context, |peer| async move {
283                        peer.get_prompt_once(request).await
284                    })
285                    .await?;
286                debug!("Proxying get_prompt response");
287                Ok(result)
288            }
289            None => {
290                warn!("Server doesn't support prompts capability");
291                let messages = Vec::new();
292                Ok(rmcp::model::GetPromptResult::new(messages).into())
293            }
294        }
295    }
296
297    async fn complete(
298        &self,
299        request: rmcp::model::CompleteRequestParams,
300        context: RequestContext<RoleServer>,
301    ) -> Result<rmcp::model::CompleteResult, ErrorData> {
302        let result = self
303            .forward_backend(&context, |peer| async move { peer.complete(request).await })
304            .await?;
305        debug!("Proxying complete response");
306        Ok(result)
307    }
308
309    async fn subscribe(
310        &self,
311        request: SubscribeRequestParams,
312        context: RequestContext<RoleServer>,
313    ) -> Result<(), ErrorData> {
314        match self
315            .capabilities()
316            .resources
317            .as_ref()
318            .and_then(|r| r.subscribe)
319        {
320            Some(true) => {
321                self.forward_backend(
322                    &context,
323                    |peer| async move { peer.subscribe(request).await },
324                )
325                .await
326            }
327            _ => Err(ErrorData::method_not_found::<SubscribeRequestMethod>()),
328        }
329    }
330
331    async fn unsubscribe(
332        &self,
333        request: UnsubscribeRequestParams,
334        context: RequestContext<RoleServer>,
335    ) -> Result<(), ErrorData> {
336        match self
337            .capabilities()
338            .resources
339            .as_ref()
340            .and_then(|r| r.subscribe)
341        {
342            Some(true) => {
343                self.forward_backend(
344                    &context,
345                    |peer| async move { peer.unsubscribe(request).await },
346                )
347                .await
348            }
349            _ => Err(ErrorData::method_not_found::<UnsubscribeRequestMethod>()),
350        }
351    }
352
353    #[allow(deprecated)]
354    async fn set_level(
355        &self,
356        request: rmcp::model::SetLevelRequestParams,
357        context: RequestContext<RoleServer>,
358    ) -> Result<(), ErrorData> {
359        if self.capabilities().logging.is_none() {
360            return Err(ErrorData::method_not_found::<SetLevelRequestMethod>());
361        }
362        self.forward_backend(
363            &context,
364            |peer| async move { peer.set_level(request).await },
365        )
366        .await
367    }
368
369    async fn enqueue_task(
370        &self,
371        mut request: CallToolRequestParams,
372        context: RequestContext<RoleServer>,
373    ) -> Result<CreateTaskResult, ErrorData> {
374        let tasks = self.capabilities().tasks.clone();
375        if tasks.as_ref().is_none_or(|t| !t.supports_tools_call()) {
376            return Err(ErrorData::internal_error(
377                "Backend does not support task-based tools/call".to_string(),
378                None,
379            ));
380        }
381        merge_context_meta_into_params(&mut request, &context.meta);
382        let _progress_guard = self.maybe_progress_guard(&context);
383        let result = self
384            .forward_backend(&context, |peer| async move {
385                let result = peer
386                    .send_request(ClientRequest::CallToolRequest(CallToolRequest::new(
387                        request,
388                    )))
389                    .await?;
390                match result {
391                    ServerResult::CreateTaskResult(r) => Ok(r),
392                    _ => Err(ServiceError::UnexpectedResponse),
393                }
394            })
395            .await?;
396
397        // Shared isolation: route task notifications via registry.
398        // Per-session: bridge delivers 1:1 via notify_slot; no registry entry needed.
399        if self.isolation == BackendIsolation::Shared {
400            self.upstream_peers
401                .register_task(result.task.task_id.clone(), context.peer.clone());
402        }
403        Ok(result)
404    }
405
406    async fn list_tasks(
407        &self,
408        request: Option<PaginatedRequestParams>,
409        context: RequestContext<RoleServer>,
410    ) -> Result<ListTasksResult, ErrorData> {
411        let tasks = self.capabilities().tasks.clone();
412        if tasks.as_ref().is_none_or(|t| !t.supports_list()) {
413            return Err(ErrorData::method_not_found::<rmcp::model::ListTasksMethod>());
414        }
415        self.forward_backend(&context, |peer| async move {
416            let result = peer
417                .send_request(ClientRequest::ListTasksRequest(ListTasksRequest {
418                    method: Default::default(),
419                    params: request,
420                    extensions: Default::default(),
421                }))
422                .await?;
423            match result {
424                ServerResult::ListTasksResult(r) => Ok(r),
425                _ => Err(ServiceError::UnexpectedResponse),
426            }
427        })
428        .await
429    }
430
431    async fn get_task_info(
432        &self,
433        request: GetTaskParams,
434        context: RequestContext<RoleServer>,
435    ) -> Result<GetTaskResult, ErrorData> {
436        if self.capabilities().tasks.is_none() {
437            return Err(ErrorData::method_not_found::<rmcp::model::GetTaskMethod>());
438        }
439        self.forward_backend(&context, |peer| async move {
440            let result = peer
441                .send_request(ClientRequest::GetTaskRequest(GetTaskRequest::new(request)))
442                .await?;
443            match result {
444                ServerResult::GetTaskResult(r) => Ok(r),
445                _ => Err(ServiceError::UnexpectedResponse),
446            }
447        })
448        .await
449    }
450
451    async fn get_task_result(
452        &self,
453        request: GetTaskPayloadParams,
454        context: RequestContext<RoleServer>,
455    ) -> Result<rmcp::model::GetTaskPayloadResult, ErrorData> {
456        if self.capabilities().tasks.is_none() {
457            return Err(ErrorData::method_not_found::<
458                rmcp::model::GetTaskPayloadMethod,
459            >());
460        }
461        self.forward_backend(&context, |peer| async move {
462            let result = peer
463                .send_request(ClientRequest::GetTaskPayloadRequest(
464                    GetTaskPayloadRequest::new(request),
465                ))
466                .await?;
467            match result {
468                ServerResult::GetTaskPayloadResult(r) => Ok(r),
469                _ => Err(ServiceError::UnexpectedResponse),
470            }
471        })
472        .await
473    }
474
475    async fn cancel_task(
476        &self,
477        request: CancelTaskParams,
478        context: RequestContext<RoleServer>,
479    ) -> Result<rmcp::model::CancelTaskResult, ErrorData> {
480        let tasks = self.capabilities().tasks.clone();
481        if tasks.as_ref().is_none_or(|t| !t.supports_cancel()) {
482            return Err(ErrorData::method_not_found::<rmcp::model::CancelTaskMethod>());
483        }
484        let task_id = request.task_id.clone();
485        let result = self
486            .forward_backend(&context, |peer| async move {
487                let result = peer
488                    .send_request(ClientRequest::CancelTaskRequest(CancelTaskRequest::new(
489                        request,
490                    )))
491                    .await?;
492                match result {
493                    ServerResult::CancelTaskResult(r) => Ok(r),
494                    _ => Err(ServiceError::UnexpectedResponse),
495                }
496            })
497            .await?;
498        self.upstream_peers.unregister_task(&task_id);
499        Ok(result)
500    }
501
502    async fn on_initialized(&self, context: NotificationContext<RoleServer>) {
503        if let Some(session) = &self.per_session {
504            // Refresh slot (initialize already connected); keep idempotent ensure.
505            session.notify_slot.set(context.peer.clone());
506            if let Err(err) = self.ensure_connected().await {
507                error!(
508                    "[on_initialized] Per-session backend ensure failed - MCP ID: {}, err: {:?}",
509                    self.mcp_id, err
510                );
511            }
512            return;
513        }
514
515        let id = self.upstream_peers.register(context.peer.clone());
516        info!(
517            "[on_initialized] Registered upstream peer {} - MCP ID: {}, peers: {}",
518            id,
519            self.mcp_id,
520            self.upstream_peers.len()
521        );
522    }
523
524    async fn on_progress(
525        &self,
526        notification: rmcp::model::ProgressNotificationParam,
527        _context: NotificationContext<RoleServer>,
528    ) {
529        let peer = {
530            let inner_guard = self.peer.load();
531            let inner = match inner_guard.as_ref() {
532                Some(inner) => inner,
533                None => {
534                    error!(
535                        "Backend connection is not available, cannot forward progress notification"
536                    );
537                    return;
538                }
539            };
540            if inner.peer.is_transport_closed() {
541                error!("Backend transport is closed, cannot forward progress notification");
542                return;
543            }
544            inner.peer.clone()
545        };
546
547        match peer.notify_progress(notification).await {
548            Ok(_) => {
549                debug!("Proxying progress notification");
550            }
551            Err(err) => {
552                error!("Error notifying progress: {:?}", err);
553            }
554        }
555    }
556
557    async fn on_cancelled(
558        &self,
559        notification: rmcp::model::CancelledNotificationParam,
560        _context: NotificationContext<RoleServer>,
561    ) {
562        let peer = {
563            let inner_guard = self.peer.load();
564            let inner = match inner_guard.as_ref() {
565                Some(inner) => inner,
566                None => {
567                    error!(
568                        "Backend connection is not available, cannot forward cancelled notification"
569                    );
570                    return;
571                }
572            };
573            if inner.peer.is_transport_closed() {
574                error!("Backend transport is closed, cannot forward cancelled notification");
575                return;
576            }
577            inner.peer.clone()
578        };
579
580        match peer.notify_cancelled(notification).await {
581            Ok(_) => {
582                debug!("Proxying cancelled notification");
583            }
584            Err(err) => {
585                error!("Error notifying cancelled: {:?}", err);
586            }
587        }
588    }
589}
mcp_streamable_proxy/proxy_handler/server_handler.rs

mcp_streamable_proxy/proxy_handler/
server_handler.rs
