1use std::future::Future;
4use std::pin::Pin;
5use std::sync::Arc;
6use std::time::SystemTime;
7
8use rmcp::ErrorData;
9use rmcp::handler::server::ServerHandler;
10use rmcp::model::{
11 CallToolRequestParams, CallToolResponse, CallToolResult, CancelTaskParams, ContentBlock,
12 CreateTaskResult, DetailedTask, GetTaskParams, GetTaskResult, Implementation, JsonObject,
13 ListResourceTemplatesResult, ListResourcesResult, ListToolsResult, MetaObject,
14 PaginatedRequestParams, ReadResourceRequestParams, ReadResourceResponse, ReadResourceResult,
15 Resource, ResourceContents, ResourceTemplate, ServerCapabilities, ServerConfig, Task,
16 TaskPayload, TaskStatus, ToolAnnotations,
17};
18use rmcp::model::{
19 ClientCapabilities, ClientResult, CompleteRequestParams, CompleteResult, CompletionInfo,
20 ElicitRequest, ElicitRequestParams, ElicitationSchema, GetPromptRequestParams,
21 GetPromptResponse, GetPromptResult, InputRequest, InputRequests, InputRequiredResult,
22 InputResponses, ListPromptsResult, ProgressNotificationParam, ProgressToken, PromptMessage,
23 ProtocolVersion, ServerRequest, UpdateTaskParams,
24};
25use rmcp::model::{
26 ResourceUpdatedNotification, ResourceUpdatedNotificationParam, ResourcesCapability,
27 ServerNotification, SubscriptionFilter,
28};
29use rmcp::service::{RequestContext, RoleServer, SubscriptionContext};
30use serde_json::Value;
31use tokio_util::sync::CancellationToken;
32use tracing::Instrument;
33
34use crate::kit::Reply;
35use crate::kit::{CallError, Caller, Kit};
36use crate::module::Module;
37use crate::ops::{APP_MIME, Ask, Body, Completing, Content, Form, Hints, Looks, Progress, Role};
38use crate::seams::{Entitlements, Identity, Metering, Space};
39use crate::tasks::{StoredTask, TASK_TTL, TaskState, TaskStore};
40use crate::trace;
41
42fn ui(tool: &crate::ops::Tool) -> Option<JsonObject> {
46 let mut ui = JsonObject::new();
47 if let Some(uri) = tool.ui {
48 ui.insert("resourceUri".into(), uri.into());
49 }
50 if tool.app_only {
51 ui.insert("visibility".into(), serde_json::json!(["app"]));
52 }
53 (!ui.is_empty()).then(|| JsonObject::from_iter([("ui".to_owned(), Value::Object(ui))]))
54}
55
56fn annotated(tool: rmcp::model::Tool, hints: Hints) -> rmcp::model::Tool {
58 let title = tool.title.clone();
59 tool.annotate(ToolAnnotations::from_raw(
60 title,
61 Some(hints.read_only),
62 Some(hints.destructive),
63 Some(hints.idempotent),
64 Some(hints.open_world),
65 ))
66}
67
68macro_rules! shown {
71 ($item:expr, $looks:expr) => {{
72 let looks: &Looks = $looks;
73 let item = $item;
74 let item = match looks.title {
75 Some(title) => item.with_title(title),
76 None => item,
77 };
78 match icons(looks) {
79 Some(icons) => item.with_icons(icons),
80 None => item,
81 }
82 }};
83}
84
85macro_rules! shown_resource {
87 ($item:expr, $looks:expr) => {{
88 let looks: &Looks = $looks;
89 let item = shown!($item, looks);
90 let item = match looks.mime {
91 Some(mime) => item.with_mime_type(mime),
92 None => item,
93 };
94 match annotations(looks) {
95 Some(annotations) => item.with_annotations(annotations),
96 None => item,
97 }
98 }};
99}
100
101fn icons(looks: &Looks) -> Option<Vec<rmcp::model::Icon>> {
102 (!looks.icons.is_empty()).then(|| {
103 looks
104 .icons
105 .iter()
106 .map(|icon| {
107 let mut shown = rmcp::model::Icon::new(icon.src);
108 shown.mime_type = icon.mime.map(str::to_owned);
109 shown.sizes = (!icon.sizes.is_empty())
110 .then(|| icon.sizes.iter().map(|size| (*size).to_owned()).collect());
111 shown
112 })
113 .collect()
114 })
115}
116
117fn annotations(looks: &Looks) -> Option<rmcp::model::Annotations> {
118 if looks.audience.is_empty() && looks.priority.is_none() {
119 return None;
120 }
121 let mut annotations = rmcp::model::Annotations::default();
122 annotations.audience =
123 (!looks.audience.is_empty()).then(|| looks.audience.iter().copied().map(role).collect());
124 annotations.priority = looks.priority;
125 Some(annotations)
126}
127
128fn role(role: Role) -> rmcp::model::Role {
129 match role {
130 Role::User => rmcp::model::Role::User,
131 Role::Assistant => rmcp::model::Role::Assistant,
132 }
133}
134
135const POLL_MS: u64 = 500;
137
138pub(crate) struct Server<M, I, E, T, S> {
140 kit: Arc<Kit<M, I, E, T, S>>,
141 tools: Arc<[rmcp::model::Tool]>,
143 stopping: CancellationToken,
146}
147
148impl<M, I, E, T, S> Clone for Server<M, I, E, T, S> {
149 fn clone(&self) -> Self {
150 Self {
151 kit: self.kit.clone(),
152 tools: self.tools.clone(),
153 stopping: self.stopping.clone(),
154 }
155 }
156}
157
158impl<M: Module, I: Identity, E: Entitlements, T: Metering, S: TaskStore> Server<M, I, E, T, S> {
159 pub(crate) fn new(kit: Kit<M, I, E, T, S>) -> Self {
160 let tools = kit
161 .tools()
162 .map(|tool| {
163 let listed = rmcp::model::Tool::new_with_raw(
164 tool.name,
165 tool.description.clone().map(Into::into),
166 tool.input.clone(),
167 );
168 let listed = match &tool.output {
169 Some(output) => listed.with_raw_output_schema(output.clone()),
170 None => listed,
171 };
172 let listed = shown!(listed, &tool.looks);
173 let listed = match ui(tool) {
174 Some(meta) => listed.with_meta(MetaObject(meta)),
175 None => listed,
176 };
177 match tool.hints {
178 Some(hints) => annotated(listed, hints),
179 None => listed,
180 }
181 })
182 .collect();
183 Self {
184 kit: Arc::new(kit),
185 tools,
186 stopping: CancellationToken::new(),
187 }
188 }
189}
190
191impl<M: Module, I: Identity, E: Entitlements, T: Metering, S: TaskStore> ServerHandler
192 for Server<M, I, E, T, S>
193{
194 fn get_info(&self) -> ServerConfig {
195 let manifest = self.kit.manifest();
196 let mut capabilities = ServerCapabilities::builder().enable_tools().build();
197 if !self.kit.apps().is_empty() || self.kit.resources().next().is_some() {
198 let mut resources = ResourcesCapability::default();
199 resources.subscribe = self.kit.resources().next().is_some().then_some(true);
201 capabilities.resources = Some(resources);
202 }
203 if self.kit.prompts().next().is_some() {
204 capabilities.prompts = Some(Default::default());
205 }
206 if self.kit.has_completions() {
207 capabilities.completions = Some(JsonObject::new());
208 }
209 if self.kit.tools().any(|tool| tool.long) {
210 capabilities.extensions = ServerCapabilities::builder()
211 .enable_tasks()
212 .build()
213 .extensions;
214 }
215 let looks = self.kit.looks();
216 let mut implementation = shown!(
217 Implementation::new(manifest.name, manifest.contract_version),
218 looks
219 );
220 implementation.description = looks.about.map(str::to_owned);
221 implementation.website_url = looks.website.map(str::to_owned);
222 let info = ServerConfig::new(capabilities).with_server_info(implementation);
223 match self.kit.instructions() {
224 Some(text) => info.with_instructions(text),
225 None => info,
226 }
227 }
228
229 fn get_tool(&self, name: &str) -> Option<rmcp::model::Tool> {
232 self.tools.iter().find(|tool| tool.name == name).cloned()
233 }
234
235 async fn list_tools(
236 &self,
237 _request: Option<PaginatedRequestParams>,
238 _context: RequestContext<RoleServer>,
239 ) -> Result<ListToolsResult, ErrorData> {
240 Ok(ListToolsResult::with_all_items(self.tools.to_vec()))
241 }
242
243 fn accepted_subscription_filter(
244 &self,
245 requested: &SubscriptionFilter,
246 ) -> Option<SubscriptionFilter> {
247 self.kit.resources().next()?;
250 let mut accepted = SubscriptionFilter::default();
251 accepted.resource_subscriptions = requested.resource_subscriptions.clone();
252 Some(accepted)
253 }
254
255 async fn listen(&self, subscription: SubscriptionContext) -> Result<(), ErrorData> {
256 let request = subscription.request_context();
257 let space = self
258 .kit
259 .identified(&caller(request, None))
260 .await
261 .or_else(refusal)?;
262 let mut changes = self.kit.changes();
263 loop {
264 let changed = tokio::select! {
265 () = subscription.cancelled() => return Ok(()),
266 () = self.stopping.cancelled() => return Ok(()),
269 changed = changes.recv() => changed,
270 };
271 let changed = match changed {
272 Ok(changed) => changed,
273 Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
276 Err(tokio::sync::broadcast::error::RecvError::Closed) => return Ok(()),
277 };
278 if changed.space != space {
280 continue;
281 }
282 let wanted = subscription
283 .sink()
284 .accepted()
285 .resource_subscriptions
286 .as_ref()
287 .is_some_and(|uris| uris.contains(&changed.uri));
288 if wanted {
289 let note = ResourceUpdatedNotification::new(ResourceUpdatedNotificationParam::new(
290 changed.uri,
291 ));
292 if subscription
293 .sink()
294 .send(ServerNotification::ResourceUpdatedNotification(note))
295 .await
296 .is_err()
297 {
298 return Ok(());
299 }
300 }
301 }
302 }
303
304 async fn list_resources(
305 &self,
306 _request: Option<PaginatedRequestParams>,
307 _context: RequestContext<RoleServer>,
308 ) -> Result<ListResourcesResult, ErrorData> {
309 let pages = self
310 .kit
311 .apps()
312 .iter()
313 .map(|app| Resource::new(app.uri, app.uri).with_mime_type(APP_MIME));
314 let resources = self
315 .kit
316 .resources()
317 .filter(|resource| !resource.is_template())
318 .map(|resource| {
319 shown_resource!(Resource::new(resource.uri, resource.name), &resource.looks)
320 });
321 Ok(ListResourcesResult::with_all_items(
322 pages.chain(resources).collect(),
323 ))
324 }
325
326 async fn list_resource_templates(
327 &self,
328 _request: Option<PaginatedRequestParams>,
329 _context: RequestContext<RoleServer>,
330 ) -> Result<ListResourceTemplatesResult, ErrorData> {
331 let templates = self
332 .kit
333 .resources()
334 .filter(|resource| resource.is_template())
335 .map(|resource| {
336 let template = shown_resource!(
337 ResourceTemplate::new(resource.uri, resource.name),
338 &resource.looks
339 );
340 match &resource.description {
341 Some(text) => template.with_description(text),
342 None => template,
343 }
344 });
345 Ok(ListResourceTemplatesResult::with_all_items(
346 templates.collect(),
347 ))
348 }
349
350 async fn read_resource(
351 &self,
352 request: ReadResourceRequestParams,
353 context: RequestContext<RoleServer>,
354 ) -> Result<ReadResourceResponse, ErrorData> {
355 let span = trace::request("resources/read", &request.uri, &context);
356 let result = answered(&context, request.input_responses, |answers| {
357 self.kit.read_as(&request.uri, caller(&context, answers))
358 })
359 .instrument(span.clone())
360 .await;
361 let result = match result {
362 Err(CallError::Unknown(uri)) => Err(ErrorData::resource_not_found(uri, None)),
363 result => respond(result, |body| {
364 ReadResourceResult::new(vec![contents(&request.uri, body)])
365 }),
366 };
367 trace::outcome(&span, &result);
368 result
369 }
370
371 async fn list_prompts(
372 &self,
373 _request: Option<PaginatedRequestParams>,
374 _context: RequestContext<RoleServer>,
375 ) -> Result<ListPromptsResult, ErrorData> {
376 let prompts = self.kit.prompts().map(|prompt| {
377 let arguments = prompt
378 .arguments
379 .iter()
380 .map(|argument| {
381 let listed = rmcp::model::PromptArgument::new(&argument.name)
382 .with_required(argument.required);
383 match &argument.description {
384 Some(text) => listed.with_description(text),
385 None => listed,
386 }
387 })
388 .collect();
389 shown!(
390 rmcp::model::Prompt::new(
391 prompt.name,
392 prompt.description.as_deref(),
393 Some(arguments)
394 ),
395 &prompt.looks
396 )
397 });
398 Ok(ListPromptsResult::with_all_items(prompts.collect()))
399 }
400
401 async fn get_prompt(
402 &self,
403 request: GetPromptRequestParams,
404 context: RequestContext<RoleServer>,
405 ) -> Result<GetPromptResponse, ErrorData> {
406 let span = trace::request("prompts/get", &request.name, &context);
407 let arguments = request.arguments.unwrap_or_default();
408 let result = answered(&context, request.input_responses, |answers| {
409 self.kit
410 .prompt_as(&request.name, caller(&context, answers), arguments.clone())
411 })
412 .instrument(span.clone())
413 .await;
414 let result = respond(result, |messages| {
415 let messages = messages
416 .into_iter()
417 .map(|message| PromptMessage::new(role(message.role), block(message.content)))
418 .collect();
419 GetPromptResult::new(messages)
420 });
421 trace::outcome(&span, &result);
422 result
423 }
424
425 async fn complete(
426 &self,
427 request: CompleteRequestParams,
428 context: RequestContext<RoleServer>,
429 ) -> Result<CompleteResult, ErrorData> {
430 let reference = &request.r#ref;
431 let Some(target) = reference
433 .as_prompt_name()
434 .or_else(|| reference.as_resource_uri())
435 else {
436 return Ok(completion(Vec::new()));
437 };
438 let span = trace::request("completion/complete", target, &context);
439 let result = self
440 .kit
441 .complete_as(
442 target,
443 &request.argument.name,
444 Completing {
445 typed: request.argument.value.clone(),
446 given: request
447 .context
448 .and_then(|context| context.arguments)
449 .unwrap_or_default(),
450 },
451 caller(&context, None),
452 )
453 .instrument(span.clone())
454 .await
455 .map(completion)
456 .or_else(refusal);
457 trace::outcome(&span, &result);
458 result
459 }
460
461 async fn call_tool(
462 &self,
463 request: CallToolRequestParams,
464 context: RequestContext<RoleServer>,
465 ) -> Result<CallToolResponse, ErrorData> {
466 let span = trace::request("tools/call", &request.name, &context);
467 let result = self.call(request, context).instrument(span.clone()).await;
468 trace::outcome(&span, &result);
469 result
470 }
471
472 async fn get_task(
473 &self,
474 request: GetTaskParams,
475 context: RequestContext<RoleServer>,
476 ) -> Result<GetTaskResult, ErrorData> {
477 let span = trace::request("tasks/get", &request.task_id, &context);
478 let result = self
479 .kit
480 .task_as(&caller(&context, None), &request.task_id)
481 .instrument(span.clone())
482 .await
483 .map(|task| GetTaskResult::new(detailed(task)))
484 .or_else(refusal);
485 trace::outcome(&span, &result);
486 result
487 }
488
489 async fn update_task(
490 &self,
491 request: UpdateTaskParams,
492 context: RequestContext<RoleServer>,
493 ) -> Result<(), ErrorData> {
494 let span = trace::request("tasks/update", &request.task_id, &context);
495 let result = self
496 .kit
497 .update_as(
498 &caller(&context, None),
499 &request.task_id,
500 request.input_responses,
501 )
502 .instrument(span.clone())
503 .await
504 .map(drop)
505 .or_else(refusal);
506 trace::outcome(&span, &result);
507 result
508 }
509
510 async fn cancel_task(
511 &self,
512 request: CancelTaskParams,
513 context: RequestContext<RoleServer>,
514 ) -> Result<(), ErrorData> {
515 let span = trace::request("tasks/cancel", &request.task_id, &context);
516 let result = self
517 .kit
518 .cancel_as(&caller(&context, None), &request.task_id)
519 .instrument(span.clone())
520 .await
521 .map(drop)
522 .or_else(refusal);
523 trace::outcome(&span, &result);
524 result
525 }
526}
527
528impl<M: Module, I: Identity, E: Entitlements, T: Metering, S: TaskStore> Server<M, I, E, T, S> {
529 async fn call(
530 &self,
531 request: CallToolRequestParams,
532 context: RequestContext<RoleServer>,
533 ) -> Result<CallToolResponse, ErrorData> {
534 let arguments = request.arguments.unwrap_or_default();
535 let long = self.kit.tool(&request.name).is_some_and(|tool| tool.long);
536 if long && supports_tasks(&context) {
537 let task = self
538 .kit
539 .start_as(&request.name, &caller(&context, None), arguments)
540 .await
541 .or_else(refusal)?;
542 return Ok(CallToolResponse::Task(CreateTaskResult::new(listed(&task))));
543 }
544 let result = answered(&context, request.input_responses, |answers| {
545 self.kit
546 .call_as(&request.name, caller(&context, answers), arguments.clone())
547 })
548 .await;
549 match result {
550 Err(error @ (CallError::NotEntitled(_) | CallError::Failed(_))) => {
551 Ok(tool_error(error).into())
552 }
553 result => respond(result, replied),
554 }
555 }
556}
557
558fn respond<V, D, R>(result: Result<V, CallError>, done: impl FnOnce(V) -> D) -> Result<R, ErrorData>
561where
562 D: Into<R>,
563 R: From<InputRequiredResult>,
564{
565 match result {
566 Ok(value) => Ok(done(value).into()),
567 Err(CallError::InputRequired(asked)) => Ok(questions(asked)?.into()),
568 Err(error) => refusal(error),
569 }
570}
571
572fn questions(asked: Vec<Ask>) -> Result<InputRequiredResult, ErrorData> {
574 let requests = asked
575 .into_iter()
576 .map(|ask| {
577 let (key, form) = form(ask)?;
578 Ok((key, InputRequest::Elicitation(ElicitRequest::new(form))))
579 })
580 .collect::<Result<InputRequests, ErrorData>>()?;
581 Ok(InputRequiredResult::from_input_requests(requests))
582}
583
584fn form(ask: Ask) -> Result<(String, ElicitRequestParams), ErrorData> {
587 let form = match ask.form {
588 Form::Fields(schema) => {
589 let schema = ElicitationSchema::from_json_schema(schema).map_err(|error| {
590 ErrorData::internal_error(format!("cannot ask {}: {error}", ask.key), None)
593 })?;
594 ElicitRequestParams::FormElicitationParams {
595 meta: None,
596 message: ask.message,
597 requested_schema: schema,
598 }
599 }
600 Form::Url(url) => ElicitRequestParams::UrlElicitationParams {
601 meta: None,
602 message: ask.message,
603 url,
604 elicitation_id: uuid::Uuid::new_v4().to_string(),
607 },
608 };
609 Ok((ask.key, form))
610}
611
612fn askable(context: &RequestContext<RoleServer>, asked: &[Ask]) -> bool {
615 let urls = capabilities(context)
616 .and_then(|capabilities| capabilities.elicitation)
617 .is_some_and(|elicitation| elicitation.url.is_some());
618 urls || asked.iter().all(|ask| matches!(ask.form, Form::Fields(_)))
619}
620
621fn completion(mut values: Vec<String>) -> CompleteResult {
624 let total = values.len();
625 values.truncate(CompletionInfo::MAX_VALUES);
626 let mut info = CompletionInfo::new(values).expect("cut to the maximum");
627 if total > CompletionInfo::MAX_VALUES {
628 info.total = u32::try_from(total).ok();
629 info.has_more = Some(true);
630 }
631 CompleteResult::new(info)
632}
633
634fn replied(reply: Reply) -> CallToolResult {
638 let mut result = CallToolResult::structured(reply.value);
639 result.content.extend(reply.content.into_iter().map(block));
640 result
641}
642
643fn block(content: Content) -> ContentBlock {
645 use base64::Engine;
646 let encode = |data: Vec<u8>| base64::engine::general_purpose::STANDARD.encode(data);
647 match content {
648 Content::Text(text) => ContentBlock::text(text),
649 Content::Image { mime, data } => ContentBlock::image(encode(data), mime),
650 Content::Audio { mime, data } => ContentBlock::audio(encode(data), mime),
651 Content::Link { uri, name, mime } => {
652 let link = Resource::new(uri, name);
653 ContentBlock::resource_link(match mime {
654 Some(mime) => link.with_mime_type(mime),
655 None => link,
656 })
657 }
658 Content::Resource { uri, body } => ContentBlock::resource(contents(&uri, body)),
659 }
660}
661
662fn contents(uri: &str, body: Body) -> ResourceContents {
664 match body {
665 Body::Text { mime, text } => ResourceContents::text(text, uri).with_mime_type(mime),
666 Body::Bytes { mime, bytes } => {
667 use base64::Engine;
668 let blob = base64::engine::general_purpose::STANDARD.encode(bytes);
669 ResourceContents::blob(blob, uri).with_mime_type(mime)
670 }
671 }
672}
673
674fn caller(context: &RequestContext<RoleServer>, answers: Option<InputResponses>) -> Caller<'_> {
676 Caller {
677 space: identified(context),
678 answers,
679 bearer: bearer(context),
680 cancel: Some(context.ct.clone()),
681 progress: context.meta.get_progress_token().map(|token| {
682 Arc::new(PeerProgress {
683 peer: context.peer.clone(),
684 token,
685 }) as Arc<dyn Progress>
686 }),
687 }
688}
689
690struct PeerProgress {
692 peer: rmcp::service::Peer<RoleServer>,
693 token: ProgressToken,
694}
695
696impl Progress for PeerProgress {
697 fn report(
698 &self,
699 done: f64,
700 total: Option<f64>,
701 message: Option<String>,
702 ) -> Pin<Box<dyn Future<Output = ()> + Send + '_>> {
703 let mut note = ProgressNotificationParam::new(self.token.clone(), done);
704 note.total = total;
705 note.message = message;
706 Box::pin(async move {
707 let _ = self.peer.notify_progress(note).await;
709 })
710 }
711}
712
713fn tool_error(error: CallError) -> CallToolResult {
715 trace::tool_error();
716 CallToolResult::error(vec![ContentBlock::text(error.to_string())])
717}
718
719fn refusal<V>(error: CallError) -> Result<V, ErrorData> {
722 let message = error.to_string();
723 Err(match error {
724 CallError::Unidentified => ErrorData::invalid_request(message, None),
725 CallError::Unavailable(_) => ErrorData::internal_error(message, None),
726 _ => ErrorData::invalid_params(message, None),
727 })
728}
729
730fn supports_tasks(context: &RequestContext<RoleServer>) -> bool {
733 capabilities(context).is_some_and(|capabilities| capabilities.supports_tasks())
734}
735
736fn capabilities(context: &RequestContext<RoleServer>) -> Option<ClientCapabilities> {
739 context.meta.client_capabilities().or_else(|| {
740 context
741 .peer
742 .peer_info()
743 .map(|info| info.capabilities.clone())
744 })
745}
746
747const MAX_ROUNDS: usize = 10;
750
751async fn answered<V, Fut>(
757 context: &RequestContext<RoleServer>,
758 mut answers: Option<InputResponses>,
759 attempt: impl Fn(Option<InputResponses>) -> Fut,
760) -> Result<V, CallError>
761where
762 Fut: Future<Output = Result<V, CallError>>,
763{
764 let mrtr = context
765 .protocol_version()
766 .is_some_and(|version| version >= ProtocolVersion::V_2026_07_28);
767 let can_ask = capabilities(context).is_some_and(|c| c.elicitation.is_some());
768 for _ in 0..MAX_ROUNDS {
769 let asked = match attempt(answers.clone()).await {
770 Err(CallError::InputRequired(asked)) if !askable(context, &asked) => {
771 return Err(CallError::Failed(
772 "the operation needs the user to open a page, which this client cannot do"
773 .into(),
774 ));
775 }
776 Err(CallError::InputRequired(asked)) if !mrtr && can_ask => asked,
777 result => return result,
778 };
779 let given = answers.get_or_insert_with(InputResponses::new);
780 for ask in asked {
781 let (key, form) = form(ask).map_err(|error| CallError::Failed(error.message.into()))?;
782 let asking = ServerRequest::ElicitRequest(ElicitRequest::new(form));
785 let answer = match context.peer.send_request(asking).await {
786 Ok(ClientResult::ElicitResult(answer)) => answer,
787 Ok(_) => {
788 return Err(CallError::Failed(
789 "the client answered something else".into(),
790 ));
791 }
792 Err(error) => {
793 return Err(CallError::Failed(format!(
794 "asking the user failed: {error}"
795 )));
796 }
797 };
798 let answer = serde_json::to_value(answer)
799 .map_err(|error| CallError::Failed(error.to_string()))?;
800 given.insert(key, answer);
801 }
802 }
803 Err(CallError::Failed(format!(
804 "the operation asked the user more than {MAX_ROUNDS} rounds of questions"
805 )))
806}
807
808fn listed(task: &StoredTask) -> Task {
809 let status = match task.state {
810 TaskState::Working => TaskStatus::Working,
811 TaskState::InputRequired(_) => TaskStatus::InputRequired,
812 TaskState::Succeeded(_) | TaskState::Errored(_) => TaskStatus::Completed,
813 TaskState::Cancelled => TaskStatus::Cancelled,
814 };
815 let listed = Task::new(&task.id, status, iso(task.created_at), iso(task.updated_at))
816 .with_ttl_ms(TASK_TTL.as_millis() as u64)
817 .with_poll_interval_ms(POLL_MS);
818 match &task.status {
819 Some(status) => listed.with_status_message(status),
820 None => listed,
821 }
822}
823
824fn detailed(task: StoredTask) -> DetailedTask {
825 let listed = listed(&task);
826 let result = match task.state {
827 TaskState::Working => return DetailedTask::new(listed, TaskPayload::Working),
828 TaskState::Cancelled => return DetailedTask::new(listed, TaskPayload::Cancelled),
829 TaskState::InputRequired(asked) => {
830 let payload = match questions(asked) {
831 Ok(result) => TaskPayload::InputRequired {
832 input_requests: result.input_requests.unwrap_or_default(),
833 },
834 Err(error) => {
837 let Ok(Value::Object(error)) = serde_json::to_value(error) else {
838 unreachable!("an error serializes to an object")
839 };
840 TaskPayload::Failed { error }
841 }
842 };
843 return DetailedTask::new(listed, payload);
844 }
845 TaskState::Succeeded(reply) => replied(reply),
846 TaskState::Errored(message) => CallToolResult::error(vec![ContentBlock::text(message)]),
849 };
850 let Ok(Value::Object(result)) = serde_json::to_value(result) else {
851 unreachable!("a tool result serializes to an object")
852 };
853 DetailedTask::new(listed, TaskPayload::Completed { result })
854}
855
856fn iso(time: SystemTime) -> String {
857 chrono::DateTime::<chrono::Utc>::from(time).to_rfc3339()
858}
859
860#[cfg(feature = "http")]
862fn bearer(context: &RequestContext<RoleServer>) -> Option<&str> {
863 context
864 .extensions
865 .get::<http::request::Parts>()?
866 .headers
867 .get(http::header::AUTHORIZATION)?
868 .to_str()
869 .ok()?
870 .strip_prefix("Bearer ")
871}
872
873#[cfg(not(feature = "http"))]
874fn bearer(_context: &RequestContext<RoleServer>) -> Option<&str> {
875 None
876}
877
878#[cfg(feature = "http")]
881fn identified(context: &RequestContext<RoleServer>) -> Option<Arc<Space>> {
882 context
883 .extensions
884 .get::<http::request::Parts>()?
885 .extensions
886 .get::<Arc<Space>>()
887 .cloned()
888}
889
890#[cfg(not(feature = "http"))]
891fn identified(_context: &RequestContext<RoleServer>) -> Option<Arc<Space>> {
892 None
893}
894
895impl<M: Module, I: Identity, E: Entitlements, T: Metering, S: TaskStore> Kit<M, I, E, T, S> {
896 pub async fn serve_stdio(self) -> std::io::Result<()> {
898 use rmcp::ServiceExt;
899 let running = Server::new(self)
900 .serve(rmcp::transport::stdio())
901 .await
902 .map_err(std::io::Error::other)?;
903 running.waiting().await.map_err(std::io::Error::other)?;
904 Ok(())
905 }
906}
907
908#[cfg(feature = "http")]
910#[derive(Clone, Debug, Default)]
911pub struct HttpOptions {
912 pub allowed_hosts: Vec<String>,
915 pub protected: Option<ProtectedResource>,
919}
920
921#[cfg(feature = "http")]
923#[derive(Clone, Debug)]
924pub struct ProtectedResource {
925 pub resource: String,
927 pub authorization_servers: Vec<String>,
929}
930
931#[cfg(feature = "http")]
933#[derive(Clone, serde::Serialize)]
934struct ResourceMetadata {
935 resource: String,
936 authorization_servers: Vec<String>,
937 scopes_supported: Vec<&'static str>,
938 bearer_methods_supported: [&'static str; 1],
939 resource_name: &'static str,
940}
941
942#[cfg(feature = "http")]
943impl<M: Module, I: Identity, E: Entitlements, T: Metering, S: TaskStore> Kit<M, I, E, T, S> {
944 pub async fn serve_http(
950 self,
951 listener: tokio::net::TcpListener,
952 options: HttpOptions,
953 ) -> std::io::Result<()> {
954 self.serve_http_with_shutdown(listener, options, std::future::pending())
955 .await
956 }
957
958 pub async fn serve_http_with_shutdown(
975 self,
976 listener: tokio::net::TcpListener,
977 options: HttpOptions,
978 signal: impl Future<Output = ()> + Send + 'static,
979 ) -> std::io::Result<()> {
980 use rmcp::transport::streamable_http_server::session::never::NeverSessionManager;
981 use rmcp::transport::streamable_http_server::{
982 StreamableHttpServerConfig, StreamableHttpService,
983 };
984 let server = Server::new(self);
985 let kit = server.kit.clone();
986 let stopping = server.stopping.clone();
987 let mut config = StreamableHttpServerConfig::default()
988 .with_legacy_session_mode(false)
989 .with_json_response(true);
990 if !options.allowed_hosts.is_empty() {
991 config = config.with_allowed_hosts(options.allowed_hosts);
992 }
993 let service = StreamableHttpService::new(
994 move || Ok(server.clone()),
995 Arc::new(NeverSessionManager::default()),
996 config,
997 );
998 let mut router = axum::Router::new();
999 match options.protected {
1000 None => router = router.nest_service("/mcp", service),
1001 Some(protected) => {
1002 let (path, metadata) = published(&protected, kit.manifest())?;
1003 let challenge = format!(
1004 "Bearer resource_metadata=\"{}\"",
1005 protected_url(&protected.resource, &path)?
1006 );
1007 let mcp = axum::Router::new().nest_service("/mcp", service).layer(
1009 axum::middleware::from_fn(move |request, next| {
1010 identify_request(kit.clone(), challenge.clone(), request, next)
1011 }),
1012 );
1013 router = router
1014 .route(
1015 &path,
1016 axum::routing::get(move || {
1017 std::future::ready(axum::Json(metadata.clone()))
1018 }),
1019 )
1020 .merge(mcp);
1021 }
1022 }
1023 axum::serve(listener, router)
1024 .with_graceful_shutdown(async move {
1025 signal.await;
1026 stopping.cancel();
1027 })
1028 .await
1029 }
1030}
1031
1032#[cfg(feature = "http")]
1036async fn identify_request<M, I, E, T, S>(
1037 kit: Arc<Kit<M, I, E, T, S>>,
1038 challenge: String,
1039 mut request: axum::extract::Request,
1040 next: axum::middleware::Next,
1041) -> axum::response::Response
1042where
1043 M: Module,
1044 I: Identity,
1045 E: Entitlements,
1046 T: Metering,
1047 S: TaskStore,
1048{
1049 let bearer = request
1050 .headers()
1051 .get(http::header::AUTHORIZATION)
1052 .and_then(|value| value.to_str().ok())
1053 .and_then(|value| value.strip_prefix("Bearer "));
1054 match kit.identify(bearer).await {
1055 Ok(space) => {
1056 request.extensions_mut().insert(space);
1057 next.run(request).await
1058 }
1059 Err(_) => axum::response::IntoResponse::into_response((
1060 http::StatusCode::UNAUTHORIZED,
1061 [(http::header::WWW_AUTHENTICATE, challenge)],
1062 )),
1063 }
1064}
1065
1066#[cfg(feature = "http")]
1069fn published(
1070 protected: &ProtectedResource,
1071 manifest: &crate::module::Manifest,
1072) -> std::io::Result<(String, ResourceMetadata)> {
1073 let resource: http::Uri = protected.resource.parse().map_err(std::io::Error::other)?;
1074 let path = format!(
1075 "/.well-known/oauth-protected-resource{}",
1076 resource.path().trim_end_matches('/')
1077 );
1078 let metadata = ResourceMetadata {
1079 resource: protected.resource.clone(),
1080 authorization_servers: protected.authorization_servers.clone(),
1081 scopes_supported: manifest.scopes.clone(),
1082 bearer_methods_supported: ["header"],
1083 resource_name: manifest.name,
1084 };
1085 Ok((path, metadata))
1086}
1087
1088#[cfg(feature = "http")]
1090fn protected_url(resource: &str, path: &str) -> std::io::Result<String> {
1091 let resource: http::Uri = resource.parse().map_err(std::io::Error::other)?;
1092 let (Some(scheme), Some(authority)) = (resource.scheme_str(), resource.authority()) else {
1093 return Err(std::io::Error::other(format!(
1094 "the protected resource {resource} is not an absolute URL"
1095 )));
1096 };
1097 Ok(format!("{scheme}://{authority}{path}"))
1098}