agent_engine/mcp/lease.rs
1//! Task 19 (Commit B) — session-scoped exact MCP runtime leases (spec §7.4).
2//!
3//! Deferred descriptor-backed MCP tools execute ONLY through this manager,
4//! AFTER the `ExecutionGate` has already authorized the exact call. The
5//! manager owns every live MCP child for the runtime:
6//!
7//! - a lease is keyed by (session, server) and PINNED to the config
8//! fingerprint it was acquired under; per-key SINGLE-FLIGHT acquisition
9//! (a `Starting` placeholder holding a `watch` receiver, whose stored
10//! state makes wakeups lost-proof) guarantees concurrent first calls
11//! never spawn duplicate children, and the manager map lock is NEVER
12//! held across process/pipe I/O;
13//! - acquisition starts exactly one child (`kill_on_drop`), initializes
14//! and lists ONCE, and pins the bounded listing; every call validates
15//! the selected exact runtime name and canonical schema digest against
16//! the cached expectation BEFORE `tools/call` — a mismatch terminates
17//! the poisoned lease without calling;
18//! - revocation/termination semantics (exact and honest): revocation
19//! marks the lease cancelled under the per-lease call gate. A call that
20//! already passed the gate check (taken while holding the connection
21//! lock) MAY still complete; NO call passes the gate after revocation
22//! returns, and a completing call that observes cancellation kills the
23//! child before releasing the connection. Teardown is guaranteed by
24//! `kill_on_drop` once the last `Arc` drops (map removal plus bailing
25//! callers make that inevitable) with an immediate best-effort
26//! `start_kill` accelerator — the manager retains nothing;
27//! - idle leases are reaped opportunistically with a CONSTANT per-pass
28//! scan cap; session end is the drop of the LAST owner of the shared
29//! durable [`McpSessionEndGuard`] scope (runtime clones + in-flight
30//! streams hold it; no per-turn guard exists). No PID signalling, no
31//! long-lived task, no unbounded channel;
32//! - a server configured `shared: true` (daemon mode, opt-in) is keyed
33//! under [`SHARED_SESSION_KEY`] instead of the session: one child serves
34//! every session, `terminate_session` spares it, `reap_idle` and
35//! `terminate_all` still reap it. Every call still validates exact
36//! name + digest; sharing widens the CHILD, never the grant.
37
38use std::collections::HashMap;
39use std::sync::Arc;
40use std::time::{Duration, Instant};
41
42use serde_json::Value;
43use thiserror::Error;
44
45use super::connection::McpConnection;
46use super::descriptors::{self, server_config_fingerprint};
47use super::McpServerConfig;
48use crate::tools::activation::SessionId;
49use crate::tools::catalog::SchemaDigest;
50
51/// Injectable current-config lookup: production reads `mcp.json` per
52/// acquisition (bounded local file read); tests inject a mutable map.
53pub type ConfigSource = Arc<dyn Fn(&str) -> Option<McpServerConfig> + Send + Sync>;
54
55/// Production config source: profile-resolved `mcp.json` lookup.
56pub fn config_source_from_disk() -> ConfigSource {
57 Arc::new(|server: &str| {
58 super::load_mcp_config().and_then(|c| c.mcp_servers.get(server).cloned())
59 })
60}
61
62/// Lease-map session key for `shared: true` servers (one child per
63/// process, not per session). Never a real session's key: `terminate_session`
64/// skips it explicitly.
65pub const SHARED_SESSION_KEY: &str = "*";
66
67/// Session component of the lease key for `server` under `config`.
68pub fn lease_session_key(config: &McpServerConfig, session: &SessionId) -> String {
69 if config.shared {
70 SHARED_SESSION_KEY.to_string()
71 } else {
72 session.as_str().to_string()
73 }
74}
75
76/// Default idle bound for opportunistic reaping.
77pub const DEFAULT_IDLE_MAX: Duration = Duration::from_secs(300);
78/// Hard cap on simultaneously live leases per manager, so termination
79/// fanout (and cleanup-task count) is always bounded.
80pub const MAX_LIVE_LEASES: usize = 64;
81/// Reap passes scan up to the FULL capped map (equal to the live-lease
82/// cap), so idle entries can never starve behind an active HashMap
83/// iteration prefix.
84pub const REAP_SCAN_MAX: usize = MAX_LIVE_LEASES;
85/// Outer total bound on one lease cleanup task (lock wait + shutdown).
86const CLEANUP_TOTAL_TIMEOUT: Duration = Duration::from_secs(3);
87
88/// Typed lease failures. Provider-controlled content never appears here —
89/// only static text, local identities, and numeric metadata.
90#[derive(Debug, Error)]
91pub enum McpLeaseError {
92 #[error("MCP server '{0}' is no longer configured; lease revoked")]
93 ServerNotConfigured(String),
94 #[error(
95 "MCP server '{0}' config fingerprint changed; cached descriptors and lease invalidated"
96 )]
97 FingerprintDrift(String),
98 #[error("MCP server '{0}' did not list expected tool '{1}'; poisoned lease terminated without calling")]
99 NameNotListed(String, String),
100 #[error("MCP tool '{1}' on server '{0}' listed a schema that does not match the pinned descriptor digest; poisoned lease terminated without calling")]
101 SchemaMismatch(String, String),
102 #[error("MCP lease for server '{0}' was revoked before the call could start")]
103 Revoked(String),
104 #[error("MCP transport failure for server '{0}': {1}")]
105 Transport(String, String),
106 #[error(
107 "MCP lease capacity of {0} live leases reached; try again after idle leases are reaped"
108 )]
109 Capacity(usize),
110}
111
112struct LeaseInner {
113 fingerprint: String,
114 connection: tokio::sync::Mutex<McpConnection>,
115 /// Bounded pinned listing from the single `tools/list`.
116 listing: HashMap<String, SchemaDigest>,
117 /// Per-lease call/revocation gate: `true` = cancelled. Held only for
118 /// flag reads/writes — never across I/O.
119 gate: std::sync::Mutex<bool>,
120 last_used: std::sync::Mutex<Instant>,
121 /// Test-observable ownership token: a `Weak` to this proves when the
122 /// manager (and any bounded shutdown task) released the lease.
123 liveness: Arc<()>,
124}
125
126impl LeaseInner {
127 fn cancelled(&self) -> bool {
128 *self
129 .gate
130 .lock()
131 .unwrap_or_else(std::sync::PoisonError::into_inner)
132 }
133
134 /// Mark cancelled: no NEW call passes the gate after this returns.
135 fn set_cancelled(&self) {
136 *self
137 .gate
138 .lock()
139 .unwrap_or_else(std::sync::PoisonError::into_inner) = true;
140 }
141}
142
143enum LeaseState {
144 /// Single-flight placeholder. The `watch` channel stores completion
145 /// state, so a follower that clones the receiver under the map lock and
146 /// awaits after releasing it can never miss the wakeup (`changed()`
147 /// resolves immediately once the starter sent or dropped the sender).
148 Starting {
149 token: u64,
150 done: tokio::sync::watch::Receiver<bool>,
151 },
152 Ready(Arc<LeaseInner>),
153}
154
155/// Session-scoped exact MCP runtime lease manager. See module docs.
156pub struct McpRuntimeManager {
157 leases: std::sync::Mutex<HashMap<(String, String), LeaseState>>,
158 config_source: ConfigSource,
159 idle_max: Duration,
160 next_token: std::sync::atomic::AtomicU64,
161 /// Descriptor-cache path for write-back after `tools/list`
162 /// (production: profile-resolved default; tests inject a temp path).
163 cache_path: Arc<dyn Fn() -> std::path::PathBuf + Send + Sync>,
164}
165
166impl McpRuntimeManager {
167 pub fn new(config_source: ConfigSource, idle_max: Duration) -> Self {
168 Self {
169 leases: std::sync::Mutex::new(HashMap::new()),
170 config_source,
171 idle_max,
172 next_token: std::sync::atomic::AtomicU64::new(1),
173 cache_path: Arc::new(descriptors::default_cache_path),
174 }
175 }
176
177 /// Override the descriptor-cache write-back location (tests, or a
178 /// host that keeps its cache elsewhere).
179 pub fn with_cache_path(mut self, path: std::path::PathBuf) -> Self {
180 self.cache_path = Arc::new(move || path.clone());
181 self
182 }
183
184 /// Lease-map key for (session, server), honouring `shared: true`.
185 /// Reads the config source (bounded local read); an unconfigured
186 /// server keys per session (it has no lease to find anyway).
187 fn lease_key(&self, session: &SessionId, server: &str) -> (String, String) {
188 let session_key = match (self.config_source)(server) {
189 Some(config) => lease_session_key(&config, session),
190 None => session.as_str().to_string(),
191 };
192 (session_key, server.to_string())
193 }
194
195 /// Sync RAII-safe termination: mark cancelled, then — when the
196 /// connection is not in-flight — close stdin immediately (graceful EOF)
197 /// and hand the Arc to one BOUNDED spawned task that waits ~250ms for
198 /// the child to exit before `start_kill` + bounded reap. Outside a
199 /// runtime the fallback is an immediate `start_kill`. If the connection
200 /// IS in-flight, the completing call observes cancellation and runs the
201 /// async graceful shutdown itself. In every path the last `Arc` drop
202 /// triggers `kill_on_drop` as the hard backstop.
203 fn cancel_state(state: LeaseState) {
204 match state {
205 LeaseState::Ready(inner) => {
206 inner.set_cancelled();
207 match inner.connection.try_lock() {
208 Ok(mut conn) => {
209 conn.close_stdin();
210 drop(conn);
211 Self::spawn_bounded_cleanup(Arc::clone(&inner));
212 }
213 // In-flight: the completion path observes cancellation
214 // and shuts down; the bounded cleanup task below is a
215 // second, outer-timeout-bounded chance to reap sooner.
216 Err(_) => Self::spawn_bounded_cleanup(Arc::clone(&inner)),
217 }
218 }
219 // The in-flight starter detects its placeholder is gone and
220 // kills its own child instead of publishing the lease; dropping
221 // the receiver here wakes nobody wrongly (watch state persists).
222 LeaseState::Starting { .. } => {}
223 }
224 }
225
226 /// One BOUNDED one-shot cleanup task per terminated lease: strict
227 /// outer timeout around (connection lock wait + graceful 250ms EOF +
228 /// kill/1s reap). Finite input, no producer — not a long-lived task.
229 /// If the outer timeout is exceeded (a call still in flight), the
230 /// completion path sees the cancelled gate and shuts the child down;
231 /// `kill_on_drop` on the final `Arc` drop remains the hard backstop.
232 /// Without a runtime handle: sync close_stdin + start_kill, then drop.
233 fn spawn_bounded_cleanup(inner: Arc<LeaseInner>) {
234 match tokio::runtime::Handle::try_current() {
235 Ok(handle) => {
236 handle.spawn(async move {
237 let _ = tokio::time::timeout(CLEANUP_TOTAL_TIMEOUT, async {
238 let mut conn = inner.connection.lock().await;
239 conn.finish_shutdown().await;
240 })
241 .await;
242 });
243 }
244 Err(_) => {
245 if let Ok(mut conn) = inner.connection.try_lock() {
246 conn.start_kill();
247 }
248 }
249 }
250 }
251
252 /// Number of live (Ready or Starting) leases.
253 pub fn lease_count(&self) -> usize {
254 self.leases
255 .lock()
256 .unwrap_or_else(std::sync::PoisonError::into_inner)
257 .len()
258 }
259
260 /// Terminate one session's runtime lease on one SERVER (the lease unit:
261 /// one child process shared by that session's exact tools of that
262 /// server).
263 pub fn revoke_server_lease(&self, session: &SessionId, server: &str) {
264 let key = self.lease_key(session, server);
265 let removed = {
266 let mut map = self
267 .leases
268 .lock()
269 .unwrap_or_else(std::sync::PoisonError::into_inner);
270 map.remove(&key)
271 };
272 if let Some(state) = removed {
273 Self::cancel_state(state);
274 }
275 }
276
277 /// Revoke one EXACT tool grant's runtime backing. Grant revocation
278 /// itself is gate-level (the `SessionToolSet`); at the runtime layer the
279 /// lease unit is the server child, so this conservatively terminates
280 /// the shared server lease — the child loses all session access rather
281 /// than retaining any after an exact revocation.
282 pub fn revoke_exact_tool(&self, session: &SessionId, server: &str, server_tool_name: &str) {
283 let _ = server_tool_name; // exactness lives in the gate; lease unit is the server child
284 self.revoke_server_lease(session, server);
285 }
286
287 /// Terminate every lease held by one session. Returns how many.
288 /// Shared (`"*"`) leases are never a session's to terminate.
289 pub fn terminate_session(&self, session: &SessionId) -> usize {
290 let removed: Vec<LeaseState> = {
291 let mut map = self
292 .leases
293 .lock()
294 .unwrap_or_else(std::sync::PoisonError::into_inner);
295 let keys: Vec<(String, String)> = map
296 .keys()
297 .filter(|(s, _)| s != SHARED_SESSION_KEY && s == session.as_str())
298 .cloned()
299 .collect();
300 keys.into_iter().filter_map(|k| map.remove(&k)).collect()
301 };
302 let count = removed.len();
303 for state in removed {
304 Self::cancel_state(state);
305 }
306 count
307 }
308
309 /// Terminate every lease (runtime shutdown).
310 pub fn terminate_all(&self) {
311 let removed: Vec<LeaseState> = {
312 let mut map = self
313 .leases
314 .lock()
315 .unwrap_or_else(std::sync::PoisonError::into_inner);
316 map.drain().map(|(_, v)| v).collect()
317 };
318 for state in removed {
319 Self::cancel_state(state);
320 }
321 }
322
323 /// Reap Ready leases idle past the configured bound. Work per pass is
324 /// bounded by [`REAP_SCAN_MAX`] examined keys.
325 pub fn reap_idle(&self) {
326 let removed: Vec<LeaseState> = {
327 let mut map = self
328 .leases
329 .lock()
330 .unwrap_or_else(std::sync::PoisonError::into_inner);
331 let idle: Vec<(String, String)> = map
332 .iter()
333 .take(REAP_SCAN_MAX)
334 .filter(|(_, state)| match state {
335 LeaseState::Ready(inner) => {
336 inner
337 .last_used
338 .lock()
339 .unwrap_or_else(std::sync::PoisonError::into_inner)
340 .elapsed()
341 >= self.idle_max
342 }
343 LeaseState::Starting { .. } => false,
344 })
345 .map(|(k, _)| k.clone())
346 .collect();
347 idle.into_iter().filter_map(|k| map.remove(&k)).collect()
348 };
349 for state in removed {
350 Self::cancel_state(state);
351 }
352 }
353
354 /// Execute one EXACT already-gate-authorized MCP tool call. See module
355 /// docs for the lease/single-flight/validation/revocation contract.
356 pub async fn call_exact(
357 &self,
358 session: &SessionId,
359 server: &str,
360 expected_fingerprint: &str,
361 server_tool_name: &str,
362 expected_digest: &SchemaDigest,
363 params: Value,
364 ) -> Result<String, McpLeaseError> {
365 // 1. Current config + fingerprint (pure local; no lock held).
366 let Some(config) = (self.config_source)(server) else {
367 self.revoke_server_lease(session, server);
368 return Err(McpLeaseError::ServerNotConfigured(server.to_string()));
369 };
370 if server_config_fingerprint(&config) != expected_fingerprint {
371 self.revoke_server_lease(session, server);
372 return Err(McpLeaseError::FingerprintDrift(server.to_string()));
373 }
374 self.reap_idle();
375
376 let key = (lease_session_key(&config, session), server.to_string());
377 let inner: Arc<LeaseInner> = loop {
378 // 2. Single-flight acquisition. The guard-scoped decision below
379 // returns an owned action, so the map lock is NEVER held across
380 // I/O or awaits (also keeps this future Send).
381 enum Acquire {
382 Use(Arc<LeaseInner>),
383 CancelStale(LeaseState),
384 Wait(tokio::sync::watch::Receiver<bool>),
385 Start {
386 token: u64,
387 done_tx: tokio::sync::watch::Sender<bool>,
388 },
389 AtCapacity,
390 }
391 let action = {
392 let mut map = self
393 .leases
394 .lock()
395 .unwrap_or_else(std::sync::PoisonError::into_inner);
396 match map.get(&key) {
397 Some(LeaseState::Ready(inner)) => {
398 if inner.cancelled() || inner.fingerprint != expected_fingerprint {
399 match map.remove(&key) {
400 Some(state) => Acquire::CancelStale(state),
401 None => continue,
402 }
403 } else {
404 Acquire::Use(Arc::clone(inner))
405 }
406 }
407 Some(LeaseState::Starting { done, .. }) => Acquire::Wait(done.clone()),
408 None => {
409 if map.len() >= MAX_LIVE_LEASES {
410 Acquire::AtCapacity
411 } else {
412 let token = self
413 .next_token
414 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
415 let (done_tx, done_rx) = tokio::sync::watch::channel(false);
416 map.insert(
417 key.clone(),
418 LeaseState::Starting {
419 token,
420 done: done_rx,
421 },
422 );
423 Acquire::Start { token, done_tx }
424 }
425 }
426 }
427 };
428 match action {
429 Acquire::Use(inner) => break inner,
430 Acquire::CancelStale(state) => {
431 Self::cancel_state(state);
432 continue;
433 }
434 Acquire::AtCapacity => {
435 self.reap_idle();
436 let still_full = self
437 .leases
438 .lock()
439 .unwrap_or_else(std::sync::PoisonError::into_inner)
440 .len()
441 >= MAX_LIVE_LEASES;
442 if still_full {
443 return Err(McpLeaseError::Capacity(MAX_LIVE_LEASES));
444 }
445 continue;
446 }
447 Acquire::Wait(mut follower) => {
448 // Lost-wakeup-proof: the watch channel stores state, so
449 // if the starter finished between our map release and
450 // this await, `changed()` resolves immediately; sender
451 // drop also wakes us.
452 if !*follower.borrow() {
453 let _ = follower.changed().await;
454 }
455 continue;
456 }
457 Acquire::Start { token, done_tx } => {
458 let started = self
459 .start_lease(server, expected_fingerprint, &config)
460 .await
461 .map_err(|e| McpLeaseError::Transport(server.to_string(), e));
462 let ours = {
463 let mut map = self
464 .leases
465 .lock()
466 .unwrap_or_else(std::sync::PoisonError::into_inner);
467 let ours = matches!(
468 map.get(&key),
469 Some(LeaseState::Starting { token: t, .. }) if *t == token
470 );
471 match &started {
472 Ok(inner) if ours => {
473 map.insert(key.clone(), LeaseState::Ready(Arc::clone(inner)));
474 }
475 Err(_) if ours => {
476 map.remove(&key);
477 }
478 _ => {}
479 }
480 ours
481 };
482 match started {
483 Ok(inner) => {
484 if !ours {
485 // Terminated while starting (session end or
486 // revocation): never resurrect.
487 inner.set_cancelled();
488 if let Ok(mut conn) = inner.connection.try_lock() {
489 conn.shutdown().await;
490 }
491 let _ = done_tx.send(true);
492 return Err(McpLeaseError::Revoked(server.to_string()));
493 }
494 let _ = done_tx.send(true);
495 break inner;
496 }
497 Err(err) => {
498 let _ = done_tx.send(true);
499 return Err(err);
500 }
501 }
502 }
503 }
504 };
505
506 // 3. Exact validation against the pinned listing BEFORE any call.
507 match inner.listing.get(server_tool_name) {
508 None => {
509 self.revoke_server_lease(session, server);
510 return Err(McpLeaseError::NameNotListed(
511 server.to_string(),
512 server_tool_name.to_string(),
513 ));
514 }
515 Some(listed) if listed != expected_digest => {
516 self.revoke_server_lease(session, server);
517 return Err(McpLeaseError::SchemaMismatch(
518 server.to_string(),
519 server_tool_name.to_string(),
520 ));
521 }
522 Some(_) => {}
523 }
524
525 // 4. Call under the connection lock; the revocation gate check is
526 // taken while HOLDING that lock, so after `cancel()` returns no new
527 // call can pass. A call that already passed the gate may complete;
528 // if it then observes cancellation it kills the child before
529 // releasing the connection.
530 let mut conn = inner.connection.lock().await;
531 if inner.cancelled() {
532 return Err(McpLeaseError::Revoked(server.to_string()));
533 }
534 let result = conn.call_tool(server_tool_name, params).await;
535 if inner.cancelled() {
536 // Revoked while this already-started call ran: shut the child
537 // down (graceful EOF, bounded kill fallback) before releasing.
538 conn.shutdown().await;
539 }
540 drop(conn);
541 *inner
542 .last_used
543 .lock()
544 .unwrap_or_else(std::sync::PoisonError::into_inner) = Instant::now();
545 result.map_err(|e| McpLeaseError::Transport(server.to_string(), e))
546 }
547
548 /// Spawn + initialize + single bounded tools/list for one server. No
549 /// manager lock is held. Any failure kills the child before returning.
550 /// On success the listing is written back to the descriptor cache
551 /// (`SYNAPS_MCP_CACHE_WRITEBACK=0` disables) so the next boot advertises
552 /// the same dormant tools without spawning.
553 async fn start_lease(
554 &self,
555 server: &str,
556 fingerprint: &str,
557 config: &McpServerConfig,
558 ) -> Result<Arc<LeaseInner>, String> {
559 let mut conn = McpConnection::start(config).await?;
560 let defs = match conn.list_tools().await {
561 Ok(defs) => defs,
562 Err(e) => {
563 conn.start_kill();
564 return Err(e);
565 }
566 };
567 let mut listing = HashMap::new();
568 for def in &defs {
569 listing.insert(def.name.clone(), SchemaDigest::of_schema(&def.input_schema));
570 }
571 if descriptors::cache_writeback_enabled() {
572 let descriptors: Vec<descriptors::CachedToolDescriptor> = defs
573 .into_iter()
574 .map(|def| descriptors::CachedToolDescriptor {
575 name: def.name,
576 description: def.description,
577 input_schema: def.input_schema,
578 })
579 .collect();
580 let path = (self.cache_path)();
581 let server_owned = server.to_string();
582 let fingerprint_owned = fingerprint.to_string();
583 // Blocking file I/O (flock + atomic rename) off the reactor;
584 // failure is logged, never surfaced — the lease is already good.
585 let _ = tokio::task::spawn_blocking(move || {
586 match descriptors::record_server_listing(
587 &path,
588 &server_owned,
589 &fingerprint_owned,
590 &descriptors,
591 ) {
592 Ok(count) => tracing::debug!(server = %server_owned, tools = count,
593 "MCP descriptor cache written back after tools/list"),
594 Err(err) => tracing::warn!(server = %server_owned, error = %err,
595 "MCP descriptor cache write-back failed"),
596 }
597 })
598 .await;
599 }
600 Ok(Arc::new(LeaseInner {
601 fingerprint: fingerprint.to_string(),
602 connection: tokio::sync::Mutex::new(conn),
603 listing,
604 gate: std::sync::Mutex::new(false),
605 last_used: std::sync::Mutex::new(Instant::now()),
606 liveness: Arc::new(()),
607 }))
608 }
609
610 /// Test seam: the live lease keys `(session_key, server)`; shared
611 /// servers show [`SHARED_SESSION_KEY`].
612 #[doc(hidden)]
613 pub fn lease_keys_for_tests(&self) -> Vec<(String, String)> {
614 let map = self
615 .leases
616 .lock()
617 .unwrap_or_else(std::sync::PoisonError::into_inner);
618 let mut keys: Vec<_> = map.keys().cloned().collect();
619 keys.sort();
620 keys
621 }
622
623 /// Test seam: `Weak` ownership token of one lease. `upgrade() == None`
624 /// proves the manager (and any bounded shutdown task) released it.
625 #[doc(hidden)]
626 pub fn lease_liveness_for_tests(
627 &self,
628 session: &SessionId,
629 server: &str,
630 ) -> Option<std::sync::Weak<()>> {
631 let key = self.lease_key(session, server);
632 let map = self
633 .leases
634 .lock()
635 .unwrap_or_else(std::sync::PoisonError::into_inner);
636 match map.get(&key) {
637 Some(LeaseState::Ready(inner)) => Some(Arc::downgrade(&inner.liveness)),
638 _ => None,
639 }
640 }
641}
642
643impl Drop for McpRuntimeManager {
644 fn drop(&mut self) {
645 self.terminate_all();
646 }
647}
648
649impl std::fmt::Debug for McpRuntimeManager {
650 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
651 f.debug_struct("McpRuntimeManager")
652 .field("leases", &self.lease_count())
653 .finish_non_exhaustive()
654 }
655}
656
657/// Typed per-session lease capability handed to tool contexts: the exact
658/// session identity plus the shared manager. Never widens to a server
659/// grant — every call re-validates the exact tool and fingerprint.
660#[derive(Clone)]
661pub struct McpLeaseCapability {
662 session: SessionId,
663 manager: Arc<McpRuntimeManager>,
664}
665
666impl McpLeaseCapability {
667 pub fn new(session: SessionId, manager: Arc<McpRuntimeManager>) -> Self {
668 Self { session, manager }
669 }
670
671 /// The exact session identity this capability is scoped to.
672 pub fn session(&self) -> &SessionId {
673 &self.session
674 }
675
676 pub async fn call_exact(
677 &self,
678 server: &str,
679 expected_fingerprint: &str,
680 server_tool_name: &str,
681 expected_digest: &SchemaDigest,
682 params: Value,
683 ) -> Result<String, McpLeaseError> {
684 self.manager
685 .call_exact(
686 &self.session,
687 server,
688 expected_fingerprint,
689 server_tool_name,
690 expected_digest,
691 params,
692 )
693 .await
694 }
695}
696
697impl std::fmt::Debug for McpLeaseCapability {
698 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
699 f.debug_struct("McpLeaseCapability").finish_non_exhaustive()
700 }
701}
702
703/// Durable last-owner session scope. The Runtime mints ONE of these per
704/// tool session (wrapped in an `Arc` shared by every runtime clone and
705/// in-flight stream); streams HOLD it, never construct or drop their own.
706/// Leases therefore survive across provider turns, and only the drop of
707/// the LAST owner — true session end — terminates the session's leases.
708pub struct McpSessionEndGuard {
709 session: SessionId,
710 manager: Arc<McpRuntimeManager>,
711}
712
713impl McpSessionEndGuard {
714 pub fn new(session: SessionId, manager: Arc<McpRuntimeManager>) -> Self {
715 Self { session, manager }
716 }
717}
718
719impl Drop for McpSessionEndGuard {
720 fn drop(&mut self) {
721 let terminated = self.manager.terminate_session(&self.session);
722 if terminated > 0 {
723 tracing::debug!(
724 leases = terminated,
725 "MCP session leases terminated at stream end"
726 );
727 }
728 }
729}