Skip to main content

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(&note);
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, &note);
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}