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 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 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 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 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 let (record, validated) = self.record_and_validate(session, plugin)?;
720 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 let inner = self
742 .acquire_ready(session, plugin, &record, &validated)
743 .await?;
744 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 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 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 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 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 #[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#[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 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 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
977pub 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
1006pub 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 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 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 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 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 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 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 self.manager.revoke_plugin_all_sessions(&self.plugin);
1202 }
1203
1204 async fn health(&self) -> super::runtime::ExtensionHealth {
1205 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}