johnandersen777.bsky.social/knotless
git that lives in your Atmosphere account
Forked from chadtmiller.com/knotless
0
0
//! MCP over streamable HTTP, for assistants that reach a server rather than
//! start one, such as claude.ai's connectors: at `/mcp` on a public
//! `knotless-serve`, for any account.
//!
//! The server is its own OAuth authorization server, as the MCP
//! authorization spec has it: an assistant registers (RFC 7591), sends its
//! user here with PKCE, and gets Bearer tokens for `/mcp`. The user signs in
//! with atproto OAuth, this server a confidential client of their PDS, and
//! the server keeps that session to act as them. Any account gets a token
//! but one the operator refuses, such as one taken down.
//!
//! What it hands out is signed rather than stored: a registered client is
//! its own ID, and a token holds whom it is for. Refusing an account stops
//! its tokens at their next use. Each account's
//! sessions and private mirrors live in a directory of their own, so no one
//! reads as anyone else, forgotten once its refresh tokens have all run
//! out; public repositories are read from the server's own mirrors. Each
//! account has a share of the server: so many tool calls at once, so many
//! of them waiting, and so many a minute.
use crate::Server;
use axum::extract::{Form, Query, Request, State};
use axum::http::request::Parts;
use axum::http::{StatusCode, header};
use axum::middleware::{self, Next};
use axum::response::{Html, IntoResponse, Redirect, Response};
use axum::routing::{get, post};
use axum::{Json, Router};
use base64::Engine;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use knotless::crypto::{self, P256Key};
use knotless::env::Ctx;
use knotless::session::{self, Session};
use knotless::telemetry::{self, Fields, Level};
use knotless::{identity, keys, oauth};
use rmcp::model::{
CallToolRequestParams, CallToolResponse, CallToolResult, ContentBlock,
ListResourceTemplatesResult, ListToolsResult, PaginatedRequestParams,
ReadResourceRequestParams, ReadResourceResponse, ServerConfig, Tool,
};
use rmcp::service::RequestContext;
use rmcp::transport::streamable_http_server::session::never::NeverSessionManager;
use rmcp::transport::streamable_http_server::{StreamableHttpServerConfig, StreamableHttpService};
use rmcp::{ErrorData as McpError, RoleServer, ServerHandler};
use serde::de::DeserializeOwned;
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant, SystemTime};
const MCP: &str = "/mcp";
const METADATA: &str = "/mcp/oauth-client-metadata.json";
const CALLBACK: &str = "/mcp/oauth/callback";
const ACCESS_TTL: u64 = 60 * 60;
const REFRESH_TTL: u64 = 90 * 24 * 60 * 60;
/// A registered client lasts as long as the key that signed it.
const CLIENT_TTL: u64 = 100 * 365 * 24 * 60 * 60;
/// How long a sign-in may take, from the assistant sending its user here to
/// the user coming back from their PDS.
const SIGN_IN: Duration = Duration::from_secs(15 * 60);
/// How long a code waits to be traded for tokens.
const CODE: Duration = Duration::from_secs(5 * 60);
const MINUTE: Duration = Duration::from_secs(60);
/// How often accounts gone quiet are looked for.
const FORGET: Duration = Duration::from_secs(60 * 60);
/// The file whose time says when a token was last made for an account.
const USED: &str = "used";
/// How much of the server sign-ins and each account may take.
pub struct Limits {
/// Sign-ins gone to a PDS and not back yet, from everyone.
pub signing_in: usize,
/// Tool calls an account may have running at once.
pub running: usize,
/// Of those, the ones that may wait for checks or a merge.
pub waiting: usize,
/// Tool calls an account may start in a minute.
pub per_minute: usize,
}
impl Default for Limits {
fn default() -> Limits {
Limits {
signing_in: 1000,
running: 8,
waiting: 2,
per_minute: 120,
}
}
}
/// The MCP server over HTTP, and the authorization server in front of it.
pub struct Mcp {
/// Config and cache under `mcp/`, apart from the rest of the server's.
ctx: Ctx,
/// The public URL, the issuer, with no trailing slash.
url: String,
limits: Limits,
/// Signs what this server hands out, and its assertions as a client of
/// PDSes.
key: P256Key,
/// Sign-ins gone to a PDS, by the state of their atproto request.
pending: Mutex<HashMap<String, Pending>>,
codes: Mutex<HashMap<String, Code>>,
/// Told of each sign-in and tool call, for the operator's counts.
events: std::sync::OnceLock<Events>,
/// Says which accounts the operator turns away, listed or not.
refused: std::sync::OnceLock<Refused>,
/// The tool calls each account has running, and made lately.
calls: Mutex<HashMap<String, Calls>>,
/// When accounts gone quiet were last looked for.
looked: Mutex<Option<Instant>>,
}
/// Whether the operator turns an account away, by its DID.
pub type Refused = Arc<dyn Fn(&str) -> bool + Send + Sync>;
/// An account's tool calls: running now, of them waiting, and made since
/// `since`, a minute at most ago.
struct Calls {
running: usize,
waiting: usize,
made: usize,
since: Instant,
}
/// A tool call an account has running, given back when it ends, however
/// it ends.
struct Turn<'a> {
mcp: &'a Mcp,
did: &'a str,
waits: bool,
}
impl Drop for Turn<'_> {
fn drop(&mut self) {
let mut calls = self.mcp.calls.lock().unwrap();
let waits = usize::from(self.waits);
calls.entry(self.did.to_string()).and_modify(|c| {
c.running -= 1;
c.waiting -= waits;
});
}
}
/// What the server tells its operator's counts of: `mcp.signin`, `mcp.tool`
/// with the tool's name, or `mcp.resource` with the name of the tool that
/// answers the resource read; whether it worked, how long it took, and the
/// account.
pub struct Event<'a> {
pub name: &'static str,
pub subject: &'a str,
pub ok: bool,
pub took: Duration,
pub did: Option<&'a str>,
/// Why it failed, as it answered, if it did and said.
pub error: Option<&'a str>,
}
/// Where events go.
pub type Events = Arc<dyn Fn(Event) + Send + Sync>;
/// An assistant's authorization request, carried through the sign-in.
#[derive(Serialize, Deserialize)]
struct Ask {
client_id: String,
name: String,
redirect_uri: String,
state: Option<String>,
challenge: String,
}
/// A registered assistant, which its client ID holds.
#[derive(Serialize, Deserialize)]
struct Client {
redirect_uris: Vec<String>,
name: String,
/// Whether it proves itself with a secret.
secret: bool,
}
/// Whom a token is for.
#[derive(Serialize, Deserialize)]
struct Grant {
sub: String,
client_id: String,
}
struct Pending {
ask: Ask,
request: oauth::Request,
started: Instant,
}
struct Code {
ask: Ask,
did: String,
issued: Instant,
}
/// The account a request's token is for, as the MCP service sees it.
#[derive(Clone)]
struct Caller(String);
/// An OAuth error, as JSON.
struct Refusal(StatusCode, &'static str, String);
impl IntoResponse for Refusal {
fn into_response(self) -> Response {
(
self.0,
[(header::CACHE_CONTROL, "no-store")],
Json(json!({ "error": self.1, "error_description": self.2 })),
)
.into_response()
}
}
fn invalid_grant(why: &str) -> Refusal {
Refusal(StatusCode::BAD_REQUEST, "invalid_grant", why.into())
}
impl Mcp {
/// The server at `url`, within `limits`, keeping what it
/// needs under `ctx`'s config and cache directories.
pub fn new(ctx: &Ctx, url: &str, limits: Limits) -> anyhow::Result<Arc<Mcp>> {
let mut env = ctx.env.clone();
env.config_dir = env.config_dir.join("mcp");
// Public repositories are read from the server's own mirrors, one
// for everyone, and private ones each account keeps.
env.public_mirrors = Some(env.cache_dir.join("mirrors"));
env.cache_dir = env.cache_dir.join("mcp");
let ctx = Ctx {
env,
http: ctx.http.clone(),
};
let key = oauth::client_key(&ctx)?;
Ok(Arc::new(Mcp {
ctx,
url: url.trim_end_matches('/').to_string(),
limits,
key,
pending: Mutex::default(),
codes: Mutex::default(),
events: std::sync::OnceLock::new(),
refused: std::sync::OnceLock::new(),
calls: Mutex::default(),
looked: Mutex::default(),
}))
}
/// Turns away the accounts `refused` says from now on.
pub fn refuse(&self, refused: Refused) {
let _ = self.refused.set(refused);
}
/// Whether `did` may use this server: any account not turned away.
fn lets_in(&self, did: &str) -> bool {
!self.refused.get().is_some_and(|refused| refused(did))
}
/// A turn for `did` to call a tool at `now`, one that `waits` or not,
/// or why it may not: it has as many running, or waiting, as it may, or
/// has made as many this minute.
fn turn<'a>(&'a self, did: &'a str, waits: bool, now: Instant) -> Result<Turn<'a>, String> {
let mut calls = self.calls.lock().unwrap();
calls.retain(|_, c| c.running > 0 || now.duration_since(c.since) < MINUTE);
let fresh = Calls {
running: 0,
waiting: 0,
made: 0,
since: now,
};
let c = calls.entry(did.to_string()).or_insert(fresh);
if now.duration_since(c.since) >= MINUTE {
(c.made, c.since) = (0, now);
}
let Limits {
running,
waiting,
per_minute,
..
} = self.limits;
if c.running >= running {
return Err(format!(
"This account has {running} tool calls running, as many as it may: try again once one ends"
));
}
if waits && c.waiting >= waiting {
return Err(format!(
"This account has {waiting} tool calls waiting, as many as it may: try again once one ends, or without wait"
));
}
if c.made >= per_minute {
return Err(format!(
"This account has made {per_minute} tool calls this minute, as many as it may: try again in a minute"
));
}
c.running += 1;
c.waiting += usize::from(waits);
c.made += 1;
Ok(Turn {
mcp: self,
did,
waits,
})
}
/// Tells `events` of each sign-in and tool call from now on.
pub fn observe(&self, events: Events) {
let _ = self.events.set(events);
}
fn tell(&self, event: Event) {
self.events.get().inspect(|events| events(event));
}
fn resource(&self) -> String {
format!("{}{MCP}", self.url)
}
/// Acting as `did`: its own config, sessions and mirrors, and never a
/// password.
fn ctx_for(&self, did: &str) -> Ctx {
let mut env = self.ctx.env.clone();
let dir = did.replace(':', "_");
env.config_dir = env.config_dir.join(&dir);
env.cache_dir = env.cache_dir.join(&dir);
env.identifier = Some(did.to_string());
env.password = None;
Ctx {
env,
http: self.ctx.http.clone(),
}
}
/// Notes that a token was made for `did` now, and forgets the accounts
/// none has been made for since their last refresh token ran out.
fn used(&self, did: &str) {
let _ = std::fs::write(self.ctx_for(did).env.config_dir.join(USED), "");
let mut looked = self.looked.lock().unwrap();
if looked.is_none_or(|at| at.elapsed() >= FORGET) {
*looked = Some(Instant::now());
self.forget(SystemTime::now());
}
}
/// Forgets each account no token was made for in the [`REFRESH_TTL`]
/// before `now`, whose tokens have all run out: its session and its
/// mirrors. One with no note of it is dated by its directory.
fn forget(&self, now: SystemTime) {
let env = &self.ctx.env;
let dirs = std::fs::read_dir(&env.config_dir).into_iter().flatten();
let accounts = dirs
.flatten()
.filter(|d| d.file_name().to_string_lossy().starts_with("did_"));
for dir in accounts {
let path = dir.path();
let used = std::fs::metadata(path.join(USED)).or_else(|_| dir.metadata());
let at = used.and_then(|m| m.modified()).unwrap_or(now);
let idle = now.duration_since(at).unwrap_or_default();
if idle >= Duration::from_secs(REFRESH_TTL) {
let _ = std::fs::remove_dir_all(&path);
let _ = std::fs::remove_dir_all(env.cache_dir.join(dir.file_name()));
}
}
}
/// Whether another sign-in may go to a PDS, with those that took too
/// long forgotten.
fn room(&self) -> bool {
let mut pending = self.pending.lock().unwrap();
pending.retain(|_, p| p.started.elapsed() < SIGN_IN);
pending.len() < self.limits.signing_in
}
/// The client this server is to PDSes.
fn atproto(&self) -> oauth::Client {
let key = P256Key::from_bytes(&self.key.to_bytes()).expect("a key reads back");
oauth::Client::confidential(&format!("{}{METADATA}", self.url), key)
}
fn sign(&self, kind: &str, ttl: u64, body: &impl Serialize) -> String {
let mut claims = serde_json::to_value(body).expect("claims serialize");
claims["kind"] = kind.into();
claims["iss"] = self.url.clone().into();
claims["exp"] = (keys::now_micros() / 1_000_000 + ttl).into();
self.key.sign_jwt(&json!({ "typ": "JWT" }), &claims)
}
/// What this server signed as `kind`, unexpired.
fn open<T: DeserializeOwned>(&self, kind: &str, token: &str) -> Option<T> {
let claims: Value = crypto::verify_jwt(token, &self.key.public()).ok()?;
(claims["kind"] == kind).then_some(())?;
serde_json::from_value(claims).ok()
}
fn client(&self, id: &str) -> Option<Client> {
self.open("client", id)
}
/// `redirect`, with `values`, the state and who answers.
fn back(&self, redirect: &str, state: Option<&str>, values: &[(&str, &str)]) -> String {
let mut url = reqwest::Url::parse(redirect).expect("a registered redirect is a URL");
{
let mut query = url.query_pairs_mut();
query.extend_pairs(values);
query.extend_pairs(state.map(|s| ("state", s)));
query.append_pair("iss", &self.url);
}
url.to_string()
}
/// Keeps the session to act as its account with.
async fn keep(&self, mut session: Session) -> anyhow::Result<()> {
let ctx = self.ctx_for(&session.did);
// A refresh signs as this client, with the key in the session's
// config directory.
oauth::set_client_key(&ctx, &self.key)?;
session.handle = identity::handle(&self.ctx, &session.did)
.await
.unwrap_or_default();
session::save(&ctx, &session)
}
/// The PDS's answer to a sign-in, traded for a code for the assistant,
/// and whose it is.
async fn finish(
&self,
pending: Pending,
params: &HashMap<String, String>,
) -> anyhow::Result<(String, String)> {
let session = pending
.request
.redeem(&self.ctx, &self.atproto(), params)
.await?;
let did = session.did.clone();
self.keep(session).await?;
Ok((self.code(pending.ask, did.clone()), did))
}
fn code(&self, ask: Ask, did: String) -> String {
let code = oauth::random(32);
let mut codes = self.codes.lock().unwrap();
codes.retain(|_, c| c.issued.elapsed() < CODE);
codes.insert(
code.clone(),
Code {
ask,
did,
issued: Instant::now(),
},
);
code
}
/// The client a token request is from, with its secret if it has one.
fn authenticate(&self, form: &HashMap<String, String>) -> Result<String, Refusal> {
let id = form.get("client_id").cloned().unwrap_or_default();
let unknown = || {
Refusal(
StatusCode::UNAUTHORIZED,
"invalid_client",
"Unknown client".into(),
)
};
let client = self.client(&id).ok_or_else(unknown)?;
let proved = !client.secret
|| form
.get("client_secret")
.and_then(|s| URL_SAFE_NO_PAD.decode(s).ok())
.is_some_and(|sig| self.key.public().verify(id.as_bytes(), &sig));
proved.then_some(id).ok_or_else(unknown)
}
fn redeem(&self, client_id: &str, form: &HashMap<String, String>) -> Result<String, Refusal> {
let code = form.get("code").cloned().unwrap_or_default();
let mut codes = self.codes.lock().unwrap();
codes.retain(|_, c| c.issued.elapsed() < CODE);
let code = codes
.remove(&code)
.filter(|c| c.ask.client_id == client_id)
.ok_or_else(|| invalid_grant("Invalid or expired code"))?;
let ok = form
.get("redirect_uri")
.is_none_or(|r| *r == code.ask.redirect_uri)
&& form
.get("code_verifier")
.is_some_and(|v| oauth::sha256(v) == code.ask.challenge);
ok.then_some(code.did)
.ok_or_else(|| invalid_grant("The redirect_uri or code_verifier does not match"))
}
/// The account `form`'s refresh token is `client_id`'s for. One whose
/// login with its PDS is over gets no more tokens, so that its assistant
/// has its user sign in again.
async fn refreshing(
&self,
client_id: &str,
form: &HashMap<String, String>,
) -> Result<String, Refusal> {
let did = form
.get("refresh_token")
.and_then(|t| self.open::<Grant>("refresh", t))
.filter(|g| g.client_id == client_id)
.map(|g| g.sub)
.ok_or_else(|| invalid_grant("Invalid refresh token"))?;
if session::ended(&self.ctx_for(&did), &did).await {
return Err(invalid_grant("The login with the account's PDS is over"));
}
Ok(did)
}
async fn tokens(&self, form: &HashMap<String, String>) -> Result<Value, Refusal> {
let client_id = self.authenticate(form)?;
let did = match form.get("grant_type").map(String::as_str) {
Some("authorization_code") => self.redeem(&client_id, form)?,
Some("refresh_token") => self.refreshing(&client_id, form).await?,
other => {
return Err(Refusal(
StatusCode::BAD_REQUEST,
"unsupported_grant_type",
format!("Unsupported grant_type {other:?}"),
));
}
};
if !self.lets_in(&did) {
return Err(invalid_grant(&format!(
"{did} may no longer use this server"
)));
}
self.used(&did);
let grant = Grant {
sub: did,
client_id,
};
let mut access = serde_json::to_value(&grant).expect("a grant serializes");
access["aud"] = self.resource().into();
Ok(json!({
"access_token": self.sign("access", ACCESS_TTL, &access),
"token_type": "Bearer",
"expires_in": ACCESS_TTL,
"refresh_token": self.sign("refresh", REFRESH_TTL, &grant),
}))
}
}
/// The routes: `/mcp`, and the OAuth around it.
pub fn router(mcp: Arc<Mcp>) -> Router {
let tools = mcp.clone();
let service = StreamableHttpService::new(
move || {
Ok(Remote {
tools: Server::new(tools.ctx.clone()),
mcp: tools.clone(),
})
},
Arc::new(NeverSessionManager::default()),
// Every request carries a token, which a rebound DNS name cannot.
StreamableHttpServerConfig::default()
.with_legacy_session_mode(false)
.with_json_response(true)
.disable_allowed_hosts(),
);
let guarded = Router::new()
.route_service(MCP, service)
.layer(middleware::from_fn_with_state(mcp.clone(), bearer));
Router::new()
.route("/.well-known/oauth-protected-resource", get(resource))
.route("/.well-known/oauth-protected-resource/mcp", get(resource))
.route(
"/.well-known/oauth-authorization-server",
get(server_metadata),
)
.route(METADATA, get(client_metadata))
.route("/mcp/register", post(register))
.route("/mcp/authorize", get(authorize).post(sign_in))
.route(CALLBACK, get(callback))
.route("/mcp/token", post(token))
.with_state(mcp)
.merge(guarded)
}
/// Lets in a request with a token for an account still on the list.
async fn bearer(State(mcp): State<Arc<Mcp>>, mut req: Request, next: Next) -> Response {
let grant = req
.headers()
.get(header::AUTHORIZATION)
.and_then(|v| v.to_str().ok())
.and_then(|v| v.strip_prefix("Bearer "))
.and_then(|t| mcp.open::<Grant>("access", t))
.filter(|g| mcp.lets_in(&g.sub));
let ended = match &grant {
Some(g) => session::ended(&mcp.ctx_for(&g.sub), &g.sub).await,
None => false,
};
// With no login left, the tools would read as no one, and a private
// repository would seem not to exist: this has the user sign in again.
let Some(grant) = grant.filter(|_| !ended) else {
let metadata = format!("{}/.well-known/oauth-protected-resource{MCP}", mcp.url);
return (
StatusCode::UNAUTHORIZED,
[(
header::WWW_AUTHENTICATE,
format!("Bearer resource_metadata=\"{metadata}\""),
)],
)
.into_response();
};
req.extensions_mut().insert(Caller(grant.sub));
next.run(req).await
}
/// The tools and resources, each call made as the account its token is for.
#[derive(Clone)]
struct Remote {
mcp: Arc<Mcp>,
tools: Server,
}
/// How a call of one kind is logged and counted: its event's name, and
/// its log line's when it worked and when it failed.
struct Kind {
event: &'static str,
said: [&'static str; 2],
}
const TOOL: Kind = Kind {
event: "mcp.tool",
said: ["mcp tool", "mcp tool failed"],
};
const RESOURCE: Kind = Kind {
event: "mcp.resource",
said: ["mcp resource", "mcp resource failed"],
};
/// Whom a call is made as, and the request's trace, handed on: rmcp
/// answers in a task of its own, and what the call does is in the trace.
struct As {
did: String,
trace: String,
}
impl As {
fn of(context: &RequestContext<RoleServer>) -> As {
let parts = context.extensions.get::<Parts>();
let did = parts
.and_then(|p| p.extensions.get::<Caller>())
.map(|c| c.0.clone())
.unwrap_or_default();
let trace = parts.and_then(|p| p.extensions.get::<telemetry::Trace>());
let trace = telemetry::Trace::or_new(trace);
As { did, trace }
}
}
impl Remote {
/// A line a call, named for its tool, in the request's trace: at warn
/// when it failed, or was slow and not asked to wait; and the
/// operator's counts told of it.
async fn told(
&self,
kind: &Kind,
who: As,
tool: &str,
started: Instant,
waits: bool,
failure: Option<String>,
) {
let took = started.elapsed();
let f = Fields {
did: Some(&who.did),
tool: Some(tool),
error: failure.clone(),
duration_ms: Some(telemetry::millis(took)),
..Fields::default()
};
let warn = failure.is_some() || (telemetry::slow(took) && !waits);
let level = [Level::INFO, Level::WARN][usize::from(warn)];
let message = kind.said[usize::from(failure.is_some())];
telemetry::traced(who.trace, async { telemetry::note(level, message, &f) }).await;
self.mcp.tell(Event {
name: kind.event,
subject: tool,
ok: failure.is_none(),
took,
did: Some(&who.did),
error: failure.as_deref(),
});
}
}
impl ServerHandler for Remote {
fn get_info(&self) -> ServerConfig {
self.tools.get_info()
}
fn get_tool(&self, name: &str) -> Option<Tool> {
self.tools.get_tool(name)
}
async fn list_tools(
&self,
request: Option<PaginatedRequestParams>,
context: RequestContext<RoleServer>,
) -> Result<ListToolsResult, McpError> {
self.tools.list_tools(request, context).await
}
async fn call_tool(
&self,
request: CallToolRequestParams,
context: RequestContext<RoleServer>,
) -> Result<CallToolResponse, McpError> {
let who = As::of(&context);
let server = Server::new(self.mcp.ctx_for(&who.did));
let (started, tool) = (Instant::now(), request.name.to_string());
let waits = request
.arguments
.as_ref()
.and_then(|a| a.get("wait"))
.and_then(serde_json::Value::as_u64)
.unwrap_or_default()
> 0;
let res = match self.mcp.turn(&who.did, waits, started) {
Ok(_turn) => {
let called = server.call_tool(request, context);
telemetry::traced(who.trace.clone(), called).await
}
Err(why) => Ok(CallToolResponse::Complete(CallToolResult::error(vec![
ContentBlock::text(why),
]))),
};
// A tool that fails answers as a result that says so; a call no
// tool takes, as an error.
let failure = match &res {
Ok(CallToolResponse::Complete(r)) if r.is_error == Some(true) => {
let text = r.content.iter().filter_map(|c| c.as_text());
Some(text.map(|t| t.text.as_str()).collect::<Vec<_>>().join("\n"))
}
Ok(_) => None,
Err(err) => Some(err.message.to_string()),
};
self.told(&TOOL, who, &tool, started, waits, failure).await;
res
}
async fn list_resource_templates(
&self,
request: Option<PaginatedRequestParams>,
context: RequestContext<RoleServer>,
) -> Result<ListResourceTemplatesResult, McpError> {
self.tools.list_resource_templates(request, context).await
}
/// A resource read as a tool call is made, and counted as one, named
/// for the tool that answers it.
async fn read_resource(
&self,
request: ReadResourceRequestParams,
context: RequestContext<RoleServer>,
) -> Result<ReadResourceResponse, McpError> {
let who = As::of(&context);
let server = Server::new(self.mcp.ctx_for(&who.did));
let started = Instant::now();
let tool = crate::resources::parse(&request.uri).map_or("", |(tool, _)| tool);
let res = match self.mcp.turn(&who.did, false, started) {
Ok(_turn) => {
let read = server.read_resource(request, context);
telemetry::traced(who.trace.clone(), read).await
}
Err(why) => Err(McpError::invalid_request(why, None)),
};
let failure = res.as_ref().err().map(|err| err.message.to_string());
self.told(&RESOURCE, who, tool, started, false, failure)
.await;
res
}
}
async fn resource(State(mcp): State<Arc<Mcp>>) -> Json<Value> {
Json(json!({
"resource": mcp.resource(),
"authorization_servers": [mcp.url],
"bearer_methods_supported": ["header"],
"resource_name": "knotless",
}))
}
async fn server_metadata(State(mcp): State<Arc<Mcp>>) -> Json<Value> {
let url = &mcp.url;
Json(json!({
"issuer": url,
"authorization_endpoint": format!("{url}/mcp/authorize"),
"token_endpoint": format!("{url}/mcp/token"),
"registration_endpoint": format!("{url}/mcp/register"),
"response_types_supported": ["code"],
"grant_types_supported": ["authorization_code", "refresh_token"],
"code_challenge_methods_supported": ["S256"],
"token_endpoint_auth_methods_supported": ["none", "client_secret_post"],
"authorization_response_iss_parameter_supported": true,
}))
}
/// This server as a client of PDSes, which its client ID names.
async fn client_metadata(State(mcp): State<Arc<Mcp>>) -> Json<Value> {
let client = mcp.atproto();
Json(json!({
"client_id": client.id,
"client_name": "knotless MCP",
"application_type": "web",
"grant_types": ["authorization_code", "refresh_token"],
"response_types": ["code"],
"redirect_uris": [format!("{}{CALLBACK}", mcp.url)],
"scope": oauth::SCOPE.as_str(),
"token_endpoint_auth_method": "private_key_jwt",
"token_endpoint_auth_signing_alg": "ES256",
"dpop_bound_access_tokens": true,
"jwks": client.jwks(),
}))
}
/// Where an assistant may be sent back to: https, or the machine it runs on.
fn redirect_allowed(uri: &str) -> bool {
reqwest::Url::parse(uri).is_ok_and(|u| {
u.fragment().is_none()
&& (u.scheme() == "https"
|| u.scheme() == "http"
&& matches!(u.host_str(), Some("localhost" | "127.0.0.1" | "[::1]")))
})
}
/// Registers an assistant (RFC 7591). A client that asks to prove itself
/// gets a secret, sent in the token request's body.
async fn register(State(mcp): State<Arc<Mcp>>, Json(m): Json<Value>) -> Response {
let uris: Vec<String> = m["redirect_uris"]
.as_array()
.into_iter()
.flatten()
.filter_map(|u| u.as_str().map(String::from))
.collect();
if uris.is_empty() || !uris.iter().all(|u| redirect_allowed(u)) {
return Refusal(
StatusCode::BAD_REQUEST,
"invalid_redirect_uri",
"Each redirect_uri must be https, or http on localhost".into(),
)
.into_response();
}
let name = m["client_name"].as_str().unwrap_or("An assistant").trim();
let client = Client {
redirect_uris: uris,
name: name.chars().take(60).collect(),
secret: m["token_endpoint_auth_method"].as_str().unwrap_or("none") != "none",
};
let id = mcp.sign("client", CLIENT_TTL, &client);
let mut out = json!({
"client_id": id,
"client_id_issued_at": keys::now_micros() / 1_000_000,
"client_name": client.name,
"redirect_uris": client.redirect_uris,
"grant_types": ["authorization_code", "refresh_token"],
"response_types": ["code"],
"token_endpoint_auth_method": if client.secret { "client_secret_post" } else { "none" },
});
if client.secret {
out["client_secret"] = URL_SAFE_NO_PAD.encode(mcp.key.sign(id.as_bytes())).into();
out["client_secret_expires_at"] = 0.into();
}
(StatusCode::CREATED, Json(out)).into_response()
}
/// Where the assistant sends its user: checks what it asks, then asks the
/// user who they are.
async fn authorize(
State(mcp): State<Arc<Mcp>>,
Query(q): Query<HashMap<String, String>>,
) -> Response {
let Some(client) = q.get("client_id").and_then(|id| mcp.client(id)) else {
return page(
StatusCode::BAD_REQUEST,
"<p>This assistant didn't say who it is, or isn't registered here. Go back to it and connect again.",
);
};
let redirect = q
.get("redirect_uri")
.or(client
.redirect_uris
.first()
.filter(|_| client.redirect_uris.len() == 1))
.filter(|r| client.redirect_uris.contains(r));
let Some(redirect) = redirect else {
return page(
StatusCode::BAD_REQUEST,
"<p>This assistant asked to be sent back somewhere it didn't register. Go back to it and connect again.",
);
};
let state = q.get("state").cloned();
let refuse = |error: &str, why: &str| {
Redirect::to(&mcp.back(
redirect,
state.as_deref(),
&[("error", error), ("error_description", why)],
))
.into_response()
};
if q.get("response_type").map(String::as_str) != Some("code") {
return refuse("unsupported_response_type", "Only code");
}
let challenge = q
.get("code_challenge")
.filter(|_| q.get("code_challenge_method").map(String::as_str) == Some("S256"));
let Some(challenge) = challenge else {
return refuse("invalid_request", "PKCE with S256 is required");
};
if q.get("resource")
.is_some_and(|r| r.trim_end_matches('/') != mcp.resource())
{
return refuse(
"invalid_target",
&format!("This server is {}", mcp.resource()),
);
}
let ask = Ask {
client_id: q["client_id"].clone(),
name: client.name,
redirect_uri: redirect.clone(),
state,
challenge: challenge.clone(),
};
let signed = mcp.sign("ask", SIGN_IN.as_secs(), &ask);
page(StatusCode::OK, &form(&ask.name, &signed, None))
}
#[derive(Deserialize)]
struct SignIn {
ask: String,
handle: String,
}
/// The user's handle: sends them to their PDS to sign in, if the account
/// may use this server.
async fn sign_in(State(mcp): State<Arc<Mcp>>, Form(f): Form<SignIn>) -> Response {
let Some(ask) = mcp.open::<Ask>("ask", &f.ask) else {
return page(
StatusCode::BAD_REQUEST,
"<p>This sign-in took too long. Go back to the assistant and connect again.",
);
};
if !mcp.room() {
return page(
StatusCode::SERVICE_UNAVAILABLE,
"<p>Too many sign-ins are under way here. Try again in a few minutes.",
);
}
let handle = f.handle.trim().trim_start_matches('@').to_lowercase();
let Ok(id) = identity::resolve(&mcp.ctx, &handle).await else {
let problem = format!("No account answers to {handle}.");
return page(
StatusCode::BAD_REQUEST,
&form(&ask.name, &f.ask, Some(&problem)),
);
};
if !mcp.lets_in(&id.did) {
let why = format!(
"<p>{} can't use this server's MCP: it is taken down here.",
escape(&handle)
);
return page(StatusCode::FORBIDDEN, &why);
}
let callback = format!("{}{CALLBACK}", mcp.url);
let request = oauth::authorize(
&mcp.ctx,
&mcp.atproto(),
id,
&handle,
&callback,
&oauth::SCOPE,
)
.await;
let request = match request {
Ok(request) => request,
Err(err) => {
let why = format!(
"<p>{}'s server didn't take the sign-in: {}",
escape(&handle),
escape(&format!("{err:#}"))
);
return page(StatusCode::BAD_GATEWAY, &why);
}
};
let url = request.url.clone();
let mut pending = mcp.pending.lock().unwrap();
pending.insert(
request.state.clone(),
Pending {
ask,
request,
started: Instant::now(),
},
);
Redirect::to(&url).into_response()
}
/// Where the user's PDS sends them back: keeps the session, and sends them
/// on to the assistant with a code, or why there is none.
async fn callback(
State(mcp): State<Arc<Mcp>>,
Query(params): Query<HashMap<String, String>>,
) -> Response {
let state = params.get("state").cloned().unwrap_or_default();
let Some(pending) = mcp.pending.lock().unwrap().remove(&state) else {
return page(
StatusCode::BAD_REQUEST,
"<p>No sign-in is waiting for this. Go back to the assistant and connect again.",
);
};
let redirect = pending.ask.redirect_uri.clone();
let state = pending.ask.state.clone();
let started = pending.started;
let finished = mcp.finish(pending, ¶ms).await;
let did = finished.as_ref().ok().map(|(_, did)| did.as_str());
mcp.tell(Event {
name: "mcp.signin",
subject: "",
ok: did.is_some(),
took: started.elapsed(),
did,
error: None,
});
let to = match finished {
Ok((code, _)) => mcp.back(&redirect, state.as_deref(), &[("code", &code)]),
Err(err) => mcp.back(
&redirect,
state.as_deref(),
&[
("error", "access_denied"),
("error_description", &format!("{err:#}")),
],
),
};
Redirect::to(&to).into_response()
}
async fn token(State(mcp): State<Arc<Mcp>>, Form(form): Form<HashMap<String, String>>) -> Response {
match mcp.tokens(&form).await {
Ok(tokens) => ([(header::CACHE_CONTROL, "no-store")], Json(tokens)).into_response(),
Err(refusal) => refusal.into_response(),
}
}
fn escape(text: &str) -> String {
text.replace('&', "&")
.replace('<', "<")
.replace('>', ">")
.replace('"', """)
}
/// Asks who the user is, for the assistant `name`.
fn form(name: &str, ask: &str, problem: Option<&str>) -> String {
let name = escape(name);
let problem = problem
.map(|p| format!("<p class=problem>{}", escape(p)))
.unwrap_or_default();
format!(
"<h1>Connect {name} to knotless</h1>\
<p>{name} asks to work on your repositories as you. Sign in with your atproto account to let it.\
{problem}\
<form method=post action=/mcp/authorize>\
<input type=hidden name=ask value=\"{ask}\">\
<label>Handle <input name=handle autocomplete=username placeholder=you.bsky.social required autofocus></label>\
<button>Sign in</button></form>"
)
}
fn page(status: StatusCode, body: &str) -> Response {
let html = format!(
"<!doctype html><meta charset=utf-8><meta name=viewport content=\"width=device-width\">\
<title>knotless</title><style>\
body{{font:16px/1.5 system-ui,sans-serif;max-width:32rem;margin:3rem auto;padding:0 1rem;color:#222;background:#fff}}\
@media(prefers-color-scheme:dark){{body{{color:#ddd;background:#161616}}}}\
label,input,button{{display:block;font:inherit;margin:.5rem 0}}input{{width:100%;padding:.4rem}}\
.problem{{color:#b00}}</style>{body}"
);
(status, Html(html)).into_response()
}
#[cfg(test)]
mod tests {
use super::*;
fn mcp(dir: &std::path::Path, limits: Limits) -> Arc<Mcp> {
let home = dir.to_string_lossy().to_string();
let ctx = Ctx::new(knotless::env::Env::from_vars(|name| {
(name == "HOME").then(|| home.clone())
}));
Mcp::new(&ctx, "https://mcp.test", limits).unwrap()
}
#[test]
fn each_account_has_its_share() {
let dir = tempfile::tempdir().unwrap();
let limits = Limits {
signing_in: 1,
running: 3,
waiting: 1,
per_minute: 5,
};
let mcp = mcp(dir.path(), limits);
let now = Instant::now();
// So many waiting, and as many at once as it may; others still
// have theirs.
let waiting = mcp.turn("did:a", true, now).unwrap();
let more = mcp.turn("did:a", true, now).err().unwrap();
assert!(more.contains("1 tool calls waiting"), "{more}");
let (one, two) = (mcp.turn("did:a", false, now), mcp.turn("did:a", false, now));
let busy = mcp.turn("did:a", false, now).err().unwrap();
assert!(busy.contains("3 tool calls running"), "{busy}");
assert!(mcp.turn("did:b", true, now).is_ok());
// One that ends gives its turn back, but not the minute's count:
// four made, and those refused not counted, so a fifth.
drop((waiting, one, two));
drop(mcp.turn("did:a", true, now).unwrap());
drop(mcp.turn("did:a", false, now).unwrap());
let spent = mcp.turn("did:a", false, now).err().unwrap();
assert!(spent.contains("5 tool calls this minute"), "{spent}");
// A minute on, the count starts again, and what has gone quiet is
// forgotten.
let later = now + MINUTE;
let held = mcp.turn("did:a", false, later).unwrap();
assert_eq!(mcp.calls.lock().unwrap().len(), 1);
// One still running is kept, its count started again.
let after = later + MINUTE;
drop((mcp.turn("did:a", false, after).unwrap(), held));
assert_eq!(mcp.calls.lock().unwrap()["did:a"].made, 1);
}
#[test]
fn accounts_gone_quiet_are_forgotten() {
let dir = tempfile::tempdir().unwrap();
let mcp = mcp(dir.path(), Limits::default());
let (config, cache) = (&mcp.ctx.env.config_dir, &mcp.ctx.env.cache_dir);
for did in ["did:quiet", "did:new", "did:unnoted"] {
let dir = mcp.ctx_for(did).env;
std::fs::create_dir_all(&dir.config_dir).unwrap();
std::fs::create_dir_all(dir.cache_dir.join("mirrors")).unwrap();
}
std::fs::create_dir_all(config.join("other")).unwrap();
mcp.used("did:quiet");
mcp.used("did:new");
assert!(config.join("did_new").join(USED).exists());
// Its last token's refresh run out, an account is forgotten, and
// one never noted is dated by its directory.
let ttl = Duration::from_secs(REFRESH_TTL);
let old = SystemTime::now() - ttl - MINUTE;
let quiet = std::fs::File::open(config.join("did_quiet").join(USED)).unwrap();
quiet.set_modified(old).unwrap();
mcp.forget(SystemTime::now());
assert!(!config.join("did_quiet").exists());
assert!(!cache.join("did_quiet").exists());
assert!(config.join("did_new").exists() && cache.join("did_new").exists());
assert!(config.join("did_unnoted").exists() && config.join("other").exists());
// A season on, every one is.
mcp.forget(SystemTime::now() + ttl);
assert!(!config.join("did_unnoted").exists() && !config.join("did_new").exists());
assert!(config.join("other").exists());
// Looked for as tokens are made, once an hour.
for did in ["did_new", "did_unnoted"] {
std::fs::create_dir_all(config.join(did)).unwrap();
}
let unnoted = std::fs::File::open(config.join("did_unnoted")).unwrap();
unnoted.set_modified(old).unwrap();
*mcp.looked.lock().unwrap() = None;
mcp.used("did:new");
assert!(config.join("did_new").exists());
assert!(!config.join("did_unnoted").exists());
}
}