1use 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
59pub const DEFAULT_IDLE_MAX: Duration = Duration::from_secs(300);
61pub const MAX_LIVE_LEASES: usize = 64;
63pub const REAP_SCAN_MAX: usize = MAX_LIVE_LEASES;
66const CLEANUP_TOTAL_TIMEOUT: Duration = Duration::from_secs(5);
68const TRANSPORT_DETAIL_MAX: usize = 200;
70const EXTENSION_REPORTED_PREFIX: &str = "Extension error:";
73
74#[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 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
124fn 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
138fn 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 call_permit: tokio::sync::Mutex<()>,
154 listing: HashMap<String, SchemaDigest>,
157 gate: std::sync::Mutex<bool>,
159 last_used: std::sync::Mutex<Instant>,
160 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 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 Starting {
203 token: u64,
204 done: tokio::sync::watch::Receiver<bool>,
205 },
206 Ready(Arc<LeaseInner>),
207}
208
209pub struct ExtensionRuntimeManager {
211 leases: std::sync::Mutex<HashMap<(String, String), LeaseState>>,
212 records: SharedDeferredRecords,
215 context_catalog: SharedContextProviderCatalog,
220 idle_max: Duration,
221 next_token: std::sync::atomic::AtomicU64,
222 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 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 pub fn bind_host_scope(&self, session: SessionId) {
271 let _ = self.host_scope.set(session);
272 }
273
274 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 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 fn cancel_state(state: LeaseState) {
318 match state {
319 LeaseState::Ready(inner) => {
320 inner.interrupt();
321 Self::spawn_bounded_cleanup(inner);
322 }
323 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 let _permit = inner.call_permit.lock().await;
336 inner.handler.shutdown().await;
337 })
338 .await;
339 });
340 }
341 }
344
345 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 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 pub fn revoke_exact_tool(&self, session: &SessionId, plugin: &str, tool_name: &str) {
372 let _ = tool_name; self.revoke_plugin_lease(session, plugin);
374 }
375
376 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 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 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 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 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 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 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 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 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 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 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 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 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 let (record, validated) = self.record_and_validate(session, plugin)?;
681 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 let inner = self
703 .acquire_ready(session, plugin, &record, &validated)
704 .await?;
705 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 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 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 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 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_memory_policy(
771 plugin,
772 &manifest.command,
773 &manifest.args,
774 record.cwd.clone(),
775 record.exclusive_memory,
776 )
777 .await
778 .map_err(|e| ExtensionLeaseError::Transport(plugin.to_string(), bound_detail(&e)))?;
779 process.set_permissions(validated.permissions.clone()).await;
780 let caps = match process
781 .initialize(record.cwd.clone(), record.config.clone())
782 .await
783 {
784 Ok(caps) => caps,
785 Err(e) => {
786 process.shutdown().await;
787 let err = if let Some(rest) = e.strip_prefix(EXTENSION_REPORTED_PREFIX) {
788 ExtensionLeaseError::ExtensionReported {
789 plugin: plugin.to_string(),
790 length: rest.trim().len(),
791 }
792 } else {
793 ExtensionLeaseError::Transport(plugin.to_string(), bound_detail(&e))
794 };
795 return Err(err);
796 }
797 };
798 let deferred = manifest.deferred.as_ref();
799 let declared_tools = deferred.map(|d| d.tools.as_slice()).unwrap_or(&[]);
800 let declared_providers = deferred.map(|d| d.providers.as_slice()).unwrap_or(&[]);
801 if let Err(reason) = validate_runtime_tool_declarations(declared_tools, &caps.tools) {
802 process.shutdown().await;
803 return Err(ExtensionLeaseError::DeclarationMismatch(
804 plugin.to_string(),
805 reason,
806 ));
807 }
808 if let Err(reason) =
809 validate_runtime_provider_declarations(declared_providers, &caps.providers)
810 {
811 process.shutdown().await;
812 return Err(ExtensionLeaseError::DeclarationMismatch(
813 plugin.to_string(),
814 reason,
815 ));
816 }
817 for decl in &caps.capabilities {
818 if super::runtime::process::validate_capability(decl, &validated.permissions).is_err() {
819 process.shutdown().await;
820 return Err(ExtensionLeaseError::DeclarationMismatch(
821 plugin.to_string(),
822 "capability_declaration_rejected",
823 ));
824 }
825 }
826 let mut listing = HashMap::new();
827 for spec in &caps.tools {
828 listing.insert(
829 spec.name.clone(),
830 SchemaDigest::of_schema(&spec.input_schema),
831 );
832 }
833 Ok(Arc::new(LeaseInner {
834 fingerprint,
835 handler: Arc::new(process),
836 call_permit: tokio::sync::Mutex::new(()),
837 listing,
838 gate: std::sync::Mutex::new(false),
839 last_used: std::sync::Mutex::new(Instant::now()),
840 liveness: Arc::new(()),
841 }))
842 }
843
844 #[doc(hidden)]
847 pub fn lease_liveness_for_tests(
848 &self,
849 session: &SessionId,
850 plugin: &str,
851 ) -> Option<std::sync::Weak<()>> {
852 let map = self
853 .leases
854 .lock()
855 .unwrap_or_else(std::sync::PoisonError::into_inner);
856 match map.get(&(session.as_str().to_string(), plugin.to_string())) {
857 Some(LeaseState::Ready(inner)) => Some(Arc::downgrade(&inner.liveness)),
858 _ => None,
859 }
860 }
861}
862
863impl Drop for ExtensionRuntimeManager {
864 fn drop(&mut self) {
865 self.terminate_all();
866 }
867}
868
869impl std::fmt::Debug for ExtensionRuntimeManager {
870 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
871 f.debug_struct("ExtensionRuntimeManager")
872 .field("leases", &self.lease_count())
873 .finish_non_exhaustive()
874 }
875}
876
877#[derive(Clone)]
881pub struct ExtensionLeaseCapability {
882 session: SessionId,
883 manager: Arc<ExtensionRuntimeManager>,
884}
885
886impl ExtensionLeaseCapability {
887 pub fn new(session: SessionId, manager: Arc<ExtensionRuntimeManager>) -> Self {
888 Self { session, manager }
889 }
890
891 pub fn session(&self) -> &SessionId {
893 &self.session
894 }
895
896 pub async fn call_exact(
897 &self,
898 plugin: &str,
899 tool_name: &str,
900 expected_digest: &SchemaDigest,
901 params: Value,
902 ) -> Result<Value, ExtensionLeaseError> {
903 self.manager
904 .call_exact(&self.session, plugin, tool_name, expected_digest, params)
905 .await
906 }
907}
908
909impl std::fmt::Debug for ExtensionLeaseCapability {
910 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
911 f.debug_struct("ExtensionLeaseCapability")
912 .finish_non_exhaustive()
913 }
914}
915
916pub struct ExtensionSessionEndGuard {
922 session: SessionId,
923 manager: Arc<ExtensionRuntimeManager>,
924}
925
926impl ExtensionSessionEndGuard {
927 pub fn new(session: SessionId, manager: Arc<ExtensionRuntimeManager>) -> Self {
928 Self { session, manager }
929 }
930}
931
932impl Drop for ExtensionSessionEndGuard {
933 fn drop(&mut self) {
934 self.manager.terminate_session(&self.session);
935 }
936}
937
938impl std::fmt::Debug for ExtensionSessionEndGuard {
939 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
940 f.debug_struct("ExtensionSessionEndGuard")
941 .finish_non_exhaustive()
942 }
943}
944
945pub struct LazyExtensionHandler {
954 plugin: String,
955 manager: Arc<ExtensionRuntimeManager>,
956}
957
958impl LazyExtensionHandler {
959 pub fn new(plugin: &str, manager: Arc<ExtensionRuntimeManager>) -> Self {
960 Self {
961 plugin: plugin.to_string(),
962 manager,
963 }
964 }
965
966 async fn live(&self) -> Result<Arc<LeaseInner>, String> {
969 self.manager
970 .acquire_handler_scope(&self.plugin)
971 .await
972 .map_err(|e| e.to_string())
973 }
974
975 fn touch(inner: &LeaseInner) {
976 *inner
977 .last_used
978 .lock()
979 .unwrap_or_else(std::sync::PoisonError::into_inner) = Instant::now();
980 }
981}
982
983#[async_trait::async_trait]
984impl ExtensionHandler for LazyExtensionHandler {
985 fn id(&self) -> &str {
986 &self.plugin
987 }
988
989 async fn handle(
990 &self,
991 event: &crate::extensions::hooks::events::HookEvent,
992 ) -> crate::extensions::hooks::events::HookResult {
993 match self.live().await {
998 Ok(inner) => {
999 let result = inner.handler.handle(event).await;
1000 Self::touch(&inner);
1001 result
1002 }
1003 Err(error) => {
1004 tracing::warn!(
1005 extension = %self.plugin,
1006 error = %error,
1007 "Deferred extension could not start for hook event — continuing",
1008 );
1009 crate::extensions::hooks::events::HookResult::Continue
1010 }
1011 }
1012 }
1013
1014 async fn call_tool(&self, _name: &str, _input: Value) -> Result<Value, String> {
1015 Err(format!(
1021 "extension '{}' is deferred: tool calls are only served through the \
1022 gate-authorized exact activation path",
1023 self.plugin
1024 ))
1025 }
1026
1027 async fn provider_complete(
1028 &self,
1029 params: super::runtime::process::ProviderCompleteParams,
1030 ) -> Result<super::runtime::process::ProviderCompleteResult, String> {
1031 let inner = self.live().await?;
1032 let result = inner.handler.provider_complete(params).await;
1033 Self::touch(&inner);
1034 result
1035 }
1036
1037 async fn provider_stream(
1038 &self,
1039 params: super::runtime::process::ProviderCompleteParams,
1040 sink: tokio::sync::mpsc::Sender<super::runtime::process::ProviderStreamEvent>,
1041 ) -> Result<super::runtime::process::ProviderCompleteResult, String> {
1042 let inner = self.live().await?;
1043 let result = inner.handler.provider_stream(params, sink).await;
1044 Self::touch(&inner);
1045 result
1046 }
1047
1048 async fn invoke_command(
1049 &self,
1050 command: &str,
1051 args: Vec<String>,
1052 request_id: &str,
1053 sink: crate::extensions::invoke_output::InvokeEventSink,
1054 ) -> Result<Value, String> {
1055 let inner = self.live().await?;
1056 let result = inner
1057 .handler
1058 .invoke_command(command, args, request_id, sink)
1059 .await;
1060 Self::touch(&inner);
1061 result
1062 }
1063
1064 async fn get_info(&self) -> Result<crate::extensions::info::PluginInfo, String> {
1065 Err("deferred extension info is unavailable without activation".to_string())
1067 }
1068
1069 async fn sidecar_spawn_args(&self) -> Result<crate::sidecar::spawn::SidecarSpawnArgs, String> {
1070 let inner = self.live().await?;
1072 let result = inner.handler.sidecar_spawn_args().await;
1073 Self::touch(&inner);
1074 result
1075 }
1076
1077 async fn settings_editor_open(&self, category: &str, field: &str) -> Result<Value, String> {
1078 let inner = self.live().await?;
1079 let result = inner.handler.settings_editor_open(category, field).await;
1080 Self::touch(&inner);
1081 result
1082 }
1083
1084 async fn settings_editor_key(
1085 &self,
1086 category: &str,
1087 field: &str,
1088 key: &str,
1089 ) -> Result<Value, String> {
1090 let inner = self.live().await?;
1091 let result = inner
1092 .handler
1093 .settings_editor_key(category, field, key)
1094 .await;
1095 Self::touch(&inner);
1096 result
1097 }
1098
1099 async fn settings_editor_commit(
1100 &self,
1101 category: &str,
1102 field: &str,
1103 value: Value,
1104 ) -> Result<Value, String> {
1105 let inner = self.live().await?;
1106 let result = inner
1107 .handler
1108 .settings_editor_commit(category, field, value)
1109 .await;
1110 Self::touch(&inner);
1111 result
1112 }
1113
1114 async fn shutdown(&self) {
1115 self.manager.revoke_plugin_all_sessions(&self.plugin);
1117 }
1118
1119 async fn health(&self) -> super::runtime::ExtensionHealth {
1120 if self
1122 .manager
1123 .has_ready_lease(&self.manager.host_scope(), &self.plugin)
1124 {
1125 super::runtime::ExtensionHealth::Running
1126 } else {
1127 super::runtime::ExtensionHealth::Loaded
1128 }
1129 }
1130}
1131
1132impl std::fmt::Debug for LazyExtensionHandler {
1133 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1134 f.debug_struct("LazyExtensionHandler")
1135 .field("plugin", &self.plugin)
1136 .finish_non_exhaustive()
1137 }
1138}