mcp/server.rs
1// SPDX-License-Identifier: AGPL-3.0-only
2//! The reusable **MCP server** base: transport, framing, connection handling, the
3//! lifecycle/version machinery, and the resource-subscription registry — agentd's
4//! served self-MCP (and any other embedder's server) builds its domain surface on
5//! top by implementing [`Handler`].
6//!
7//! The split mirrors the client: this module owns the *protocol* (how bytes become
8//! requests, how `initialize` / `server/discover` / `ping` are answered across
9//! both eras, how a subscriber is pushed a `notifications/resources/updated`),
10//! while the embedder owns the *domain* (which tools exist, which resources are
11//! readable, who may subscribe to what). One [`Handler`] trait is the seam.
12//!
13//! Transport is deliberately minimal and dependency-light: a
14//! blocking listener, one thread per connection, speaking the same NDJSON JSON-RPC
15//! codec ([`crate::rpc::frame`]) as the client. No async, no mio. [`ServeStream`]
16//! type-erases unix vs. vsock so the framing, threading, and dispatch are entirely
17//! transport-agnostic ("the unix server with the socket type swapped").
18
19use crate::rpc::{Incoming, Notification, Request, Response, frame};
20use crate::wire::method;
21use serde_json::{Value, json};
22use std::collections::HashMap;
23use std::io::{self, BufReader, Read, Write};
24use std::os::unix::net::{UnixListener, UnixStream};
25use std::sync::atomic::{AtomicU64, Ordering};
26use std::sync::{Arc, Mutex};
27use std::thread;
28use std::time::Duration;
29
30/// Which transport a connection arrived on, and therefore its trust domain. The
31/// framework only carries this two-domain model and hands it to the [`Handler`];
32/// the embedder assigns meaning, and is the layer that must refuse a management
33/// peer anything reserved for the in-process caller:
34/// * [`Stdio`](PeerOrigin::Stdio) — an in-process / same-trust caller (agentd's
35/// own driving harness over the process stdio).
36/// * [`Management`](PeerOrigin::Management) — a peer that dialed a listener (unix
37/// socket / vsock), i.e. the management trust domain.
38#[derive(Debug, Clone, Copy, PartialEq, Eq)]
39pub enum PeerOrigin {
40 /// The process's own stdio / an in-process caller (the driving harness).
41 Stdio,
42 /// A peer on a listener (unix / vsock) — the management trust domain.
43 Management,
44}
45
46impl PeerOrigin {
47 /// Stable lowercase label for logs/metrics.
48 pub fn as_str(self) -> &'static str {
49 match self {
50 PeerOrigin::Stdio => "stdio",
51 PeerOrigin::Management => "management",
52 }
53 }
54}
55
56/// The served-MCP transport, type-erased to one concrete enum so the connection
57/// registry ([`SharedWriter`], [`Subscriber`]) stays monomorphic across transports
58/// while the *same* connection code serves each. The socket variants are
59/// `Read + Write` with a [`try_clone`](ServeStream::try_clone) (their write half is
60/// shared with the threads that push NDJSON notifications). The [`Http`](ServeStream::Http)
61/// variant is a **write-only SSE sink**: an HTTP subscription stream's write half,
62/// so an HTTP subscriber registers in the SAME [`SubRegistry`] and the embedder's
63/// existing `notify_*` calls reach it transparently — the framing (NDJSON vs SSE
64/// `data:` events) is chosen per-variant in [`write_notification`](ServeStream::write_notification).
65pub enum ServeStream {
66 /// A unix-domain-socket peer.
67 Unix(UnixStream),
68 /// An AF_VSOCK peer (host↔guest management transport).
69 #[cfg(feature = "vsock")]
70 Vsock(vsock::VsockStream),
71 /// The write half of an HTTP subscription (SSE) stream — a push-only sink.
72 /// Never read from, never a reply channel; notifications are framed as SSE
73 /// `data:` events. Boxed so it spans plain TCP and (feature `tls`) TLS.
74 Http(Box<dyn Write + Send>),
75}
76
77impl ServeStream {
78 /// Clone the handle (a second fd onto the same connection) for the shared write
79 /// half. Mirrors `UnixStream::try_clone`. The [`Http`](ServeStream::Http) sink
80 /// is single-owner (its SharedWriter is built directly from the stream half),
81 /// so it is never cloned — attempting to is an error.
82 pub fn try_clone(&self) -> io::Result<ServeStream> {
83 match self {
84 ServeStream::Unix(s) => s.try_clone().map(ServeStream::Unix),
85 #[cfg(feature = "vsock")]
86 ServeStream::Vsock(s) => s.try_clone().map(ServeStream::Vsock),
87 ServeStream::Http(_) => Err(io::Error::new(
88 io::ErrorKind::Unsupported,
89 "http SSE sink is not clonable",
90 )),
91 }
92 }
93
94 /// Bound a stalled-but-alive peer so it can't pin the writer Mutex forever.
95 /// The HTTP sink sets its timeout on the underlying stream before boxing.
96 pub fn set_write_timeout(&self, dur: Option<Duration>) -> io::Result<()> {
97 match self {
98 ServeStream::Unix(s) => s.set_write_timeout(dur),
99 #[cfg(feature = "vsock")]
100 ServeStream::Vsock(s) => s.set_write_timeout(dur),
101 ServeStream::Http(_) => Ok(()),
102 }
103 }
104
105 /// Push one notification, framed for this transport: NDJSON (one JSON object +
106 /// `\n`) on the socket variants, an SSE `data:` event on [`Http`](ServeStream::Http).
107 /// This is the single seam the `notify_*` push helpers write through.
108 pub fn write_notification(&mut self, note: &Notification) -> io::Result<()> {
109 if let ServeStream::Http(sink) = self {
110 let json = serde_json::to_string(note).map_err(io::Error::other)?;
111 sink.write_all(format!("data: {json}\n\n").as_bytes())?;
112 return sink.flush();
113 }
114 frame::write_line(self, note)
115 }
116
117 /// Write one JSON-RPC RESPONSE frame with per-transport framing — the
118 /// server-streaming twin of [`write_notification`]: an HTTP (SSE) sink gets
119 /// a `data:` event, a socket peer gets an NDJSON line. Streaming method
120 /// handlers (A2A `StreamResponse` frames) push through this so one handler
121 /// serves every transport.
122 pub fn write_response(&mut self, resp: &Response) -> io::Result<()> {
123 if let ServeStream::Http(sink) = self {
124 let json = serde_json::to_string(resp).map_err(io::Error::other)?;
125 sink.write_all(format!("data: {json}\n\n").as_bytes())?;
126 return sink.flush();
127 }
128 frame::write_line(self, resp)
129 }
130}
131
132impl Read for ServeStream {
133 fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
134 match self {
135 ServeStream::Unix(s) => s.read(buf),
136 #[cfg(feature = "vsock")]
137 ServeStream::Vsock(s) => s.read(buf),
138 // A push-only sink; a reader loop never runs over it.
139 ServeStream::Http(_) => Ok(0),
140 }
141 }
142}
143
144impl Write for ServeStream {
145 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
146 match self {
147 ServeStream::Unix(s) => s.write(buf),
148 #[cfg(feature = "vsock")]
149 ServeStream::Vsock(s) => s.write(buf),
150 ServeStream::Http(s) => s.write(buf),
151 }
152 }
153 fn flush(&mut self) -> io::Result<()> {
154 match self {
155 ServeStream::Unix(s) => s.flush(),
156 #[cfg(feature = "vsock")]
157 ServeStream::Vsock(s) => s.flush(),
158 ServeStream::Http(s) => s.flush(),
159 }
160 }
161}
162
163/// A connection's shared write half — both replies and pushed notifications go
164/// through it, serialized by the Mutex (a reply and a notification can't interleave
165/// bytes). The [`ServeStream`] enum keeps this one type across unix + vsock peers.
166pub type SharedWriter = Arc<Mutex<ServeStream>>;
167
168/// A peer subscribed to a resource: which connection, and the writer to push a
169/// `notifications/resources/updated` to. Opaque — fields are private; construct +
170/// mutate a registry through [`register_subscriber`] / [`drop_subscription`] /
171/// [`remove_conn_subscriptions`] and fire pushes through the `notify_*` helpers.
172pub struct Subscriber {
173 conn: u64,
174 writer: SharedWriter,
175}
176
177/// `uri` → its subscribers. Pushed when a resource changes. `Arc`-shared with the
178/// background threads that mutate resource state (a run reaching a terminal status,
179/// a reload landing, an event-ring growth).
180pub type SubRegistry = Arc<Mutex<HashMap<String, Vec<Subscriber>>>>;
181
182/// Register `conn` (with its `writer`) as a subscriber of `uri`, idempotently — a
183/// second subscribe from the same connection is a no-op rather than a duplicate
184/// push target. The embedder does its own gating (which URIs are subscribable, who
185/// may subscribe) *before* calling this.
186pub fn register_subscriber(subs: &SubRegistry, uri: &str, conn: u64, writer: &SharedWriter) {
187 let mut g = subs.lock().unwrap_or_else(|e| e.into_inner());
188 let list = g.entry(uri.to_string()).or_default();
189 if !list.iter().any(|s| s.conn == conn) {
190 list.push(Subscriber {
191 conn,
192 writer: Arc::clone(writer),
193 });
194 }
195}
196
197/// Drop `conn`'s subscription to a single `uri` (the `resources/unsubscribe` path).
198/// Prunes the uri entry entirely once its last subscriber leaves.
199pub fn drop_subscription(subs: &SubRegistry, uri: &str, conn: u64) {
200 let mut g = subs.lock().unwrap_or_else(|e| e.into_inner());
201 if let Some(list) = g.get_mut(uri) {
202 list.retain(|s| s.conn != conn);
203 if list.is_empty() {
204 g.remove(uri);
205 }
206 }
207}
208
209/// Drop every subscription held by a (now-closed) connection — called when a
210/// connection's reader loop ends so pushes never target a dead socket.
211pub fn remove_conn_subscriptions(subs: &SubRegistry, conn: u64) {
212 let mut g = subs.lock().unwrap_or_else(|e| e.into_inner());
213 g.retain(|_uri, list| {
214 list.retain(|s| s.conn != conn);
215 !list.is_empty()
216 });
217}
218
219/// Push `notifications/resources/updated{uri}` to every current subscriber of
220/// `uri`, **consuming** the subscription list (the resource changes exactly once —
221/// e.g. a subagent run reaching its terminal status — so no entry should linger
222/// after its one event). Best-effort: a write to a dead peer fails and is cleaned
223/// up when that connection's reader loop ends. The lock is released before writing,
224/// so a slow/blocked peer can't stall other notifications.
225pub fn notify_resource_updated(subs: &SubRegistry, uri: &str) {
226 let writers: Vec<SharedWriter> = {
227 let mut g = subs.lock().unwrap_or_else(|e| e.into_inner());
228 match g.remove(uri) {
229 Some(list) => list.into_iter().map(|s| s.writer).collect(),
230 None => return,
231 }
232 };
233 push_updated(&writers, uri);
234}
235
236/// Like [`notify_resource_updated`] but **keeps** the subscriber list — for
237/// resources that change REPEATEDLY (a run aggregate on each spawn, a warm session
238/// on each turn boundary, `config/effective` on each reload, an event ring on each
239/// batch). Cloning the writers under the lock (then releasing it before writing)
240/// keeps the entry intact for the next emission. Dead peers are pruned when their
241/// reader loop ends ([`remove_conn_subscriptions`]).
242pub fn notify_resource_updated_keep(subs: &SubRegistry, uri: &str) {
243 let writers: Vec<SharedWriter> = {
244 let g = subs.lock().unwrap_or_else(|e| e.into_inner());
245 match g.get(uri) {
246 Some(list) => list.iter().map(|s| Arc::clone(&s.writer)).collect(),
247 None => return,
248 }
249 };
250 push_updated(&writers, uri);
251}
252
253fn push_updated(writers: &[SharedWriter], uri: &str) {
254 let note = Notification::new(
255 method::NOTIFY_RESOURCES_UPDATED,
256 Some(json!({ "uri": uri })),
257 );
258 for w in writers {
259 if let Ok(mut wl) = w.lock() {
260 let _ = wl.write_notification(¬e);
261 }
262 }
263}
264
265/// Broadcast a payload-free `note` to every DISTINCT writer currently in the
266/// registry — for connection-scoped notifications that aren't tied to a single uri
267/// (e.g. `notifications/tools/list_changed` after a hot reload changed the tool
268/// set). A connection subscribed to several resources is written to once. Dead
269/// writers are pruned by their own reader loop.
270pub fn broadcast_distinct(subs: &SubRegistry, note: &Notification) {
271 let writers: Vec<SharedWriter> = {
272 let g = subs.lock().unwrap_or_else(|e| e.into_inner());
273 let mut seen: Vec<*const Mutex<ServeStream>> = Vec::new();
274 let mut out: Vec<SharedWriter> = Vec::new();
275 for list in g.values() {
276 for s in list {
277 let ptr = Arc::as_ptr(&s.writer);
278 if !seen.contains(&ptr) {
279 seen.push(ptr);
280 out.push(Arc::clone(&s.writer));
281 }
282 }
283 }
284 out
285 };
286 for w in writers {
287 if let Ok(mut wl) = w.lock() {
288 let _ = wl.write_notification(note);
289 }
290 }
291}
292
293// ---------------------------------------------------------------------------
294// The connection framework: the lifecycle/version machinery, the `Handler` seam,
295// and the blocking thread-per-connection listeners. An embedder implements
296// [`Handler`] for its domain surface and calls [`serve_unix`] / [`serve_vsock`];
297// everything below is transport- and domain-agnostic.
298// ---------------------------------------------------------------------------
299
300/// The embedder's domain seam. The framework owns the transport, the framing, the
301/// connection lifecycle, and the subscription registry; the `Handler` supplies the
302/// *meaning* — which tools/resources exist, who may call/read/subscribe to what.
303///
304/// Lifecycle (`initialize` / `server/discover` / `ping`) is NOT routed here: a
305/// handler answers it once, version-aware, by calling [`lifecycle_response`] at the
306/// top of its [`dispatch`](Handler::dispatch) (so the multi-version negotiation
307/// lives in one place). Everything else — `tools/*`, `resources/*`, and any custom
308/// method — flows through `dispatch`.
309pub trait Handler: Send + Sync + 'static {
310 /// Route one request to a response. `origin` is the caller's trust domain;
311 /// `writer`/`conn` identify the connection so a `resources/subscribe` can
312 /// register a push target via [`register_subscriber`]. Called from the
313 /// connection's own thread — implementations do their own locking.
314 fn dispatch(
315 &self,
316 req: Request,
317 origin: PeerOrigin,
318 writer: &SharedWriter,
319 conn: u64,
320 ) -> Response;
321
322 /// Called once when a connection is accepted (before its first request), for
323 /// logging/metrics. Default: nothing.
324 fn on_connect(&self, _origin: PeerOrigin, _conn: u64) {}
325
326 /// Whether `method` responds as a SERVER STREAM (several frames pushed
327 /// through the dispatch `writer`, the returned `Response` being the final
328 /// frame). The HTTP transport upgrades such a request to a `text/event-stream`
329 /// response instead of a unary `application/json` one. Default: nothing
330 /// streams.
331 fn streams(&self, _method: &str) -> bool {
332 false
333 }
334
335 /// Called once when a connection's reader loop ends. The framework has already
336 /// dropped the connection's subscriptions; this is for logging/metrics or any
337 /// embedder-side per-connection cleanup. Default: nothing.
338 fn on_disconnect(&self, _origin: PeerOrigin, _conn: u64) {}
339}
340
341/// Answer the three lifecycle methods every MCP server must handle, in ONE place,
342/// version-aware across both eras — the server-side mirror of the client's version
343/// negotiation. Returns `Some(response)` for `initialize` / `server/discover` /
344/// `ping`, or `None` if `req.method` is a domain method the [`Handler`] must route.
345///
346/// * `initialize` (legacy handshake): negotiate the protocol version — echo the
347/// peer's requested version when it's [supported](crate::version::is_supported_version),
348/// else fall back to our latest legacy [`PROTOCOL_VERSION`](crate::version::PROTOCOL_VERSION).
349/// * `server/discover` (modern, stateless): advertise the full
350/// [`SUPPORTED_PROTOCOL_VERSIONS`](crate::version::SUPPORTED_PROTOCOL_VERSIONS)
351/// list + capabilities in one call, so a modern client needn't fall back to the
352/// legacy handshake. Answering both is what makes the embedder a *dual-era
353/// server*: a peer of either era completes its lifecycle unaided.
354/// * `ping`: an empty result.
355///
356/// `server_info` is the `{name, version}` object and `capabilities` the advertised
357/// capability object; both are echoed verbatim into the two lifecycle replies. When
358/// the crate gains support for a new protocol version, both replies pick it up here
359/// without the embedder changing anything.
360pub fn lifecycle_response(
361 req: &Request,
362 server_info: &Value,
363 capabilities: &Value,
364) -> Option<Response> {
365 match req.method.as_str() {
366 "initialize" => {
367 let requested = req
368 .params
369 .as_ref()
370 .and_then(|p| p.get("protocolVersion"))
371 .and_then(Value::as_str);
372 let version = match requested {
373 Some(v) if crate::version::is_supported_version(v) => v,
374 _ => crate::version::PROTOCOL_VERSION,
375 };
376 Some(Response::ok(
377 req.id.clone(),
378 json!({
379 "protocolVersion": version,
380 "capabilities": capabilities,
381 "serverInfo": server_info,
382 }),
383 ))
384 }
385 method::SERVER_DISCOVER => Some(Response::ok(
386 req.id.clone(),
387 json!({
388 "resultType": "complete",
389 "supportedVersions": crate::version::SUPPORTED_PROTOCOL_VERSIONS,
390 "capabilities": capabilities,
391 "serverInfo": server_info,
392 }),
393 )),
394 "ping" => Some(Response::ok(req.id.clone(), json!({}))),
395 _ => None,
396 }
397}
398
399/// Serve one accepted connection to completion: the blocking NDJSON read loop.
400/// Requests get a reply (through the shared writer, which a background thread may
401/// also push notifications on — the Mutex serializes them); notifications
402/// (`initialized`, …) are read and dropped. On EOF/hangup the connection's
403/// subscriptions are dropped so no push ever targets a dead socket. A write timeout
404/// bounds a stalled-but-alive peer so it can't pin the writer Mutex (and a pushing
405/// thread) forever.
406pub fn handle_conn(
407 stream: ServeStream,
408 origin: PeerOrigin,
409 handler: &Arc<dyn Handler>,
410 subs: &SubRegistry,
411 conn_counter: &AtomicU64,
412 write_timeout: Duration,
413) {
414 let writer: SharedWriter = match stream.try_clone() {
415 Ok(w) => {
416 let _ = w.set_write_timeout(Some(write_timeout));
417 Arc::new(Mutex::new(w))
418 }
419 Err(_) => return,
420 };
421 let conn = conn_counter.fetch_add(1, Ordering::Relaxed);
422 handler.on_connect(origin, conn);
423 let mut reader = BufReader::new(stream);
424 while let Ok(Some(bytes)) = frame::read_line(&mut reader) {
425 if let Ok(Incoming::Request(req)) = serde_json::from_slice::<Incoming>(&bytes) {
426 let resp = handler.dispatch(req, origin, &writer, conn);
427 let wrote = writer
428 .lock()
429 .is_ok_and(|mut w| frame::write_line(&mut *w, &resp).is_ok());
430 if !wrote {
431 break; // peer hung up mid-reply
432 }
433 }
434 }
435 remove_conn_subscriptions(subs, conn); // don't push to a dead socket
436 handler.on_disconnect(origin, conn);
437}
438
439/// Bind a unix socket for serving, clearing any stale socket file first. Returned
440/// separately from the accept loop so the caller can log/act on a successful bind
441/// (or propagate the bind error) before the accept thread starts.
442pub fn bind_unix(path: &str) -> io::Result<UnixListener> {
443 // A stale socket from a crashed prior run would block the bind; clear it.
444 let _ = std::fs::remove_file(path);
445 UnixListener::bind(path)
446}
447
448/// Spawn the background accept thread for `listener`: one blocking thread per
449/// connection, each running [`handle_conn`] against `handler`. Peers arrive in the
450/// [`PeerOrigin::Management`] trust domain (they dialed a listener). Returns once
451/// the accept thread is spawned; a thread-spawn failure is surfaced as the error.
452pub fn spawn_accept_unix(
453 listener: UnixListener,
454 handler: Arc<dyn Handler>,
455 subs: SubRegistry,
456 conn_counter: Arc<AtomicU64>,
457 write_timeout: Duration,
458) -> io::Result<()> {
459 thread::Builder::new()
460 .name("serve-mcp".into())
461 .spawn(move || {
462 for stream in listener.incoming().flatten() {
463 let handler = Arc::clone(&handler);
464 let subs = Arc::clone(&subs);
465 let conn_counter = Arc::clone(&conn_counter);
466 thread::Builder::new()
467 .name("serve-mcp-conn".into())
468 .spawn(move || {
469 handle_conn(
470 ServeStream::Unix(stream),
471 PeerOrigin::Management,
472 &handler,
473 &subs,
474 &conn_counter,
475 write_timeout,
476 )
477 })
478 .ok();
479 }
480 })
481 .map(|_| ())
482}
483
484/// Bind an AF_VSOCK `(cid, port)` for serving — the management transport for a
485/// host↔guest boundary. The vsock counterpart of [`bind_unix`].
486#[cfg(feature = "vsock")]
487pub fn bind_vsock(cid: u32, port: u32) -> io::Result<vsock::VsockListener> {
488 vsock::VsockListener::bind_with_cid_port(cid, port)
489}
490
491/// Spawn the background accept thread for a vsock `listener` — byte-for-byte
492/// [`spawn_accept_unix`] with the socket type swapped (the same [`handle_conn`], no
493/// new framing). Peers arrive in [`PeerOrigin::Management`].
494#[cfg(feature = "vsock")]
495pub fn spawn_accept_vsock(
496 listener: vsock::VsockListener,
497 handler: Arc<dyn Handler>,
498 subs: SubRegistry,
499 conn_counter: Arc<AtomicU64>,
500 write_timeout: Duration,
501) -> io::Result<()> {
502 thread::Builder::new()
503 .name("serve-mcp-vsock".into())
504 .spawn(move || {
505 for stream in listener.incoming().flatten() {
506 let handler = Arc::clone(&handler);
507 let subs = Arc::clone(&subs);
508 let conn_counter = Arc::clone(&conn_counter);
509 thread::Builder::new()
510 .name("serve-mcp-conn".into())
511 .spawn(move || {
512 handle_conn(
513 ServeStream::Vsock(stream),
514 PeerOrigin::Management,
515 &handler,
516 &subs,
517 &conn_counter,
518 write_timeout,
519 )
520 })
521 .ok();
522 }
523 })
524 .map(|_| ())
525}
526
527#[cfg(test)]
528mod tests {
529 use super::*;
530 use crate::rpc::Request;
531 use std::io::BufReader;
532 use std::os::unix::net::UnixStream;
533
534 // --- lifecycle / version negotiation ---------------------------------
535
536 fn info() -> Value {
537 json!({"name": "test-server", "version": "9.9.9"})
538 }
539 fn caps() -> Value {
540 json!({"tools": {}, "resources": {"subscribe": true}})
541 }
542
543 #[test]
544 fn initialize_echoes_a_supported_requested_version() {
545 // A legacy client requesting a version we support gets it echoed back.
546 let want = crate::version::SUPPORTED_PROTOCOL_VERSIONS[1]; // a non-latest supported one
547 let req = Request::new(1, "initialize", Some(json!({"protocolVersion": want})));
548 let resp = lifecycle_response(&req, &info(), &caps()).expect("lifecycle handled");
549 let r = resp.result.expect("ok");
550 assert_eq!(r["protocolVersion"], want);
551 assert_eq!(r["serverInfo"]["name"], "test-server");
552 assert!(
553 r["capabilities"]["resources"]["subscribe"]
554 .as_bool()
555 .unwrap()
556 );
557 }
558
559 #[test]
560 fn initialize_falls_back_to_latest_legacy_for_an_unsupported_version() {
561 // An unknown/too-old version isn't echoed — we answer with our own latest.
562 let req = Request::new(
563 1,
564 "initialize",
565 Some(json!({"protocolVersion": "1999-01-01"})),
566 );
567 let resp = lifecycle_response(&req, &info(), &caps()).expect("handled");
568 let r = resp.result.expect("ok");
569 assert_eq!(r["protocolVersion"], crate::version::PROTOCOL_VERSION);
570 }
571
572 #[test]
573 fn initialize_defaults_when_no_version_is_requested() {
574 let req = Request::new(1, "initialize", Some(json!({})));
575 let resp = lifecycle_response(&req, &info(), &caps()).expect("handled");
576 assert_eq!(
577 resp.result.expect("ok")["protocolVersion"],
578 crate::version::PROTOCOL_VERSION
579 );
580 }
581
582 #[test]
583 fn server_discover_advertises_every_supported_version() {
584 // The modern stateless probe learns our full version list + caps in one call.
585 let req = Request::new(7, method::SERVER_DISCOVER, None);
586 let resp = lifecycle_response(&req, &info(), &caps()).expect("handled");
587 let r = resp.result.expect("ok");
588 assert_eq!(r["resultType"], "complete");
589 let listed = r["supportedVersions"].as_array().expect("array");
590 assert_eq!(
591 listed.len(),
592 crate::version::SUPPORTED_PROTOCOL_VERSIONS.len()
593 );
594 assert!(
595 listed
596 .iter()
597 .any(|v| v == crate::version::FIRST_MODERN_VERSION)
598 );
599 assert!(listed.iter().any(|v| v == crate::version::PROTOCOL_VERSION));
600 assert_eq!(r["serverInfo"]["name"], "test-server");
601 }
602
603 #[test]
604 fn ping_is_an_empty_ok_and_domain_methods_fall_through() {
605 let ping = Request::new(1, "ping", None);
606 assert_eq!(
607 lifecycle_response(&ping, &info(), &caps())
608 .expect("handled")
609 .result,
610 Some(json!({}))
611 );
612 // A non-lifecycle method is the handler's job — the helper declines it.
613 let dom = Request::new(2, "tools/call", None);
614 assert!(lifecycle_response(&dom, &info(), &caps()).is_none());
615 }
616
617 // --- subscription registry ------------------------------------------
618
619 /// A writer whose pushes can be read back off its peer end.
620 fn wired() -> (SharedWriter, BufReader<UnixStream>) {
621 let (tx, rx) = UnixStream::pair().unwrap();
622 // Bound the read so a "should push nothing" assertion can't hang.
623 rx.set_read_timeout(Some(Duration::from_millis(250)))
624 .unwrap();
625 (
626 Arc::new(Mutex::new(ServeStream::Unix(tx))),
627 BufReader::new(rx),
628 )
629 }
630
631 /// Read one pushed notification and return `(method, uri-or-empty)`.
632 fn read_note(rx: &mut BufReader<UnixStream>) -> (String, String) {
633 let bytes = frame::read_line(rx).expect("read").expect("a frame");
634 let v: Value = serde_json::from_slice(&bytes).expect("json");
635 let method = v["method"].as_str().unwrap_or_default().to_string();
636 let uri = v["params"]["uri"].as_str().unwrap_or_default().to_string();
637 (method, uri)
638 }
639
640 /// Assert nothing more was pushed (the bounded read times out / hits EOF).
641 fn assert_silent(rx: &mut BufReader<UnixStream>) {
642 assert!(
643 !matches!(frame::read_line(rx), Ok(Some(_))),
644 "expected no further push"
645 );
646 }
647
648 #[test]
649 fn register_is_idempotent_per_connection() {
650 let subs: SubRegistry = Arc::new(Mutex::new(HashMap::new()));
651 let (w, _rx) = wired();
652 register_subscriber(&subs, "res://a", 1, &w);
653 register_subscriber(&subs, "res://a", 1, &w); // same conn again — no dup
654 let g = subs.lock().unwrap();
655 assert_eq!(g.get("res://a").unwrap().len(), 1);
656 }
657
658 #[test]
659 fn notify_updated_consumes_but_keep_retains() {
660 let subs: SubRegistry = Arc::new(Mutex::new(HashMap::new()));
661 let (w, mut rx) = wired();
662 register_subscriber(&subs, "res://run", 1, &w);
663
664 // keep-variant fires and leaves the subscription in place …
665 notify_resource_updated_keep(&subs, "res://run");
666 let (m, uri) = read_note(&mut rx);
667 assert_eq!(m, method::NOTIFY_RESOURCES_UPDATED);
668 assert_eq!(uri, "res://run");
669 assert!(subs.lock().unwrap().contains_key("res://run"));
670
671 // … the consume-variant fires once and drops the entry.
672 notify_resource_updated(&subs, "res://run");
673 let (_m, uri2) = read_note(&mut rx);
674 assert_eq!(uri2, "res://run");
675 assert!(!subs.lock().unwrap().contains_key("res://run"));
676 }
677
678 #[test]
679 fn drop_and_conn_cleanup_remove_subscriptions() {
680 let subs: SubRegistry = Arc::new(Mutex::new(HashMap::new()));
681 let (w, _rx) = wired();
682 register_subscriber(&subs, "res://a", 1, &w);
683 register_subscriber(&subs, "res://b", 1, &w);
684
685 drop_subscription(&subs, "res://a", 1);
686 assert!(!subs.lock().unwrap().contains_key("res://a"));
687 assert!(subs.lock().unwrap().contains_key("res://b"));
688
689 remove_conn_subscriptions(&subs, 1);
690 assert!(subs.lock().unwrap().is_empty());
691 }
692
693 #[test]
694 fn broadcast_distinct_writes_once_per_connection() {
695 let subs: SubRegistry = Arc::new(Mutex::new(HashMap::new()));
696 let (w, mut rx) = wired();
697 // Same connection subscribed to two resources → one distinct writer.
698 register_subscriber(&subs, "res://a", 1, &w);
699 register_subscriber(&subs, "res://b", 1, &w);
700
701 let note = Notification::new(method::NOTIFY_TOOLS_LIST_CHANGED, None);
702 broadcast_distinct(&subs, ¬e);
703
704 let (m, _) = read_note(&mut rx);
705 assert_eq!(m, method::NOTIFY_TOOLS_LIST_CHANGED);
706 assert_silent(&mut rx); // not written twice
707 }
708}