use std::collections::{BTreeMap, BTreeSet};
use std::sync::{Mutex, PoisonError};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use base64::Engine as _;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use hmac::{Hmac, KeyInit as _, Mac as _};
use schemars::JsonSchema;
use serde::de::DeserializeOwned;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value, json};
use sha2::{Digest as _, Sha256};
use crate::context::Ctx;
use crate::error::McpError;
use crate::spec::{self, methods};
use crate::tool::{ToolError, schema_for};
type HmacSha256 = Hmac<Sha256>;
pub const MIN_REQUEST_STATE_KEY_BYTES: usize = 32;
const STATE_VERSION: u8 = 1;
const SIGNATURE_CONTEXT: &[u8] = b"flymcprust/requestState/v1\0";
pub(crate) struct StateCodec {
key: Vec<u8>,
ttl: Duration,
}
impl std::fmt::Debug for StateCodec {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("StateCodec")
.field("ttl", &self.ttl)
.finish_non_exhaustive()
}
}
#[derive(Serialize, Deserialize)]
struct Payload {
v: u8,
m: String,
t: String,
a: String,
e: u64,
n: u32,
r: BTreeMap<String, Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
p: Option<String>,
}
impl StateCodec {
pub(crate) fn new(key: Vec<u8>, ttl: Duration) -> Self {
Self { key, ttl }
}
fn mac(&self, context: &[u8], data: &[u8]) -> Option<HmacSha256> {
let mut mac = HmacSha256::new_from_slice(&self.key).ok()?;
mac.update(context);
mac.update(data);
Some(mac)
}
pub(crate) fn seal_for<T: Serialize>(&self, context: &[u8], payload: &T) -> Option<String> {
let json = serde_json::to_vec(payload).ok()?;
let body = URL_SAFE_NO_PAD.encode(json);
let tag = self.mac(context, body.as_bytes())?.finalize().into_bytes();
Some(format!("{body}.{}", URL_SAFE_NO_PAD.encode(tag)))
}
pub(crate) fn open_for<T: DeserializeOwned>(&self, context: &[u8], token: &str) -> Option<T> {
let (body, tag) = token.rsplit_once('.')?;
let tag = URL_SAFE_NO_PAD.decode(tag).ok()?;
self.mac(context, body.as_bytes())?
.verify_slice(&tag)
.ok()?;
let json = URL_SAFE_NO_PAD.decode(body).ok()?;
serde_json::from_slice(&json).ok()
}
fn seal(&self, payload: &Payload) -> Result<String, McpError> {
self.seal_for(SIGNATURE_CONTEXT, payload)
.ok_or_else(McpError::internal)
}
fn open(&self, token: &str, round: &RoundKey<'_>) -> Result<Payload, McpError> {
let digest = round.args_digest();
let invalid = || McpError::invalid_params("Invalid or expired requestState");
let payload: Payload = self
.open_for(SIGNATURE_CONTEXT, token)
.ok_or_else(invalid)?;
if payload.v != STATE_VERSION
|| payload.m != round.method
|| payload.t != round.target
|| payload.a != digest
|| payload.p.as_deref() != round.principal
|| payload.e < now_secs()
{
return Err(invalid());
}
Ok(payload)
}
}
pub(crate) fn now_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map_or(0, |d| d.as_secs())
}
pub(crate) struct RoundKey<'a> {
pub(crate) method: &'static str,
pub(crate) target: &'a str,
pub(crate) arguments: Args<'a>,
pub(crate) principal: Option<&'a str>,
}
#[derive(Clone, Copy)]
pub(crate) enum Args<'a> {
None,
Raw(&'a str),
Map(&'a BTreeMap<String, String>),
}
impl RoundKey<'_> {
fn args_digest(&self) -> String {
let canonical = match self.arguments {
Args::None => Vec::new(),
Args::Raw(raw) => serde_json::from_str::<Value>(raw)
.ok()
.and_then(|mut v| {
v.sort_all_objects();
serde_json::to_vec(&v).ok()
})
.unwrap_or_default(),
Args::Map(map) => serde_json::to_vec(map).unwrap_or_default(),
};
URL_SAFE_NO_PAD.encode(Sha256::digest(&canonical))
}
}
#[derive(Debug, Default)]
pub(crate) struct Round {
number: u32,
responses: BTreeMap<String, Value>,
pending: Mutex<Vec<(String, Value)>>,
used: Mutex<BTreeSet<String>>,
}
impl Round {
pub(crate) fn start(
codec: &StateCodec,
key: &RoundKey<'_>,
input_responses: Option<Value>,
request_state: Option<&str>,
) -> Result<Self, McpError> {
let (number, mut responses) = match request_state {
Some(token) => {
let payload = codec.open(token, key)?;
(payload.n, payload.r)
}
None => (0, BTreeMap::new()),
};
match input_responses {
None | Some(Value::Null) => {}
Some(Value::Object(map)) => {
for (k, v) in map {
responses.entry(k).or_insert(v);
}
}
Some(_) => {
return Err(McpError::invalid_params(
"`inputResponses` must be an object",
));
}
}
Ok(Self {
number,
responses,
pending: Mutex::new(Vec::new()),
used: Mutex::new(BTreeSet::new()),
})
}
pub(crate) fn input_required(
&self,
codec: &StateCodec,
key: &RoundKey<'_>,
meta: Option<&spec::ResultMetaObject>,
) -> Result<Value, McpError> {
let pending =
std::mem::take(&mut *self.pending.lock().unwrap_or_else(PoisonError::into_inner));
let expiry = now_secs().saturating_add(codec.ttl.as_secs());
let state = codec.seal(&Payload {
v: STATE_VERSION,
m: key.method.to_owned(),
t: key.target.to_owned(),
a: key.args_digest(),
e: expiry,
p: key.principal.map(str::to_owned),
n: self.number.saturating_add(1),
r: {
let used = self.used.lock().unwrap_or_else(PoisonError::into_inner);
self.responses
.iter()
.filter(|(k, _)| used.contains(k.as_str()))
.map(|(k, v)| (k.clone(), v.clone()))
.collect()
},
})?;
let mut result = Map::new();
result.insert("resultType".to_owned(), Value::from("input_required"));
if !pending.is_empty() {
result.insert(
"inputRequests".to_owned(),
Value::Object(pending.into_iter().collect()),
);
}
result.insert("requestState".to_owned(), Value::from(state));
if let Some(meta) = meta {
let meta = serde_json::to_value(meta).map_err(|_| McpError::internal())?;
result.insert("_meta".to_owned(), meta);
}
Ok(Value::Object(result))
}
fn response(&self, key: &str) -> Option<&Value> {
let value = self.responses.get(key)?;
self.used
.lock()
.unwrap_or_else(PoisonError::into_inner)
.insert(key.to_owned());
Some(value)
}
fn push(&self, key: &str, request: Value) {
let mut pending = self.pending.lock().unwrap_or_else(PoisonError::into_inner);
if !pending.iter().any(|(k, _)| k == key) {
pending.push((key.to_owned(), request));
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Elicitation<T> {
Accepted(T),
Declined,
Cancelled,
}
#[derive(Deserialize)]
struct ElicitAnswer {
action: String,
#[serde(default)]
content: Option<Value>,
}
#[derive(Clone, Copy)]
enum Needs {
Elicitation,
Sampling,
Roots,
}
impl Ctx {
#[must_use]
pub fn input_round(&self) -> u32 {
self.round().map_or(0, |r| r.number)
}
pub fn input_response<T: DeserializeOwned>(&self, key: &str) -> Result<Option<T>, ToolError> {
let Some(value) = self.round().and_then(|r| r.response(key)) else {
return Ok(None);
};
T::deserialize(value)
.map(Some)
.map_err(|_| invalid_response(key))
}
pub async fn elicit<T: DeserializeOwned + JsonSchema>(
&self,
key: &str,
message: impl Into<String>,
) -> Result<T, ToolError> {
match self.elicit_with_action::<T>(key, message).await? {
Elicitation::Accepted(value) => Ok(value),
Elicitation::Declined => Err(ToolError::user(
"The user declined to provide the requested information.",
)),
Elicitation::Cancelled => Err(ToolError::user("The user cancelled the request.")),
}
}
pub async fn elicit_with_action<T: DeserializeOwned + JsonSchema>(
&self,
key: &str,
message: impl Into<String>,
) -> Result<Elicitation<T>, ToolError> {
if let Some(answer) = self.input_response::<ElicitAnswer>(key)? {
return match answer.action.as_str() {
"accept" => {
let content = answer.content.unwrap_or_else(|| json!({}));
T::deserialize(&content)
.map(Elicitation::Accepted)
.map_err(|_| invalid_response(key))
}
"decline" => Ok(Elicitation::Declined),
"cancel" => Ok(Elicitation::Cancelled),
_ => Err(invalid_response(key)),
};
}
let schema = elicitation_schema::<T>()?;
let request = json!({
"method": methods::ELICITATION_CREATE,
"params": { "message": message.into(), "requestedSchema": schema },
});
Err(self.ask(key, request, Needs::Elicitation))
}
pub async fn elicit_with_schema(
&self,
key: &str,
message: impl Into<String>,
requested_schema: Value,
) -> Result<Elicitation<Map<String, Value>>, ToolError> {
if let Some(answer) = self.input_response::<ElicitAnswer>(key)? {
return match answer.action.as_str() {
"accept" => match answer.content.unwrap_or_else(|| json!({})) {
Value::Object(content) => Ok(Elicitation::Accepted(content)),
_ => Err(invalid_response(key)),
},
"decline" => Ok(Elicitation::Declined),
"cancel" => Ok(Elicitation::Cancelled),
_ => Err(invalid_response(key)),
};
}
let request = json!({
"method": methods::ELICITATION_CREATE,
"params": { "message": message.into(), "requestedSchema": requested_schema },
});
Err(self.ask(key, request, Needs::Elicitation))
}
#[deprecated(note = "Sampling is deprecated as of protocol version 2026-07-28 (SEP-2577).")]
#[allow(deprecated)]
pub async fn sample(
&self,
key: &str,
params: spec::CreateMessageRequestParams,
) -> Result<spec::CreateMessageResult, ToolError> {
if let Some(answer) = self.input_response(key)? {
return Ok(answer);
}
let request = json!({ "method": methods::SAMPLING_CREATE_MESSAGE, "params": params });
Err(self.ask(key, request, Needs::Sampling))
}
#[deprecated(note = "Roots are deprecated as of protocol version 2026-07-28 (SEP-2577).")]
#[allow(deprecated)]
pub async fn list_roots(&self, key: &str) -> Result<spec::ListRootsResult, ToolError> {
if let Some(answer) = self.input_response(key)? {
return Ok(answer);
}
let request = json!({ "method": methods::ROOTS_LIST, "params": {} });
Err(self.ask(key, request, Needs::Roots))
}
fn ask(&self, key: &str, request: Value, needs: Needs) -> ToolError {
let Some(round) = self.round() else {
return ToolError::internal(std::io::Error::other(
"input requests are only available in tools/call, resources/read and prompts/get",
));
};
#[allow(deprecated)]
let caps = self.client_capabilities();
#[allow(deprecated)]
let (supported, required) = match needs {
Needs::Elicitation => (
caps.elicitation
.as_ref()
.is_some_and(|e| e.form.is_some() || e.url.is_none()),
spec::ClientCapabilities {
elicitation: Some(spec::ClientCapabilitiesElicitation {
form: Some(Map::new()),
url: None,
}),
..spec::ClientCapabilities::default()
},
),
Needs::Sampling => (
caps.sampling.is_some(),
spec::ClientCapabilities {
sampling: Some(spec::ClientCapabilitiesSampling::default()),
..spec::ClientCapabilities::default()
},
),
Needs::Roots => (
caps.roots.is_some(),
spec::ClientCapabilities {
roots: Some(spec::ClientCapabilitiesRoots::default()),
..spec::ClientCapabilities::default()
},
),
};
if !supported {
static HINT: crate::hints::Once = crate::hints::Once::new();
crate::hints::warn_once(
&HINT,
format_args!(
"a client that did not declare the needed capability called a handler \
that asks it for input; it received -32021"
),
"check ctx.client_capabilities() before asking if the handler can work \
without the answer",
);
return ToolError::missing_client_capability(&required);
}
round.push(key, request);
ToolError::input_required()
}
}
fn invalid_response(key: &str) -> ToolError {
ToolError::protocol(McpError::invalid_params(format!(
"Invalid input response for `{}`",
crate::error::truncate(key)
)))
}
fn elicitation_schema<T: JsonSchema>() -> Result<Value, ToolError> {
let mut schema = schema_for::<T>();
let unsupported = || {
ToolError::internal(std::io::Error::other(format!(
"`{}` cannot be used with ctx.elicit\n = help: use a struct whose fields are strings, \
numbers, integers, booleans or enums (optionally in Option<..>); forms cannot nest",
std::any::type_name::<T>()
)))
};
if schema.get("type").and_then(Value::as_str) != Some("object") {
return Err(unsupported());
}
schema.remove("additionalProperties");
let Some(Value::Object(properties)) = schema.get_mut("properties") else {
return Err(unsupported());
};
for property in properties.values_mut() {
let Value::Object(property) = property else {
return Err(unsupported());
};
if let Some(Value::Array(types)) = property.get("type") {
let non_null: Vec<&Value> = types
.iter()
.filter(|t| t.as_str() != Some("null"))
.collect();
if let [single] = non_null.as_slice() {
let single = (*single).clone();
property.insert("type".to_owned(), single);
}
}
let primitive = matches!(
property.get("type").and_then(Value::as_str),
Some("string" | "number" | "integer" | "boolean" | "array")
) || property.contains_key("enum")
|| property.contains_key("oneOf")
|| property.contains_key("anyOf");
if !primitive {
return Err(unsupported());
}
}
Ok(Value::Object(schema))
}
#[cfg(test)]
mod tests {
use super::*;
fn codec() -> StateCodec {
StateCodec::new(vec![7; 32], Duration::from_secs(600))
}
fn key(target: &str) -> RoundKey<'_> {
RoundKey {
method: "tools/call",
target,
arguments: Args::Raw(r#"{"a": 1}"#),
principal: None,
}
}
fn round_trip(codec: &StateCodec, round: &Round) -> String {
round.input_required(codec, &key("t"), None).unwrap()["requestState"]
.as_str()
.unwrap()
.to_owned()
}
#[test]
fn state_round_trips_and_accumulates_answers() {
let codec = codec();
let first = Round::start(
&codec,
&key("t"),
Some(json!({ "q1": 1, "unread": 0 })),
None,
)
.unwrap();
assert_eq!(first.response("q1"), Some(&json!(1)));
let token = round_trip(&codec, &first);
let second =
Round::start(&codec, &key("t"), Some(json!({ "q2": 2 })), Some(&token)).unwrap();
assert_eq!(second.number, 1);
assert_eq!(second.response("q1"), Some(&json!(1)));
assert_eq!(second.response("q2"), Some(&json!(2)));
assert_eq!(second.response("unread"), None);
}
#[test]
fn tampered_or_rebound_state_is_rejected() {
let codec = codec();
let round = Round::start(&codec, &key("t"), None, None).unwrap();
let token = round_trip(&codec, &round);
let reject = |token: &str, key: RoundKey<'_>| {
Round::start(&codec, &key, None, Some(token))
.unwrap_err()
.code()
};
assert_eq!(reject(&format!("{token}-TAMPERED"), key("t")), -32602);
assert_eq!(reject("garbage", key("t")), -32602);
assert_eq!(reject(&token, key("other")), -32602);
let mut other_args = key("t");
other_args.arguments = Args::Raw(r#"{"a":2}"#);
assert_eq!(reject(&token, other_args), -32602);
let mut other_user = key("t");
other_user.principal = Some("someone-else");
assert_eq!(reject(&token, other_user), -32602);
let other_key = StateCodec::new(vec![8; 32], Duration::from_secs(600));
assert!(Round::start(&other_key, &key("t"), None, Some(&token)).is_err());
}
#[test]
fn expired_state_is_rejected() {
let codec = StateCodec::new(vec![7; 32], Duration::ZERO);
let round = Round::start(&codec, &key("t"), None, None).unwrap();
let mut payload: Payload = {
let token = round_trip(&codec, &round);
let body = token.split('.').next().unwrap();
serde_json::from_slice(&URL_SAFE_NO_PAD.decode(body).unwrap()).unwrap()
};
payload.e = now_secs() - 1;
let token = codec.seal(&payload).unwrap();
assert!(Round::start(&codec, &key("t"), None, Some(&token)).is_err());
}
#[test]
fn argument_digest_ignores_whitespace() {
let digest = |raw| {
RoundKey {
method: "tools/call",
target: "t",
arguments: Args::Raw(raw),
principal: None,
}
.args_digest()
};
assert_eq!(digest(r#"{"a":1}"#), digest(r#"{ "a" : 1 }"#));
assert_eq!(
digest(r#"{"a":1,"b":{"y":2,"x":1}}"#),
digest(r#"{"b":{"x":1,"y":2},"a":1}"#)
);
assert_ne!(digest(r#"{"a":1}"#), digest(r#"{"a":2}"#));
}
#[test]
fn invalid_input_responses_are_rejected() {
let err = Round::start(&codec(), &key("t"), Some(json!(5)), None).unwrap_err();
assert_eq!(err.code(), -32602);
}
#[derive(JsonSchema)]
#[allow(dead_code)]
struct Form {
name: String,
age: Option<u32>,
subscribe: bool,
}
#[test]
fn elicitation_schema_is_flat() {
let schema = elicitation_schema::<Form>().unwrap();
assert_eq!(
schema,
json!({
"type": "object",
"properties": {
"name": { "type": "string" },
"age": { "type": "integer", "minimum": 0 },
"subscribe": { "type": "boolean" },
},
"required": ["name", "subscribe"],
})
);
#[derive(JsonSchema)]
#[allow(dead_code)]
struct Nested {
inner: Form,
}
assert!(elicitation_schema::<Nested>().is_err());
}
}