1use std::collections::HashMap;
40use std::path::Path;
41use std::sync::Arc;
42use std::time::{Duration, Instant};
43
44use serde_json::Value;
45use thiserror::Error;
46
47use super::lifecycle::{
48 validate_runtime_provider_declarations, validate_runtime_tool_declarations,
49};
50use super::manager::{DeferredExtensionRecord, SharedDeferredRecords};
51use super::runtime::process::ProcessExtension;
52use super::runtime::ExtensionHandler;
53use crate::tools::activation::SessionId;
54use crate::tools::catalog::SchemaDigest;
55use agent_core::memory::author::Author;
56
57pub const DEFAULT_IDLE_MAX: Duration = Duration::from_secs(300);
59pub const MAX_LIVE_LEASES: usize = 64;
61pub const REAP_SCAN_MAX: usize = MAX_LIVE_LEASES;
64const CLEANUP_TOTAL_TIMEOUT: Duration = Duration::from_secs(5);
66const TRANSPORT_DETAIL_MAX: usize = 200;
68const EXTENSION_REPORTED_PREFIX: &str = "Extension error:";
71
72#[derive(Debug, Error)]
75pub enum ExtensionLeaseError {
76 #[error("extension '{0}' has no retained deferred launch record; lease revoked")]
77 NotDeferred(String),
78 #[error("extension '{0}' manifest failed re-validation at lease acquisition; lease revoked")]
79 ManifestInvalid(String),
80 #[error(
81 "extension '{0}' launch record changed since the lease was acquired; lease invalidated"
82 )]
83 LaunchDrift(String),
84 #[error("extension '{0}' no longer declares tool '{1}' with the pinned schema; lease revoked")]
85 DeclarationDrift(String, String),
86 #[error("extension '{0}' runtime declarations do not match the manifest ({1}); child terminated without any call")]
87 DeclarationMismatch(String, &'static str),
88 #[error("extension '{0}' did not register expected tool '{1}'; poisoned lease terminated without calling")]
89 NameNotListed(String, String),
90 #[error("extension '{0}' registered a schema for tool '{1}' that does not match the pinned digest; poisoned lease terminated without calling")]
91 SchemaMismatch(String, String),
92 #[error("extension lease for '{0}' was revoked before the call could start")]
93 Revoked(String),
94 #[error("extension lease capacity of {0} live leases reached; try again after idle leases are reaped")]
95 Capacity(usize),
96 #[error("extension '{0}' transport failure: {1}")]
97 Transport(String, String),
98 #[error("extension '{plugin}' reported a tool error ({length} bytes, content withheld)")]
99 ExtensionReported { plugin: String, length: usize },
100}
101
102impl ExtensionLeaseError {
103 pub fn revokes_exact_grant(&self) -> bool {
109 matches!(
110 self,
111 Self::NotDeferred(_)
112 | Self::ManifestInvalid(_)
113 | Self::LaunchDrift(_)
114 | Self::DeclarationDrift(_, _)
115 | Self::DeclarationMismatch(_, _)
116 | Self::NameNotListed(_, _)
117 | Self::SchemaMismatch(_, _)
118 )
119 }
120}
121
122fn bound_detail(detail: &str) -> String {
126 if detail.len() <= TRANSPORT_DETAIL_MAX {
127 return detail.to_string();
128 }
129 let mut cut = TRANSPORT_DETAIL_MAX;
130 while cut > 0 && !detail.is_char_boundary(cut) {
131 cut -= 1;
132 }
133 format!("{}… ({} bytes total)", &detail[..cut], detail.len())
134}
135
136fn record_fingerprint(record: &DeferredExtensionRecord) -> SchemaDigest {
139 SchemaDigest::of_schema(&serde_json::json!({
140 "manifest": serde_json::to_value(&record.manifest).unwrap_or(Value::Null),
141 "cwd": record.cwd.as_ref().map(|p| p.display().to_string()),
142 "config": record.config,
143 }))
144}
145
146struct LeaseInner {
147 fingerprint: SchemaDigest,
148 handler: Arc<ProcessExtension>,
149 call_permit: tokio::sync::Mutex<()>,
152 listing: HashMap<String, SchemaDigest>,
155 gate: std::sync::Mutex<bool>,
157 last_used: std::sync::Mutex<Instant>,
158 liveness: Arc<()>,
160}
161
162impl LeaseInner {
163 fn cancelled(&self) -> bool {
164 *self
165 .gate
166 .lock()
167 .unwrap_or_else(std::sync::PoisonError::into_inner)
168 }
169
170 fn set_cancelled(&self) {
171 *self
172 .gate
173 .lock()
174 .unwrap_or_else(std::sync::PoisonError::into_inner) = true;
175 }
176
177 fn interrupt(&self) {
178 self.set_cancelled();
179
180 if self.call_permit.try_lock().is_err() {
187 let handler = Arc::clone(&self.handler);
188 if let Ok(handle) = tokio::runtime::Handle::try_current() {
189 handle.spawn(async move {
190 handler.force_shutdown().await;
191 });
192 }
193 }
194 }
195}
196
197enum LeaseState {
198 Starting {
201 token: u64,
202 done: tokio::sync::watch::Receiver<bool>,
203 },
204 Ready(Arc<LeaseInner>),
205}
206
207pub struct ExtensionRuntimeManager {
209 leases: std::sync::Mutex<HashMap<(String, String), LeaseState>>,
210 records: SharedDeferredRecords,
213 idle_max: Duration,
214 next_token: std::sync::atomic::AtomicU64,
215 host_scope: std::sync::OnceLock<SessionId>,
222}
223
224impl ExtensionRuntimeManager {
225 pub(crate) fn new(
226 records: SharedDeferredRecords,
227 idle_max: Duration,
228 ) -> Self {
229 Self {
230 leases: std::sync::Mutex::new(HashMap::new()),
231 records,
232 idle_max,
233 next_token: std::sync::atomic::AtomicU64::new(1),
234 host_scope: std::sync::OnceLock::new(),
235 }
236 }
237
238 pub fn bind_host_scope(&self, session: SessionId) {
242 let _ = self.host_scope.set(session);
243 }
244
245 pub fn host_scope(&self) -> SessionId {
247 self.host_scope
248 .get_or_init(|| {
249 SessionId::parse("extension-host-scope")
250 .expect("static host scope identity is valid")
251 })
252 .clone()
253 }
254
255 fn current_record(&self, plugin: &str) -> Option<DeferredExtensionRecord> {
256 self.records
257 .lock()
258 .unwrap_or_else(std::sync::PoisonError::into_inner)
259 .get(plugin)
260 .cloned()
261 }
262
263 fn cancel_state(state: LeaseState) {
268 match state {
269 LeaseState::Ready(inner) => {
270 inner.interrupt();
271 Self::spawn_bounded_cleanup(inner);
272 }
273 LeaseState::Starting { .. } => {}
276 }
277 }
278
279 fn spawn_bounded_cleanup(inner: Arc<LeaseInner>) {
280 if let Ok(handle) = tokio::runtime::Handle::try_current() {
281 handle.spawn(async move {
282 let _ = tokio::time::timeout(CLEANUP_TOTAL_TIMEOUT, async {
283 let _permit = inner.call_permit.lock().await;
286 inner.handler.shutdown().await;
287 })
288 .await;
289 });
290 }
291 }
294
295 pub fn lease_count(&self) -> usize {
297 self.leases
298 .lock()
299 .unwrap_or_else(std::sync::PoisonError::into_inner)
300 .len()
301 }
302
303 pub fn revoke_plugin_lease(&self, session: &SessionId, plugin: &str) {
306 let removed = {
307 let mut map = self
308 .leases
309 .lock()
310 .unwrap_or_else(std::sync::PoisonError::into_inner);
311 map.remove(&(session.as_str().to_string(), plugin.to_string()))
312 };
313 if let Some(state) = removed {
314 Self::cancel_state(state);
315 }
316 }
317
318 pub fn revoke_exact_tool(&self, session: &SessionId, plugin: &str, tool_name: &str) {
322 let _ = tool_name; self.revoke_plugin_lease(session, plugin);
324 }
325
326 pub fn revoke_plugin_all_sessions(&self, plugin: &str) -> usize {
329 let removed: Vec<LeaseState> = {
330 let mut map = self
331 .leases
332 .lock()
333 .unwrap_or_else(std::sync::PoisonError::into_inner);
334 let keys: Vec<(String, String)> =
335 map.keys().filter(|(_, p)| p == plugin).cloned().collect();
336 keys.into_iter().filter_map(|k| map.remove(&k)).collect()
337 };
338 let count = removed.len();
339 for state in removed {
340 Self::cancel_state(state);
341 }
342 count
343 }
344
345 pub fn terminate_session(&self, session: &SessionId) -> usize {
347 let removed: Vec<LeaseState> = {
348 let mut map = self
349 .leases
350 .lock()
351 .unwrap_or_else(std::sync::PoisonError::into_inner);
352 let keys: Vec<(String, String)> = map
353 .keys()
354 .filter(|(s, _)| s == session.as_str())
355 .cloned()
356 .collect();
357 keys.into_iter().filter_map(|k| map.remove(&k)).collect()
358 };
359 let count = removed.len();
360 for state in removed {
361 Self::cancel_state(state);
362 }
363 count
364 }
365
366 pub fn terminate_all(&self) {
368 let removed: Vec<LeaseState> = {
369 let mut map = self
370 .leases
371 .lock()
372 .unwrap_or_else(std::sync::PoisonError::into_inner);
373 map.drain().map(|(_, v)| v).collect()
374 };
375 for state in removed {
376 Self::cancel_state(state);
377 }
378 }
379
380 pub fn reap_idle(&self) {
383 let removed: Vec<LeaseState> = {
384 let mut map = self
385 .leases
386 .lock()
387 .unwrap_or_else(std::sync::PoisonError::into_inner);
388 let idle: Vec<(String, String)> = map
389 .iter()
390 .take(REAP_SCAN_MAX)
391 .filter(|(_, state)| match state {
392 LeaseState::Ready(inner) => {
393 inner
394 .last_used
395 .lock()
396 .unwrap_or_else(std::sync::PoisonError::into_inner)
397 .elapsed()
398 >= self.idle_max
399 }
400 LeaseState::Starting { .. } => false,
401 })
402 .map(|(k, _)| k.clone())
403 .collect();
404 idle.into_iter().filter_map(|k| map.remove(&k)).collect()
405 };
406 for state in removed {
407 Self::cancel_state(state);
408 }
409 }
410
411 async fn acquire_ready(
418 &self,
419 session: &SessionId,
420 plugin: &str,
421 record: &DeferredExtensionRecord,
422 validated: &super::manifest::ValidatedExtensionManifest,
423 ) -> Result<Arc<LeaseInner>, ExtensionLeaseError> {
424 let expected_fingerprint = record_fingerprint(record);
425 let key = (session.as_str().to_string(), plugin.to_string());
426 let inner: Arc<LeaseInner> = loop {
427 enum Acquire {
431 Use(Arc<LeaseInner>),
432 CancelStale(LeaseState),
433 Wait(tokio::sync::watch::Receiver<bool>),
434 Start {
435 token: u64,
436 done_tx: tokio::sync::watch::Sender<bool>,
437 },
438 AtCapacity,
439 }
440 let action = {
441 let mut map = self
442 .leases
443 .lock()
444 .unwrap_or_else(std::sync::PoisonError::into_inner);
445 match map.get(&key) {
446 Some(LeaseState::Ready(inner)) => {
447 if inner.cancelled() || inner.fingerprint != expected_fingerprint {
448 match map.remove(&key) {
449 Some(state) => Acquire::CancelStale(state),
450 None => continue,
451 }
452 } else {
453 Acquire::Use(Arc::clone(inner))
454 }
455 }
456 Some(LeaseState::Starting { done, .. }) => Acquire::Wait(done.clone()),
457 None => {
458 if map.len() >= MAX_LIVE_LEASES {
459 Acquire::AtCapacity
460 } else {
461 let token = self
462 .next_token
463 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
464 let (done_tx, done_rx) = tokio::sync::watch::channel(false);
465 map.insert(
466 key.clone(),
467 LeaseState::Starting {
468 token,
469 done: done_rx,
470 },
471 );
472 Acquire::Start { token, done_tx }
473 }
474 }
475 }
476 };
477 match action {
478 Acquire::Use(inner) => break inner,
479 Acquire::CancelStale(state) => {
480 let was_drift = matches!(
484 &state,
485 LeaseState::Ready(inner) if inner.fingerprint != expected_fingerprint
486 && !inner.cancelled()
487 );
488 Self::cancel_state(state);
489 if was_drift {
490 return Err(ExtensionLeaseError::LaunchDrift(plugin.to_string()));
491 }
492 continue;
493 }
494 Acquire::AtCapacity => {
495 self.reap_idle();
496 let still_full = self
497 .leases
498 .lock()
499 .unwrap_or_else(std::sync::PoisonError::into_inner)
500 .len()
501 >= MAX_LIVE_LEASES;
502 if still_full {
503 return Err(ExtensionLeaseError::Capacity(MAX_LIVE_LEASES));
504 }
505 continue;
506 }
507 Acquire::Wait(mut follower) => {
508 if !*follower.borrow() {
510 let _ = follower.changed().await;
511 }
512 continue;
513 }
514 Acquire::Start { token, done_tx } => {
515 let started = self
516 .start_lease(plugin, record, validated, expected_fingerprint.clone())
517 .await;
518 let ours = {
519 let mut map = self
520 .leases
521 .lock()
522 .unwrap_or_else(std::sync::PoisonError::into_inner);
523 let ours = matches!(
524 map.get(&key),
525 Some(LeaseState::Starting { token: t, .. }) if *t == token
526 );
527 match &started {
528 Ok(inner) if ours => {
529 map.insert(key.clone(), LeaseState::Ready(Arc::clone(inner)));
530 }
531 Err(_) if ours => {
532 map.remove(&key);
533 }
534 _ => {}
535 }
536 ours
537 };
538 match started {
539 Ok(inner) => {
540 if !ours {
541 inner.set_cancelled();
544 inner.handler.shutdown().await;
545 let _ = done_tx.send(true);
546 return Err(ExtensionLeaseError::Revoked(plugin.to_string()));
547 }
548 let _ = done_tx.send(true);
549 break inner;
550 }
551 Err(err) => {
552 let _ = done_tx.send(true);
553 return Err(err);
554 }
555 }
556 }
557 }
558 };
559
560 Ok(inner)
561 }
562
563 async fn acquire_handler_scope(
571 &self,
572 plugin: &str,
573 ) -> Result<Arc<LeaseInner>, ExtensionLeaseError> {
574 let session = self.host_scope();
575 let (record, validated) = self.record_and_validate(&session, plugin)?;
576 self.reap_idle();
577 self.acquire_ready(&session, plugin, &record, &validated)
578 .await
579 }
580
581 fn has_ready_lease(&self, session: &SessionId, plugin: &str) -> bool {
583 let map = self
584 .leases
585 .lock()
586 .unwrap_or_else(std::sync::PoisonError::into_inner);
587 matches!(
588 map.get(&(session.as_str().to_string(), plugin.to_string())),
589 Some(LeaseState::Ready(inner)) if !inner.cancelled()
590 )
591 }
592
593 fn ready_handler_scope(&self, plugin: &str) -> Option<Arc<LeaseInner>> {
598 let session = self.host_scope();
599 let (record, _validated) = self.record_and_validate(&session, plugin).ok()?;
600 let expected_fingerprint = record_fingerprint(&record);
601 let map = self
602 .leases
603 .lock()
604 .unwrap_or_else(std::sync::PoisonError::into_inner);
605 match map.get(&(session.as_str().to_string(), plugin.to_string())) {
606 Some(LeaseState::Ready(inner))
607 if !inner.cancelled() && inner.fingerprint == expected_fingerprint =>
608 {
609 Some(Arc::clone(inner))
610 }
611 _ => None,
612 }
613 }
614
615 fn record_and_validate(
620 &self,
621 session: &SessionId,
622 plugin: &str,
623 ) -> Result<
624 (
625 DeferredExtensionRecord,
626 super::manifest::ValidatedExtensionManifest,
627 ),
628 ExtensionLeaseError,
629 > {
630 let Some(record) = self.current_record(plugin) else {
631 self.revoke_plugin_lease(session, plugin);
632 return Err(ExtensionLeaseError::NotDeferred(plugin.to_string()));
633 };
634 let Ok(validated) = record.manifest.validate(plugin) else {
635 self.revoke_plugin_lease(session, plugin);
636 return Err(ExtensionLeaseError::ManifestInvalid(plugin.to_string()));
637 };
638 Ok((record, validated))
639 }
640
641 #[cfg(any(test, feature = "testing"))]
650 pub async fn call_exact_no_chat(
651 &self,
652 session: &SessionId,
653 plugin: &str,
654 tool_name: &str,
655 expected_digest: &SchemaDigest,
656 params: Value,
657 ) -> Result<Value, ExtensionLeaseError> {
658 self.call_exact_for(session, plugin, tool_name, expected_digest, params, None, None)
659 .await
660 }
661
662 #[allow(clippy::too_many_arguments)] pub async fn call_exact_for(
673 &self,
674 session: &SessionId,
675 plugin: &str,
676 tool_name: &str,
677 expected_digest: &SchemaDigest,
678 params: Value,
679 conversation_id: Option<&str>,
680 cwd: Option<&Path>,
681 ) -> Result<Value, ExtensionLeaseError> {
682 self.call_exact_as(
683 session,
684 plugin,
685 tool_name,
686 expected_digest,
687 params,
688 conversation_id,
689 cwd,
690 None,
691 )
692 .await
693 }
694
695 #[allow(clippy::too_many_arguments)] pub async fn call_exact_as(
701 &self,
702 session: &SessionId,
703 plugin: &str,
704 tool_name: &str,
705 expected_digest: &SchemaDigest,
706 params: Value,
707 conversation_id: Option<&str>,
708 cwd: Option<&Path>,
709 caller: Option<&Author>,
710 ) -> Result<Value, ExtensionLeaseError> {
711 let (record, validated) = self.record_and_validate(session, plugin)?;
713 let declared_tools = record
716 .manifest
717 .deferred
718 .as_ref()
719 .map(|d| d.tools.as_slice())
720 .unwrap_or(&[]);
721 let still_declared = declared_tools.iter().any(|t| {
722 t.name == tool_name && &SchemaDigest::of_schema(&t.input_schema) == expected_digest
723 });
724 if !still_declared {
725 self.revoke_plugin_lease(session, plugin);
726 return Err(ExtensionLeaseError::DeclarationDrift(
727 plugin.to_string(),
728 tool_name.to_string(),
729 ));
730 }
731 self.reap_idle();
732
733 let inner = self
735 .acquire_ready(session, plugin, &record, &validated)
736 .await?;
737 match inner.listing.get(tool_name) {
740 None => {
741 self.revoke_plugin_lease(session, plugin);
742 return Err(ExtensionLeaseError::NameNotListed(
743 plugin.to_string(),
744 tool_name.to_string(),
745 ));
746 }
747 Some(listed) if listed != expected_digest => {
748 self.revoke_plugin_lease(session, plugin);
749 return Err(ExtensionLeaseError::SchemaMismatch(
750 plugin.to_string(),
751 tool_name.to_string(),
752 ));
753 }
754 Some(_) => {}
755 }
756
757 let permit = inner.call_permit.lock().await;
760 if inner.cancelled() {
761 return Err(ExtensionLeaseError::Revoked(plugin.to_string()));
762 }
763 let result = inner
764 .handler
765 .call_tool_once_no_spawn(tool_name, params, conversation_id, cwd, caller)
766 .await;
767 if inner.cancelled() {
768 inner.handler.shutdown().await;
771 }
772 drop(permit);
773 *inner
774 .last_used
775 .lock()
776 .unwrap_or_else(std::sync::PoisonError::into_inner) = Instant::now();
777 result.map_err(|e| {
778 if let Some(rest) = e.strip_prefix(EXTENSION_REPORTED_PREFIX) {
779 ExtensionLeaseError::ExtensionReported {
781 plugin: plugin.to_string(),
782 length: rest.trim().len(),
783 }
784 } else {
785 ExtensionLeaseError::Transport(plugin.to_string(), bound_detail(&e))
786 }
787 })
788 }
789
790 async fn start_lease(
795 &self,
796 plugin: &str,
797 record: &DeferredExtensionRecord,
798 validated: &super::manifest::ValidatedExtensionManifest,
799 fingerprint: SchemaDigest,
800 ) -> Result<Arc<LeaseInner>, ExtensionLeaseError> {
801 let manifest = &record.manifest;
802 let process = ProcessExtension::spawn_with_cwd(
803 plugin,
804 &manifest.command,
805 &manifest.args,
806 record.cwd.clone(),
807 )
808 .await
809 .map_err(|e| ExtensionLeaseError::Transport(plugin.to_string(), bound_detail(&e)))?;
810 process.set_permissions(validated.permissions.clone()).await;
811 let caps = match process
812 .initialize(record.cwd.clone(), record.config.clone())
813 .await
814 {
815 Ok(caps) => caps,
816 Err(e) => {
817 process.shutdown().await;
818 let err = if let Some(rest) = e.strip_prefix(EXTENSION_REPORTED_PREFIX) {
819 ExtensionLeaseError::ExtensionReported {
820 plugin: plugin.to_string(),
821 length: rest.trim().len(),
822 }
823 } else {
824 ExtensionLeaseError::Transport(plugin.to_string(), bound_detail(&e))
825 };
826 return Err(err);
827 }
828 };
829 let deferred = manifest.deferred.as_ref();
830 let declared_tools = deferred.map(|d| d.tools.as_slice()).unwrap_or(&[]);
831 let declared_providers = deferred.map(|d| d.providers.as_slice()).unwrap_or(&[]);
832 if let Err(reason) = validate_runtime_tool_declarations(declared_tools, &caps.tools) {
833 process.shutdown().await;
834 return Err(ExtensionLeaseError::DeclarationMismatch(
835 plugin.to_string(),
836 reason,
837 ));
838 }
839 if let Err(reason) =
840 validate_runtime_provider_declarations(declared_providers, &caps.providers)
841 {
842 process.shutdown().await;
843 return Err(ExtensionLeaseError::DeclarationMismatch(
844 plugin.to_string(),
845 reason,
846 ));
847 }
848 for decl in &caps.capabilities {
849 if super::runtime::process::validate_capability(decl, &validated.permissions).is_err() {
850 process.shutdown().await;
851 return Err(ExtensionLeaseError::DeclarationMismatch(
852 plugin.to_string(),
853 "capability_declaration_rejected",
854 ));
855 }
856 }
857 let mut listing = HashMap::new();
858 for spec in &caps.tools {
859 listing.insert(
860 spec.name.clone(),
861 SchemaDigest::of_schema(&spec.input_schema),
862 );
863 }
864 Ok(Arc::new(LeaseInner {
865 fingerprint,
866 handler: Arc::new(process),
867 call_permit: tokio::sync::Mutex::new(()),
868 listing,
869 gate: std::sync::Mutex::new(false),
870 last_used: std::sync::Mutex::new(Instant::now()),
871 liveness: Arc::new(()),
872 }))
873 }
874
875 #[doc(hidden)]
878 pub fn lease_liveness_for_tests(
879 &self,
880 session: &SessionId,
881 plugin: &str,
882 ) -> Option<std::sync::Weak<()>> {
883 let map = self
884 .leases
885 .lock()
886 .unwrap_or_else(std::sync::PoisonError::into_inner);
887 match map.get(&(session.as_str().to_string(), plugin.to_string())) {
888 Some(LeaseState::Ready(inner)) => Some(Arc::downgrade(&inner.liveness)),
889 _ => None,
890 }
891 }
892}
893
894impl Drop for ExtensionRuntimeManager {
895 fn drop(&mut self) {
896 self.terminate_all();
897 }
898}
899
900impl std::fmt::Debug for ExtensionRuntimeManager {
901 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
902 f.debug_struct("ExtensionRuntimeManager")
903 .field("leases", &self.lease_count())
904 .finish_non_exhaustive()
905 }
906}
907
908#[derive(Clone)]
912pub struct ExtensionLeaseCapability {
913 session: SessionId,
914 manager: Arc<ExtensionRuntimeManager>,
915}
916
917impl ExtensionLeaseCapability {
918 pub fn new(session: SessionId, manager: Arc<ExtensionRuntimeManager>) -> Self {
919 Self { session, manager }
920 }
921
922 pub fn session(&self) -> &SessionId {
924 &self.session
925 }
926
927 #[cfg(any(test, feature = "testing"))]
930 pub async fn call_exact_no_chat(
931 &self,
932 plugin: &str,
933 tool_name: &str,
934 expected_digest: &SchemaDigest,
935 params: Value,
936 ) -> Result<Value, ExtensionLeaseError> {
937 self.manager
938 .call_exact_no_chat(&self.session, plugin, tool_name, expected_digest, params)
939 .await
940 }
941
942 pub async fn call_exact_for(
945 &self,
946 plugin: &str,
947 tool_name: &str,
948 expected_digest: &SchemaDigest,
949 params: Value,
950 conversation_id: Option<&str>,
951 cwd: Option<&Path>,
952 ) -> Result<Value, ExtensionLeaseError> {
953 self.manager
954 .call_exact_for(
955 &self.session,
956 plugin,
957 tool_name,
958 expected_digest,
959 params,
960 conversation_id,
961 cwd,
962 )
963 .await
964 }
965
966 #[allow(clippy::too_many_arguments)] pub async fn call_exact_as(
970 &self,
971 plugin: &str,
972 tool_name: &str,
973 expected_digest: &SchemaDigest,
974 params: Value,
975 conversation_id: Option<&str>,
976 cwd: Option<&Path>,
977 caller: Option<&Author>,
978 ) -> Result<Value, ExtensionLeaseError> {
979 self.manager
980 .call_exact_as(
981 &self.session,
982 plugin,
983 tool_name,
984 expected_digest,
985 params,
986 conversation_id,
987 cwd,
988 caller,
989 )
990 .await
991 }
992}
993
994impl std::fmt::Debug for ExtensionLeaseCapability {
995 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
996 f.debug_struct("ExtensionLeaseCapability")
997 .finish_non_exhaustive()
998 }
999}
1000
1001pub struct ExtensionSessionEndGuard {
1007 session: SessionId,
1008 manager: Arc<ExtensionRuntimeManager>,
1009}
1010
1011impl ExtensionSessionEndGuard {
1012 pub fn new(session: SessionId, manager: Arc<ExtensionRuntimeManager>) -> Self {
1013 Self { session, manager }
1014 }
1015}
1016
1017impl Drop for ExtensionSessionEndGuard {
1018 fn drop(&mut self) {
1019 self.manager.terminate_session(&self.session);
1020 }
1021}
1022
1023impl std::fmt::Debug for ExtensionSessionEndGuard {
1024 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1025 f.debug_struct("ExtensionSessionEndGuard")
1026 .finish_non_exhaustive()
1027 }
1028}
1029
1030pub struct LazyExtensionHandler {
1039 plugin: String,
1040 manager: Arc<ExtensionRuntimeManager>,
1041}
1042
1043impl LazyExtensionHandler {
1044 pub fn new(plugin: &str, manager: Arc<ExtensionRuntimeManager>) -> Self {
1045 Self {
1046 plugin: plugin.to_string(),
1047 manager,
1048 }
1049 }
1050
1051 async fn live(&self) -> Result<Arc<LeaseInner>, String> {
1054 self.manager
1055 .acquire_handler_scope(&self.plugin)
1056 .await
1057 .map_err(|e| e.to_string())
1058 }
1059
1060 fn touch(inner: &LeaseInner) {
1061 *inner
1062 .last_used
1063 .lock()
1064 .unwrap_or_else(std::sync::PoisonError::into_inner) = Instant::now();
1065 }
1066}
1067
1068#[async_trait::async_trait]
1069impl ExtensionHandler for LazyExtensionHandler {
1070 fn id(&self) -> &str {
1071 &self.plugin
1072 }
1073
1074 async fn handle(
1075 &self,
1076 event: &crate::extensions::hooks::events::HookEvent,
1077 ) -> crate::extensions::hooks::events::HookResult {
1078 match self.live().await {
1083 Ok(inner) => {
1084 let result = inner.handler.handle(event).await;
1085 Self::touch(&inner);
1086 result
1087 }
1088 Err(error) => {
1089 tracing::warn!(
1090 extension = %self.plugin,
1091 error = %error,
1092 "Deferred extension could not start for hook event — continuing",
1093 );
1094 crate::extensions::hooks::events::HookResult::Continue
1095 }
1096 }
1097 }
1098
1099 async fn handle_no_spawn(
1104 &self,
1105 event: &crate::extensions::hooks::events::HookEvent,
1106 ) -> crate::extensions::hooks::events::HookResult {
1107 match self.manager.ready_handler_scope(&self.plugin) {
1108 Some(inner) => {
1109 let result = inner.handler.handle_no_spawn(event).await;
1110 Self::touch(&inner);
1111 result
1112 }
1113 None => {
1114 tracing::debug!(
1115 extension = %self.plugin,
1116 "Deferred extension not running; worker hook event passes through",
1117 );
1118 crate::extensions::hooks::events::HookResult::Continue
1119 }
1120 }
1121 }
1122
1123 async fn call_tool_for(
1124 &self,
1125 _name: &str,
1126 _input: Value,
1127 _session_id: Option<&str>,
1128 _cwd: Option<&Path>,
1129 ) -> Result<Value, String> {
1130 Err(format!(
1136 "extension '{}' is deferred: tool calls are only served through the \
1137 gate-authorized exact activation path",
1138 self.plugin
1139 ))
1140 }
1141
1142 async fn provider_complete(
1143 &self,
1144 params: super::runtime::process::ProviderCompleteParams,
1145 ) -> Result<super::runtime::process::ProviderCompleteResult, String> {
1146 let inner = self.live().await?;
1147 let result = inner.handler.provider_complete(params).await;
1148 Self::touch(&inner);
1149 result
1150 }
1151
1152 async fn provider_stream(
1153 &self,
1154 params: super::runtime::process::ProviderCompleteParams,
1155 sink: tokio::sync::mpsc::Sender<super::runtime::process::ProviderStreamEvent>,
1156 ) -> Result<super::runtime::process::ProviderCompleteResult, String> {
1157 let inner = self.live().await?;
1158 let result = inner.handler.provider_stream(params, sink).await;
1159 Self::touch(&inner);
1160 result
1161 }
1162
1163 async fn invoke_command(
1164 &self,
1165 command: &str,
1166 args: Vec<String>,
1167 request_id: &str,
1168 sink: crate::extensions::invoke_output::InvokeEventSink,
1169 ) -> Result<Value, String> {
1170 let inner = self.live().await?;
1171 let result = inner
1172 .handler
1173 .invoke_command(command, args, request_id, sink)
1174 .await;
1175 Self::touch(&inner);
1176 result
1177 }
1178
1179 async fn get_info(&self) -> Result<crate::extensions::info::PluginInfo, String> {
1180 Err("deferred extension info is unavailable without activation".to_string())
1182 }
1183
1184 async fn settings_editor_open(&self, category: &str, field: &str) -> Result<Value, String> {
1185 let inner = self.live().await?;
1186 let result = inner.handler.settings_editor_open(category, field).await;
1187 Self::touch(&inner);
1188 result
1189 }
1190
1191 async fn settings_editor_key(
1192 &self,
1193 category: &str,
1194 field: &str,
1195 key: &str,
1196 ) -> Result<Value, String> {
1197 let inner = self.live().await?;
1198 let result = inner
1199 .handler
1200 .settings_editor_key(category, field, key)
1201 .await;
1202 Self::touch(&inner);
1203 result
1204 }
1205
1206 async fn settings_editor_commit(
1207 &self,
1208 category: &str,
1209 field: &str,
1210 value: Value,
1211 ) -> Result<Value, String> {
1212 let inner = self.live().await?;
1213 let result = inner
1214 .handler
1215 .settings_editor_commit(category, field, value)
1216 .await;
1217 Self::touch(&inner);
1218 result
1219 }
1220
1221 async fn shutdown(&self) {
1222 self.manager.revoke_plugin_all_sessions(&self.plugin);
1224 }
1225
1226 async fn health(&self) -> super::runtime::ExtensionHealth {
1227 if self
1229 .manager
1230 .has_ready_lease(&self.manager.host_scope(), &self.plugin)
1231 {
1232 super::runtime::ExtensionHealth::Running
1233 } else {
1234 super::runtime::ExtensionHealth::Loaded
1235 }
1236 }
1237}
1238
1239impl std::fmt::Debug for LazyExtensionHandler {
1240 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1241 f.debug_struct("LazyExtensionHandler")
1242 .field("plugin", &self.plugin)
1243 .finish_non_exhaustive()
1244 }
1245}