Skip to main content

agent_engine/extensions/
lease.rs

1//! Task 20 (Commit B) — session-scoped exact extension runtime leases
2//! (spec §7.5, mirroring the §7.4 MCP lease discipline).
3//!
4//! Deferred manifest-declared extension tools execute ONLY through this
5//! manager, strictly AFTER the `ExecutionGate` has authorized the exact
6//! call. One `ExtensionRuntimeManager` is shared by the `ExtensionManager`
7//! (which owns the retained internal launch records) and the `Runtime`
8//! (which mints per-session capabilities and the durable session-end
9//! scope):
10//!
11//! - a lease is keyed by (session, plugin) and PINNED to the launch-record
12//!   fingerprint (manifest + cwd + resolved config) it was acquired under;
13//!   per-key SINGLE-FLIGHT acquisition (a `Starting` placeholder holding a
14//!   `watch` receiver whose stored state makes wakeups lost-proof)
15//!   guarantees concurrent first calls never spawn duplicate children, and
16//!   the manager map lock is NEVER held across process/pipe I/O;
17//! - acquisition re-reads the retained launch record, RE-VALIDATES the
18//!   manifest (permissions included), starts exactly one child
19//!   (`kill_on_drop`), initializes ONCE, and requires the runtime's
20//!   registered declarations to match the manifest's passive declarations
21//!   EXACTLY (tool names, descriptions, canonical schema digests — and
22//!   provider declarations when present) BEFORE any call; a mismatch shuts
23//!   the child down and fails closed;
24//! - every call validates the selected exact tool name and canonical
25//!   schema digest against the pinned validated listing BEFORE `tool.call`;
26//! - revocation/termination semantics are exact and honest: revocation
27//!   marks the lease cancelled under the per-lease call permit; a call
28//!   that already passed the gate check MAY complete and then shuts the
29//!   child down; NO call passes after revocation returns. `kill_on_drop`
30//!   on the child is the hard backstop once the last `Arc` drops;
31//! - idle leases are reaped opportunistically with a CONSTANT per-pass
32//!   scan cap; session end is the drop of the LAST owner of the shared
33//!   durable [`ExtensionSessionEndGuard`] scope. No PID signalling, no
34//!   long-lived task, no unbounded channel;
35//! - lease failures carry ONLY static text, local identities, and numeric
36//!   metadata — extension-reported error content is withheld (length
37//!   only), and no launch-record (secret) config ever appears anywhere.
38
39use std::collections::HashMap;
40use std::sync::Arc;
41use std::time::{Duration, Instant};
42
43use serde_json::Value;
44use thiserror::Error;
45
46use super::context_provider::{
47    resolve_context_provider, ContextProviderId, ContextProviderLookupError,
48    RegisteredContextProviderDescriptor, SharedContextProviderCatalog,
49};
50use super::lifecycle::{
51    validate_runtime_provider_declarations, validate_runtime_tool_declarations,
52};
53use super::manager::{DeferredExtensionRecord, SharedDeferredRecords};
54use super::runtime::process::ProcessExtension;
55use super::runtime::ExtensionHandler;
56use crate::tools::activation::SessionId;
57use crate::tools::catalog::SchemaDigest;
58
59/// Default idle bound for opportunistic reaping (parity with MCP leases).
60pub const DEFAULT_IDLE_MAX: Duration = Duration::from_secs(300);
61/// Hard cap on simultaneously live leases per manager.
62pub const MAX_LIVE_LEASES: usize = 64;
63/// Reap passes scan up to the full capped map so idle entries can never
64/// starve behind an active HashMap iteration prefix.
65pub const REAP_SCAN_MAX: usize = MAX_LIVE_LEASES;
66/// Outer total bound on one lease cleanup task.
67const CLEANUP_TOTAL_TIMEOUT: Duration = Duration::from_secs(5);
68/// Bound on locally-generated transport detail strings surfaced in errors.
69const TRANSPORT_DETAIL_MAX: usize = 200;
70/// Stable prefix `ProcessExtension` puts on extension-REPORTED JSON-RPC
71/// error messages; content behind it is extension-controlled and withheld.
72const EXTENSION_REPORTED_PREFIX: &str = "Extension error:";
73
74/// Typed lease failures. Static text, local identities, and numeric
75/// metadata only — never extension-controlled content, never config.
76#[derive(Debug, Error)]
77pub enum ExtensionLeaseError {
78    #[error("extension '{0}' has no retained deferred launch record; lease revoked")]
79    NotDeferred(String),
80    #[error("extension '{0}' manifest failed re-validation at lease acquisition; lease revoked")]
81    ManifestInvalid(String),
82    #[error(
83        "extension '{0}' launch record changed since the lease was acquired; lease invalidated"
84    )]
85    LaunchDrift(String),
86    #[error("extension '{0}' no longer declares tool '{1}' with the pinned schema; lease revoked")]
87    DeclarationDrift(String, String),
88    #[error("extension '{0}' runtime declarations do not match the manifest ({1}); child terminated without any call")]
89    DeclarationMismatch(String, &'static str),
90    #[error("extension '{0}' did not register expected tool '{1}'; poisoned lease terminated without calling")]
91    NameNotListed(String, String),
92    #[error("extension '{0}' registered a schema for tool '{1}' that does not match the pinned digest; poisoned lease terminated without calling")]
93    SchemaMismatch(String, String),
94    #[error("extension lease for '{0}' was revoked before the call could start")]
95    Revoked(String),
96    #[error("extension lease capacity of {0} live leases reached; try again after idle leases are reaped")]
97    Capacity(usize),
98    #[error("extension '{0}' transport failure: {1}")]
99    Transport(String, String),
100    #[error("extension '{plugin}' reported a tool error ({length} bytes, content withheld)")]
101    ExtensionReported { plugin: String, length: usize },
102}
103
104impl ExtensionLeaseError {
105    /// Whether this failure class poisons the pinned declaration/grant
106    /// (spec §7.5): record removal, permission/config/catalog drift, and
107    /// runtime declaration mismatches invalidate the EXACT session grant.
108    /// Transport/capacity/revocation-race and extension-reported tool
109    /// errors are transient and must NOT revoke.
110    pub fn revokes_exact_grant(&self) -> bool {
111        matches!(
112            self,
113            Self::NotDeferred(_)
114                | Self::ManifestInvalid(_)
115                | Self::LaunchDrift(_)
116                | Self::DeclarationDrift(_, _)
117                | Self::DeclarationMismatch(_, _)
118                | Self::NameNotListed(_, _)
119                | Self::SchemaMismatch(_, _)
120        )
121    }
122}
123
124/// Truncate a locally-generated detail string to a hard byte bound (valid
125/// UTF-8 preserved). Extension-REPORTED content never reaches this — it is
126/// withheld entirely by the caller.
127fn bound_detail(detail: &str) -> String {
128    if detail.len() <= TRANSPORT_DETAIL_MAX {
129        return detail.to_string();
130    }
131    let mut cut = TRANSPORT_DETAIL_MAX;
132    while cut > 0 && !detail.is_char_boundary(cut) {
133        cut -= 1;
134    }
135    format!("{}… ({} bytes total)", &detail[..cut], detail.len())
136}
137
138/// Canonical fingerprint of a retained launch record: manifest identity,
139/// process cwd, and RESOLVED config. Computed locally; never logged.
140fn record_fingerprint(record: &DeferredExtensionRecord) -> SchemaDigest {
141    SchemaDigest::of_schema(&serde_json::json!({
142        "manifest": serde_json::to_value(&record.manifest).unwrap_or(Value::Null),
143        "cwd": record.cwd.as_ref().map(|p| p.display().to_string()),
144        "config": record.config,
145    }))
146}
147
148struct LeaseInner {
149    fingerprint: SchemaDigest,
150    handler: Arc<ProcessExtension>,
151    /// Per-lease call/revocation permit: the cancellation gate is read
152    /// while HOLDING this, so no new call passes after `cancel` returns.
153    call_permit: tokio::sync::Mutex<()>,
154    /// Pinned validated listing from the single initialize (tool name →
155    /// canonical schema digest).
156    listing: HashMap<String, SchemaDigest>,
157    /// `true` = cancelled. Held only for flag reads/writes — never I/O.
158    gate: std::sync::Mutex<bool>,
159    last_used: std::sync::Mutex<Instant>,
160    /// Test-observable ownership token.
161    liveness: Arc<()>,
162}
163
164impl LeaseInner {
165    fn cancelled(&self) -> bool {
166        *self
167            .gate
168            .lock()
169            .unwrap_or_else(std::sync::PoisonError::into_inner)
170    }
171
172    fn set_cancelled(&self) {
173        *self
174            .gate
175            .lock()
176            .unwrap_or_else(std::sync::PoisonError::into_inner) = true;
177    }
178
179    fn interrupt(&self) {
180        self.set_cancelled();
181
182        // An idle lease can take the normal bounded cleanup path and send a
183        // graceful shutdown frame. Only an in-flight call needs the hard
184        // interrupt to drain its pending responder; otherwise racing an
185        // unconditional force shutdown against graceful cleanup removes the
186        // child first and used to make shutdown's restart-capable call spawn
187        // a replacement process during teardown.
188        if self.call_permit.try_lock().is_err() {
189            let handler = Arc::clone(&self.handler);
190            if let Ok(handle) = tokio::runtime::Handle::try_current() {
191                handle.spawn(async move {
192                    handler.force_shutdown().await;
193                });
194            }
195        }
196    }
197}
198
199enum LeaseState {
200    /// Single-flight placeholder; the `watch` channel stores completion
201    /// state so followers can never miss the wakeup.
202    Starting {
203        token: u64,
204        done: tokio::sync::watch::Receiver<bool>,
205    },
206    Ready(Arc<LeaseInner>),
207}
208
209/// Session-scoped exact extension runtime lease manager. See module docs.
210pub struct ExtensionRuntimeManager {
211    leases: std::sync::Mutex<HashMap<(String, String), LeaseState>>,
212    /// The `ExtensionManager`'s retained internal launch records (shared
213    /// handle; sync lock held only for map operations, never across I/O).
214    records: SharedDeferredRecords,
215    /// The `ExtensionManager`'s published context-provider catalog
216    /// snapshot (task A6; same sharing discipline as `records`). Read by
217    /// [`Self::resolve_context_provider`] for enable-time memory-context
218    /// provider validation — a PURE catalog read that never spawns.
219    context_catalog: SharedContextProviderCatalog,
220    idle_max: Duration,
221    next_token: std::sync::atomic::AtomicU64,
222    /// Host-scope session identity for HANDLER-based acquisition (hook
223    /// events, provider routing, user sidecar/command APIs). Bound at
224    /// engine boot to the Runtime's durable tool session, so a MIXED
225    /// extension's tool leases and handler leases share ONE key — one
226    /// shared child process per plugin. Unbound (tests/manual) falls back
227    /// to a fixed private scope.
228    host_scope: std::sync::OnceLock<SessionId>,
229}
230
231impl ExtensionRuntimeManager {
232    pub(crate) fn new(
233        records: SharedDeferredRecords,
234        context_catalog: SharedContextProviderCatalog,
235        idle_max: Duration,
236    ) -> Self {
237        Self {
238            leases: std::sync::Mutex::new(HashMap::new()),
239            records,
240            context_catalog,
241            idle_max,
242            next_token: std::sync::atomic::AtomicU64::new(1),
243            host_scope: std::sync::OnceLock::new(),
244        }
245    }
246
247    /// Task A6: resolve one memory-context provider request against the
248    /// published catalog of loaded extensions' declared context providers.
249    /// PURE CATALOG READ — never spawns a process, never touches lease
250    /// state (consistent with "discovery/status spawn zero processes").
251    /// Exact-scope, fail-closed matching per
252    /// [`resolve_context_provider`]: `Ok` only when exactly one loaded
253    /// declaration matches; zero matches fail `NotRegistered` and
254    /// overlapping declarations fail `Ambiguous`. Returns an owned clone
255    /// of the dormant descriptor (bounded metadata only).
256    pub fn resolve_context_provider(
257        &self,
258        requested: Option<&ContextProviderId>,
259    ) -> Result<RegisteredContextProviderDescriptor, ContextProviderLookupError> {
260        let catalog = self
261            .context_catalog
262            .lock()
263            .unwrap_or_else(std::sync::PoisonError::into_inner);
264        resolve_context_provider(&catalog, requested).cloned()
265    }
266
267    /// Bind the host-scope session identity used by handler-based
268    /// acquisition (first binder wins; the engine binds the Runtime's
269    /// durable tool session at boot so Mixed extensions share one child).
270    pub fn bind_host_scope(&self, session: SessionId) {
271        let _ = self.host_scope.set(session);
272    }
273
274    /// The host-scope session identity (bound, or the private fallback).
275    pub fn host_scope(&self) -> SessionId {
276        self.host_scope
277            .get_or_init(|| {
278                SessionId::parse("extension-host-scope")
279                    .expect("static host scope identity is valid")
280            })
281            .clone()
282    }
283
284    fn current_record(&self, plugin: &str) -> Option<DeferredExtensionRecord> {
285        self.records
286            .lock()
287            .unwrap_or_else(std::sync::PoisonError::into_inner)
288            .get(plugin)
289            .cloned()
290    }
291
292    /// Task B4: pinned schema digest of one plugin's CURRENTLY declared
293    /// deferred tool — a pure launch-record read that never spawns. The
294    /// memory recall dispatch uses it to address the provider's declared
295    /// recall tool through the same exact-digest [`Self::call_exact`] gate
296    /// as every other extension call (drift still re-checked there).
297    pub(crate) fn declared_tool_digest(
298        &self,
299        plugin: &str,
300        tool_name: &str,
301    ) -> Option<SchemaDigest> {
302        let record = self.current_record(plugin)?;
303        record
304            .manifest
305            .deferred
306            .as_ref()?
307            .tools
308            .iter()
309            .find(|tool| tool.name == tool_name)
310            .map(|tool| SchemaDigest::of_schema(&tool.input_schema))
311    }
312
313    /// Mark cancelled and hand the lease to ONE bounded cleanup task
314    /// (graceful `shutdown` RPC + kill under an outer timeout). Without a
315    /// runtime handle the last `Arc` drop triggers `kill_on_drop` — the
316    /// hard backstop in every path.
317    fn cancel_state(state: LeaseState) {
318        match state {
319            LeaseState::Ready(inner) => {
320                inner.interrupt();
321                Self::spawn_bounded_cleanup(inner);
322            }
323            // The in-flight starter detects its placeholder is gone and
324            // shuts its own child down instead of publishing the lease.
325            LeaseState::Starting { .. } => {}
326        }
327    }
328
329    fn spawn_bounded_cleanup(inner: Arc<LeaseInner>) {
330        if let Ok(handle) = tokio::runtime::Handle::try_current() {
331            handle.spawn(async move {
332                let _ = tokio::time::timeout(CLEANUP_TOTAL_TIMEOUT, async {
333                    // Wait for any in-flight call to release the permit so
334                    // shutdown never races an active tool.call frame.
335                    let _permit = inner.call_permit.lock().await;
336                    inner.handler.shutdown().await;
337                })
338                .await;
339            });
340        }
341        // No runtime: dropping the Arc (here or when bailing callers
342        // release theirs) kills the child via kill_on_drop.
343    }
344
345    /// Number of live (Ready or Starting) leases.
346    pub fn lease_count(&self) -> usize {
347        self.leases
348            .lock()
349            .unwrap_or_else(std::sync::PoisonError::into_inner)
350            .len()
351    }
352
353    /// Terminate one session's runtime lease for one PLUGIN (the lease
354    /// unit: one child shared by that session's exact tools of the plugin).
355    pub fn revoke_plugin_lease(&self, session: &SessionId, plugin: &str) {
356        let removed = {
357            let mut map = self
358                .leases
359                .lock()
360                .unwrap_or_else(std::sync::PoisonError::into_inner);
361            map.remove(&(session.as_str().to_string(), plugin.to_string()))
362        };
363        if let Some(state) = removed {
364            Self::cancel_state(state);
365        }
366    }
367
368    /// Revoke one EXACT tool grant's runtime backing. Grant revocation
369    /// itself is gate-level; at the runtime layer the lease unit is the
370    /// plugin child, so this conservatively terminates the shared lease.
371    pub fn revoke_exact_tool(&self, session: &SessionId, plugin: &str, tool_name: &str) {
372        let _ = tool_name; // exactness lives in the gate; lease unit is the plugin child
373        self.revoke_plugin_lease(session, plugin);
374    }
375
376    /// Terminate every session's lease for one plugin (manager unload/
377    /// reload path). Returns how many were terminated.
378    pub fn revoke_plugin_all_sessions(&self, plugin: &str) -> usize {
379        let removed: Vec<LeaseState> = {
380            let mut map = self
381                .leases
382                .lock()
383                .unwrap_or_else(std::sync::PoisonError::into_inner);
384            let keys: Vec<(String, String)> =
385                map.keys().filter(|(_, p)| p == plugin).cloned().collect();
386            keys.into_iter().filter_map(|k| map.remove(&k)).collect()
387        };
388        let count = removed.len();
389        for state in removed {
390            Self::cancel_state(state);
391        }
392        count
393    }
394
395    /// Terminate every lease held by one session. Returns how many.
396    pub fn terminate_session(&self, session: &SessionId) -> usize {
397        let removed: Vec<LeaseState> = {
398            let mut map = self
399                .leases
400                .lock()
401                .unwrap_or_else(std::sync::PoisonError::into_inner);
402            let keys: Vec<(String, String)> = map
403                .keys()
404                .filter(|(s, _)| s == session.as_str())
405                .cloned()
406                .collect();
407            keys.into_iter().filter_map(|k| map.remove(&k)).collect()
408        };
409        let count = removed.len();
410        for state in removed {
411            Self::cancel_state(state);
412        }
413        count
414    }
415
416    /// Terminate every lease (runtime shutdown).
417    pub fn terminate_all(&self) {
418        let removed: Vec<LeaseState> = {
419            let mut map = self
420                .leases
421                .lock()
422                .unwrap_or_else(std::sync::PoisonError::into_inner);
423            map.drain().map(|(_, v)| v).collect()
424        };
425        for state in removed {
426            Self::cancel_state(state);
427        }
428    }
429
430    /// Reap Ready leases idle past the configured bound. Work per pass is
431    /// bounded by [`REAP_SCAN_MAX`] examined keys.
432    pub fn reap_idle(&self) {
433        let removed: Vec<LeaseState> = {
434            let mut map = self
435                .leases
436                .lock()
437                .unwrap_or_else(std::sync::PoisonError::into_inner);
438            let idle: Vec<(String, String)> = map
439                .iter()
440                .take(REAP_SCAN_MAX)
441                .filter(|(_, state)| match state {
442                    LeaseState::Ready(inner) => {
443                        inner
444                            .last_used
445                            .lock()
446                            .unwrap_or_else(std::sync::PoisonError::into_inner)
447                            .elapsed()
448                            >= self.idle_max
449                    }
450                    LeaseState::Starting { .. } => false,
451                })
452                .map(|(k, _)| k.clone())
453                .collect();
454            idle.into_iter().filter_map(|k| map.remove(&k)).collect()
455        };
456        for state in removed {
457            Self::cancel_state(state);
458        }
459    }
460
461    /// Single-flight acquisition of the validated Ready lease for one
462    /// (session, plugin) key. Shared by exact tool calls and handler-based
463    /// (hook/provider/user) acquisition — one child per key. The map lock
464    /// is never held across I/O; a live lease whose pinned launch
465    /// fingerprint no longer matches the CURRENT record is launch drift
466    /// and fails closed after terminating the stale child.
467    async fn acquire_ready(
468        &self,
469        session: &SessionId,
470        plugin: &str,
471        record: &DeferredExtensionRecord,
472        validated: &super::manifest::ValidatedExtensionManifest,
473    ) -> Result<Arc<LeaseInner>, ExtensionLeaseError> {
474        let expected_fingerprint = record_fingerprint(record);
475        let key = (session.as_str().to_string(), plugin.to_string());
476        let inner: Arc<LeaseInner> = loop {
477            // 4. Single-flight acquisition: the guard-scoped decision
478            // returns an owned action, so the map lock is NEVER held
479            // across I/O or awaits.
480            enum Acquire {
481                Use(Arc<LeaseInner>),
482                CancelStale(LeaseState),
483                Wait(tokio::sync::watch::Receiver<bool>),
484                Start {
485                    token: u64,
486                    done_tx: tokio::sync::watch::Sender<bool>,
487                },
488                AtCapacity,
489            }
490            let action = {
491                let mut map = self
492                    .leases
493                    .lock()
494                    .unwrap_or_else(std::sync::PoisonError::into_inner);
495                match map.get(&key) {
496                    Some(LeaseState::Ready(inner)) => {
497                        if inner.cancelled() || inner.fingerprint != expected_fingerprint {
498                            match map.remove(&key) {
499                                Some(state) => Acquire::CancelStale(state),
500                                None => continue,
501                            }
502                        } else {
503                            Acquire::Use(Arc::clone(inner))
504                        }
505                    }
506                    Some(LeaseState::Starting { done, .. }) => Acquire::Wait(done.clone()),
507                    None => {
508                        if map.len() >= MAX_LIVE_LEASES {
509                            Acquire::AtCapacity
510                        } else {
511                            let token = self
512                                .next_token
513                                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
514                            let (done_tx, done_rx) = tokio::sync::watch::channel(false);
515                            map.insert(
516                                key.clone(),
517                                LeaseState::Starting {
518                                    token,
519                                    done: done_rx,
520                                },
521                            );
522                            Acquire::Start { token, done_tx }
523                        }
524                    }
525                }
526            };
527            match action {
528                Acquire::Use(inner) => break inner,
529                Acquire::CancelStale(state) => {
530                    // A live lease whose pinned fingerprint no longer
531                    // matches the current record is launch drift: the
532                    // pinned grant falls with it (spec §7.5).
533                    let was_drift = matches!(
534                        &state,
535                        LeaseState::Ready(inner) if inner.fingerprint != expected_fingerprint
536                            && !inner.cancelled()
537                    );
538                    Self::cancel_state(state);
539                    if was_drift {
540                        return Err(ExtensionLeaseError::LaunchDrift(plugin.to_string()));
541                    }
542                    continue;
543                }
544                Acquire::AtCapacity => {
545                    self.reap_idle();
546                    let still_full = self
547                        .leases
548                        .lock()
549                        .unwrap_or_else(std::sync::PoisonError::into_inner)
550                        .len()
551                        >= MAX_LIVE_LEASES;
552                    if still_full {
553                        return Err(ExtensionLeaseError::Capacity(MAX_LIVE_LEASES));
554                    }
555                    continue;
556                }
557                Acquire::Wait(mut follower) => {
558                    // Lost-wakeup-proof: the watch channel stores state.
559                    if !*follower.borrow() {
560                        let _ = follower.changed().await;
561                    }
562                    continue;
563                }
564                Acquire::Start { token, done_tx } => {
565                    let started = self
566                        .start_lease(plugin, record, validated, expected_fingerprint.clone())
567                        .await;
568                    let ours = {
569                        let mut map = self
570                            .leases
571                            .lock()
572                            .unwrap_or_else(std::sync::PoisonError::into_inner);
573                        let ours = matches!(
574                            map.get(&key),
575                            Some(LeaseState::Starting { token: t, .. }) if *t == token
576                        );
577                        match &started {
578                            Ok(inner) if ours => {
579                                map.insert(key.clone(), LeaseState::Ready(Arc::clone(inner)));
580                            }
581                            Err(_) if ours => {
582                                map.remove(&key);
583                            }
584                            _ => {}
585                        }
586                        ours
587                    };
588                    match started {
589                        Ok(inner) => {
590                            if !ours {
591                                // Terminated while starting (session end or
592                                // revocation): never resurrect.
593                                inner.set_cancelled();
594                                inner.handler.shutdown().await;
595                                let _ = done_tx.send(true);
596                                return Err(ExtensionLeaseError::Revoked(plugin.to_string()));
597                            }
598                            let _ = done_tx.send(true);
599                            break inner;
600                        }
601                        Err(err) => {
602                            let _ = done_tx.send(true);
603                            return Err(err);
604                        }
605                    }
606                }
607            }
608        };
609
610        Ok(inner)
611    }
612
613    /// Acquire the shared validated per-plugin lease under the HOST scope
614    /// for handler-based capabilities (hook events, provider routing, user
615    /// sidecar/command/settings APIs). Same record re-validation, exact
616    /// runtime declaration matching, single-flight, and teardown rules as
617    /// exact tool calls — and, when the host scope is bound to the
618    /// Runtime's tool session (engine boot), the SAME lease key: a Mixed
619    /// extension runs ONE shared child.
620    async fn acquire_handler_scope(
621        &self,
622        plugin: &str,
623    ) -> Result<Arc<LeaseInner>, ExtensionLeaseError> {
624        let session = self.host_scope();
625        let (record, validated) = self.record_and_validate(&session, plugin)?;
626        self.reap_idle();
627        self.acquire_ready(&session, plugin, &record, &validated)
628            .await
629    }
630
631    /// Whether a Ready (non-cancelled) lease exists for one key.
632    fn has_ready_lease(&self, session: &SessionId, plugin: &str) -> bool {
633        let map = self
634            .leases
635            .lock()
636            .unwrap_or_else(std::sync::PoisonError::into_inner);
637        matches!(
638            map.get(&(session.as_str().to_string(), plugin.to_string())),
639            Some(LeaseState::Ready(inner)) if !inner.cancelled()
640        )
641    }
642
643    /// Steps shared by every acquisition path: re-read the retained
644    /// internal launch record and RE-VALIDATE the manifest (permissions
645    /// included) before any spawn decision. Failure revokes the session's
646    /// plugin lease and fails closed.
647    fn record_and_validate(
648        &self,
649        session: &SessionId,
650        plugin: &str,
651    ) -> Result<
652        (
653            DeferredExtensionRecord,
654            super::manifest::ValidatedExtensionManifest,
655        ),
656        ExtensionLeaseError,
657    > {
658        let Some(record) = self.current_record(plugin) else {
659            self.revoke_plugin_lease(session, plugin);
660            return Err(ExtensionLeaseError::NotDeferred(plugin.to_string()));
661        };
662        let Ok(validated) = record.manifest.validate(plugin) else {
663            self.revoke_plugin_lease(session, plugin);
664            return Err(ExtensionLeaseError::ManifestInvalid(plugin.to_string()));
665        };
666        Ok((record, validated))
667    }
668
669    /// Execute one EXACT already-gate-authorized extension tool call. See
670    /// module docs for the lease/single-flight/validation contract.
671    pub async fn call_exact(
672        &self,
673        session: &SessionId,
674        plugin: &str,
675        tool_name: &str,
676        expected_digest: &SchemaDigest,
677        params: Value,
678    ) -> Result<Value, ExtensionLeaseError> {
679        // 1.-2. Record re-read + manifest re-validation (pure local).
680        let (record, validated) = self.record_and_validate(session, plugin)?;
681        // 3. Catalog-drift check: the CURRENT record must still declare
682        // this exact tool with the digest pinned in the catalog entry.
683        let declared_tools = record
684            .manifest
685            .deferred
686            .as_ref()
687            .map(|d| d.tools.as_slice())
688            .unwrap_or(&[]);
689        let still_declared = declared_tools.iter().any(|t| {
690            t.name == tool_name && &SchemaDigest::of_schema(&t.input_schema) == expected_digest
691        });
692        if !still_declared {
693            self.revoke_plugin_lease(session, plugin);
694            return Err(ExtensionLeaseError::DeclarationDrift(
695                plugin.to_string(),
696                tool_name.to_string(),
697            ));
698        }
699        self.reap_idle();
700
701        // 4. Single-flight acquisition of the validated per-plugin lease.
702        let inner = self
703            .acquire_ready(session, plugin, &record, &validated)
704            .await?;
705        // 5. Exact validation against the pinned validated listing BEFORE
706        // any call.
707        match inner.listing.get(tool_name) {
708            None => {
709                self.revoke_plugin_lease(session, plugin);
710                return Err(ExtensionLeaseError::NameNotListed(
711                    plugin.to_string(),
712                    tool_name.to_string(),
713                ));
714            }
715            Some(listed) if listed != expected_digest => {
716                self.revoke_plugin_lease(session, plugin);
717                return Err(ExtensionLeaseError::SchemaMismatch(
718                    plugin.to_string(),
719                    tool_name.to_string(),
720                ));
721            }
722            Some(_) => {}
723        }
724
725        // 6. Call under the per-lease permit; the revocation gate check is
726        // taken while HOLDING it, so after cancellation no new call passes.
727        let permit = inner.call_permit.lock().await;
728        if inner.cancelled() {
729            return Err(ExtensionLeaseError::Revoked(plugin.to_string()));
730        }
731        let result = inner
732            .handler
733            .call_tool_once_no_spawn(tool_name, params)
734            .await;
735        if inner.cancelled() {
736            // Revoked while this already-started call ran: shut the child
737            // down before releasing the permit.
738            inner.handler.shutdown().await;
739        }
740        drop(permit);
741        *inner
742            .last_used
743            .lock()
744            .unwrap_or_else(std::sync::PoisonError::into_inner) = Instant::now();
745        result.map_err(|e| {
746            if let Some(rest) = e.strip_prefix(EXTENSION_REPORTED_PREFIX) {
747                // Extension-controlled content: withheld, length only.
748                ExtensionLeaseError::ExtensionReported {
749                    plugin: plugin.to_string(),
750                    length: rest.trim().len(),
751                }
752            } else {
753                ExtensionLeaseError::Transport(plugin.to_string(), bound_detail(&e))
754            }
755        })
756    }
757
758    /// Spawn + set permissions + initialize ONCE for one plugin, then
759    /// require the runtime's declarations to match the manifest EXACTLY.
760    /// No manager lock is held. Any failure shuts the child down before
761    /// returning (`kill_on_drop` backs the local paths).
762    async fn start_lease(
763        &self,
764        plugin: &str,
765        record: &DeferredExtensionRecord,
766        validated: &super::manifest::ValidatedExtensionManifest,
767        fingerprint: SchemaDigest,
768    ) -> Result<Arc<LeaseInner>, ExtensionLeaseError> {
769        let manifest = &record.manifest;
770        let process = ProcessExtension::spawn_with_cwd(
771            plugin,
772            &manifest.command,
773            &manifest.args,
774            record.cwd.clone(),
775        )
776        .await
777        .map_err(|e| ExtensionLeaseError::Transport(plugin.to_string(), bound_detail(&e)))?;
778        process.set_permissions(validated.permissions.clone()).await;
779        let caps = match process
780            .initialize(record.cwd.clone(), record.config.clone())
781            .await
782        {
783            Ok(caps) => caps,
784            Err(e) => {
785                process.shutdown().await;
786                let err = if let Some(rest) = e.strip_prefix(EXTENSION_REPORTED_PREFIX) {
787                    ExtensionLeaseError::ExtensionReported {
788                        plugin: plugin.to_string(),
789                        length: rest.trim().len(),
790                    }
791                } else {
792                    ExtensionLeaseError::Transport(plugin.to_string(), bound_detail(&e))
793                };
794                return Err(err);
795            }
796        };
797        let deferred = manifest.deferred.as_ref();
798        let declared_tools = deferred.map(|d| d.tools.as_slice()).unwrap_or(&[]);
799        let declared_providers = deferred.map(|d| d.providers.as_slice()).unwrap_or(&[]);
800        if let Err(reason) = validate_runtime_tool_declarations(declared_tools, &caps.tools) {
801            process.shutdown().await;
802            return Err(ExtensionLeaseError::DeclarationMismatch(
803                plugin.to_string(),
804                reason,
805            ));
806        }
807        if let Err(reason) =
808            validate_runtime_provider_declarations(declared_providers, &caps.providers)
809        {
810            process.shutdown().await;
811            return Err(ExtensionLeaseError::DeclarationMismatch(
812                plugin.to_string(),
813                reason,
814            ));
815        }
816        for decl in &caps.capabilities {
817            if super::runtime::process::validate_capability(decl, &validated.permissions).is_err() {
818                process.shutdown().await;
819                return Err(ExtensionLeaseError::DeclarationMismatch(
820                    plugin.to_string(),
821                    "capability_declaration_rejected",
822                ));
823            }
824        }
825        let mut listing = HashMap::new();
826        for spec in &caps.tools {
827            listing.insert(
828                spec.name.clone(),
829                SchemaDigest::of_schema(&spec.input_schema),
830            );
831        }
832        Ok(Arc::new(LeaseInner {
833            fingerprint,
834            handler: Arc::new(process),
835            call_permit: tokio::sync::Mutex::new(()),
836            listing,
837            gate: std::sync::Mutex::new(false),
838            last_used: std::sync::Mutex::new(Instant::now()),
839            liveness: Arc::new(()),
840        }))
841    }
842
843    /// Test seam: `Weak` ownership token of one lease. `upgrade() == None`
844    /// proves the manager (and any bounded shutdown task) released it.
845    #[doc(hidden)]
846    pub fn lease_liveness_for_tests(
847        &self,
848        session: &SessionId,
849        plugin: &str,
850    ) -> Option<std::sync::Weak<()>> {
851        let map = self
852            .leases
853            .lock()
854            .unwrap_or_else(std::sync::PoisonError::into_inner);
855        match map.get(&(session.as_str().to_string(), plugin.to_string())) {
856            Some(LeaseState::Ready(inner)) => Some(Arc::downgrade(&inner.liveness)),
857            _ => None,
858        }
859    }
860}
861
862impl Drop for ExtensionRuntimeManager {
863    fn drop(&mut self) {
864        self.terminate_all();
865    }
866}
867
868impl std::fmt::Debug for ExtensionRuntimeManager {
869    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
870        f.debug_struct("ExtensionRuntimeManager")
871            .field("leases", &self.lease_count())
872            .finish_non_exhaustive()
873    }
874}
875
876/// Typed per-session lease capability handed to tool contexts: the exact
877/// session identity plus the shared manager. Never widens to a plugin
878/// grant — every call re-validates the exact tool and launch record.
879#[derive(Clone)]
880pub struct ExtensionLeaseCapability {
881    session: SessionId,
882    manager: Arc<ExtensionRuntimeManager>,
883}
884
885impl ExtensionLeaseCapability {
886    pub fn new(session: SessionId, manager: Arc<ExtensionRuntimeManager>) -> Self {
887        Self { session, manager }
888    }
889
890    /// The exact session identity this capability is scoped to.
891    pub fn session(&self) -> &SessionId {
892        &self.session
893    }
894
895    pub async fn call_exact(
896        &self,
897        plugin: &str,
898        tool_name: &str,
899        expected_digest: &SchemaDigest,
900        params: Value,
901    ) -> Result<Value, ExtensionLeaseError> {
902        self.manager
903            .call_exact(&self.session, plugin, tool_name, expected_digest, params)
904            .await
905    }
906}
907
908impl std::fmt::Debug for ExtensionLeaseCapability {
909    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
910        f.debug_struct("ExtensionLeaseCapability")
911            .finish_non_exhaustive()
912    }
913}
914
915/// Durable last-owner session scope. The Runtime mints ONE of these per
916/// tool session (wrapped in an `Arc` shared by every runtime clone and
917/// in-flight stream); streams HOLD it, never construct or drop their own.
918/// Leases therefore survive across provider turns, and only the drop of
919/// the LAST owner — true session end — terminates the session's leases.
920pub struct ExtensionSessionEndGuard {
921    session: SessionId,
922    manager: Arc<ExtensionRuntimeManager>,
923}
924
925impl ExtensionSessionEndGuard {
926    pub fn new(session: SessionId, manager: Arc<ExtensionRuntimeManager>) -> Self {
927        Self { session, manager }
928    }
929}
930
931impl Drop for ExtensionSessionEndGuard {
932    fn drop(&mut self) {
933        self.manager.terminate_session(&self.session);
934    }
935}
936
937impl std::fmt::Debug for ExtensionSessionEndGuard {
938    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
939        f.debug_struct("ExtensionSessionEndGuard")
940            .finish_non_exhaustive()
941    }
942}
943
944/// Lazy [`ExtensionHandler`] for deferred hook/provider/user capabilities
945/// (Task 20 Commit C). Registered/subscribed at load WITHOUT a process;
946/// the first AUTHORIZED use (a matching hook event delivered by the
947/// permission-checked `HookBus` subscription, a selected provider
948/// complete/stream, or an explicit user sidecar/command/settings action)
949/// single-flights through the SAME per-plugin lease manager — starting
950/// and exact-validating the child once. Discovery, load, search, and
951/// diagnostics never acquire.
952pub struct LazyExtensionHandler {
953    plugin: String,
954    manager: Arc<ExtensionRuntimeManager>,
955}
956
957impl LazyExtensionHandler {
958    pub fn new(plugin: &str, manager: Arc<ExtensionRuntimeManager>) -> Self {
959        Self {
960            plugin: plugin.to_string(),
961            manager,
962        }
963    }
964
965    /// Acquire the shared validated lease (host scope) and hand back the
966    /// live inner. Static-only failure text.
967    async fn live(&self) -> Result<Arc<LeaseInner>, String> {
968        self.manager
969            .acquire_handler_scope(&self.plugin)
970            .await
971            .map_err(|e| e.to_string())
972    }
973
974    fn touch(inner: &LeaseInner) {
975        *inner
976            .last_used
977            .lock()
978            .unwrap_or_else(std::sync::PoisonError::into_inner) = Instant::now();
979    }
980}
981
982#[async_trait::async_trait]
983impl ExtensionHandler for LazyExtensionHandler {
984    fn id(&self) -> &str {
985        &self.plugin
986    }
987
988    async fn handle(
989        &self,
990        event: &crate::extensions::hooks::events::HookEvent,
991    ) -> crate::extensions::hooks::events::HookResult {
992        // First AUTHORIZED matching event (the HookBus checked the
993        // permission at subscribe time and matched the filter) starts the
994        // child. Acquisition failure follows the documented eager hook
995        // transport policy: warn + Continue (never a silent block).
996        match self.live().await {
997            Ok(inner) => {
998                let result = inner.handler.handle(event).await;
999                Self::touch(&inner);
1000                result
1001            }
1002            Err(error) => {
1003                tracing::warn!(
1004                    extension = %self.plugin,
1005                    error = %error,
1006                    "Deferred extension could not start for hook event — continuing",
1007                );
1008                crate::extensions::hooks::events::HookResult::Continue
1009            }
1010        }
1011    }
1012
1013    async fn call_tool(&self, _name: &str, _input: Value) -> Result<Value, String> {
1014        // FAIL CLOSED: deferred extension tool execution flows ONLY
1015        // through the gate-authorized exact lease path
1016        // (`ExtensionRuntimeManager::call_exact`, which pins and checks
1017        // the declared name + canonical schema digest). A generic handler
1018        // call with an arbitrary tool name would bypass those checks.
1019        Err(format!(
1020            "extension '{}' is deferred: tool calls are only served through the \
1021             gate-authorized exact activation path",
1022            self.plugin
1023        ))
1024    }
1025
1026    async fn provider_complete(
1027        &self,
1028        params: super::runtime::process::ProviderCompleteParams,
1029    ) -> Result<super::runtime::process::ProviderCompleteResult, String> {
1030        let inner = self.live().await?;
1031        let result = inner.handler.provider_complete(params).await;
1032        Self::touch(&inner);
1033        result
1034    }
1035
1036    async fn provider_stream(
1037        &self,
1038        params: super::runtime::process::ProviderCompleteParams,
1039        sink: tokio::sync::mpsc::Sender<super::runtime::process::ProviderStreamEvent>,
1040    ) -> Result<super::runtime::process::ProviderCompleteResult, String> {
1041        let inner = self.live().await?;
1042        let result = inner.handler.provider_stream(params, sink).await;
1043        Self::touch(&inner);
1044        result
1045    }
1046
1047    async fn invoke_command(
1048        &self,
1049        command: &str,
1050        args: Vec<String>,
1051        request_id: &str,
1052        sink: crate::extensions::invoke_output::InvokeEventSink,
1053    ) -> Result<Value, String> {
1054        let inner = self.live().await?;
1055        let result = inner
1056            .handler
1057            .invoke_command(command, args, request_id, sink)
1058            .await;
1059        Self::touch(&inner);
1060        result
1061    }
1062
1063    async fn get_info(&self) -> Result<crate::extensions::info::PluginInfo, String> {
1064        // Diagnostics must never spawn a deferred extension.
1065        Err("deferred extension info is unavailable without activation".to_string())
1066    }
1067
1068    async fn sidecar_spawn_args(&self) -> Result<crate::sidecar::spawn::SidecarSpawnArgs, String> {
1069        // Explicit user sidecar action: a legitimate acquisition trigger.
1070        let inner = self.live().await?;
1071        let result = inner.handler.sidecar_spawn_args().await;
1072        Self::touch(&inner);
1073        result
1074    }
1075
1076    async fn settings_editor_open(&self, category: &str, field: &str) -> Result<Value, String> {
1077        let inner = self.live().await?;
1078        let result = inner.handler.settings_editor_open(category, field).await;
1079        Self::touch(&inner);
1080        result
1081    }
1082
1083    async fn settings_editor_key(
1084        &self,
1085        category: &str,
1086        field: &str,
1087        key: &str,
1088    ) -> Result<Value, String> {
1089        let inner = self.live().await?;
1090        let result = inner
1091            .handler
1092            .settings_editor_key(category, field, key)
1093            .await;
1094        Self::touch(&inner);
1095        result
1096    }
1097
1098    async fn settings_editor_commit(
1099        &self,
1100        category: &str,
1101        field: &str,
1102        value: Value,
1103    ) -> Result<Value, String> {
1104        let inner = self.live().await?;
1105        let result = inner
1106            .handler
1107            .settings_editor_commit(category, field, value)
1108            .await;
1109        Self::touch(&inner);
1110        result
1111    }
1112
1113    async fn shutdown(&self) {
1114        // Terminates every session's lease for this plugin; never spawns.
1115        self.manager.revoke_plugin_all_sessions(&self.plugin);
1116    }
1117
1118    async fn health(&self) -> super::runtime::ExtensionHealth {
1119        // Dormant reporting only — health checks never spawn.
1120        if self
1121            .manager
1122            .has_ready_lease(&self.manager.host_scope(), &self.plugin)
1123        {
1124            super::runtime::ExtensionHealth::Running
1125        } else {
1126            super::runtime::ExtensionHealth::Loaded
1127        }
1128    }
1129}
1130
1131impl std::fmt::Debug for LazyExtensionHandler {
1132    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1133        f.debug_struct("LazyExtensionHandler")
1134            .field("plugin", &self.plugin)
1135            .finish_non_exhaustive()
1136    }
1137}