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    /// The host-scope lease [`Self::acquire_handler_scope`] would use, only
644    /// if it is already Ready, live and launched from the current record.
645    /// Never starts, waits for or replaces a child: the subagent-worker
646    /// hook path ([`ExtensionHandler::handle_no_spawn`]).
647    fn ready_handler_scope(&self, plugin: &str) -> Option<Arc<LeaseInner>> {
648        let session = self.host_scope();
649        let (record, _validated) = self.record_and_validate(&session, plugin).ok()?;
650        let expected_fingerprint = record_fingerprint(&record);
651        let map = self
652            .leases
653            .lock()
654            .unwrap_or_else(std::sync::PoisonError::into_inner);
655        match map.get(&(session.as_str().to_string(), plugin.to_string())) {
656            Some(LeaseState::Ready(inner))
657                if !inner.cancelled() && inner.fingerprint == expected_fingerprint =>
658            {
659                Some(Arc::clone(inner))
660            }
661            _ => None,
662        }
663    }
664
665    /// Steps shared by every acquisition path: re-read the retained
666    /// internal launch record and RE-VALIDATE the manifest (permissions
667    /// included) before any spawn decision. Failure revokes the session's
668    /// plugin lease and fails closed.
669    fn record_and_validate(
670        &self,
671        session: &SessionId,
672        plugin: &str,
673    ) -> Result<
674        (
675            DeferredExtensionRecord,
676            super::manifest::ValidatedExtensionManifest,
677        ),
678        ExtensionLeaseError,
679    > {
680        let Some(record) = self.current_record(plugin) else {
681            self.revoke_plugin_lease(session, plugin);
682            return Err(ExtensionLeaseError::NotDeferred(plugin.to_string()));
683        };
684        let Ok(validated) = record.manifest.validate(plugin) else {
685            self.revoke_plugin_lease(session, plugin);
686            return Err(ExtensionLeaseError::ManifestInvalid(plugin.to_string()));
687        };
688        Ok((record, validated))
689    }
690
691    /// Execute one EXACT already-gate-authorized extension tool call. See
692    /// module docs for the lease/single-flight/validation contract.
693    /// Host-originated: the wire `session_id` is `null`.
694    pub async fn call_exact(
695        &self,
696        session: &SessionId,
697        plugin: &str,
698        tool_name: &str,
699        expected_digest: &SchemaDigest,
700        params: Value,
701    ) -> Result<Value, ExtensionLeaseError> {
702        self.call_exact_for(session, plugin, tool_name, expected_digest, params, None)
703            .await
704    }
705
706    /// [`Self::call_exact`] on behalf of a conversation: `conversation_id`
707    /// is forwarded as `tool.call` `params.session_id` (the owning
708    /// conversation, not the lease scope `session`).
709    pub async fn call_exact_for(
710        &self,
711        session: &SessionId,
712        plugin: &str,
713        tool_name: &str,
714        expected_digest: &SchemaDigest,
715        params: Value,
716        conversation_id: Option<&str>,
717    ) -> Result<Value, ExtensionLeaseError> {
718        // 1.-2. Record re-read + manifest re-validation (pure local).
719        let (record, validated) = self.record_and_validate(session, plugin)?;
720        // 3. Catalog-drift check: the CURRENT record must still declare
721        // this exact tool with the digest pinned in the catalog entry.
722        let declared_tools = record
723            .manifest
724            .deferred
725            .as_ref()
726            .map(|d| d.tools.as_slice())
727            .unwrap_or(&[]);
728        let still_declared = declared_tools.iter().any(|t| {
729            t.name == tool_name && &SchemaDigest::of_schema(&t.input_schema) == expected_digest
730        });
731        if !still_declared {
732            self.revoke_plugin_lease(session, plugin);
733            return Err(ExtensionLeaseError::DeclarationDrift(
734                plugin.to_string(),
735                tool_name.to_string(),
736            ));
737        }
738        self.reap_idle();
739
740        // 4. Single-flight acquisition of the validated per-plugin lease.
741        let inner = self
742            .acquire_ready(session, plugin, &record, &validated)
743            .await?;
744        // 5. Exact validation against the pinned validated listing BEFORE
745        // any call.
746        match inner.listing.get(tool_name) {
747            None => {
748                self.revoke_plugin_lease(session, plugin);
749                return Err(ExtensionLeaseError::NameNotListed(
750                    plugin.to_string(),
751                    tool_name.to_string(),
752                ));
753            }
754            Some(listed) if listed != expected_digest => {
755                self.revoke_plugin_lease(session, plugin);
756                return Err(ExtensionLeaseError::SchemaMismatch(
757                    plugin.to_string(),
758                    tool_name.to_string(),
759                ));
760            }
761            Some(_) => {}
762        }
763
764        // 6. Call under the per-lease permit; the revocation gate check is
765        // taken while HOLDING it, so after cancellation no new call passes.
766        let permit = inner.call_permit.lock().await;
767        if inner.cancelled() {
768            return Err(ExtensionLeaseError::Revoked(plugin.to_string()));
769        }
770        let result = inner
771            .handler
772            .call_tool_once_no_spawn(tool_name, params, conversation_id)
773            .await;
774        if inner.cancelled() {
775            // Revoked while this already-started call ran: shut the child
776            // down before releasing the permit.
777            inner.handler.shutdown().await;
778        }
779        drop(permit);
780        *inner
781            .last_used
782            .lock()
783            .unwrap_or_else(std::sync::PoisonError::into_inner) = Instant::now();
784        result.map_err(|e| {
785            if let Some(rest) = e.strip_prefix(EXTENSION_REPORTED_PREFIX) {
786                // Extension-controlled content: withheld, length only.
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        })
795    }
796
797    /// Spawn + set permissions + initialize ONCE for one plugin, then
798    /// require the runtime's declarations to match the manifest EXACTLY.
799    /// No manager lock is held. Any failure shuts the child down before
800    /// returning (`kill_on_drop` backs the local paths).
801    async fn start_lease(
802        &self,
803        plugin: &str,
804        record: &DeferredExtensionRecord,
805        validated: &super::manifest::ValidatedExtensionManifest,
806        fingerprint: SchemaDigest,
807    ) -> Result<Arc<LeaseInner>, ExtensionLeaseError> {
808        let manifest = &record.manifest;
809        let process = ProcessExtension::spawn_with_memory_policy(
810            plugin,
811            &manifest.command,
812            &manifest.args,
813            record.cwd.clone(),
814            record.exclusive_memory,
815        )
816        .await
817        .map_err(|e| ExtensionLeaseError::Transport(plugin.to_string(), bound_detail(&e)))?;
818        process.set_permissions(validated.permissions.clone()).await;
819        let caps = match process
820            .initialize(record.cwd.clone(), record.config.clone())
821            .await
822        {
823            Ok(caps) => caps,
824            Err(e) => {
825                process.shutdown().await;
826                let err = if let Some(rest) = e.strip_prefix(EXTENSION_REPORTED_PREFIX) {
827                    ExtensionLeaseError::ExtensionReported {
828                        plugin: plugin.to_string(),
829                        length: rest.trim().len(),
830                    }
831                } else {
832                    ExtensionLeaseError::Transport(plugin.to_string(), bound_detail(&e))
833                };
834                return Err(err);
835            }
836        };
837        let deferred = manifest.deferred.as_ref();
838        let declared_tools = deferred.map(|d| d.tools.as_slice()).unwrap_or(&[]);
839        let declared_providers = deferred.map(|d| d.providers.as_slice()).unwrap_or(&[]);
840        if let Err(reason) = validate_runtime_tool_declarations(declared_tools, &caps.tools) {
841            process.shutdown().await;
842            return Err(ExtensionLeaseError::DeclarationMismatch(
843                plugin.to_string(),
844                reason,
845            ));
846        }
847        if let Err(reason) =
848            validate_runtime_provider_declarations(declared_providers, &caps.providers)
849        {
850            process.shutdown().await;
851            return Err(ExtensionLeaseError::DeclarationMismatch(
852                plugin.to_string(),
853                reason,
854            ));
855        }
856        for decl in &caps.capabilities {
857            if super::runtime::process::validate_capability(decl, &validated.permissions).is_err() {
858                process.shutdown().await;
859                return Err(ExtensionLeaseError::DeclarationMismatch(
860                    plugin.to_string(),
861                    "capability_declaration_rejected",
862                ));
863            }
864        }
865        let mut listing = HashMap::new();
866        for spec in &caps.tools {
867            listing.insert(
868                spec.name.clone(),
869                SchemaDigest::of_schema(&spec.input_schema),
870            );
871        }
872        Ok(Arc::new(LeaseInner {
873            fingerprint,
874            handler: Arc::new(process),
875            call_permit: tokio::sync::Mutex::new(()),
876            listing,
877            gate: std::sync::Mutex::new(false),
878            last_used: std::sync::Mutex::new(Instant::now()),
879            liveness: Arc::new(()),
880        }))
881    }
882
883    /// Test seam: `Weak` ownership token of one lease. `upgrade() == None`
884    /// proves the manager (and any bounded shutdown task) released it.
885    #[doc(hidden)]
886    pub fn lease_liveness_for_tests(
887        &self,
888        session: &SessionId,
889        plugin: &str,
890    ) -> Option<std::sync::Weak<()>> {
891        let map = self
892            .leases
893            .lock()
894            .unwrap_or_else(std::sync::PoisonError::into_inner);
895        match map.get(&(session.as_str().to_string(), plugin.to_string())) {
896            Some(LeaseState::Ready(inner)) => Some(Arc::downgrade(&inner.liveness)),
897            _ => None,
898        }
899    }
900}
901
902impl Drop for ExtensionRuntimeManager {
903    fn drop(&mut self) {
904        self.terminate_all();
905    }
906}
907
908impl std::fmt::Debug for ExtensionRuntimeManager {
909    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
910        f.debug_struct("ExtensionRuntimeManager")
911            .field("leases", &self.lease_count())
912            .finish_non_exhaustive()
913    }
914}
915
916/// Typed per-session lease capability handed to tool contexts: the exact
917/// session identity plus the shared manager. Never widens to a plugin
918/// grant — every call re-validates the exact tool and launch record.
919#[derive(Clone)]
920pub struct ExtensionLeaseCapability {
921    session: SessionId,
922    manager: Arc<ExtensionRuntimeManager>,
923}
924
925impl ExtensionLeaseCapability {
926    pub fn new(session: SessionId, manager: Arc<ExtensionRuntimeManager>) -> Self {
927        Self { session, manager }
928    }
929
930    /// The exact session identity this capability is scoped to.
931    pub fn session(&self) -> &SessionId {
932        &self.session
933    }
934
935    pub async fn call_exact(
936        &self,
937        plugin: &str,
938        tool_name: &str,
939        expected_digest: &SchemaDigest,
940        params: Value,
941    ) -> Result<Value, ExtensionLeaseError> {
942        self.manager
943            .call_exact(&self.session, plugin, tool_name, expected_digest, params)
944            .await
945    }
946
947    /// [`Self::call_exact`] on behalf of a conversation; see
948    /// [`ExtensionRuntimeManager::call_exact_for`].
949    pub async fn call_exact_for(
950        &self,
951        plugin: &str,
952        tool_name: &str,
953        expected_digest: &SchemaDigest,
954        params: Value,
955        conversation_id: Option<&str>,
956    ) -> Result<Value, ExtensionLeaseError> {
957        self.manager
958            .call_exact_for(
959                &self.session,
960                plugin,
961                tool_name,
962                expected_digest,
963                params,
964                conversation_id,
965            )
966            .await
967    }
968}
969
970impl std::fmt::Debug for ExtensionLeaseCapability {
971    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
972        f.debug_struct("ExtensionLeaseCapability")
973            .finish_non_exhaustive()
974    }
975}
976
977/// Durable last-owner session scope. The Runtime mints ONE of these per
978/// tool session (wrapped in an `Arc` shared by every runtime clone and
979/// in-flight stream); streams HOLD it, never construct or drop their own.
980/// Leases therefore survive across provider turns, and only the drop of
981/// the LAST owner — true session end — terminates the session's leases.
982pub struct ExtensionSessionEndGuard {
983    session: SessionId,
984    manager: Arc<ExtensionRuntimeManager>,
985}
986
987impl ExtensionSessionEndGuard {
988    pub fn new(session: SessionId, manager: Arc<ExtensionRuntimeManager>) -> Self {
989        Self { session, manager }
990    }
991}
992
993impl Drop for ExtensionSessionEndGuard {
994    fn drop(&mut self) {
995        self.manager.terminate_session(&self.session);
996    }
997}
998
999impl std::fmt::Debug for ExtensionSessionEndGuard {
1000    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1001        f.debug_struct("ExtensionSessionEndGuard")
1002            .finish_non_exhaustive()
1003    }
1004}
1005
1006/// Lazy [`ExtensionHandler`] for deferred hook/provider/user capabilities
1007/// (Task 20 Commit C). Registered/subscribed at load WITHOUT a process;
1008/// the first AUTHORIZED use (a matching hook event delivered by the
1009/// permission-checked `HookBus` subscription, a selected provider
1010/// complete/stream, or an explicit user sidecar/command/settings action)
1011/// single-flights through the SAME per-plugin lease manager — starting
1012/// and exact-validating the child once. Discovery, load, search, and
1013/// diagnostics never acquire.
1014pub struct LazyExtensionHandler {
1015    plugin: String,
1016    manager: Arc<ExtensionRuntimeManager>,
1017}
1018
1019impl LazyExtensionHandler {
1020    pub fn new(plugin: &str, manager: Arc<ExtensionRuntimeManager>) -> Self {
1021        Self {
1022            plugin: plugin.to_string(),
1023            manager,
1024        }
1025    }
1026
1027    /// Acquire the shared validated lease (host scope) and hand back the
1028    /// live inner. Static-only failure text.
1029    async fn live(&self) -> Result<Arc<LeaseInner>, String> {
1030        self.manager
1031            .acquire_handler_scope(&self.plugin)
1032            .await
1033            .map_err(|e| e.to_string())
1034    }
1035
1036    fn touch(inner: &LeaseInner) {
1037        *inner
1038            .last_used
1039            .lock()
1040            .unwrap_or_else(std::sync::PoisonError::into_inner) = Instant::now();
1041    }
1042}
1043
1044#[async_trait::async_trait]
1045impl ExtensionHandler for LazyExtensionHandler {
1046    fn id(&self) -> &str {
1047        &self.plugin
1048    }
1049
1050    async fn handle(
1051        &self,
1052        event: &crate::extensions::hooks::events::HookEvent,
1053    ) -> crate::extensions::hooks::events::HookResult {
1054        // First AUTHORIZED matching event (the HookBus checked the
1055        // permission at subscribe time and matched the filter) starts the
1056        // child. Acquisition failure follows the documented eager hook
1057        // transport policy: warn + Continue (never a silent block).
1058        match self.live().await {
1059            Ok(inner) => {
1060                let result = inner.handler.handle(event).await;
1061                Self::touch(&inner);
1062                result
1063            }
1064            Err(error) => {
1065                tracing::warn!(
1066                    extension = %self.plugin,
1067                    error = %error,
1068                    "Deferred extension could not start for hook event — continuing",
1069                );
1070                crate::extensions::hooks::events::HookResult::Continue
1071            }
1072        }
1073    }
1074
1075    /// Worker hook path: only an already-running lease serves it; a
1076    /// deferred extension no foreground event has started (or one being
1077    /// restarted) passes through instead of being started from a worker's
1078    /// throwaway runtime.
1079    async fn handle_no_spawn(
1080        &self,
1081        event: &crate::extensions::hooks::events::HookEvent,
1082    ) -> crate::extensions::hooks::events::HookResult {
1083        match self.manager.ready_handler_scope(&self.plugin) {
1084            Some(inner) => {
1085                let result = inner.handler.handle_no_spawn(event).await;
1086                Self::touch(&inner);
1087                result
1088            }
1089            None => {
1090                tracing::debug!(
1091                    extension = %self.plugin,
1092                    "Deferred extension not running; worker hook event passes through",
1093                );
1094                crate::extensions::hooks::events::HookResult::Continue
1095            }
1096        }
1097    }
1098
1099    async fn call_tool(&self, _name: &str, _input: Value) -> Result<Value, String> {
1100        // FAIL CLOSED: deferred extension tool execution flows ONLY
1101        // through the gate-authorized exact lease path
1102        // (`ExtensionRuntimeManager::call_exact`, which pins and checks
1103        // the declared name + canonical schema digest). A generic handler
1104        // call with an arbitrary tool name would bypass those checks.
1105        Err(format!(
1106            "extension '{}' is deferred: tool calls are only served through the \
1107             gate-authorized exact activation path",
1108            self.plugin
1109        ))
1110    }
1111
1112    async fn provider_complete(
1113        &self,
1114        params: super::runtime::process::ProviderCompleteParams,
1115    ) -> Result<super::runtime::process::ProviderCompleteResult, String> {
1116        let inner = self.live().await?;
1117        let result = inner.handler.provider_complete(params).await;
1118        Self::touch(&inner);
1119        result
1120    }
1121
1122    async fn provider_stream(
1123        &self,
1124        params: super::runtime::process::ProviderCompleteParams,
1125        sink: tokio::sync::mpsc::Sender<super::runtime::process::ProviderStreamEvent>,
1126    ) -> Result<super::runtime::process::ProviderCompleteResult, String> {
1127        let inner = self.live().await?;
1128        let result = inner.handler.provider_stream(params, sink).await;
1129        Self::touch(&inner);
1130        result
1131    }
1132
1133    async fn invoke_command(
1134        &self,
1135        command: &str,
1136        args: Vec<String>,
1137        request_id: &str,
1138        sink: crate::extensions::invoke_output::InvokeEventSink,
1139    ) -> Result<Value, String> {
1140        let inner = self.live().await?;
1141        let result = inner
1142            .handler
1143            .invoke_command(command, args, request_id, sink)
1144            .await;
1145        Self::touch(&inner);
1146        result
1147    }
1148
1149    async fn get_info(&self) -> Result<crate::extensions::info::PluginInfo, String> {
1150        // Diagnostics must never spawn a deferred extension.
1151        Err("deferred extension info is unavailable without activation".to_string())
1152    }
1153
1154    async fn sidecar_spawn_args(&self) -> Result<crate::sidecar::spawn::SidecarSpawnArgs, String> {
1155        // Explicit user sidecar action: a legitimate acquisition trigger.
1156        let inner = self.live().await?;
1157        let result = inner.handler.sidecar_spawn_args().await;
1158        Self::touch(&inner);
1159        result
1160    }
1161
1162    async fn settings_editor_open(&self, category: &str, field: &str) -> Result<Value, String> {
1163        let inner = self.live().await?;
1164        let result = inner.handler.settings_editor_open(category, field).await;
1165        Self::touch(&inner);
1166        result
1167    }
1168
1169    async fn settings_editor_key(
1170        &self,
1171        category: &str,
1172        field: &str,
1173        key: &str,
1174    ) -> Result<Value, String> {
1175        let inner = self.live().await?;
1176        let result = inner
1177            .handler
1178            .settings_editor_key(category, field, key)
1179            .await;
1180        Self::touch(&inner);
1181        result
1182    }
1183
1184    async fn settings_editor_commit(
1185        &self,
1186        category: &str,
1187        field: &str,
1188        value: Value,
1189    ) -> Result<Value, String> {
1190        let inner = self.live().await?;
1191        let result = inner
1192            .handler
1193            .settings_editor_commit(category, field, value)
1194            .await;
1195        Self::touch(&inner);
1196        result
1197    }
1198
1199    async fn shutdown(&self) {
1200        // Terminates every session's lease for this plugin; never spawns.
1201        self.manager.revoke_plugin_all_sessions(&self.plugin);
1202    }
1203
1204    async fn health(&self) -> super::runtime::ExtensionHealth {
1205        // Dormant reporting only — health checks never spawn.
1206        if self
1207            .manager
1208            .has_ready_lease(&self.manager.host_scope(), &self.plugin)
1209        {
1210            super::runtime::ExtensionHealth::Running
1211        } else {
1212            super::runtime::ExtensionHealth::Loaded
1213        }
1214    }
1215}
1216
1217impl std::fmt::Debug for LazyExtensionHandler {
1218    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1219        f.debug_struct("LazyExtensionHandler")
1220            .field("plugin", &self.plugin)
1221            .finish_non_exhaustive()
1222    }
1223}