Skip to main content

pamin_kit/
serve.rs

1//! Serving a [`Kit`] as MCP. The only module in the system that names `rmcp`.
2
3use 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
42/// A tool's `_meta` naming the MCP App that shows its results.
43/// A tool's MCP Apps `_meta`: the page that shows its results, and whether
44/// only an app may call it. `None` when it has neither.
45fn 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
56/// A tool with its hints as MCP annotations.
57fn 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
68/// Applies the parts of `looks` that are set, with the builder each MCP
69/// type has for them: the one place looks become MCP.
70macro_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
85/// A resource or template shown with its looks, annotations included.
86macro_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
135/// How often a caller should poll a task, in milliseconds.
136const POLL_MS: u64 = 500;
137
138/// A kit as an MCP server handler. Cheap to clone: one per HTTP request.
139pub(crate) struct Server<M, I, E, T, S> {
140    kit: Arc<Kit<M, I, E, T, S>>,
141    /// Built once: listing tools is the first thing every client does.
142    tools: Arc<[rmcp::model::Tool]>,
143    /// Cancelled when the server shuts down, which ends every open
144    /// subscription so it does not hold the shutdown open.
145    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            // Module resources can change; app pages cannot.
200            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    /// A tool by name, which rmcp's HTTP transport checks `Mcp-Param-*`
230    /// headers (SEP-2243) against.
231    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        // Only resources change: the lists of tools, prompts and resources
248        // are fixed when the kit is built.
249        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                // Ending it sends the final result that tells the client the
267                // server is going away.
268                () = self.stopping.cancelled() => return Ok(()),
269                changed = changes.recv() => changed,
270            };
271            let changed = match changed {
272                Ok(changed) => changed,
273                // Fell behind: the changes skipped are lost to this listener,
274                // which rereads what it cares about on the next one.
275                Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
276                Err(tokio::sync::broadcast::error::RecvError::Closed) => return Ok(()),
277            };
278            // Another space's change is not this listener's to see.
279            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        // A kind of reference newer than this kit has nothing to complete.
432        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
558/// A kit result as the MCP response `R`: the result, the questions the
559/// operation asks the user (MRTR), or a refusal.
560fn 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
572/// What an operation asks the user, as MCP elicitation requests.
573fn 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
584/// One question as an elicitation form, with the key its answer comes back
585/// under.
586fn 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                // The module asked with a type an elicitation form cannot
591                // show: a mistake in the module.
592                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            // Unique per question asked, as MCP requires; the key ties the
605            // answer to the question.
606            elicitation_id: uuid::Uuid::new_v4().to_string(),
607        },
608    };
609    Ok((ask.key, form))
610}
611
612/// Whether the client can answer every question asked: a URL question
613/// only when the client declared URL elicitation.
614fn 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
621/// Suggestions as MCP completion, cut to the most MCP allows with the rest
622/// counted.
623fn 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
634/// A tool's reply as an MCP tool result: the structured result (with its
635/// JSON as text, for clients that read only content), then the content the
636/// operation attached.
637fn 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
643/// Kit content as an MCP content block.
644fn 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
662/// A resource's body as MCP resource contents.
663fn 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
674/// What the transport knows about a request, for the kit.
675fn 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
690/// Progress sent to the client as `notifications/progress`.
691struct 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            // A client that has gone away cannot be told; the call goes on.
708            let _ = self.peer.notify_progress(note).await;
709        })
710    }
711}
712
713/// A call the model should read and act on: not entitled, or failed.
714fn tool_error(error: CallError) -> CallToolResult {
715    trace::tool_error();
716    CallToolResult::error(vec![ContentBlock::text(error.to_string())])
717}
718
719/// What no caller should retry as it is, as a protocol error. When starting
720/// a task, any error is one: there is no task to report it in.
721fn 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
730/// Whether this request's client declared the tasks extension: per request
731/// over stateless HTTP, at initialization over stdio.
732fn supports_tasks(context: &RequestContext<RoleServer>) -> bool {
733    capabilities(context).is_some_and(|capabilities| capabilities.supports_tasks())
734}
735
736/// What the client declared it can do: per request over stateless HTTP, at
737/// initialization over stdio.
738fn 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
747/// The most rounds of questions one request may take, as rmcp's client
748/// allows.
749const MAX_ROUNDS: usize = 10;
750
751/// Runs `attempt` with `answers`, and when the operation asks the user and
752/// the client speaks a protocol before MRTR (2026-07-28), asks the client
753/// itself (`elicitation/create`) and runs it again with the answers. A newer
754/// client gets the questions back as MRTR and calls again on its own; a
755/// client that cannot ask at all gets the questions as a refusal.
756async 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            // The request rmcp's `elicitation` feature would send, without the
783            // URL parser it brings for nothing the kit uses.
784            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                // The module asked with a form MCP cannot show: the task
835                // cannot go on, and the poller is told why.
836                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        // An operation's error is a completed call whose result is an error,
847        // as MCP has it; `failed` is for protocol errors, and there are none.
848        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/// The request's bearer token: present only over HTTP.
861#[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/// The space the HTTP layer already identified for this request, so the
879/// identity seam runs once.
880#[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    /// Serves this kit over stdin and stdout until the client disconnects.
897    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/// How [`Kit::serve_http`] serves.
909#[cfg(feature = "http")]
910#[derive(Clone, Debug, Default)]
911pub struct HttpOptions {
912    /// `Host` headers accepted; empty accepts loopback hosts only, which
913    /// stops DNS rebinding against a local server.
914    pub allowed_hosts: Vec<String>,
915    /// When set, the server is an OAuth protected resource: it publishes
916    /// its metadata (RFC 9728) and answers an unidentified request with 401,
917    /// which is how an MCP client learns where to sign in.
918    pub protected: Option<ProtectedResource>,
919}
920
921/// An MCP server as an OAuth protected resource (RFC 9728).
922#[cfg(feature = "http")]
923#[derive(Clone, Debug)]
924pub struct ProtectedResource {
925    /// The server's public URL, such as `https://notes.example.com/mcp`.
926    pub resource: String,
927    /// The authorization servers that issue its tokens.
928    pub authorization_servers: Vec<String>,
929}
930
931/// The metadata a protected resource publishes (RFC 9728, section 2).
932#[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    /// Serves this kit as Streamable HTTP at `/mcp` on `listener`, until the
945    /// listener fails.
946    ///
947    /// Stateless: no session outlives a request, so any instance serves any
948    /// request. [`Kit::serve_http_with_shutdown`] serves until a signal.
949    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    /// Serves as [`Kit::serve_http`] does, and shuts down gracefully once
959    /// `signal` completes, such as on SIGTERM.
960    ///
961    /// When `signal` completes, the server stops accepting connections and
962    /// closes the listener. Requests already being served, including one
963    /// whose headers are still arriving, run to completion and get their
964    /// responses, each on a connection that then closes. A connection
965    /// between requests, such as an idle keep-alive one, closes at once. The
966    /// returned future resolves once every connection has closed.
967    ///
968    /// An open `subscriptions/listen` stream ends with its final result, as
969    /// MCP has a server that is going away do.
970    ///
971    /// Shutdown waits for in-flight requests without a limit: bound the wait
972    /// with a timeout if the process must exit by a deadline. Tasks (the MCP
973    /// tasks extension) run on in the background and are not waited for.
974    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                // Only /mcp needs an identity; the metadata is public.
1008                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/// Identifies an HTTP request once, before MCP sees it: an unidentified one
1033/// gets 401 with where to sign in; an identified one carries its space to
1034/// the kit, which does not ask the identity seam again.
1035#[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/// Where the metadata is published and what it says. RFC 9728 puts it at
1067/// `/.well-known/oauth-protected-resource` followed by the resource's path.
1068#[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/// The metadata's absolute URL: the resource's origin and `path`.
1089#[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}