clover's git

home-infra

1	use crate::*;
2	use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
3	use futures::SinkExt;
4	use mcp::{delete, get, hash, list, put, secret};
5	use rmcp::{
6	ErrorData, RoleServer, ServerHandler,
7	model::{
8	CacheScope, CallToolRequestParams, CallToolResponse, CallToolResult, ContentBlock,
9	ListToolsResult, PaginatedRequestParams, ServerCapabilities, ServerConfig, Tool,
10	ToolAnnotations,
11	},
12	service::RequestContext,
13	};
14	use tokio::sync::{mpsc, oneshot};
15
16	pub struct Broker {
17	connections: Mutex<HashMap<String, Arc<Connection>>>,
18	pending_slots: Semaphore,
19	pub(crate) changes: watch::Sender<()>,
20	}
21	impl Default for Broker {
22	fn default() -> Self {
23	Self {
24	connections: Mutex::new(HashMap::new()),
25	pending_slots: Semaphore::new(128),
26	changes: watch::channel(()).0,
27	}
28	}
29	}
30	struct Connection {
31	messages: mpsc::Sender<Message>,
32	pending: Mutex<HashMap<String, oneshot::Sender<Result<Value>>>>,
33	closed: watch::Sender<bool>,
34	}
35	struct Pending {
36	connection: Arc<Connection>,
37	id: String,
38	}
39	impl Drop for Pending {
40	fn drop(&mut self) {
41	self.connection.pending.lock().unwrap().remove(&self.id);
42	}
43	}
44	fn unknown() -> Error {
45	Error::new(
46	409,
47	"Command outcome is unknown. Read the thread before retrying.",
48	)
49	}
50	impl Broker {
51	pub fn disconnect(&self, id: &str) {
52	if let Some(connection) = self.connections.lock().unwrap().remove(id) {
53	self.changes.send_replace(());
54	connection.closed.send_replace(true);
55	for (_, response) in connection.pending.lock().unwrap().drain() {
56	let _ = response.send(Err(unknown()));
57	}
58	}
59	}
60	fn remove(&self, id: &str, connection: &Arc<Connection>) {
61	let mut connections = self.connections.lock().unwrap();
62	if connections
63	.get(id)
64	.is_some_and(|current| Arc::ptr_eq(current, connection))
65	{
66	connections.remove(id);
67	self.changes.send_replace(());
68	}
69	connection.closed.send_replace(true);
70	for (_, response) in connection.pending.lock().unwrap().drain() {
71	let _ = response.send(Err(unknown()));
72	}
73	}
74	pub fn view(&self, machines: Vec<Value>, grant: Option<&Value>) -> Vec<Value> {
75	let connections = self.connections.lock().unwrap();
76	machines
77	.into_iter()
78	.filter(|machine| grant.is_none_or(|g| array(&g["resources"]).contains(&machine["id"])))
79	.map(|mut machine| {
80	machine.as_object_mut().unwrap().remove("tokenHash");
81	machine.as_object_mut().unwrap().remove("user");
82	machine["online"] = json!(connections.contains_key(string(&machine["id"])));
83	if let Some(grant) = grant {
84	machine["selected"] = json!(array(&grant["targets"]).contains(&machine["id"]));
85	}
86	machine
87	})
88	.collect()
89	}
90	async fn dispatch(
91	&self,
92	store: &mcp::Store,
93	grant_id: &str,
94	machine_id: &str,
95	method: &str,
96	params: Value,
97	deadline: Duration,
98	) -> Result<Value> {
99	let params = command(method, params)?;
100	let (pending, receive, slot) = {
101	let db = store.db.lock().unwrap();
102	let grant = get(&db, &format!("grant:{grant_id}"))?;
103	let machine = get(&db, &format!("machine:{machine_id}"))?;
104	if grant.is_null()
105	|| grant["resource"] != store.resource("agents")
106	|| machine.is_null()
107	|| machine["user"] != grant["user"]
108	|| !array(&grant["resources"]).iter().any(|id| id == machine_id)
109	{
110	return Err(Error::new(
111	403,
112	"Choose a machine granted to this connection.",
113	));
114	}
115	let scope = if matches!(method, "send_message" | "interrupt_thread" | "start_thread") {
116	"sessions:write"
117	} else {
118	"sessions:read"
119	};
120	if !array(&grant["scopes"]).iter().any(|s| s == scope) {
121	return Err(Error::new(
122	403,
123	"This connection has read access only. Connect again to request control.",
124	));
125	}
126	let slot = self.pending_slots.try_acquire().map_err(|_| {
127	Error::new(
128	429,
129	"The relay has too many pending commands. Try again shortly.",
130	)
131	})?;
132	let connection = self
133	.connections
134	.lock()
135	.unwrap()
136	.get(machine_id)
137	.cloned()
138	.ok_or_else(|| Error::new(409, "Machine is offline. Start its local agent."))?;
139	let (send, receive) = oneshot::channel();
140	let id = uuid::Uuid::new_v4().to_string();
141	{
142	let mut pending = connection.pending.lock().unwrap();
143	if *connection.closed.borrow() {
144	return Err(Error::new(
145	409,
146	"Machine is offline. Start its local agent.",
147	));
148	}
149	if pending.len() >= 8 {
150	return Err(Error::new(
151	429,
152	"Machine has eight pending commands. Wait for one to finish.",
153	));
154	}
155	pending.insert(id.clone(), send);
156	}
157	let pending = Pending {
158	connection: connection.clone(),
159	id: id.clone(),
160	};
161	let frame = json!({"id":id,"method":method,"params":params}).to_string();
162	if frame.len() > 256 * 1024 {
163	return Err(Error::new(
164	413,
165	"Send a shorter message. Local agent commands are limited to 256 KB.",
166	));
167	}
168	let frame = Message::Text(frame.into());
169	if connection.messages.try_send(frame).is_err() {
170	return Err(unknown());
171	}
172	(pending, receive, slot)
173	};
174	let result = tokio::time::timeout(deadline, receive)
175	.await
176	.map_err(|_| unknown())?
177	.map_err(|_| unknown())?;
178	drop(pending);
179	drop(slot);
180	result
181	}
182	}
183	pub(crate) fn machines(db: &rusqlite::Connection, user: &str) -> Result<Vec<Value>> {
184	Ok(list(db, "machine:")?
185	.into_iter()
186	.filter(|m| m["user"] == user)
187	.collect())
188	}
189	fn device(db: &rusqlite::Connection, headers: &HeaderMap) -> Result<Value> {
190	let token = headers
191	.get("authorization")
192	.and_then(|v| v.to_str().ok())
193	.and_then(|s| s.strip_prefix("Bearer "))
194	.filter(|s| !s.is_empty() && s.len() <= 256)
195	.ok_or_else(|| Error::new(401, "Pair this machine again."))?;
196	let fingerprint = hash(token);
197	list(db, "machine:")?
198	.into_iter()
199	.find(|machine| {
200	string(&machine["tokenHash"])
201	.as_bytes()
202	.ct_eq(fingerprint.as_bytes())
203	.unwrap_u8()
204	== 1
205	})
206	.ok_or_else(|| Error::new(401, "Pair this machine again."))
207	}
208	fn name(value: &Value) -> Result<&str> {
209	value
210	.as_str()
211	.filter(|s| !s.trim().is_empty() && s.encode_utf16().count() <= 100)
212	.ok_or_else(|| Error::new(400, "Enter a name up to 100 characters."))
213	}
214	pub(crate) fn manage(
215	app: &App,
216	db: &rusqlite::Connection,
217	parts: &[&str],
218	method: &Method,
219	owner: &str,
220	body: &Value,
221	) -> Result<Value> {
222	match parts {
223	["pair"] if method == Method::POST => {
224	let code = body["code"]
225	.as_str()
226	.filter(|s| !s.is_empty() && s.len() <= 40)
227	.ok_or_else(|| Error::new(400, "Enter the code from your local agent."))?;
228	let code: String = code
229	.chars()
230	.filter(|c| *c != '-' && !c.is_whitespace())
231	.flat_map(char::to_uppercase)
232	.collect();
233	let key = format!("pair:{}", hash(&code));
234	let pairing = get(db, &key)?;
235	if pairing.is_null() {
236	return Err(Error::new(
237	410,
238	"This code expired or was used. Start pairing again.",
239	));
240	}
241	if machines(db, owner)?.len() >= 128 {
242	return Err(Error::new(
243	409,
244	"Unlink an unused machine before adding another.",
245	));
246	}
247	let id = uuid::Uuid::new_v4().to_string();
248	let machine = json!({"id":id,"user":owner,"name":pairing["name"],"platform":pairing["platform"],"tokenHash":pairing["tokenHash"]});
249	put(db, &format!("machine:{id}"), &machine, 0)?;
250	delete(db, &key)?;
251	Ok(json!(app.relay.view(vec![machine], None).remove(0)))
252	}
253	["machines", id] if method == Method::DELETE || method == Method::PATCH => {
254	let mut machine = get(db, &format!("machine:{id}"))?;
255	if machine["user"] != owner {
256	return Err(Error::new(404, "No linked machine with that ID."));
257	}
258	if method == Method::DELETE {
259	delete(db, &format!("machine:{id}"))?;
260	app.relay.disconnect(id);
261	Ok(Value::Null)
262	} else {
263	machine["name"] = json!(name(&body["name"])?);
264	put(db, &format!("machine:{id}"), &machine, 0)?;
265	Ok(json!(app.relay.view(vec![machine], None).remove(0)))
266	}
267	}
268	["keys"] if method == Method::POST => {
269	let available = machines(db, owner)?;
270	let resources = selection(
271	&body["resources"],
272	&available
273	.iter()
274	.map(|m| m["id"].clone())
275	.collect::<Vec<_>>(),
276	)?;
277	let name = name(&body["name"])?;
278	let write = body
279	.get("write")
280	.map(|v| {
281	v.as_bool()
282	.ok_or_else(|| Error::new(400, "Choose read access or control."))
283	})
284	.transpose()?
285	.unwrap_or(false);
286	if list(db, "grant:")?
287	.iter()
288	.filter(|g| g["user"] == owner)
289	.count()
290	>= 256
291	{
292	return Err(Error::new(
293	409,
294	"Revoke an unused connection before adding another.",
295	));
296	}
297	let id = uuid::Uuid::new_v4().to_string();
298	let grant = json!({"id":id,"user":owner,"name":name,"client":null,"resource":app.mcp.resource("agents"),"scopes":if write {json!(["sessions:read","sessions:write"])} else {json!(["sessions:read"])},"resources":resources,"targets":resources,"createdAt":now()});
299	let key = format!("ar_{}", secret());
300	put(db, &format!("grant:{id}"), &grant, 0)?;
301	put(
302	db,
303	&format!("access:{}", hash(&key)),
304	&json!({"grant":id,"resource":grant["resource"]}),
305	0,
306	)?;
307	Ok(json!({"key":key,"id":id}))
308	}
309	_ => Err(Error::new(404, "No endpoint here.")),
310	}
311	}
312	fn selection(value: &Value, allowed: &[Value]) -> Result<Vec<Value>> {
313	let ids = value
314	.as_array()
315	.filter(|ids| !ids.is_empty() && ids.len() <= 128)
316	.ok_or_else(|| Error::new(400, "Choose at least one linked machine."))?;
317	if ids
318	.iter()
319	.enumerate()
320	.any(|(i, id)| !allowed.contains(id) || ids[..i].contains(id))
321	{
322	return Err(Error::new(
323	403,
324	"Choose machines granted to this connection once each.",
325	));
326	}
327	Ok(ids.clone())
328	}
329	fn view(app: &App, grant_id: &str, targets: Option<&Value>) -> Result<Vec<Value>> {
330	let mut db = app.mcp.db.lock().unwrap();
331	let tx = db.transaction()?;
332	let mut grant = get(&tx, &format!("grant:{grant_id}"))?;
333	if grant.is_null() || grant["resource"] != app.mcp.resource("agents") {
334	return Err(Error::new(
335	401,
336	"This connection was revoked. Connect again.",
337	));
338	}
339	let machines = machines(&tx, string(&grant["user"]))?;
340	if let Some(targets) = targets {
341	let allowed: Vec<_> = machines
342	.iter()
343	.filter(|m| array(&grant["resources"]).contains(&m["id"]))
344	.map(|m| m["id"].clone())
345	.collect();
346	grant["targets"] = json!(selection(targets, &allowed)?);
347	put(&tx, &format!("grant:{grant_id}"), &grant, 0)?;
348	}
349	let result = app.relay.view(machines, Some(&grant));
350	tx.commit()?;
351	Ok(result)
352	}
353	fn schema(method: &str) -> Option<Value> {
354	let provider = json!({"type":"string","enum":["codex","claude"]});
355	let id = json!({"type":"string","format":"uuid"});
356	let message = json!({"type":"string","minLength":1,"maxLength":100000});
357	let (properties, required) = match method {
358	"list_threads" => (
359	json!({"provider":provider,"limit":{"type":"integer","minimum":1,"maximum":100,"default":30}}),
360	json!([]),
361	),
362	"read_thread" => (
363	json!({"provider":provider,"thread_id":id,"cursor":{"type":"string","minLength":1,"maxLength":512},"limit":{"type":"integer","minimum":1,"maximum":100,"default":20}}),
364	json!(["provider", "thread_id"]),
365	),
366	"send_message" => (
367	json!({"provider":provider,"thread_id":id,"message":message,"expected_turn_id":id}),
368	json!(["provider", "thread_id", "message"]),
369	),
370	"interrupt_thread" => (
371	json!({"provider":provider,"thread_id":id,"expected_turn_id":id}),
372	json!(["provider", "thread_id"]),
373	),
374	"start_thread" => (
375	json!({"provider":provider,"cwd":{"type":"string","minLength":1,"maxLength":4096},"message":message}),
376	json!(["provider", "cwd", "message"]),
377	),
378	_ => return None,
379	};
380	Some(
381	json!({"type":"object","properties":properties,"required":required,"additionalProperties":false}),
382	)
383	}
384	fn command(method: &str, value: Value) -> Result<Value> {
385	let schema = schema(method)
386	.ok_or_else(|| Error::new(404, "Choose a session tool listed by this connector."))?;
387	let mut params = value
388	.as_object()
389	.cloned()
390	.ok_or_else(|| Error::new(400, "Use the fields listed for this tool."))?;
391	let properties = schema["properties"].as_object().unwrap();
392	if params.keys().any(|key| !properties.contains_key(key))
393	|| array(&schema["required"])
394	.iter()
395	.any(|key| !params.contains_key(string(key)))
396	{
397	return Err(Error::new(400, "Use the fields listed for this tool."));
398	}
399	for (key, rule) in properties {
400	if !params.contains_key(key) && rule.get("default").is_some() {
401	params.insert(key.clone(), rule["default"].clone());
402	}
403	let Some(value) = params.get(key) else {
404	continue;
405	};
406	let valid = match string(&rule["type"]) {
407	"integer" => value.as_i64().is_some_and(|n| {
408	n >= rule["minimum"].as_i64().unwrap() && n <= rule["maximum"].as_i64().unwrap()
409	}),
410	"string" => value.as_str().is_some_and(|s| {
411	let length = s.encode_utf16().count() as u64;
412	rule["minLength"].as_u64().is_none_or(|n| length >= n)
413	&& rule["maxLength"].as_u64().is_none_or(|n| length <= n)
414	&& rule
415	.get("enum")
416	.is_none_or(|values| array(values).contains(value))
417	&& (rule["format"] != "uuid"
418	|| uuid::Uuid::parse_str(s)
419	.is_ok_and(|id| id.to_string().eq_ignore_ascii_case(s)))
420	}),
421	_ => false,
422	};
423	if !valid {
424	return Err(Error::new(
425	400,
426	format!("Check {key} against the tool's fields."),
427	));
428	}
429	}
430	Ok(json!(params))
431	}
432	async fn pairing(State(app): State<Arc<App>>, request: Request) -> Result<Response> {
433	let method = request.method().clone();
434	let headers = request.headers().clone();
435	let bytes = axum::body::to_bytes(request.into_body(), 4096)
436	.await
437	.map_err(|_| Error::new(400, "Enter a shorter machine name."))?;
438	let mut db = app.mcp.db.lock().unwrap();
439	let tx = db.transaction()?;
440	let response = if method == Method::POST {
441	if list(&tx, "pair:")?.len() >= 256 {
442	return Err(Error::new(
443	429,
444	"Too many pairing requests. Try again in ten minutes.",
445	));
446	}
447	let body: Value = serde_json::from_slice(&bytes)
448	.map_err(|_| Error::new(400, "Enter a machine name and platform."))?;
449	let name = name(&body["name"])?;
450	let platform = body["platform"]
451	.as_str()
452	.filter(|s| s.encode_utf16().count() <= 100)
453	.ok_or_else(|| Error::new(400, "Enter a platform up to 100 characters."))?;
454	let token = secret();
455	let code: String = rand::random::<[u8; 5]>()
456	.iter()
457	.map(|b| format!("{b:02X}"))
458	.collect();
459	put(
460	&tx,
461	&format!("pair:{}", hash(&code)),
462	&json!({"name":name,"platform":platform,"tokenHash":hash(&token)}),
463	600,
464	)?;
465	(StatusCode::CREATED, axum::Json(json!({"code":format!("{}-{}", &code[..5], &code[5..]),"token":token,"expires_in":600}))).into_response()
466	} else if method == Method::GET {
467	match device(&tx, &headers) {
468	Ok(machine) => axum::Json(json!({"machine_id":machine["id"]})).into_response(),
469	Err(_) => {
470	let token = headers
471	.get("authorization")
472	.and_then(|v| v.to_str().ok())
473	.and_then(|s| s.strip_prefix("Bearer "))
474	.filter(|s| !s.is_empty() && s.len() <= 256)
475	.ok_or_else(|| Error::new(401, "Start pairing from your local agent."))?;
476	if !list(&tx, "pair:")?
477	.iter()
478	.any(|p| p["tokenHash"] == hash(token))
479	{
480	return Err(Error::new(
481	410,
482	"This pairing expired. Start pairing again.",
483	));
484	}
485	(StatusCode::ACCEPTED, axum::Json(json!({"pending":true}))).into_response()
486	}
487	}
488	} else {
489	return Err(Error::new(405, "Use GET or POST for pairing."));
490	};
491	tx.commit()?;
492	Ok(response)
493	}
494	async fn connect(
495	State(app): State<Arc<App>>,
496	headers: HeaderMap,
497	ws: WebSocketUpgrade,
498	) -> Result<Response> {
499	if headers.contains_key("origin") {
500	return Err(Error::new(401, "Connect from your local agent."));
501	}
502	let machine = device(&app.mcp.db.lock().unwrap(), &headers)?;
503	let id = string(&machine["id"]).to_owned();
504	let (send, receive) = mpsc::channel(16);
505	let (closed, _) = watch::channel(false);
506	let connection = Arc::new(Connection {
507	messages: send,
508	pending: Mutex::new(HashMap::new()),
509	closed,
510	});
511	{
512	let mut connections = app.relay.connections.lock().unwrap();
513	if connections.contains_key(&id) {
514	return Err(Error::new(
515	409,
516	"Another agent is connected. Stop it before starting a second copy.",
517	));
518	}
519	if connections.len() >= 256 {
520	return Err(Error::new(
521	503,
522	"Too many agents are connected. Try again later.",
523	));
524	}
525	connections.insert(id.clone(), connection.clone());
526	app.relay.changes.send_replace(());
527	}
528	let (failure_app, failure_id, failure_connection) =
529	(app.clone(), id.clone(), connection.clone());
530	Ok(ws
531	.max_frame_size(4 * 1024 * 1024)
532	.max_message_size(4 * 1024 * 1024)
533	.read_buffer_size(16 * 1024)
534	.write_buffer_size(0)
535	.max_write_buffer_size(256 * 1024)
536	.on_failed_upgrade(move |_| failure_app.relay.remove(&failure_id, &failure_connection))
537	.on_upgrade(move |socket| session(app, id, connection, receive, socket)))
538	}
539	async fn session(
540	app: Arc<App>,
541	id: String,
542	connection: Arc<Connection>,
543	mut messages: mpsc::Receiver<Message>,
544	mut socket: WebSocket,
545	) {
546	let mut closed = connection.closed.subscribe();
547	let connected = Message::Text(
548	json!({"type":"connected","machine_id":id})
549	.to_string()
550	.into(),
551	);
552	if tokio::time::timeout(Duration::from_secs(5), socket.send(connected))
553	.await
554	.is_ok_and(|r| r.is_ok())
555	{
556	let mut heartbeat = tokio::time::interval_at(
557	tokio::time::Instant::now() + Duration::from_secs(20),
558	Duration::from_secs(20),
559	);
560	heartbeat.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
561	let mut alive = true;
562	loop {
563	if *closed.borrow()
564	|| !get(&app.mcp.db.lock().unwrap(), &format!("machine:{id}"))
565	.is_ok_and(|m| !m.is_null())
566	{
567	break;
568	}
569	let outgoing = tokio::select! {
570	_ = closed.changed() => break,
571	outgoing = messages.recv() => match outgoing {Some(frame) => frame, None => break},
572	_ = heartbeat.tick() => {
573	if !alive {break;}
574	alive = false;
575	Message::Ping(Bytes::new())
576	},
577	incoming = socket.recv() => {
578	match incoming {
579	Some(Ok(Message::Pong(_))) => alive = true,
580	Some(Ok(Message::Ping(bytes))) => {
581	if !tokio::time::timeout(Duration::from_secs(5), socket.send(Message::Pong(bytes))).await.is_ok_and(|r| r.is_ok()) {break;}
582	},
583	Some(Ok(Message::Text(text))) => {
584	let reply = serde_json::from_str::<Value>(&text);
585	let Ok(reply) = reply else {break};
586	let valid = reply.as_object().is_some_and(|fields| fields.keys().all(|k| ["id","result","error"].contains(&k.as_str())))
587	&& reply["id"].as_str().is_some_and(|s| uuid::Uuid::parse_str(s).is_ok())
588	&& (reply.get("result").is_some() != reply.get("error").is_some())
589	&& reply.get("error").is_none_or(|error| error.as_str().is_some_and(|s| s.encode_utf16().count() <= 2000));
590	if !valid {break;}
591	if let Some(send) = connection.pending.lock().unwrap().remove(string(&reply["id"])) {
592	let result = if let Some(error) = reply["error"].as_str() {Err(Error::new(400, error))} else {Ok(reply["result"].clone())};
593	let _ = send.send(result);
594	}
595	},
596	_ => break,
597	}
598	continue;
599	}
600	};
601	if !tokio::time::timeout(Duration::from_secs(5), socket.send(outgoing))
602	.await
603	.is_ok_and(|r| r.is_ok())
604	{
605	break;
606	}
607	}
608	}
609	app.relay.remove(&id, &connection);
610	let _ = tokio::time::timeout(Duration::from_secs(1), socket.close()).await;
611	}
612	async fn rest(State(app): State<Arc<App>>, request: Request) -> Response {
613	let result: Result<Response> = async {
614	if request.headers().get("origin").is_some_and(|v| {
615	v.to_str().ok() != Some(app.mcp.origin.origin().ascii_serialization().as_str())
616	}) {
617	return Err(Error::new(403, "Use the relay from its own origin."));
618	}
619	let grant = app
620	.mcp
621	.authenticate(request.headers(), &app.mcp.resource("agents"))?;
622	if !mcp::active_owner(&app, &grant)? {
623	return Err(Error::new(401, "This connection's account is no longer authorized."));
624	}
625	let id = string(&grant["id"]);
626	let method = request.method().clone();
627	let path = request
628	.uri()
629	.path()
630	.trim_start_matches("/api/v1/")
631	.to_owned();
632	let bytes = axum::body::to_bytes(request.into_body(), 256 * 1024)
633	.await
634	.map_err(|_| Error::new(400, "Send a shorter command."))?;
635	let body: Value = if bytes.is_empty() {
636	Value::Null
637	} else {
638	serde_json::from_slice(&bytes).map_err(|_| Error::new(400, "Use JSON command fields."))?
639	};
640	let value = match path.split('/').collect::<Vec<_>>().as_slice() {
641	["machines"] if method == Method::GET => json!(view(&app, id, None)?),
642	["targets"] if method == Method::PUT => json!(view(&app, id, Some(&body["machine_ids"]))?),
643	["machines", machine, "commands"] if method == Method::POST => {
644	let method = body["method"]
645	.as_str()
646	.ok_or_else(|| Error::new(400, "Choose a session command."))?;
647	json!({"result":app.relay.dispatch(&app.mcp, id, machine, method, body["params"].clone(), Duration::from_secs(30)).await?})
648	}
649	_ => return Err(Error::new(404, "No relay endpoint here.")),
650	};
651	Ok(axum::Json(value).into_response())
652	}.await;
653	match result {
654	Ok(response) => response,
655	Err(error) => (StatusCode::from_u16(error.status).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR), axum::Json(json!({"error":if error.status >= 500 {"The relay couldn't answer. Check its dashboard and retry."} else {&error.message}}))).into_response(),
656	}
657	}
658	#[derive(Clone)]
659	struct Agents(Arc<App>);
660	impl ServerHandler for Agents {
661	fn get_info(&self) -> ServerConfig {
662	ServerConfig::new(ServerCapabilities::builder().enable_tools().build())
663	}
664	async fn list_tools(
665	&self,
666	_: Option<PaginatedRequestParams>,
667	_: RequestContext<RoleServer>,
668	) -> std::result::Result<ListToolsResult, ErrorData> {
669	let tools = [
670	("list_machines", "List granted machines and their connection state."),
671	("set_target_machines", "Select default machines for this connection."),
672	("list_threads", "List recent Codex and Claude Code threads on selected machines."),
673	("read_thread", "Read a thread and its current status."),
674	("send_message", "Submit a message. Desktop control requires local opt-in and reopens unloaded chats, switching the selected desktop chat. Submission acknowledges delivery, not task completion."),
675	("interrupt_thread", "Interrupt an agent-owned session or an opted-in Codex desktop turn."),
676	("start_thread", "Start an agent-owned session under a locally allowed directory."),
677	].into_iter().map(|(name, description)| {
678	let mut schema = schema(name).unwrap_or_else(|| if name == "set_target_machines" {json!({"type":"object","properties":{"machine_ids":{"type":"array","items":{"type":"string","format":"uuid"},"minItems":1,"maxItems":128,"uniqueItems":true}},"required":["machine_ids"],"additionalProperties":false})} else {json!({"type":"object","properties":{},"additionalProperties":false})});
679	if !["list_machines","set_target_machines"].contains(&name) {schema["properties"]["machine_id"] = json!({"type":"string","format":"uuid"});}
680	let read = ["list_machines","list_threads","read_thread"].contains(&name);
681	Tool::new(name, description, schema.as_object().unwrap().clone()).with_annotations(ToolAnnotations::new().read_only(read).destructive(name == "interrupt_thread").idempotent(read))
682	}).collect();
683	Ok(ListToolsResult {
684	tools,
685	ttl_ms: Some(0),
686	cache_scope: Some(CacheScope::Private),
687	..Default::default()
688	})
689	}
690	async fn call_tool(
691	&self,
692	request: CallToolRequestParams,
693	context: RequestContext<RoleServer>,
694	) -> std::result::Result<CallToolResponse, ErrorData> {
695	let result: Result<CallToolResult> = async {
696	let grant = &context
697	.extensions
698	.get::<axum::http::request::Parts>()
699	.and_then(|parts| parts.extensions.get::<mcp::Grant>())
700	.ok_or_else(|| Error::new(401, "This connection expired. Connect again."))?
701	.0;
702	let mut arguments = request.arguments.unwrap_or_default();
703	let id = string(&grant["id"]);
704	match request.name.as_ref() {
705	"list_machines" if arguments.is_empty() => Ok(CallToolResult::structured(
706	json!({"machines":view(&self.0, id, None)?}),
707	)),
708	"set_target_machines"
709	if arguments.len() == 1 && arguments.contains_key("machine_ids") =>
710	{
711	Ok(CallToolResult::structured(
712	json!({"machines":view(&self.0, id, arguments.get("machine_ids"))?}),
713	))
714	}
715	method => {
716	let explicit = arguments.remove("machine_id");
717	let machines = view(&self.0, id, None)?;
718	let targets: Vec<_> = if let Some(explicit) = explicit {
719	let explicit = explicit
720	.as_str()
721	.ok_or_else(|| Error::new(400, "Choose a machine ID."))?;
722	if !machines.iter().any(|m| m["id"] == explicit) {
723	return Err(Error::new(403, "Choose a granted machine."));
724	}
725	vec![explicit.to_owned()]
726	} else {
727	machines
728	.iter()
729	.filter(|m| m["selected"] == true)
730	.map(|m| string(&m["id"]).to_owned())
731	.collect()
732	};
733	if targets.is_empty() || method != "list_threads" && targets.len() != 1 {
734	return Err(Error::new(
735	400,
736	"Select one machine or supply machine_id for this command.",
737	));
738	}
739	let params = command(method, json!(arguments))?;
740	let mut replies: futures::stream::FuturesUnordered<_> = targets.into_iter().map(|machine| {
741	let params = params.clone();
742	async move {
743	match self
744	.0
745	.relay
746	.dispatch(
747	&self.0.mcp,
748	id,
749	&machine,
750	method,
751	params,
752	Duration::from_secs(30),
753	)
754	.await
755	{
756	Ok(result) => json!({"machine_id":machine,"result":result}),
757	Err(error) => json!({"machine_id":machine,"error":error.message}),
758	}
759	}
760	}).collect();
761	let mut results = Vec::new();
762	let mut bytes = 0;
763	while let Some(reply) = futures::StreamExt::next(&mut replies).await {
764	bytes += reply.to_string().len();
765	if bytes > 16*1024*1024 {return Err(Error::new(413, "Machine replies exceed 16 MB. Select fewer machines or request fewer messages."));}
766	results.push(reply);
767	}
768	let errors = results.iter().any(|r| r.get("error").is_some());
769	let mut result = CallToolResult::structured(json!({"results":results}));
770	result.is_error = Some(errors);
771	Ok(result)
772	}
773	}
774	}
775	.await;
776	Ok(match result {
777	Ok(result) => result,
778	Err(error) => CallToolResult::error(vec![ContentBlock::text(if error.status >= 500 {
779	"The relay couldn't answer. Check its dashboard and retry.".to_owned()
780	} else {
781	error.message
782	})]),
783	}
784	.into())
785	}
786	}
787	async fn installer(State(app): State<Arc<App>>, request: Request) -> Response {
788	let origin = app.mcp.origin.origin().ascii_serialization();
789	let source = if request.uri().path().ends_with(".ps1") {
790	include_str!("../agent/install.ps1")
791	.replace("__SERVER__", &format!("'{}'", origin.replace('\'', "''")))
792	} else {
793	include_str!("../agent/install.sh").replace(
794	"__SERVER__",
795	&format!("'{}'", origin.replace('\'', "'\\''")),
796	)
797	};
798	([("content-type", "text/plain; charset=utf-8")], source).into_response()
799	}
800	pub fn router(app: Arc<App>) -> Router {
801	let state = app.clone();
802	let expected_host =
803	app.mcp.origin[url::Position::BeforeHost..url::Position::AfterPort].to_owned();
804	let agent = env("STUDIO_AGENT_DIR", "agent");
805	Router::new()
806	.route("/pairing", any(pairing))
807	.route("/agent/connect", any(connect))
808	.route("/agent/install.sh", axum::routing::get(installer))
809	.route("/agent/install.ps1", axum::routing::get(installer))
810	.route_service(
811	"/agent/setup.mjs",
812	ServeFile::new(format!("{agent}/setup.mjs")),
813	)
814	.route_service(
815	"/agent/relay.mjs",
816	ServeFile::new(format!("{agent}/relay.mjs")),
817	)
818	.route("/api/v1/{*path}", any(rest))
819	.with_state(app.clone())
820	.merge(mcp::router(app, "agents", move || {
821	Ok(Agents(state.clone()))
822	}))
823	.layer(axum::middleware::from_fn(
824	move |request: Request, next: axum::middleware::Next| {
825	let host = expected_host.clone();
826	async move {
827	if request.headers().get("host").and_then(|v| v.to_str().ok()) != Some(&host) {
828	return StatusCode::MISDIRECTED_REQUEST.into_response();
829	}
830	next.run(request).await
831	}
832	},
833	))
834	}
835
836	#[cfg(test)]
837	mod tests {
838	use super::*;
839	struct Fixture {
840	store: Arc<mcp::Store>,
841	broker: Arc<Broker>,
842	path: PathBuf,
843	machine: String,
844	grant: String,
845	}
846	impl Fixture {
847	fn new(write: bool) -> Self {
848	let path = std::env::temp_dir()
849	.canonicalize()
850	.unwrap()
851	.join(format!("studio-relay-test-{}", uuid::Uuid::new_v4()));
852	let store = Arc::new(mcp::Store::new(&path, "https://globe.studio.test").unwrap());
853	let machine = uuid::Uuid::new_v4().to_string();
854	let grant = uuid::Uuid::new_v4().to_string();
855	{
856	let db = store.db.lock().unwrap();
857	put(&db, &format!("machine:{machine}"), &json!({"id":machine,"user":"owner","name":"Fixture","tokenHash":hash("device-secret")}), 0).unwrap();
858	put(&db, &format!("grant:{grant}"), &json!({"id":grant,"user":"owner","resource":store.resource("agents"),"resources":[machine],"scopes":if write {json!(["sessions:read","sessions:write"])} else {json!(["sessions:read"])}}), 0).unwrap();
859	}
860	Self {
861	store,
862	broker: Arc::new(Broker::default()),
863	path,
864	machine,
865	grant,
866	}
867	}
868	fn connect(&self) -> (Arc<Connection>, mpsc::Receiver<Message>) {
869	let (messages, receive) = mpsc::channel(16);
870	let (closed, _) = watch::channel(false);
871	let connection = Arc::new(Connection {
872	messages,
873	closed,
874	pending: Mutex::new(HashMap::new()),
875	});
876	self.broker
877	.connections
878	.lock()
879	.unwrap()
880	.insert(self.machine.clone(), connection.clone());
881	(connection, receive)
882	}
883	fn dispatch(&self, duration: Duration) -> tokio::task::JoinHandle<Result<Value>> {
884	let (store, broker, machine, grant) = (
885	self.store.clone(),
886	self.broker.clone(),
887	self.machine.clone(),
888	self.grant.clone(),
889	);
890	tokio::spawn(async move {
891	broker
892	.dispatch(
893	&store,
894	&grant,
895	&machine,
896	"list_threads",
897	json!({}),
898	duration,
899	)
900	.await
901	})
902	}
903	}
904	impl Drop for Fixture {
905	fn drop(&mut self) {
906	std::fs::remove_dir_all(&self.path).unwrap();
907	}
908	}
909	#[test]
910	fn command_fields_defaults_and_provider_boundaries() {
911	assert_eq!(
912	command("list_threads", json!({})).unwrap(),
913	json!({"limit":30})
914	);
915	let id = uuid::Uuid::new_v4().to_string();
916	assert_eq!(
917	command("read_thread", json!({"provider":"claude","thread_id":id})).unwrap()["limit"],
918	20
919	);
920	assert_eq!(
921	command(
922	"read_thread",
923	json!({"provider":"codex","thread_id":id,"cursor":"page"})
924	)
925	.unwrap()["cursor"],
926	"page"
927	);
928	for (method, params) in [
929	(
930	"read_thread",
931	json!({"provider":"codex","thread_id":id,"cursor":""}),
932	),
933	(
934	"read_thread",
935	json!({"provider":"codex","thread_id":id,"cursor":"x".repeat(513)}),
936	),
937	(
938	"read_thread",
939	json!({"provider":"codex","thread_id":"not-a-uuid"}),
940	),
941	(
942	"read_thread",
943	json!({"provider":"native-chat","thread_id":id}),
944	),
945	("list_threads", json!({"limit":101})),
946	("list_threads", json!({"limit":1.5})),
947	("list_threads", json!({"approval_response":true})),
948	(
949	"send_message",
950	json!({"provider":"codex","thread_id":id,"message":""}),
951	),
952	("start_thread", json!({"provider":"claude","cwd":"/tmp"})),
953	(
954	"interrupt_thread",
955	json!({"provider":"codex","thread_id":id,"expected_turn_id":1}),
956	),
957	("approve_tool", json!({})),
958	] {
959	assert!(command(method, params).is_err(), "{method}");
960	}
961	}
962	#[tokio::test]
963	async fn dispatch_uses_current_owner_grant_and_control_scope() {
964	let fixture = Fixture::new(false);
965	let (connection, mut receive) = fixture.connect();
966	let task = fixture.dispatch(Duration::from_secs(2));
967	let Message::Text(frame) = receive.recv().await.unwrap() else {
968	panic!()
969	};
970	let frame: Value = serde_json::from_str(&frame).unwrap();
971	assert_eq!(frame["params"], json!({"limit":30}));
972	connection
973	.pending
974	.lock()
975	.unwrap()
976	.remove(string(&frame["id"]))
977	.unwrap()
978	.send(Ok(json!({"threads":[]})))
979	.unwrap();
980	assert_eq!(task.await.unwrap().unwrap(), json!({"threads":[]}));
981	let refused = fixture
982	.broker
983	.dispatch(
984	&fixture.store,
985	&fixture.grant,
986	&fixture.machine,
987	"start_thread",
988	json!({"provider":"codex","cwd":"/tmp","message":"owned fixture"}),
989	Duration::from_secs(2),
990	)
991	.await
992	.unwrap_err();
993	assert_eq!(refused.status, 403);
994	assert!(receive.try_recv().is_err());
995	let foreign = uuid::Uuid::new_v4().to_string();
996	{
997	let db = fixture.store.db.lock().unwrap();
998	let mut grant = get(&db, &format!("grant:{}", fixture.grant)).unwrap();
999	grant["resources"] = json!([foreign]);
1000	put(&db, &format!("grant:{}", fixture.grant), &grant, 0).unwrap();
1001	put(
1002	&db,
1003	&format!("machine:{foreign}"),
1004	&json!({"id":foreign,"user":"someone-else"}),
1005	0,
1006	)
1007	.unwrap();
1008	}
1009	assert_eq!(
1010	fixture
1011	.broker
1012	.dispatch(
1013	&fixture.store,
1014	&fixture.grant,
1015	&foreign,
1016	"list_threads",
1017	json!({}),
1018	Duration::from_secs(2)
1019	)
1020	.await
1021	.unwrap_err()
1022	.status,
1023	403
1024	);
1025	delete(
1026	&fixture.store.db.lock().unwrap(),
1027	&format!("grant:{}", fixture.grant),
1028	)
1029	.unwrap();
1030	assert_eq!(
1031	fixture
1032	.dispatch(Duration::from_secs(2))
1033	.await
1034	.unwrap()
1035	.unwrap_err()
1036	.status,
1037	403
1038	);
1039	}
1040	#[tokio::test]
1041	async fn unicode_message_cannot_exceed_local_agent_frame_budget() {
1042	let fixture = Fixture::new(true);
1043	let (connection, mut receive) = fixture.connect();
1044	let error = fixture.broker.dispatch(&fixture.store, &fixture.grant, &fixture.machine, "send_message", json!({"provider":"codex","thread_id":uuid::Uuid::new_v4().to_string(),"message":"雪".repeat(100000)}), Duration::from_secs(2)).await.unwrap_err();
1045	assert_eq!(error.status, 413);
1046	assert!(receive.try_recv().is_err());
1047	assert!(connection.pending.lock().unwrap().is_empty());
1048	assert_eq!(fixture.broker.pending_slots.available_permits(), 128);
1049	}
1050	#[tokio::test]
1051	async fn eight_pending_limit_and_cancellation_release_capacity() {
1052	let fixture = Fixture::new(true);
1053	let (connection, mut receive) = fixture.connect();
1054	let mut tasks = Vec::new();
1055	for _ in 0..8 {
1056	tasks.push(fixture.dispatch(Duration::from_secs(2)));
1057	receive.recv().await.unwrap();
1058	}
1059	assert_eq!(
1060	fixture
1061	.dispatch(Duration::from_secs(2))
1062	.await
1063	.unwrap()
1064	.unwrap_err()
1065	.status,
1066	429
1067	);
1068	tasks.remove(0).abort();
1069	tokio::task::yield_now().await;
1070	assert_eq!(connection.pending.lock().unwrap().len(), 7);
1071	tasks.push(fixture.dispatch(Duration::from_secs(2)));
1072	receive.recv().await.unwrap();
1073	assert_eq!(connection.pending.lock().unwrap().len(), 8);
1074	fixture.broker.disconnect(&fixture.machine);
1075	for task in tasks {
1076	assert!(task.await.unwrap().unwrap_err().message.contains("unknown"));
1077	}
1078	assert!(connection.pending.lock().unwrap().is_empty());
1079	}
1080	#[tokio::test]
1081	async fn timeout_disconnect_and_reconnect_do_not_replay() {
1082	let fixture = Fixture::new(true);
1083	let (old, mut receive) = fixture.connect();
1084	let task = fixture.dispatch(Duration::from_millis(5));
1085	receive.recv().await.unwrap();
1086	assert!(task.await.unwrap().unwrap_err().message.contains("unknown"));
1087	assert!(old.pending.lock().unwrap().is_empty());
1088	assert!(receive.try_recv().is_err());
1089	let task = fixture.dispatch(Duration::from_secs(2));
1090	receive.recv().await.unwrap();
1091	fixture.broker.disconnect(&fixture.machine);
1092	assert!(task.await.unwrap().unwrap_err().message.contains("unknown"));
1093	let (new, mut receive) = fixture.connect();
1094	fixture.broker.remove(&fixture.machine, &old);
1095	assert!(Arc::ptr_eq(
1096	fixture
1097	.broker
1098	.connections
1099	.lock()
1100	.unwrap()
1101	.get(&fixture.machine)
1102	.unwrap(),
1103	&new
1104	));
1105	assert!(receive.try_recv().is_err());
1106	let task = fixture.dispatch(Duration::from_secs(2));
1107	receive.recv().await.unwrap();
1108	fixture.broker.disconnect(&fixture.machine);
1109	assert!(task.await.unwrap().unwrap_err().message.contains("unknown"));
1110	}
1111	}