Skip to main content

max / makenotwork

9.2 KB · 261 lines History Blame Raw
1 //! Per-connection session state for the Streamable HTTP transport.
2 //!
3 //! A session is created on `initialize` (the server returns its id in the
4 //! `Mcp-Session-Id` header) and referenced by that id on every later request.
5 //! It holds the three things the push layer needs that a stateless request
6 //! cannot:
7 //!
8 //! - the set of resource uris this client has **subscribed** to,
9 //! - the **out** channel to this client's SSE stream (present while a GET
10 //! stream is open), down which the server pushes notifications,
11 //! - the **in-flight** `tools/call` requests, each with a trigger to cancel it.
12 //!
13 //! stdio has no sessions (one implicit peer, no `Mcp-Session-Id`), so its
14 //! dispatch passes `None` and the session-scoped methods are unavailable there.
15
16 use std::collections::{HashMap, HashSet};
17 use std::future::Future;
18 use std::sync::{Arc, Mutex};
19
20 use serde_json::Value;
21 use tokio::sync::{mpsc, oneshot};
22
23 use crate::error::{Error, Result};
24 use crate::tool::ToolCallResult;
25
26 /// JSON-RPC notification value pushed to a session's SSE stream.
27 pub(crate) type OutSender = mpsc::UnboundedSender<Value>;
28
29 struct Inflight {
30 /// Original JSON-RPC id, preserved with its type (number stays a number)
31 /// so a `notifications/cancelled` echoes the exact id the client sent.
32 id: Value,
33 cancel: oneshot::Sender<()>,
34 }
35
36 /// One client connection's mutable state. Shared (`Arc`) between the POST
37 /// handler that mutates it and the GET SSE handler that drains its out channel.
38 pub(crate) struct Session {
39 pub id: String,
40 subscriptions: Mutex<HashSet<String>>,
41 out: Mutex<Option<OutSender>>,
42 inflight: Mutex<HashMap<String, Inflight>>,
43 }
44
45 impl Session {
46 fn new(id: String) -> Self {
47 Self {
48 id,
49 subscriptions: Mutex::new(HashSet::new()),
50 out: Mutex::new(None),
51 inflight: Mutex::new(HashMap::new()),
52 }
53 }
54
55 pub(crate) fn subscribe(&self, uri: impl Into<String>) {
56 self.subscriptions.lock().unwrap().insert(uri.into());
57 }
58
59 /// Remove a subscription. Returns whether it was present.
60 pub(crate) fn unsubscribe(&self, uri: &str) -> bool {
61 self.subscriptions.lock().unwrap().remove(uri)
62 }
63
64 pub(crate) fn is_subscribed(&self, uri: &str) -> bool {
65 self.subscriptions.lock().unwrap().contains(uri)
66 }
67
68 /// Attach the SSE out channel (called when the client opens GET `/mcp`).
69 /// Replaces any prior channel — a reconnect supersedes the stale stream.
70 pub(crate) fn attach_out(&self, tx: OutSender) {
71 *self.out.lock().unwrap() = Some(tx);
72 }
73
74 /// Push a notification to the SSE stream if one is connected.
75 pub(crate) fn push(&self, notification: Value) {
76 if let Some(tx) = self.out.lock().unwrap().as_ref() {
77 let _ = tx.send(notification);
78 }
79 }
80
81 /// Run a `tools/call` future so it can be cancelled while in flight. The
82 /// call is registered under `id` until it completes; a concurrent
83 /// [`cancel_request`](Self::cancel_request) or [`cancel_all`](Self::cancel_all)
84 /// aborts it with [`Error::Cancelled`].
85 pub(crate) async fn run_cancellable<F>(&self, id: &Value, fut: F) -> Result<ToolCallResult>
86 where
87 F: Future<Output = Result<ToolCallResult>>,
88 {
89 let key = id.to_string();
90 let (tx, rx) = oneshot::channel();
91 self.inflight.lock().unwrap().insert(
92 key.clone(),
93 Inflight {
94 id: id.clone(),
95 cancel: tx,
96 },
97 );
98
99 let outcome = tokio::select! {
100 biased;
101 _ = rx => Err(Error::Cancelled(key.clone())),
102 result = fut => result,
103 };
104
105 self.inflight.lock().unwrap().remove(&key);
106 outcome
107 }
108
109 /// Cancel one in-flight request by its (stringified) id. Triggers the
110 /// select in [`run_cancellable`](Self::run_cancellable). No-op if the id is
111 /// not in flight (already finished or never existed).
112 pub(crate) fn cancel_request(&self, id_key: &str) {
113 if let Some(entry) = self.inflight.lock().unwrap().remove(id_key) {
114 let _ = entry.cancel.send(());
115 }
116 }
117
118 /// Cancel every in-flight request for this session and tell the client,
119 /// pushing a `notifications/cancelled` for each. This is the app-initiated
120 /// path — a human taking a buffer back. Returns the count cancelled.
121 pub(crate) fn cancel_all(&self, reason: &str) -> usize {
122 let drained: Vec<Inflight> = {
123 let mut map = self.inflight.lock().unwrap();
124 map.drain().map(|(_, v)| v).collect()
125 };
126 let n = drained.len();
127 for entry in drained {
128 let _ = entry.cancel.send(());
129 self.push(cancelled_notification(&entry.id, reason));
130 }
131 n
132 }
133 }
134
135 /// The set of live sessions, shared by the server's handlers and its
136 /// [`BoundServer`](super::BoundServer) cancel handle.
137 #[derive(Clone, Default)]
138 pub(crate) struct Sessions {
139 map: Arc<Mutex<HashMap<String, Arc<Session>>>>,
140 }
141
142 impl Sessions {
143 /// Create a fresh session with a random id and return it.
144 pub(super) fn create(&self) -> Arc<Session> {
145 let id = uuid::Uuid::new_v4().to_string();
146 let session = Arc::new(Session::new(id.clone()));
147 self.map.lock().unwrap().insert(id, session.clone());
148 session
149 }
150
151 pub(super) fn get(&self, id: &str) -> Option<Arc<Session>> {
152 self.map.lock().unwrap().get(id).cloned()
153 }
154
155 /// Remove a session (client sent DELETE `/mcp`). Returns whether present.
156 pub(super) fn remove(&self, id: &str) -> bool {
157 self.map.lock().unwrap().remove(id).is_some()
158 }
159
160 pub(super) fn ids(&self) -> Vec<String> {
161 self.map.lock().unwrap().keys().cloned().collect()
162 }
163
164 /// Cancel in-flight ops across every session. Returns the total cancelled.
165 pub(super) fn cancel_all(&self, reason: &str) -> usize {
166 let sessions: Vec<Arc<Session>> = self.map.lock().unwrap().values().cloned().collect();
167 sessions.iter().map(|s| s.cancel_all(reason)).sum()
168 }
169 }
170
171 /// Build a server→client `notifications/cancelled` for `id`.
172 pub(crate) fn cancelled_notification(id: &Value, reason: &str) -> Value {
173 serde_json::json!({
174 "jsonrpc": "2.0",
175 "method": "notifications/cancelled",
176 "params": { "requestId": id, "reason": reason },
177 })
178 }
179
180 /// Build a server→client `notifications/resources/updated` for `uri`.
181 pub(crate) fn resource_updated_notification(uri: &str) -> Value {
182 serde_json::json!({
183 "jsonrpc": "2.0",
184 "method": "notifications/resources/updated",
185 "params": { "uri": uri },
186 })
187 }
188
189 #[cfg(test)]
190 mod tests {
191 use super::*;
192
193 #[tokio::test]
194 async fn run_cancellable_returns_result_when_not_cancelled() {
195 let s = Session::new("s1".into());
196 let out = s
197 .run_cancellable(&serde_json::json!(1), async {
198 Ok(ToolCallResult::text("done"))
199 })
200 .await
201 .unwrap();
202 assert_eq!(out.content.len(), 1);
203 }
204
205 #[tokio::test]
206 async fn cancel_request_aborts_an_in_flight_call() {
207 let s = Arc::new(Session::new("s1".into()));
208 let id = serde_json::json!(7);
209 let s2 = s.clone();
210 // A call that would never finish on its own.
211 let never = async { std::future::pending::<Result<ToolCallResult>>().await };
212 let handle = tokio::spawn(async move { s2.run_cancellable(&id, never).await });
213 // Let the call register itself, then cancel it.
214 tokio::task::yield_now().await;
215 s.cancel_request("7");
216 let outcome = handle.await.unwrap();
217 assert!(matches!(outcome, Err(Error::Cancelled(_))));
218 }
219
220 #[tokio::test]
221 async fn cancel_all_notifies_the_client_with_original_id_type() {
222 let s = Arc::new(Session::new("s1".into()));
223 let (tx, mut rx) = mpsc::unbounded_channel();
224 s.attach_out(tx);
225
226 let id = serde_json::json!(42); // numeric id
227 let s2 = s.clone();
228 let never = async { std::future::pending::<Result<ToolCallResult>>().await };
229 let handle = tokio::spawn(async move { s2.run_cancellable(&id, never).await });
230 tokio::task::yield_now().await;
231
232 assert_eq!(s.cancel_all("human took the buffer back"), 1);
233 assert!(matches!(handle.await.unwrap(), Err(Error::Cancelled(_))));
234
235 let note = rx.recv().await.unwrap();
236 assert_eq!(note["method"], "notifications/cancelled");
237 assert_eq!(note["params"]["requestId"], 42); // stayed a JSON number
238 assert_eq!(note["params"]["reason"], "human took the buffer back");
239 }
240
241 #[test]
242 fn subscribe_unsubscribe_tracks_membership() {
243 let s = Session::new("s1".into());
244 assert!(!s.is_subscribed("buffer://a"));
245 s.subscribe("buffer://a");
246 assert!(s.is_subscribed("buffer://a"));
247 assert!(s.unsubscribe("buffer://a"));
248 assert!(!s.unsubscribe("buffer://a"));
249 }
250
251 #[test]
252 fn sessions_create_get_remove() {
253 let sessions = Sessions::default();
254 let s = sessions.create();
255 assert!(sessions.get(&s.id).is_some());
256 assert_eq!(sessions.ids().len(), 1);
257 assert!(sessions.remove(&s.id));
258 assert!(sessions.get(&s.id).is_none());
259 }
260 }
261