Skip to main content

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}