1use std::collections::{HashMap, HashSet};
7use std::path::PathBuf;
8use std::process::Stdio;
9use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
10use std::sync::Arc;
11
12use async_trait::async_trait;
13use serde::{Deserialize, Serialize};
14use serde_json::Value;
15use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
16use tokio::process::{Child, ChildStdin, ChildStdout, Command};
17use tokio::sync::{mpsc, oneshot, Mutex, RwLock};
18use tokio::task::JoinHandle;
19
20use super::{ExtensionHandler, ExtensionHealth, ExtensionLifecycle, RestartPolicy};
21use crate::extensions::hooks::events::{HookEvent, HookResult};
22use crate::extensions::manifest::CURRENT_EXTENSION_PROTOCOL_VERSION;
23
24#[derive(Serialize)]
25struct JsonRpcRequest {
26 jsonrpc: &'static str,
27 method: String,
28 params: Value,
29 id: u64,
30}
31
32#[derive(Serialize)]
33struct InitializeParams {
34 synaps_version: &'static str,
35 extension_protocol_version: u32,
36 plugin_id: String,
37 plugin_root: Option<String>,
38 config: Value,
39}
40
41#[derive(Debug, Clone, PartialEq, Deserialize)]
42pub struct RegisteredExtensionToolSpec {
43 pub name: String,
44 pub description: String,
45 pub input_schema: Value,
46}
47
48#[derive(Debug, Clone, PartialEq, Deserialize)]
49pub struct RegisteredProviderSpec {
50 pub id: String,
51 pub display_name: String,
52 pub description: String,
53 #[serde(default)]
54 pub models: Vec<RegisteredProviderModelSpec>,
55 #[serde(default)]
56 pub config_schema: Option<Value>,
57}
58
59#[derive(Debug, Clone, PartialEq, Deserialize)]
60pub struct RegisteredProviderModelSpec {
61 pub id: String,
62 #[serde(default)]
63 pub display_name: Option<String>,
64 #[serde(default)]
65 pub capabilities: Value,
66 #[serde(default)]
67 pub context_window: Option<u64>,
68}
69
70#[derive(Debug, Clone, Serialize)]
71pub struct ProviderCompleteParams {
72 pub provider_id: String,
73 pub model_id: String,
74 pub model: String,
75 pub messages: Vec<crate::SharedMessage>,
76 pub system_prompt: Option<String>,
77 pub tools: Vec<Value>,
78 pub temperature: Option<f32>,
79 pub max_tokens: Option<u32>,
80 pub thinking_budget: u32,
81}
82
83#[derive(Debug, Clone, PartialEq, Deserialize)]
84pub struct ProviderCompleteResult {
85 pub content: Vec<Value>,
86 #[serde(default)]
87 pub stop_reason: Option<String>,
88 #[serde(default)]
89 pub usage: Option<Value>,
90}
91
92#[derive(Debug, Clone, PartialEq)]
93pub struct ProviderToolUse {
94 pub id: String,
95 pub name: String,
96 pub input: Value,
97}
98
99#[derive(Debug, Clone, PartialEq)]
101pub enum ProviderStreamEvent {
102 TextDelta { text: String },
104 ThinkingDelta { text: String },
106 ToolUse {
108 id: String,
109 name: String,
110 input: Value,
111 },
112 Usage { usage: Value },
114 Error { message: String },
116 Done,
118}
119
120pub fn parse_provider_stream_event(params: &Value) -> Result<ProviderStreamEvent, String> {
126 let inner = match params.get("event") {
127 Some(ev) => ev,
128 None => params,
129 };
130 let obj = inner
131 .as_object()
132 .ok_or_else(|| "provider stream event must be a JSON object".to_string())?;
133
134 let ty = obj
135 .get("type")
136 .and_then(Value::as_str)
137 .ok_or_else(|| "provider stream event missing type".to_string())?;
138
139 match ty {
140 "text" => {
141 let text = obj
142 .get("delta")
143 .or_else(|| obj.get("text"))
144 .and_then(Value::as_str)
145 .ok_or_else(|| {
146 "provider stream text event missing 'delta' or 'text'".to_string()
147 })?;
148 Ok(ProviderStreamEvent::TextDelta {
149 text: text.to_string(),
150 })
151 }
152 "thinking" => {
153 let text = obj
154 .get("delta")
155 .or_else(|| obj.get("text"))
156 .and_then(Value::as_str)
157 .ok_or_else(|| {
158 "provider stream thinking event missing 'delta' or 'text'".to_string()
159 })?;
160 Ok(ProviderStreamEvent::ThinkingDelta {
161 text: text.to_string(),
162 })
163 }
164 "tool_use" => {
165 let id = obj
166 .get("id")
167 .and_then(Value::as_str)
168 .ok_or_else(|| "provider stream tool_use missing id".to_string())?;
169 if id.is_empty() {
170 return Err("provider stream tool_use id must be non-empty".to_string());
171 }
172 let name = obj
173 .get("name")
174 .and_then(Value::as_str)
175 .ok_or_else(|| "provider stream tool_use missing name".to_string())?;
176 if name.is_empty() {
177 return Err("provider stream tool_use name must be non-empty".to_string());
178 }
179 let input = match obj.get("input") {
180 None => Value::Object(Default::default()),
181 Some(v) if v.is_object() => v.clone(),
182 Some(_) => {
183 return Err("provider stream tool_use input must be a JSON object".to_string());
184 }
185 };
186 Ok(ProviderStreamEvent::ToolUse {
187 id: id.to_string(),
188 name: name.to_string(),
189 input,
190 })
191 }
192 "usage" => {
193 let mut clone = obj.clone();
194 clone.remove("type");
195 Ok(ProviderStreamEvent::Usage {
196 usage: Value::Object(clone),
197 })
198 }
199 "error" => {
200 let message = obj
201 .get("message")
202 .and_then(Value::as_str)
203 .ok_or_else(|| "provider stream error missing message".to_string())?;
204 if message.is_empty() {
205 return Err("provider stream error message must be non-empty".to_string());
206 }
207 Ok(ProviderStreamEvent::Error {
208 message: message.to_string(),
209 })
210 }
211 "done" => Ok(ProviderStreamEvent::Done),
212 other => Err(format!("unknown provider stream event type: {other}")),
213 }
214}
215
216pub async fn execute_provider_tool_use(
217 registry: &crate::ToolRegistry,
218 session_tools: &crate::tools::activation::SessionToolSet,
219 hook_bus: &Arc<crate::extensions::hooks::HookBus>,
220 tool_use: ProviderToolUse,
221 ctx: crate::ToolContext,
222 max_tool_output: usize,
223) -> Value {
224 let tool_id = tool_use.id;
225 let tool_name = tool_use.name;
226 let input = tool_use.input;
227
228 let authorized = match crate::tools::activation::ExecutionGate::authorize_wire_call(
240 registry,
241 session_tools,
242 &tool_name,
243 ) {
244 Ok(authorized) => authorized,
245 Err(denial) => {
246 return serde_json::json!({
247 "type": "tool_result",
248 "tool_use_id": tool_id,
249 "content": denial.to_string(),
250 "is_error": true,
251 });
252 }
253 };
254 let tool = authorized.implementation();
255
256 let runtime_name = authorized.runtime_name().to_string();
257 let input = registry.translate_input_for_api_tool(&tool_name, input);
258 let decision = crate::runtime::resolve_before_tool_call_decision(
259 input.clone(),
260 crate::runtime::emit_before_tool_call(
261 hook_bus,
262 &tool_name,
263 Some(&runtime_name),
264 input.clone(),
265 None,
268 )
269 .await,
270 ctx.capabilities.secret_prompt.as_ref(),
271 false,
272 ctx.capabilities.session_allow_all.as_deref(),
275 )
276 .await;
277
278 let crate::runtime::BeforeToolCallDecision::Continue { input } = decision else {
279 let crate::runtime::BeforeToolCallDecision::Block { reason } = decision else {
280 unreachable!()
281 };
282 return serde_json::json!({
283 "type": "tool_result",
284 "tool_use_id": tool_id,
285 "content": format!("Tool call blocked by extension: {}", reason),
286 "is_error": true,
287 });
288 };
289
290 let input_for_hook = input.clone();
291 let (result, is_error) = match tool.execute(input, ctx).await {
292 Ok(output) => (output, false),
293 Err(error) => (format!("Tool execution failed: {}", error), true),
294 };
295 let result = crate::runtime::emit_after_tool_call(
296 hook_bus,
297 &tool_name,
298 Some(&runtime_name),
299 input_for_hook,
300 result,
301 max_tool_output,
302 None,
303 )
304 .await;
305
306 let mut response = serde_json::json!({
307 "type": "tool_result",
308 "tool_use_id": tool_id,
309 "content": crate::truncate_str(&result, max_tool_output).to_string(),
310 });
311 if is_error {
312 response["is_error"] = serde_json::json!(true);
313 }
314 response
315}
316
317#[allow(clippy::too_many_arguments)] pub async fn complete_provider_with_tools<F>(
324 handler: Arc<dyn ExtensionHandler>,
325 mut params: ProviderCompleteParams,
326 registry: &crate::ToolRegistry,
327 session_tools: &crate::tools::activation::SessionToolSet,
328 hook_bus: &Arc<crate::extensions::hooks::HookBus>,
329 mut context_factory: F,
330 max_tool_output: usize,
331 max_iterations: usize,
332 tools_requested: &mut u32,
337) -> Result<ProviderCompleteResult, String>
338where
339 F: FnMut() -> crate::ToolContext,
340{
341 let max_iterations = max_iterations.max(1);
342 *tools_requested = 0;
343 for iteration in 0..max_iterations {
344 let result = handler.provider_complete(params.clone()).await?;
345 let tool_uses = extract_provider_tool_uses(&result.content)?;
346 if tool_uses.is_empty() {
347 return Ok(result);
348 }
349 *tools_requested = tools_requested.saturating_add(tool_uses.len() as u32);
352 if iteration + 1 == max_iterations {
353 return Err(format!(
354 "extension provider '{}' exceeded provider tool-use iteration limit ({})",
355 handler.id(),
356 max_iterations,
357 ));
358 }
359
360 let assistant_content = result.content.clone();
361 params.messages.push(std::sync::Arc::new(serde_json::json!({
362 "role": "assistant",
363 "content": assistant_content,
364 })));
365
366 let mut tool_results = Vec::with_capacity(tool_uses.len());
367 for tool_use in tool_uses {
368 tool_results.push(
369 execute_provider_tool_use(
370 registry,
371 session_tools,
372 hook_bus,
373 tool_use,
374 context_factory(),
375 max_tool_output,
376 )
377 .await,
378 );
379 }
380 params.messages.push(std::sync::Arc::new(serde_json::json!({
381 "role": "user",
382 "content": tool_results,
383 })));
384 }
385 Err(format!(
386 "extension provider '{}' exceeded provider tool-use iteration limit ({})",
387 handler.id(),
388 max_iterations,
389 ))
390}
391
392pub fn extract_provider_tool_uses(content: &[Value]) -> Result<Vec<ProviderToolUse>, String> {
393 let mut tool_uses = Vec::new();
394 for block in content {
395 if block.get("type").and_then(Value::as_str) != Some("tool_use") {
396 continue;
397 }
398 let id = block
399 .get("id")
400 .and_then(Value::as_str)
401 .ok_or_else(|| "provider tool_use missing id".to_string())?;
402 let name = block
403 .get("name")
404 .and_then(Value::as_str)
405 .ok_or_else(|| "provider tool_use missing name".to_string())?;
406 if id.trim().is_empty() {
407 return Err("provider tool_use id is empty".to_string());
408 }
409 if name.trim().is_empty() {
410 return Err("provider tool_use name is empty".to_string());
411 }
412 let input = block
413 .get("input")
414 .cloned()
415 .unwrap_or_else(|| serde_json::json!({}));
416 if !input.is_object() {
417 return Err(format!(
418 "provider tool_use '{}' input must be a JSON object",
419 id
420 ));
421 }
422 tool_uses.push(ProviderToolUse {
423 id: id.to_string(),
424 name: name.to_string(),
425 input,
426 });
427 }
428 Ok(tool_uses)
429}
430
431#[derive(Debug, Clone, Default, PartialEq)]
432pub struct InitializeCapabilitiesResult {
433 pub tools: Vec<RegisteredExtensionToolSpec>,
434 pub providers: Vec<RegisteredProviderSpec>,
435 pub capabilities: Vec<CapabilityDeclaration>,
441}
442
443#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
452pub struct CapabilityDeclaration {
453 pub kind: String,
457 pub name: String,
460 #[serde(default)]
465 pub permissions: Vec<String>,
466 #[serde(default, skip_serializing_if = "is_null_value")]
470 pub params: serde_json::Value,
471}
472
473fn is_null_value(v: &serde_json::Value) -> bool {
474 v.is_null()
475}
476
477pub fn validate_capability(
487 decl: &CapabilityDeclaration,
488 granted: &crate::extensions::permissions::PermissionSet,
489) -> Result<(), String> {
490 use crate::extensions::permissions::Permission;
491 if decl.kind.trim().is_empty() {
492 return Err("capability 'kind' must be non-empty".to_string());
493 }
494 if decl.name.trim().is_empty() {
495 return Err("capability 'name' must be non-empty".to_string());
496 }
497 for perm_name in &decl.permissions {
498 let parsed = Permission::parse(perm_name).ok_or_else(|| {
499 format!(
500 "capability '{}' declares unknown permission '{}'",
501 decl.kind, perm_name
502 )
503 })?;
504 if !granted.has(parsed) {
505 return Err(format!(
506 "capability '{}' requires permission '{}' but it is not granted",
507 decl.kind, perm_name
508 ));
509 }
510 }
511 Ok(())
512}
513
514#[derive(Deserialize)]
515struct InitializeResult {
516 protocol_version: u32,
517 #[serde(default)]
518 capabilities: InitializeCapabilities,
519}
520
521#[derive(Default, Deserialize)]
522struct InitializeCapabilities {
523 #[serde(default)]
524 tools: Vec<RegisteredExtensionToolSpec>,
525 #[serde(default)]
526 providers: Vec<RegisteredProviderSpec>,
527 #[serde(default)]
529 capabilities: Vec<CapabilityDeclaration>,
530}
531
532#[doc(hidden)]
538#[derive(Debug, Clone)]
539pub struct NotificationFrame {
540 pub method: String,
541 pub params: Value,
542}
543
544pub const NOTIFICATION_QUEUE_CAPACITY: usize = 8;
550pub const PROVIDER_EVENT_QUEUE_CAPACITY: usize = 8;
553
554struct Inbox {
561 pending: Mutex<HashMap<u64, oneshot::Sender<Result<Value, String>>>>,
562 notification_sinks: Mutex<Vec<(usize, mpsc::Sender<NotificationFrame>)>>,
570 next_sink_id: std::sync::atomic::AtomicUsize,
574 closed: std::sync::atomic::AtomicBool,
577 lifecycle: Mutex<Option<Arc<ProcessLifecycle>>>,
581 next_generation: AtomicU64,
584 permissions: RwLock<Option<crate::extensions::permissions::PermissionSet>>,
587 inbound_stdin: Mutex<Option<Arc<Mutex<ChildStdin>>>>,
591 extension_id: String,
593 exclusive_memory: bool,
594}
595
596impl Inbox {
597 fn new(extension_id: String) -> Self {
598 Self {
599 pending: Mutex::new(HashMap::new()),
600 notification_sinks: Mutex::new(Vec::new()),
601 next_sink_id: std::sync::atomic::AtomicUsize::new(0),
602 closed: std::sync::atomic::AtomicBool::new(false),
603 lifecycle: Mutex::new(None),
604 next_generation: AtomicU64::new(2),
605 permissions: RwLock::new(None),
606 inbound_stdin: Mutex::new(None),
607 exclusive_memory: crate::memory_backend::MemoryBinding::configured_current()
608 .exclusive(),
609 extension_id,
610 }
611 }
612
613 async fn invalidate_lifecycle(&self) {
614 if let Some(lifecycle) = self.lifecycle.lock().await.as_ref() {
615 lifecycle.invalidate();
616 }
617 }
618
619 async fn fail_all_pending(&self, reason: &str) {
622 self.invalidate_lifecycle().await;
623 self.closed
624 .store(true, std::sync::atomic::Ordering::Release);
625 let drained: Vec<_> = {
626 let mut pending = self.pending.lock().await;
627 pending.drain().collect()
628 };
629 for (_, tx) in drained {
630 let _ = tx.send(Err(reason.to_string()));
631 }
632 }
633}
634
635struct ProcessLifecycle {
638 child: Mutex<Child>,
639 generation: AtomicU64,
640 initialized: AtomicBool,
641}
642
643impl ProcessLifecycle {
644 fn invalidate(&self) {
645 self.generation.fetch_or(1, Ordering::AcqRel);
648 }
649
650 fn snapshot(&self, recovered: bool) -> Option<ExtensionLifecycle> {
651 if self.generation.load(Ordering::Acquire) & 1 == 0 {
652 let mut child = self.child.try_lock().ok()?;
653 match child.try_wait() {
654 Ok(None) => {}
655 Ok(Some(_)) | Err(_) => self.invalidate(),
656 }
657 }
658 let generation = self.generation.load(Ordering::Acquire);
659 let health = if generation & 1 != 0 {
660 ExtensionHealth::Failed
661 } else if !self.initialized.load(Ordering::Acquire) {
662 ExtensionHealth::Loaded
663 } else if recovered {
664 ExtensionHealth::Degraded
665 } else {
666 ExtensionHealth::Running
667 };
668 Some(ExtensionLifecycle { generation, health })
669 }
670}
671
672struct ProcessState {
673 lifecycle: Arc<ProcessLifecycle>,
674 stdin: Arc<Mutex<ChildStdin>>,
675 reader_handle: JoinHandle<()>,
676}
677
678impl Drop for ProcessState {
679 fn drop(&mut self) {
680 self.lifecycle.invalidate();
681 self.reader_handle.abort();
682 if let Ok(mut child) = self.lifecycle.child.try_lock() {
685 let _ = child.start_kill();
686 }
687 }
688}
689
690pub struct ProcessExtension {
692 id: String,
693 command: String,
694 args: Vec<String>,
695 cwd: Option<PathBuf>,
696 state: Arc<Mutex<Option<ProcessState>>>,
697 call_lock: Arc<Mutex<()>>,
699 next_id: AtomicU64,
700 restart_count: AtomicUsize,
701 initialize_config: Mutex<Value>,
705 total_restarts: AtomicUsize,
711 stopped: AtomicBool,
718 pub(crate) restart_policy: RestartPolicy,
720 inbox: Arc<Inbox>,
724}
725
726pub(crate) const APPLICATION_ERROR_PREFIX: &str = "Extension error: ";
730
731pub(crate) fn is_application_error(error: &str) -> bool {
735 error.starts_with(APPLICATION_ERROR_PREFIX)
736}
737
738impl ProcessExtension {
739 pub async fn force_shutdown(&self) {
743 self.inbox.invalidate_lifecycle().await;
744 let mut state_guard = self.state.lock().await;
745 if let Some(state) = state_guard.take() {
746 state.reader_handle.abort();
747 state.lifecycle.invalidate();
748 let _ = state.lifecycle.child.lock().await.kill().await;
749 }
750 drop(state_guard);
751 self.inbox.closed.store(true, Ordering::Release);
752 self.inbox.notification_sinks.lock().await.clear();
753 self.inbox
754 .fail_all_pending("transport closed: extension cancelled")
755 .await;
756 }
757
758 pub async fn spawn(id: &str, command: &str, args: &[String]) -> Result<Self, String> {
759 Self::spawn_with_cwd(id, command, args, None).await
760 }
761
762 pub async fn spawn_with_cwd(
767 id: &str,
768 command: &str,
769 args: &[String],
770 cwd: Option<PathBuf>,
771 ) -> Result<Self, String> {
772 let exclusive = crate::memory_backend::MemoryBinding::configured_current().exclusive();
773 Self::spawn_with_memory_policy(id, command, args, cwd, exclusive).await
774 }
775
776 pub(crate) async fn spawn_with_memory_policy(
777 id: &str,
778 command: &str,
779 args: &[String],
780 cwd: Option<PathBuf>,
781 exclusive: bool,
782 ) -> Result<Self, String> {
783 if exclusive && id == "axel-memory-manager" {
784 return Err(
785 "independent Axel plugin is disabled by the host-selected memory backend".into(),
786 );
787 }
788 let mut inbox = Inbox::new(id.to_string());
789 inbox.exclusive_memory = exclusive;
790 let inbox = Arc::new(inbox);
791 let state = Self::spawn_state(id, command, args, cwd.as_ref(), inbox.clone()).await?;
792 Ok(Self {
793 id: id.to_string(),
794 command: command.to_string(),
795 args: args.to_vec(),
796 cwd,
797 state: Arc::new(Mutex::new(Some(state))),
798 call_lock: Arc::new(Mutex::new(())),
799 next_id: AtomicU64::new(1),
800 restart_count: AtomicUsize::new(0),
801 initialize_config: Mutex::new(Value::Object(Default::default())),
802 total_restarts: AtomicUsize::new(0),
803 stopped: AtomicBool::new(false),
804 restart_policy: RestartPolicy::default(),
805 inbox,
806 })
807 }
808
809 pub fn with_restart_policy(mut self, policy: RestartPolicy) -> Self {
811 self.restart_policy = policy;
812 self
813 }
814
815 async fn spawn_state(
816 id: &str,
817 command: &str,
818 args: &[String],
819 cwd: Option<&PathBuf>,
820 inbox: Arc<Inbox>,
821 ) -> Result<ProcessState, String> {
822 let mut cmd = Command::new(command);
823 cmd.args(args)
824 .stdin(Stdio::piped())
825 .stdout(Stdio::piped())
826 .stderr(Stdio::piped());
827 if let Some(cwd) = cwd {
828 cmd.current_dir(cwd);
829 }
830
831 cmd.env_clear();
837 for var in &["PATH", "HOME", "LANG", "TERM", "XDG_RUNTIME_DIR"] {
838 if let Ok(val) = std::env::var(var) {
839 cmd.env(var, val);
840 }
841 }
842
843 cmd.kill_on_drop(true);
844
845 let mut child = cmd
846 .spawn()
847 .map_err(|e| format!("Failed to spawn extension '{}': {}", id, e))?;
848
849 let stdin = child
850 .stdin
851 .take()
852 .ok_or_else(|| format!("No stdin for extension '{}'", id))?;
853 let stdout = child
854 .stdout
855 .take()
856 .ok_or_else(|| format!("No stdout for extension '{}'", id))?;
857 if let Some(stderr) = child.stderr.take() {
858 let extension_id = id.to_string();
859 tokio::spawn(async move {
860 const MAX_STDERR_LINE: usize = 8 * 1024;
865 const MAX_STDERR_FORWARD_TOTAL: u64 = 128 * 1024;
866 let mut reader = BufReader::new(stderr);
867 let mut forwarded: u64 = 0;
868 let mut suppressed: u64 = 0;
869 loop {
870 let mut buf: Vec<u8> = Vec::new();
871 let mut limited =
872 tokio::io::AsyncReadExt::take(&mut reader, MAX_STDERR_LINE as u64 + 1);
873 match tokio::io::AsyncBufReadExt::read_until(&mut limited, b'\n', &mut buf)
874 .await
875 {
876 Ok(0) => break,
877 Ok(n) => {
878 if forwarded < MAX_STDERR_FORWARD_TOTAL {
879 forwarded = forwarded.saturating_add(n as u64);
880 let truncated = n > MAX_STDERR_LINE;
881 let end = buf.len().min(MAX_STDERR_LINE);
882 let line = String::from_utf8_lossy(&buf[..end]);
883 let line = line.trim_end_matches(['\r', '\n']);
884 tracing::debug!(
885 extension = %extension_id,
886 truncated,
887 stderr = %line,
888 );
889 } else {
890 suppressed = suppressed.saturating_add(n as u64);
891 }
892 }
893 Err(error) => {
894 tracing::debug!(
895 extension = %extension_id,
896 error = %error,
897 "Failed to read extension stderr",
898 );
899 break;
900 }
901 }
902 }
903 if suppressed > 0 {
904 tracing::debug!(
905 extension = %extension_id,
906 suppressed_bytes = suppressed,
907 "Extension stderr exceeded the forwarding budget; further bytes counted only",
908 );
909 }
910 });
911 }
912
913 let generation = inbox
914 .next_generation
915 .fetch_update(Ordering::AcqRel, Ordering::Acquire, |n| n.checked_add(2))
916 .map_err(|_| "extension lifecycle generation exhausted".to_string())?;
917 let lifecycle = Arc::new(ProcessLifecycle {
918 child: Mutex::new(child),
919 generation: AtomicU64::new(generation),
920 initialized: AtomicBool::new(false),
921 });
922 {
923 let mut slot = inbox.lifecycle.lock().await;
924 if let Some(old) = slot.as_ref() {
925 old.invalidate();
926 }
927 *slot = Some(lifecycle.clone());
928 }
929 inbox.closed.store(false, Ordering::Release);
931 let reader_handle =
932 Self::spawn_reader(stdout, inbox.clone(), id.to_string(), lifecycle.clone());
933
934 let stdin_arc = Arc::new(Mutex::new(stdin));
935 *inbox.inbound_stdin.lock().await = Some(stdin_arc.clone());
938
939 Ok(ProcessState {
940 lifecycle,
941 stdin: stdin_arc,
942 reader_handle,
943 })
944 }
945
946 fn spawn_reader(
951 stdout: ChildStdout,
952 inbox: Arc<Inbox>,
953 extension_id: String,
954 lifecycle: Arc<ProcessLifecycle>,
955 ) -> JoinHandle<()> {
956 tokio::spawn(async move {
957 let mut reader = BufReader::new(stdout);
958 loop {
959 match Self::read_one_frame(&mut reader, &extension_id).await {
960 Ok(Some(value)) => {
961 Self::dispatch_frame(value, &inbox, &extension_id).await;
962 }
963 Ok(None) => {
964 lifecycle.invalidate();
965 tracing::debug!(
966 extension = %extension_id,
967 "Extension stdout closed (EOF); failing pending requests",
968 );
969 inbox.fail_all_pending("transport closed: EOF").await;
970 inbox.notification_sinks.lock().await.clear();
972 return;
973 }
974 Err(error) => {
975 lifecycle.invalidate();
976 tracing::debug!(
977 extension = %extension_id,
978 error = %error,
979 "Extension transport read error",
980 );
981 inbox
982 .fail_all_pending(&format!("transport error: {}", error))
983 .await;
984 inbox.notification_sinks.lock().await.clear();
985 return;
986 }
987 }
988 }
989 })
990 }
991
992 async fn read_one_frame(
996 reader: &mut BufReader<ChildStdout>,
997 extension_id: &str,
998 ) -> Result<Option<Value>, String> {
999 let mut content_length: Option<usize> = None;
1000 let mut saw_any_header = false;
1001 loop {
1002 let mut header_line = String::new();
1003 let n = {
1006 let mut limited = tokio::io::AsyncReadExt::take(&mut *reader, 1025);
1007 limited
1008 .read_line(&mut header_line)
1009 .await
1010 .map_err(|e| format!("Read header error: {}", e))?
1011 };
1012 if n == 0 {
1013 if saw_any_header {
1014 return Err("Unexpected EOF while reading response headers".into());
1015 }
1016 return Ok(None);
1017 }
1018 saw_any_header = true;
1019 if header_line.len() > 1024 {
1020 return Err(format!(
1021 "Extension '{}' header line too long ({} bytes)",
1022 extension_id,
1023 header_line.len()
1024 ));
1025 }
1026 let trimmed = header_line.trim();
1027 if trimmed.is_empty() {
1028 break;
1029 }
1030 if let Some((name, value)) = trimmed.split_once(':') {
1031 if name.trim().eq_ignore_ascii_case("Content-Length") {
1032 content_length = Some(value.trim().parse().map_err(|_| {
1033 format!("Invalid Content-Length value: {:?}", value.trim())
1034 })?);
1035 }
1036 }
1037 }
1038 let content_length = content_length.ok_or_else(|| {
1039 format!(
1040 "Extension '{}' frame missing Content-Length header",
1041 extension_id
1042 )
1043 })?;
1044 const MAX_RESPONSE_SIZE: usize = 4 * 1024 * 1024;
1045 if content_length > MAX_RESPONSE_SIZE {
1046 return Err(format!(
1047 "Extension '{}' frame too large: {} bytes (max {})",
1048 extension_id, content_length, MAX_RESPONSE_SIZE
1049 ));
1050 }
1051 let mut buf = vec![0u8; content_length];
1052 tokio::io::AsyncReadExt::read_exact(reader, &mut buf)
1053 .await
1054 .map_err(|e| format!("Read body error: {}", e))?;
1055 let value: Value =
1056 serde_json::from_slice(&buf).map_err(|e| format!("Parse frame error: {}", e))?;
1057 Ok(Some(value))
1058 }
1059
1060 async fn dispatch_frame(value: Value, inbox: &Arc<Inbox>, extension_id: &str) {
1066 let id_field = value.get("id");
1067 let id_is_present = !matches!(id_field, None | Some(Value::Null));
1068 let method_field = value
1069 .get("method")
1070 .and_then(Value::as_str)
1071 .map(str::to_string);
1072
1073 if id_is_present && method_field.is_some() {
1074 let id = match id_field {
1086 Some(id @ (Value::String(_) | Value::Number(_))) => id.clone(),
1087 _ => {
1088 tracing::warn!(
1089 extension = %extension_id,
1090 frame = %value,
1091 "Discarding inbound request: JSON-RPC id must be a string or a number",
1092 );
1093 return;
1094 }
1095 };
1096 let Some(method) = method_field else { return };
1097 let params = value.get("params").cloned().unwrap_or(Value::Null);
1098 let inbox = inbox.clone();
1099 let extension_id = extension_id.to_string();
1100 tokio::spawn(async move {
1101 let outcome = Self::handle_inbound_request(&inbox, &method, params).await;
1102 let payload = match outcome {
1103 Ok(result) => serde_json::json!({
1104 "jsonrpc": "2.0",
1105 "id": id,
1106 "result": result,
1107 }),
1108 Err((code, message)) => serde_json::json!({
1109 "jsonrpc": "2.0",
1110 "id": id,
1111 "error": {"code": code, "message": message},
1112 }),
1113 };
1114 let stdin_handle = inbox.inbound_stdin.lock().await.clone();
1115 if let Some(stdin) = stdin_handle {
1116 let body = match serde_json::to_string(&payload) {
1117 Ok(s) => s,
1118 Err(error) => {
1119 tracing::warn!(
1120 extension = %extension_id,
1121 error = %error,
1122 "Failed to serialize inbound response",
1123 );
1124 return;
1125 }
1126 };
1127 let frame = format!("Content-Length: {}\r\n\r\n{}", body.len(), body);
1128 let mut stdin = stdin.lock().await;
1129 if let Err(error) = stdin.write_all(frame.as_bytes()).await {
1130 inbox.invalidate_lifecycle().await;
1131 tracing::warn!(
1132 extension = %extension_id,
1133 error = %error,
1134 "Failed to write inbound response",
1135 );
1136 return;
1137 }
1138 if let Err(error) = stdin.flush().await {
1139 inbox.invalidate_lifecycle().await;
1140 tracing::warn!(
1141 extension = %extension_id,
1142 error = %error,
1143 "Failed to flush inbound response",
1144 );
1145 }
1146 } else {
1147 tracing::warn!(
1148 extension = %extension_id,
1149 "No stdin available to reply to inbound request",
1150 );
1151 }
1152 });
1153 return;
1154 }
1155
1156 if id_is_present {
1157 let id = match id_field.and_then(Value::as_u64) {
1162 Some(id) => id,
1163 None => {
1164 tracing::warn!(
1165 extension = %extension_id,
1166 frame = %value,
1167 "Discarding response frame: id does not match any request \
1168 the runtime issued (runtime request ids are always numeric)",
1169 );
1170 return;
1171 }
1172 };
1173 let sender = inbox.pending.lock().await.remove(&id);
1174 match sender {
1175 Some(tx) => {
1176 let payload = if let Some(err) = value.get("error") {
1177 let message = err
1178 .get("message")
1179 .and_then(Value::as_str)
1180 .unwrap_or("unknown extension error")
1181 .to_string();
1182 Err(format!("{APPLICATION_ERROR_PREFIX}{}", message))
1183 } else {
1184 Ok(value.get("result").cloned().unwrap_or(Value::Null))
1185 };
1186 let _ = tx.send(payload);
1187 }
1188 None => {
1189 tracing::trace!(
1190 extension = %extension_id,
1191 id = id,
1192 "Response with unknown id (no pending request); dropping",
1193 );
1194 }
1195 }
1196 } else if let Some(method) = value.get("method").and_then(Value::as_str) {
1197 let params = value.get("params").cloned().unwrap_or(Value::Null);
1198 let frame = NotificationFrame {
1199 method: method.to_string(),
1200 params,
1201 };
1202 let targets: Vec<(usize, mpsc::Sender<NotificationFrame>)> = {
1206 let sinks = inbox.notification_sinks.lock().await;
1207 sinks.to_vec()
1208 };
1209 if targets.is_empty() {
1210 tracing::trace!(
1211 extension = %extension_id,
1212 method = %method,
1213 "Notification with no active subscribers; dropping",
1214 );
1215 } else {
1216 let mut dead = Vec::new();
1221 for (id, tx) in &targets {
1222 if tx.send(frame.clone()).await.is_err() {
1223 dead.push(*id);
1224 }
1225 }
1226 if !dead.is_empty() {
1227 let mut sinks = inbox.notification_sinks.lock().await;
1228 sinks.retain(|(id, _)| !dead.contains(id));
1229 }
1230 }
1231 } else {
1232 tracing::trace!(
1233 extension = %extension_id,
1234 frame = %value,
1235 "Unrecognized frame; dropping",
1236 );
1237 }
1238 }
1239
1240 pub fn restart_count(&self) -> usize {
1241 self.total_restarts.load(Ordering::Relaxed)
1244 }
1245
1246 pub async fn set_permissions(&self, perms: crate::extensions::permissions::PermissionSet) {
1249 *self.inbox.permissions.write().await = Some(perms);
1250 }
1251
1252 #[allow(clippy::doc_lazy_continuation)]
1260 async fn handle_inbound_request(
1261 inbox: &Arc<Inbox>,
1262 method: &str,
1263 params: Value,
1264 ) -> Result<Value, (i32, String)> {
1265 use crate::extensions::permissions::Permission;
1266 use crate::memory::store::{self, MemoryQuery};
1267
1268 if inbox.exclusive_memory && matches!(method, "memory.append" | "memory.query") {
1269 return Err((
1270 -32000,
1271 "legacy extension memory is disabled by the host-selected backend".into(),
1272 ));
1273 }
1274
1275 match method {
1276 "memory.append" => {
1277 Self::require_permission(inbox, Permission::MemoryWrite, "memory.write").await?;
1278 let namespace = Self::param_str(¶ms, "namespace")?;
1279 Self::require_namespace_matches(inbox, &namespace).await?;
1280 let content = Self::param_str(¶ms, "content")?;
1281 let tags = match params.get("tags") {
1282 None | Some(Value::Null) => Vec::new(),
1283 Some(Value::Array(arr)) => {
1284 let mut out = Vec::with_capacity(arr.len());
1285 for v in arr {
1286 match v.as_str() {
1287 Some(s) => out.push(s.to_string()),
1288 None => {
1289 return Err((
1290 -32602,
1291 "tags must be an array of strings".to_string(),
1292 ))
1293 }
1294 }
1295 }
1296 out
1297 }
1298 _ => return Err((-32602, "tags must be an array of strings".to_string())),
1299 };
1300 let meta = match params.get("meta") {
1301 None | Some(Value::Null) => None,
1302 Some(v) => Some(v.clone()),
1303 };
1304 let record = store::new_record(namespace, content, tags, meta);
1305 let timestamp_ms = record.timestamp_ms;
1306 store::append(&record).map_err(|e| (-32000, e.to_string()))?;
1307 Ok(serde_json::json!({"ok": true, "timestamp_ms": timestamp_ms}))
1308 }
1309 "memory.query" => {
1310 Self::require_permission(inbox, Permission::MemoryRead, "memory.read").await?;
1311 let namespace = Self::param_str(¶ms, "namespace")?;
1312 Self::require_namespace_matches(inbox, &namespace).await?;
1313 let q = MemoryQuery {
1314 content_contains: params
1315 .get("content_contains")
1316 .and_then(Value::as_str)
1317 .map(str::to_string),
1318 tag_prefix: params
1319 .get("tag_prefix")
1320 .and_then(Value::as_str)
1321 .map(str::to_string),
1322 since_ms: params.get("since_ms").and_then(Value::as_u64),
1323 until_ms: params.get("until_ms").and_then(Value::as_u64),
1324 limit: params
1325 .get("limit")
1326 .and_then(Value::as_u64)
1327 .map(|n| n as usize),
1328 };
1329 let records = store::query(&namespace, &q).map_err(|e| (-32000, e.to_string()))?;
1330 Ok(serde_json::json!({"records": records}))
1331 }
1332 "config.get" => {
1333 let key = Self::param_str(¶ms, "key")?;
1334 Self::validate_config_key(&key)?;
1335 let value =
1336 crate::extensions::config_store::read_plugin_config(&inbox.extension_id, &key);
1337 Ok(serde_json::json!({"value": value}))
1338 }
1339 "config.set" => {
1340 Self::require_permission(inbox, Permission::ConfigWrite, "config.write").await?;
1341 let key = Self::param_str(¶ms, "key")?;
1342 Self::validate_config_key(&key)?;
1343 let value = Self::param_str(¶ms, "value")?;
1344 crate::extensions::config_store::write_plugin_config(
1345 &inbox.extension_id,
1346 &key,
1347 &value,
1348 )
1349 .map_err(|e| (-32000, e.to_string()))?;
1350 Ok(serde_json::json!({"ok": true}))
1351 }
1352 "config.subscribe" => {
1353 Self::require_permission(inbox, Permission::ConfigSubscribe, "config.subscribe")
1354 .await?;
1355 Ok(serde_json::json!({"ok": true}))
1359 }
1360 other => Err((-32601, format!("method not found: {other}"))),
1361 }
1362 }
1363
1364 async fn require_permission(
1365 inbox: &Arc<Inbox>,
1366 perm: crate::extensions::permissions::Permission,
1367 wire: &str,
1368 ) -> Result<(), (i32, String)> {
1369 let guard = inbox.permissions.read().await;
1370 match guard.as_ref() {
1371 Some(set) if set.has(perm) => Ok(()),
1372 _ => Err((-32602, format!("permission denied: {wire} required"))),
1373 }
1374 }
1375
1376 async fn require_namespace_matches(
1377 inbox: &Arc<Inbox>,
1378 namespace: &str,
1379 ) -> Result<(), (i32, String)> {
1380 if namespace == inbox.extension_id {
1381 Ok(())
1382 } else {
1383 Err((
1384 -32602,
1385 format!(
1386 "namespace must equal extension id '{}' (got '{}')",
1387 inbox.extension_id, namespace
1388 ),
1389 ))
1390 }
1391 }
1392
1393 fn param_str(params: &Value, name: &str) -> Result<String, (i32, String)> {
1394 params
1395 .get(name)
1396 .and_then(Value::as_str)
1397 .map(str::to_string)
1398 .ok_or_else(|| (-32602, format!("missing or invalid '{name}' parameter")))
1399 }
1400
1401 fn validate_config_key(key: &str) -> Result<(), (i32, String)> {
1402 let trimmed = key.trim();
1403 if trimmed.is_empty() {
1404 return Err((-32602, "config key must be non-empty".to_string()));
1405 }
1406 if trimmed.contains('.') || trimmed.contains('/') || trimmed.contains(' ') {
1407 return Err((
1408 -32602,
1409 "config key must not contain dots, slashes, or spaces".to_string(),
1410 ));
1411 }
1412 Ok(())
1413 }
1414
1415 pub async fn initialize(
1416 &self,
1417 plugin_root: Option<PathBuf>,
1418 config: Value,
1419 ) -> Result<InitializeCapabilitiesResult, String> {
1420 *self.initialize_config.lock().await = config.clone();
1423 let params = InitializeParams {
1424 synaps_version: env!("CARGO_PKG_VERSION"),
1425 extension_protocol_version: CURRENT_EXTENSION_PROTOCOL_VERSION,
1426 plugin_id: self.id.clone(),
1427 plugin_root: plugin_root
1428 .or_else(|| self.cwd.clone())
1429 .map(|path| path.to_string_lossy().to_string()),
1430 config,
1431 };
1432 let value = self
1433 .call_no_restart(
1434 "initialize",
1435 serde_json::to_value(params).map_err(|e| e.to_string())?,
1436 )
1437 .await?;
1438 Self::parse_initialize_result(&self.id, value)
1439 }
1440
1441 fn parse_initialize_result(
1442 id: &str,
1443 value: Value,
1444 ) -> Result<InitializeCapabilitiesResult, String> {
1445 let result: InitializeResult = serde_json::from_value(value)
1446 .map_err(|e| format!("Invalid initialize response from extension '{}': {}", id, e))?;
1447 if result.protocol_version != CURRENT_EXTENSION_PROTOCOL_VERSION {
1448 return Err(format!(
1449 "Extension '{}' initialize returned unsupported protocol_version {} (supported: {})",
1450 id, result.protocol_version, CURRENT_EXTENSION_PROTOCOL_VERSION,
1451 ));
1452 }
1453 Self::validate_registered_tool_specs(id, &result.capabilities.tools)?;
1454 Self::validate_registered_provider_specs(id, &result.capabilities.providers)?;
1455 Ok(InitializeCapabilitiesResult {
1456 tools: result.capabilities.tools,
1457 providers: result.capabilities.providers,
1458 capabilities: result.capabilities.capabilities,
1459 })
1460 }
1461
1462 fn validate_registered_tool_specs(
1463 id: &str,
1464 tools: &[RegisteredExtensionToolSpec],
1465 ) -> Result<(), String> {
1466 use crate::extensions::validation::{validate_id_segment, IdValidationError};
1467 let mut names = HashSet::new();
1468 for tool in tools {
1469 let name = tool.name.trim();
1470 if let Err(err) = validate_id_segment(name) {
1471 return Err(match err {
1472 IdValidationError::Empty => format!(
1473 "Extension '{}' registered a tool with an empty tool name",
1474 id
1475 ),
1476 IdValidationError::ContainsReserved { ch } => format!(
1477 "Extension '{}' registered tool '{}' with invalid tool name: '{}' is reserved",
1478 id, name, ch
1479 ),
1480 IdValidationError::TooLong { len, max } => format!(
1481 "Extension '{}' registered tool '{}' with invalid tool name: must be at most {} chars (got {})",
1482 id, name, max, len
1483 ),
1484 IdValidationError::ContainsWhitespace => format!(
1485 "Extension '{}' registered tool '{}' with invalid tool name: must not contain whitespace",
1486 id, name
1487 ),
1488 IdValidationError::ContainsControl { ch } => format!(
1489 "Extension '{}' registered tool '{}' with invalid tool name: contains control character U+{:04X}",
1490 id, name, ch as u32
1491 ),
1492 });
1493 }
1494 if !names.insert(name.to_string()) {
1495 return Err(format!(
1496 "Extension '{}' registered duplicate tool name '{}'",
1497 id, name
1498 ));
1499 }
1500 if tool.description.trim().is_empty() {
1501 return Err(format!(
1502 "Extension '{}' registered tool '{}' with an empty description",
1503 id, name,
1504 ));
1505 }
1506 if !tool.input_schema.is_object() {
1507 return Err(format!(
1508 "Extension '{}' registered tool '{}' with invalid input_schema: input_schema must be a JSON object",
1509 id, name,
1510 ));
1511 }
1512 }
1513 Ok(())
1514 }
1515
1516 fn validate_registered_provider_specs(
1517 id: &str,
1518 providers: &[RegisteredProviderSpec],
1519 ) -> Result<(), String> {
1520 use crate::extensions::validation::{validate_id_segment, IdValidationError};
1521 for provider in providers {
1522 let provider_id = provider.id.trim();
1523 match validate_id_segment(provider_id) {
1524 Ok(()) => {
1525 if !Self::is_safe_provider_id(provider_id) {
1526 return Err(format!(
1527 "Extension '{}' registered provider '{}' with invalid provider id",
1528 id, provider_id
1529 ));
1530 }
1531 }
1532 Err(IdValidationError::Empty) => {
1533 return Err(format!(
1534 "Extension '{}' registered provider with empty provider id",
1535 id
1536 ));
1537 }
1538 Err(err) => {
1539 return Err(format!(
1540 "Extension '{}' registered provider '{}' with invalid provider id: {}",
1541 id, provider_id, err
1542 ));
1543 }
1544 }
1545 if provider.display_name.trim().is_empty() {
1546 return Err(format!(
1547 "Extension '{}' registered provider '{}' with empty display_name",
1548 id, provider_id,
1549 ));
1550 }
1551 if provider.description.trim().is_empty() {
1552 return Err(format!(
1553 "Extension '{}' registered provider '{}' with empty description",
1554 id, provider_id,
1555 ));
1556 }
1557 if provider.models.is_empty() {
1558 return Err(format!(
1559 "Extension '{}' registered provider '{}' must declare at least one model",
1560 id, provider_id,
1561 ));
1562 }
1563 let mut model_ids = HashSet::new();
1564 for model in &provider.models {
1565 let model_id = model.id.trim();
1566 if let Err(err) = validate_id_segment(model_id) {
1567 return Err(match err {
1568 IdValidationError::Empty => format!(
1569 "Extension '{}' registered provider '{}' with empty model id",
1570 id, provider_id
1571 ),
1572 IdValidationError::ContainsReserved { ch } => format!(
1573 "Extension '{}' registered provider '{}' with invalid model id '{}': '{}' is reserved",
1574 id, provider_id, model_id, ch
1575 ),
1576 IdValidationError::TooLong { len, max } => format!(
1577 "Extension '{}' registered provider '{}' with invalid model id '{}': must be at most {} chars (got {})",
1578 id, provider_id, model_id, max, len
1579 ),
1580 IdValidationError::ContainsWhitespace => format!(
1581 "Extension '{}' registered provider '{}' with invalid model id '{}': must not contain whitespace",
1582 id, provider_id, model_id
1583 ),
1584 IdValidationError::ContainsControl { ch } => format!(
1585 "Extension '{}' registered provider '{}' with invalid model id '{}': contains control character U+{:04X}",
1586 id, provider_id, model_id, ch as u32
1587 ),
1588 });
1589 }
1590 if !model_ids.insert(model_id.to_string()) {
1591 return Err(format!(
1592 "Extension '{}' registered provider '{}' with duplicate model id '{}'",
1593 id, provider_id, model_id,
1594 ));
1595 }
1596 }
1597 if let Some(config_schema) = &provider.config_schema {
1598 if !config_schema.is_object() {
1599 return Err(format!(
1600 "Extension '{}' registered provider '{}' with invalid config_schema: config_schema must be a JSON object",
1601 id, provider_id,
1602 ));
1603 }
1604 }
1605 }
1606 Ok(())
1607 }
1608
1609 fn is_safe_provider_id(id: &str) -> bool {
1610 !id.is_empty()
1611 && !id.contains(':')
1612 && id
1613 .chars()
1614 .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '-' || c == '_')
1615 }
1616
1617 #[doc(hidden)]
1618 pub async fn initialize_for_test(&self, plugin_root: Option<PathBuf>) -> Result<(), String> {
1619 self.initialize(plugin_root, Value::Object(Default::default()))
1620 .await
1621 .map(|_| ())
1622 }
1623
1624 fn hook_result(
1628 &self,
1629 outcome: Result<Result<Value, String>, tokio::time::error::Elapsed>,
1630 ) -> HookResult {
1631 match outcome {
1632 Ok(Ok(value)) => match serde_json::from_value(value.clone()) {
1633 Ok(result) => result,
1634 Err(error) => {
1635 tracing::warn!(
1636 extension = %self.id,
1637 error = %error,
1638 response = %value,
1639 "Extension hook handler returned invalid result",
1640 );
1641 match value.get("action").and_then(Value::as_str) {
1642 Some("modify") => HookResult::Block {
1645 reason: "Extension returned malformed modify result".to_string(),
1646 },
1647 Some("replace") => HookResult::Continue,
1652 _ => HookResult::Continue,
1653 }
1654 }
1655 },
1656 Ok(Err(e)) => {
1657 tracing::warn!(
1658 extension = %self.id,
1659 error = %e,
1660 "Extension hook handler failed — continuing",
1661 );
1662 HookResult::Continue
1663 }
1664 Err(_) => {
1665 tracing::warn!(
1666 extension = %self.id,
1667 timeout_secs = 5,
1668 "Extension hook handler timed out — continuing",
1669 );
1670 HookResult::Continue
1671 }
1672 }
1673 }
1674
1675 async fn restart_locked(&self, state: &mut Option<ProcessState>) -> Result<(), String> {
1676 if self.stopped.load(Ordering::Acquire) {
1677 return Err(format!(
1678 "Extension '{}' was stopped (unloaded or reloaded); not respawning it",
1679 self.id,
1680 ));
1681 }
1682 self.inbox.invalidate_lifecycle().await;
1683 let attempted = self.restart_count.fetch_add(1, Ordering::Relaxed) + 1;
1684 self.total_restarts.fetch_add(1, Ordering::Relaxed);
1685 let max_attempts = self.restart_policy.max_attempts;
1686 if attempted > max_attempts as usize {
1687 *state = None;
1688 return Err(format!(
1689 "Extension '{}' exceeded restart limit ({})",
1690 self.id, max_attempts,
1691 ));
1692 }
1693
1694 if let Some(old) = state.take() {
1695 old.reader_handle.abort();
1696 old.lifecycle.invalidate();
1697 let _ = old.lifecycle.child.lock().await.kill().await;
1698 }
1699 self.inbox
1701 .fail_all_pending("transport closed: process restarting")
1702 .await;
1703
1704 let delay = self
1705 .restart_policy
1706 .delay_for_attempt(attempted as u32)
1707 .unwrap_or_default();
1708
1709 tracing::warn!(
1710 extension = %self.id,
1711 attempt = attempted,
1712 max_attempts = max_attempts,
1713 delay_ms = delay.as_millis() as u64,
1714 "Restarting extension process after transport failure",
1715 );
1716
1717 if !delay.is_zero() {
1718 tokio::time::sleep(delay).await;
1719 }
1720
1721 *state = Some(
1722 Self::spawn_state(
1723 &self.id,
1724 &self.command,
1725 &self.args,
1726 self.cwd.as_ref(),
1727 self.inbox.clone(),
1728 )
1729 .await?,
1730 );
1731 self.initialize_locked(state).await?;
1732 Ok(())
1739 }
1740
1741 async fn initialize_locked(&self, state: &mut Option<ProcessState>) -> Result<(), String> {
1742 let config = self.initialize_config.lock().await.clone();
1743 let params = InitializeParams {
1744 synaps_version: env!("CARGO_PKG_VERSION"),
1745 extension_protocol_version: CURRENT_EXTENSION_PROTOCOL_VERSION,
1746 plugin_id: self.id.clone(),
1747 plugin_root: self
1748 .cwd
1749 .clone()
1750 .map(|path| path.to_string_lossy().to_string()),
1751 config,
1752 };
1753 let id = self.next_id.fetch_add(1, Ordering::Relaxed);
1754 let value = tokio::time::timeout(
1755 std::time::Duration::from_secs(10),
1756 self.call_once_locked(
1757 state.as_mut().expect("state should exist for initialize"),
1758 "initialize",
1759 serde_json::to_value(params).map_err(|e| e.to_string())?,
1760 id,
1761 ),
1762 )
1763 .await
1764 .map_err(|_| format!("Extension '{}' initialize timed out after 10s", self.id))??;
1765 let result = Self::parse_initialize_result(&self.id, value).map(|_| ());
1766 if let Some(state) = state.as_ref() {
1767 if result.is_ok() {
1768 state.lifecycle.initialized.store(true, Ordering::Release);
1769 } else {
1770 state.lifecycle.invalidate();
1771 }
1772 }
1773 result
1774 }
1775
1776 async fn call_once_locked(
1780 &self,
1781 state: &mut ProcessState,
1782 method: &str,
1783 params: Value,
1784 id: u64,
1785 ) -> Result<Value, String> {
1786 let body = serde_json::to_string(&JsonRpcRequest {
1787 jsonrpc: "2.0",
1788 method: method.to_string(),
1789 params,
1790 id,
1791 })
1792 .map_err(|e| format!("Serialize error: {}", e))?;
1793
1794 let (tx, rx) = oneshot::channel::<Result<Value, String>>();
1795 if self.inbox.closed.load(std::sync::atomic::Ordering::Acquire) {
1798 return Err("transport closed: inbox is shut down".to_string());
1799 }
1800
1801 self.inbox.pending.lock().await.insert(id, tx);
1804
1805 if self.inbox.closed.load(std::sync::atomic::Ordering::Acquire) {
1808 self.inbox.pending.lock().await.remove(&id);
1809 return Err("transport closed: inbox shut down during registration".to_string());
1810 }
1811
1812 let frame = format!("Content-Length: {}\r\n\r\n{}", body.len(), body);
1813 let write_result = {
1814 let mut stdin = state.stdin.lock().await;
1815 match stdin.write_all(frame.as_bytes()).await {
1816 Ok(()) => stdin.flush().await,
1817 Err(e) => Err(e),
1818 }
1819 };
1820 if let Err(e) = write_result {
1821 state.lifecycle.invalidate();
1822 self.inbox.pending.lock().await.remove(&id);
1824 return Err(format!("Write error: {}", e));
1825 }
1826
1827 match rx.await {
1828 Ok(payload) => payload,
1829 Err(_) => {
1830 state.lifecycle.invalidate();
1831 self.inbox.pending.lock().await.remove(&id);
1836 Err("transport closed: response channel dropped".to_string())
1837 }
1838 }
1839 }
1840
1841 async fn call_no_restart(&self, method: &str, params: Value) -> Result<Value, String> {
1842 let _call_guard = self.call_lock.lock().await;
1843 let id = self.next_id.fetch_add(1, Ordering::Relaxed);
1844 let mut state_guard = self.state.lock().await;
1845 if state_guard.is_none() {
1846 *state_guard = Some(
1847 Self::spawn_state(
1848 &self.id,
1849 &self.command,
1850 &self.args,
1851 self.cwd.as_ref(),
1852 self.inbox.clone(),
1853 )
1854 .await?,
1855 );
1856 }
1857 let state = state_guard.as_mut().expect("state should exist");
1858 let result = self.call_once_locked(state, method, params, id).await;
1859 if method == "initialize" {
1860 match result.as_ref() {
1863 Ok(value) if Self::parse_initialize_result(&self.id, value.clone()).is_ok() => {
1864 state.lifecycle.initialized.store(true, Ordering::Release);
1865 }
1866 _ => state.lifecycle.invalidate(),
1867 }
1868 }
1869 result
1870 }
1871
1872 async fn call_no_spawn(&self, method: &str, params: Value) -> Result<Value, String> {
1876 let _call_guard = self.call_lock.lock().await;
1877 let id = self.next_id.fetch_add(1, Ordering::Relaxed);
1878 let mut state_guard = self.state.lock().await;
1879 let Some(state) = state_guard.as_mut() else {
1880 return Err(format!(
1881 "Extension '{}' process is not running; '{}' refused (no respawn)",
1882 self.id, method
1883 ));
1884 };
1885 self.call_once_locked(state, method, params, id).await
1886 }
1887
1888 pub(crate) async fn call_tool_once_no_spawn(
1894 &self,
1895 name: &str,
1896 input: Value,
1897 session_id: Option<&str>,
1898 ) -> Result<Value, String> {
1899 let fut = async {
1900 let _call_guard = self.call_lock.lock().await;
1901 let id = self.next_id.fetch_add(1, Ordering::Relaxed);
1902 let mut state = {
1907 let mut state_guard = self.state.lock().await;
1908 state_guard.take().ok_or_else(|| {
1909 format!(
1910 "Extension '{}' process is not running; leased call refused (no respawn)",
1911 self.id
1912 )
1913 })?
1914 };
1915 let result = self
1916 .call_once_locked(
1917 &mut state,
1918 "tool.call",
1919 tool_call_params(name, input, session_id),
1920 id,
1921 )
1922 .await;
1923 if !self.inbox.closed.load(std::sync::atomic::Ordering::Acquire) {
1924 *self.state.lock().await = Some(state);
1925 } else {
1926 state.reader_handle.abort();
1927 state.lifecycle.invalidate();
1928 let _ = state.lifecycle.child.lock().await.kill().await;
1929 }
1930 result
1931 };
1932 match tokio::time::timeout(std::time::Duration::from_secs(120), fut).await {
1933 Ok(result) => result,
1934 Err(_) => Err(format!(
1935 "Extension '{}' leased tool.call timed out after 120s",
1936 self.id
1937 )),
1938 }
1939 }
1940
1941 async fn call(&self, method: &str, params: Value) -> Result<Value, String> {
1942 let timeout_secs = if method == "tool.call" { 120 } else { 30 };
1943 let id_str = self.id.clone();
1944 let method_str = method.to_string();
1945
1946 let result = tokio::time::timeout(
1947 std::time::Duration::from_secs(timeout_secs),
1948 self.call_inner(method, params),
1949 )
1950 .await;
1951
1952 match result {
1953 Ok(inner) => inner,
1954 Err(_) => Err(format!(
1955 "Extension '{}' method '{}' timed out after {}s",
1956 id_str, method_str, timeout_secs
1957 )),
1958 }
1959 }
1960
1961 async fn call_inner(&self, method: &str, params: Value) -> Result<Value, String> {
1962 let _call_guard = self.call_lock.lock().await;
1963 let id = self.next_id.fetch_add(1, Ordering::Relaxed);
1964 let mut state_guard = self.state.lock().await;
1965 if state_guard.is_none() {
1966 self.restart_locked(&mut state_guard).await?;
1967 }
1968
1969 let result = self
1970 .call_once_locked(
1971 state_guard
1972 .as_mut()
1973 .expect("state should exist after restart"),
1974 method,
1975 params.clone(),
1976 id,
1977 )
1978 .await;
1979
1980 match result {
1981 Ok(value) => {
1982 self.restart_count.store(0, Ordering::Relaxed);
1986 Ok(value)
1987 }
1988 Err(first_error) if is_application_error(&first_error) => {
1989 self.restart_count.store(0, Ordering::Relaxed);
1995 Err(first_error)
1996 }
1997 Err(first_error) => {
1998 self.restart_locked(&mut state_guard).await?;
1999 let retry_id = self.next_id.fetch_add(1, Ordering::Relaxed);
2000 self.call_once_locked(
2001 state_guard
2002 .as_mut()
2003 .expect("state should exist after restart"),
2004 method,
2005 params,
2006 retry_id,
2007 )
2008 .await
2009 .inspect(|_| {
2010 self.restart_count.store(0, Ordering::Relaxed);
2012 })
2013 .map_err(|retry_error| {
2014 format!(
2015 "{}; retry after restart failed: {}",
2016 first_error, retry_error
2017 )
2018 })
2019 }
2020 }
2021 }
2022
2023 #[doc(hidden)]
2041 pub async fn subscribe_notifications(&self) -> (usize, mpsc::Receiver<NotificationFrame>) {
2042 let (tx, rx) = mpsc::channel(NOTIFICATION_QUEUE_CAPACITY);
2043 let id = self
2044 .inbox
2045 .next_sink_id
2046 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2047 let mut sinks = self.inbox.notification_sinks.lock().await;
2048 sinks.retain(|(_, tx)| !tx.is_closed());
2052 sinks.push((id, tx));
2053 (id, rx)
2054 }
2055
2056 #[doc(hidden)]
2061 pub async fn unsubscribe_notifications(&self, id: usize) {
2062 let mut sinks = self.inbox.notification_sinks.lock().await;
2063 sinks.retain(|(sub_id, tx)| *sub_id != id && !tx.is_closed());
2064 }
2065
2066 pub(crate) async fn forward_invoke_command_frame(
2085 extension_id: &str,
2086 request_id: &str,
2087 sink: &crate::extensions::invoke_output::InvokeEventSink,
2088 sink_open: &mut bool,
2089 frame: NotificationFrame,
2090 ) -> bool {
2091 use crate::extensions::commands::parse_command_output;
2092 use crate::extensions::runtime::InvokeCommandEvent;
2093 use crate::extensions::tasks::{is_task_method, parse_task_event};
2094
2095 let mut saw_done = false;
2096 if frame.method == "command.output" {
2097 match parse_command_output(&frame.params) {
2098 Ok(parsed) if parsed.request_id == request_id => {
2099 if matches!(
2100 parsed.event,
2101 crate::extensions::commands::CommandOutputEvent::Done
2102 ) {
2103 saw_done = true;
2104 }
2105 if *sink_open
2106 && sink
2107 .send(InvokeCommandEvent::Output(parsed.event))
2108 .await
2109 .is_err()
2110 {
2111 *sink_open = false;
2112 }
2113 }
2114 Ok(_) => {
2115 tracing::trace!(
2117 extension = %extension_id,
2118 "Ignoring command.output for unrelated request_id",
2119 );
2120 }
2121 Err(error) => {
2122 tracing::warn!(
2123 extension = %extension_id,
2124 error = %error,
2125 params = %frame.params,
2126 "Skipping malformed command.output notification",
2127 );
2128 }
2129 }
2130 } else if is_task_method(&frame.method) {
2131 match parse_task_event(&frame.method, &frame.params) {
2132 Ok(event) => {
2133 if *sink_open && sink.send(InvokeCommandEvent::Task(event)).await.is_err() {
2134 *sink_open = false;
2135 }
2136 }
2137 Err(error) => {
2138 tracing::warn!(
2139 extension = %extension_id,
2140 method = %frame.method,
2141 error = %error,
2142 params = %frame.params,
2143 "Skipping malformed task notification",
2144 );
2145 }
2146 }
2147 } else {
2148 tracing::trace!(
2149 extension = %extension_id,
2150 method = %frame.method,
2151 "Ignoring non-command/task notification during command.invoke",
2152 );
2153 }
2154 saw_done
2155 }
2156
2157 async fn forward_provider_stream_frame(
2167 extension_id: &str,
2168 sink: &mpsc::Sender<ProviderStreamEvent>,
2169 sink_open: &mut bool,
2170 frame: NotificationFrame,
2171 ) {
2172 if frame.method != "provider.stream.event" {
2173 tracing::trace!(
2174 extension = %extension_id,
2175 method = %frame.method,
2176 "Ignoring non-stream notification during provider.stream",
2177 );
2178 return;
2179 }
2180 match parse_provider_stream_event(&frame.params) {
2181 Ok(event) => {
2182 if *sink_open && sink.send(event).await.is_err() {
2183 *sink_open = false;
2184 }
2185 }
2186 Err(error) => {
2187 tracing::warn!(
2188 extension = %extension_id,
2189 error = %error,
2190 params = %frame.params,
2191 "Skipping malformed provider.stream.event notification",
2192 );
2193 }
2194 }
2195 }
2196}
2197
2198pub(crate) fn tool_call_params(name: &str, input: Value, session_id: Option<&str>) -> Value {
2203 tool_call_params_gated(
2204 name,
2205 input,
2206 session_id,
2207 crate::extensions::hooks::events::hook_session_id_enabled(),
2208 )
2209}
2210
2211fn tool_call_params_gated(
2212 name: &str,
2213 input: Value,
2214 session_id: Option<&str>,
2215 session_id_enabled: bool,
2216) -> Value {
2217 let session_id = if session_id_enabled { session_id } else { None };
2218 serde_json::json!({
2219 "name": name,
2220 "input": input,
2221 "session_id": session_id,
2222 })
2223}
2224
2225#[async_trait]
2226impl ExtensionHandler for ProcessExtension {
2227 fn id(&self) -> &str {
2228 &self.id
2229 }
2230
2231 fn lifecycle_snapshot(&self) -> Option<ExtensionLifecycle> {
2232 let lifecycle = self.inbox.lifecycle.try_lock().ok()?.clone()?;
2233 lifecycle.snapshot(self.total_restarts.load(Ordering::Relaxed) > 0)
2234 }
2235
2236 async fn call_tool(&self, name: &str, input: Value) -> Result<Value, String> {
2237 self.call_tool_for(name, input, None).await
2238 }
2239
2240 async fn call_tool_for(
2241 &self,
2242 name: &str,
2243 input: Value,
2244 session_id: Option<&str>,
2245 ) -> Result<Value, String> {
2246 self.call("tool.call", tool_call_params(name, input, session_id))
2247 .await
2248 }
2249
2250 async fn provider_complete(
2251 &self,
2252 params: ProviderCompleteParams,
2253 ) -> Result<ProviderCompleteResult, String> {
2254 let value = tokio::time::timeout(
2255 std::time::Duration::from_secs(60),
2256 self.call(
2257 "provider.complete",
2258 serde_json::to_value(params).map_err(|e| e.to_string())?,
2259 ),
2260 )
2261 .await
2262 .map_err(|_| format!("Extension '{}' provider.complete timed out", self.id))??;
2263 let result: ProviderCompleteResult = serde_json::from_value(value).map_err(|e| {
2264 format!(
2265 "Invalid provider.complete response from extension '{}': {}",
2266 self.id, e
2267 )
2268 })?;
2269 if result.content.is_empty() {
2270 return Err(format!(
2271 "Extension '{}' provider.complete returned empty content",
2272 self.id
2273 ));
2274 }
2275 Ok(result)
2276 }
2277
2278 async fn provider_stream(
2279 &self,
2280 params: ProviderCompleteParams,
2281 sink: tokio::sync::mpsc::Sender<ProviderStreamEvent>,
2282 ) -> Result<ProviderCompleteResult, String> {
2283 let (sub_id, mut rx) = self.subscribe_notifications().await;
2286 let params_value = serde_json::to_value(params).map_err(|e| e.to_string())?;
2287
2288 let extension_id = self.id.clone();
2289 let stream_future = async {
2290 let mut call_fut = Box::pin(self.call("provider.stream", params_value));
2291 let mut sink_open = true;
2292 let response = loop {
2293 tokio::select! {
2294 response = &mut call_fut => break response,
2295 Some(frame) = rx.recv() => {
2296 Self::forward_provider_stream_frame(
2297 &extension_id, &sink, &mut sink_open, frame,
2298 ).await;
2299 }
2300 }
2301 };
2302 self.unsubscribe_notifications(sub_id).await;
2307 while let Some(frame) = rx.recv().await {
2308 Self::forward_provider_stream_frame(&extension_id, &sink, &mut sink_open, frame)
2309 .await;
2310 }
2311 response
2312 };
2313
2314 let outcome = tokio::time::timeout(std::time::Duration::from_secs(60), stream_future).await;
2315
2316 self.unsubscribe_notifications(sub_id).await;
2319
2320 let value =
2321 outcome.map_err(|_| format!("Extension '{}' provider.stream timed out", self.id))??;
2322
2323 let result: ProviderCompleteResult = serde_json::from_value(value).map_err(|e| {
2324 format!(
2325 "Invalid provider.stream response from extension '{}': {}",
2326 self.id, e
2327 )
2328 })?;
2329 Ok(result)
2332 }
2333
2334 async fn invoke_command(
2335 &self,
2336 command: &str,
2337 args: Vec<String>,
2338 request_id: &str,
2339 sink: crate::extensions::invoke_output::InvokeEventSink,
2340 ) -> Result<Value, String> {
2341 let (sub_id, mut rx) = self.subscribe_notifications().await;
2343 let params = serde_json::json!({
2344 "command": command,
2345 "args": args,
2346 "request_id": request_id,
2347 });
2348
2349 let extension_id = self.id.clone();
2350 let request_id_owned = request_id.to_string();
2351 let invoke_future = async {
2352 let mut call_fut = Box::pin(self.call("command.invoke", params));
2353 let mut sink_open = true;
2354 let response = loop {
2355 tokio::select! {
2356 response = &mut call_fut => break response,
2357 Some(frame) = rx.recv() => {
2358 let _ = Self::forward_invoke_command_frame(
2364 &extension_id, &request_id_owned, &sink, &mut sink_open, frame,
2365 ).await;
2366 }
2367 }
2368 };
2369 self.unsubscribe_notifications(sub_id).await;
2373 while let Ok(frame) = rx.try_recv() {
2374 let _ = Self::forward_invoke_command_frame(
2375 &extension_id,
2376 &request_id_owned,
2377 &sink,
2378 &mut sink_open,
2379 frame,
2380 )
2381 .await;
2382 }
2383 response
2384 };
2385
2386 let outcome =
2387 tokio::time::timeout(std::time::Duration::from_secs(120), invoke_future).await;
2388
2389 self.unsubscribe_notifications(sub_id).await;
2394
2395 outcome.map_err(|_| format!("Extension '{}' command.invoke timed out", self.id))?
2396 }
2397
2398 async fn handle(&self, event: &HookEvent) -> HookResult {
2399 let params = serde_json::to_value(event).unwrap_or(Value::Null);
2400 self.hook_result(
2401 tokio::time::timeout(
2402 std::time::Duration::from_secs(5),
2403 self.call("hook.handle", params),
2404 )
2405 .await,
2406 )
2407 }
2408
2409 async fn handle_no_spawn(&self, event: &HookEvent) -> HookResult {
2413 let params = serde_json::to_value(event).unwrap_or(Value::Null);
2414 self.hook_result(
2415 tokio::time::timeout(
2416 std::time::Duration::from_secs(5),
2417 self.call_no_spawn("hook.handle", params),
2418 )
2419 .await,
2420 )
2421 }
2422
2423 async fn get_info(&self) -> Result<crate::extensions::info::PluginInfo, String> {
2424 let value = tokio::time::timeout(
2425 std::time::Duration::from_secs(5),
2426 self.call("info.get", Value::Null),
2427 )
2428 .await
2429 .map_err(|_| format!("Extension '{}' info.get timed out", self.id))??;
2430 serde_json::from_value(value).map_err(|e| {
2431 format!(
2432 "Invalid info.get response from extension '{}': {}",
2433 self.id, e
2434 )
2435 })
2436 }
2437
2438 async fn sidecar_spawn_args(&self) -> Result<crate::sidecar::spawn::SidecarSpawnArgs, String> {
2439 let value = tokio::time::timeout(
2440 std::time::Duration::from_secs(5),
2441 self.call("sidecar.spawn_args", Value::Null),
2442 )
2443 .await
2444 .map_err(|_| format!("Extension '{}' sidecar.spawn_args timed out", self.id))??;
2445 serde_json::from_value(value).map_err(|e| {
2446 format!(
2447 "Invalid sidecar.spawn_args response from extension '{}': {}",
2448 self.id, e
2449 )
2450 })
2451 }
2452
2453 async fn settings_editor_open(&self, category: &str, field: &str) -> Result<Value, String> {
2454 let params = crate::extensions::settings_editor::SettingsEditorOpenParams {
2455 category: category.to_string(),
2456 field: field.to_string(),
2457 };
2458 tokio::time::timeout(
2459 std::time::Duration::from_secs(5),
2460 self.call(
2461 "settings.editor.open",
2462 serde_json::to_value(params).map_err(|e| e.to_string())?,
2463 ),
2464 )
2465 .await
2466 .map_err(|_| format!("Extension '{}' settings.editor.open timed out", self.id))?
2467 }
2468
2469 async fn settings_editor_key(
2470 &self,
2471 category: &str,
2472 field: &str,
2473 key: &str,
2474 ) -> Result<Value, String> {
2475 let mut params = serde_json::to_value(
2476 crate::extensions::settings_editor::SettingsEditorKeyParams {
2477 key: key.to_string(),
2478 },
2479 )
2480 .map_err(|e| e.to_string())?;
2481 if let Some(obj) = params.as_object_mut() {
2482 obj.insert("category".to_string(), Value::String(category.to_string()));
2483 obj.insert("field".to_string(), Value::String(field.to_string()));
2484 }
2485 tokio::time::timeout(
2486 std::time::Duration::from_secs(5),
2487 self.call("settings.editor.key", params),
2488 )
2489 .await
2490 .map_err(|_| format!("Extension '{}' settings.editor.key timed out", self.id))?
2491 }
2492
2493 async fn settings_editor_commit(
2494 &self,
2495 category: &str,
2496 field: &str,
2497 value: Value,
2498 ) -> Result<Value, String> {
2499 let params = serde_json::json!({
2500 "category": category,
2501 "field": field,
2502 "value": value,
2503 });
2504 tokio::time::timeout(
2505 std::time::Duration::from_secs(5),
2506 self.call("settings.editor.commit", params),
2507 )
2508 .await
2509 .map_err(|_| format!("Extension '{}' settings.editor.commit timed out", self.id))?
2510 }
2511
2512 async fn shutdown(&self) {
2513 self.stopped.store(true, Ordering::Release);
2516 self.inbox.invalidate_lifecycle().await;
2517 let _ = tokio::time::timeout(
2523 std::time::Duration::from_millis(500),
2524 self.call_no_spawn("shutdown", Value::Null),
2525 )
2526 .await;
2527
2528 tokio::time::sleep(std::time::Duration::from_millis(500)).await;
2529 let mut state_guard = self.state.lock().await;
2530 if let Some(state) = state_guard.take() {
2531 state.reader_handle.abort();
2532 state.lifecycle.invalidate();
2533 let _ = state.lifecycle.child.lock().await.kill().await;
2534 }
2535 self.inbox.notification_sinks.lock().await.clear();
2537 self.inbox
2538 .fail_all_pending("transport closed: extension shutdown")
2539 .await;
2540 }
2541
2542 async fn subscribe_notifications(
2543 &self,
2544 ) -> (usize, tokio::sync::mpsc::Receiver<NotificationFrame>) {
2545 ProcessExtension::subscribe_notifications(self).await
2546 }
2547
2548 async fn restart_count(&self) -> usize {
2549 self.restart_count()
2550 }
2551
2552 async fn health(&self) -> ExtensionHealth {
2553 let consecutive = self.restart_count.load(Ordering::Relaxed);
2557 let lifetime = self.total_restarts.load(Ordering::Relaxed);
2558 let max = self.restart_policy.max_attempts as usize;
2559 if consecutive >= max {
2560 ExtensionHealth::Failed
2561 } else if lifetime > 0 {
2562 let state_alive = self.state.try_lock().map(|g| g.is_some()).unwrap_or(true);
2566 if state_alive {
2567 ExtensionHealth::Degraded
2568 } else {
2569 ExtensionHealth::Restarting
2570 }
2571 } else {
2572 ExtensionHealth::Running
2573 }
2574 }
2575}
2576
2577#[cfg(test)]
2578mod exclusive_memory_tests {
2579 use super::*;
2580 #[tokio::test]
2581 async fn exclusive_inbox_refuses_legacy_memory_without_writes() {
2582 let mut inbox = Inbox::new("fixture".into());
2583 inbox.exclusive_memory = true;
2584 let inbox = Arc::new(inbox);
2585 for method in ["memory.append", "memory.query"] {
2586 let result = ProcessExtension::handle_inbound_request(
2587 &inbox,
2588 method,
2589 serde_json::json!({"namespace":"fixture","content":"must not persist"}),
2590 )
2591 .await;
2592 assert!(result.unwrap_err().1.contains("disabled"));
2593 }
2594 }
2595}
2596
2597#[cfg(test)]
2598mod tool_call_params_tests {
2599 use super::*;
2600 use serde_json::json;
2601
2602 #[test]
2603 fn carries_caller_session_id() {
2604 assert_eq!(
2605 tool_call_params_gated("echo", json!({"text": "hi"}), Some("sess-A"), true),
2606 json!({"name": "echo", "input": {"text": "hi"}, "session_id": "sess-A"})
2607 );
2608 }
2609
2610 #[test]
2611 fn worker_call_has_null_session_id_key() {
2612 let params = tool_call_params_gated("echo", json!({}), None, true);
2613 assert_eq!(params.get("session_id"), Some(&Value::Null));
2614 }
2615
2616 #[test]
2617 fn kill_switch_forces_null() {
2618 let params = tool_call_params_gated("echo", json!({}), Some("sess-A"), false);
2619 assert_eq!(params.get("session_id"), Some(&Value::Null));
2620 assert_eq!(params["name"], json!("echo"));
2621 }
2622}
2623
2624#[cfg(test)]
2625mod stream_event_tests {
2626 use super::*;
2627 use serde_json::json;
2628
2629 #[test]
2630 fn parses_text_delta_with_delta_key() {
2631 let v = json!({"type": "text", "delta": "hi"});
2632 assert_eq!(
2633 parse_provider_stream_event(&v).unwrap(),
2634 ProviderStreamEvent::TextDelta { text: "hi".into() }
2635 );
2636 }
2637
2638 #[test]
2639 fn parses_text_delta_with_text_key() {
2640 let v = json!({"type": "text", "text": "hi"});
2641 assert_eq!(
2642 parse_provider_stream_event(&v).unwrap(),
2643 ProviderStreamEvent::TextDelta { text: "hi".into() }
2644 );
2645 }
2646
2647 #[test]
2648 fn parses_thinking_delta() {
2649 let v = json!({"type": "thinking", "delta": "hmm"});
2650 assert_eq!(
2651 parse_provider_stream_event(&v).unwrap(),
2652 ProviderStreamEvent::ThinkingDelta { text: "hmm".into() }
2653 );
2654 let v2 = json!({"type": "thinking", "text": "hmm"});
2655 assert_eq!(
2656 parse_provider_stream_event(&v2).unwrap(),
2657 ProviderStreamEvent::ThinkingDelta { text: "hmm".into() }
2658 );
2659 }
2660
2661 #[test]
2662 fn parses_tool_use() {
2663 let v = json!({
2664 "type": "tool_use",
2665 "id": "t1",
2666 "name": "echo",
2667 "input": {"x": 1}
2668 });
2669 assert_eq!(
2670 parse_provider_stream_event(&v).unwrap(),
2671 ProviderStreamEvent::ToolUse {
2672 id: "t1".into(),
2673 name: "echo".into(),
2674 input: json!({"x": 1}),
2675 }
2676 );
2677 }
2678
2679 #[test]
2680 fn tool_use_input_defaults_to_empty_object() {
2681 let v = json!({"type": "tool_use", "id": "t1", "name": "echo"});
2682 assert_eq!(
2683 parse_provider_stream_event(&v).unwrap(),
2684 ProviderStreamEvent::ToolUse {
2685 id: "t1".into(),
2686 name: "echo".into(),
2687 input: json!({}),
2688 }
2689 );
2690 }
2691
2692 #[test]
2693 fn parses_usage_strips_type() {
2694 let v = json!({"type": "usage", "input_tokens": 5, "output_tokens": 7});
2695 assert_eq!(
2696 parse_provider_stream_event(&v).unwrap(),
2697 ProviderStreamEvent::Usage {
2698 usage: json!({"input_tokens": 5, "output_tokens": 7})
2699 }
2700 );
2701 }
2702
2703 #[test]
2704 fn parses_error() {
2705 let v = json!({"type": "error", "message": "boom"});
2706 assert_eq!(
2707 parse_provider_stream_event(&v).unwrap(),
2708 ProviderStreamEvent::Error {
2709 message: "boom".into()
2710 }
2711 );
2712 }
2713
2714 #[test]
2715 fn parses_done() {
2716 let v = json!({"type": "done"});
2717 assert_eq!(
2718 parse_provider_stream_event(&v).unwrap(),
2719 ProviderStreamEvent::Done
2720 );
2721 }
2722
2723 #[test]
2724 fn nested_event_shape_matches_flat() {
2725 let flat = json!({"type": "text", "delta": "hi"});
2726 let nested = json!({"event": {"type": "text", "delta": "hi"}});
2727 assert_eq!(
2728 parse_provider_stream_event(&flat).unwrap(),
2729 parse_provider_stream_event(&nested).unwrap()
2730 );
2731 }
2732
2733 #[test]
2734 fn missing_type_errors() {
2735 let v = json!({"delta": "hi"});
2736 let err = parse_provider_stream_event(&v).unwrap_err();
2737 assert!(err.contains("missing type"), "got: {err}");
2738 }
2739
2740 #[test]
2741 fn unknown_type_errors_with_type() {
2742 let v = json!({"type": "wat"});
2743 let err = parse_provider_stream_event(&v).unwrap_err();
2744 assert!(err.contains("wat"), "got: {err}");
2745 }
2746
2747 #[test]
2748 fn tool_use_missing_id_errors() {
2749 let v = json!({"type": "tool_use", "name": "echo"});
2750 let err = parse_provider_stream_event(&v).unwrap_err();
2751 assert!(err.contains("id"), "got: {err}");
2752 }
2753
2754 #[test]
2755 fn tool_use_missing_name_errors() {
2756 let v = json!({"type": "tool_use", "id": "t1"});
2757 let err = parse_provider_stream_event(&v).unwrap_err();
2758 assert!(err.contains("name"), "got: {err}");
2759 }
2760
2761 #[test]
2762 fn tool_use_empty_id_errors() {
2763 let v = json!({"type": "tool_use", "id": "", "name": "echo"});
2764 assert!(parse_provider_stream_event(&v).is_err());
2765 }
2766
2767 #[test]
2768 fn tool_use_empty_name_errors() {
2769 let v = json!({"type": "tool_use", "id": "t1", "name": ""});
2770 assert!(parse_provider_stream_event(&v).is_err());
2771 }
2772
2773 #[test]
2774 fn tool_use_non_object_input_errors() {
2775 let v = json!({"type": "tool_use", "id": "t1", "name": "echo", "input": "nope"});
2776 let err = parse_provider_stream_event(&v).unwrap_err();
2777 assert!(err.contains("input"), "got: {err}");
2778 }
2779
2780 #[test]
2781 fn text_missing_delta_and_text_errors() {
2782 let v = json!({"type": "text"});
2783 let err = parse_provider_stream_event(&v).unwrap_err();
2784 assert!(err.contains("delta") || err.contains("text"), "got: {err}");
2785 }
2786
2787 #[test]
2788 fn error_missing_message_errors() {
2789 let v = json!({"type": "error"});
2790 assert!(parse_provider_stream_event(&v).is_err());
2791 }
2792
2793 #[test]
2794 fn error_empty_message_errors() {
2795 let v = json!({"type": "error", "message": ""});
2796 assert!(parse_provider_stream_event(&v).is_err());
2797 }
2798}
2799
2800#[cfg(test)]
2801mod restart_policy_tests {
2802 use super::*;
2803
2804 #[tokio::test]
2805 async fn restart_policy_default_max_attempts_is_3() {
2806 let ext = ProcessExtension::spawn("policy-test", "/bin/cat", &[])
2811 .await
2812 .expect("spawn /bin/cat");
2813 assert_eq!(ext.restart_policy.max_attempts, 3);
2814 ext.shutdown().await;
2815 }
2816
2817 #[tokio::test]
2818 async fn with_restart_policy_overrides_default() {
2819 let ext = ProcessExtension::spawn("policy-test-override", "/bin/cat", &[])
2820 .await
2821 .expect("spawn /bin/cat");
2822 let custom = RestartPolicy {
2823 max_attempts: 7,
2824 ..RestartPolicy::default()
2825 };
2826 let ext = ext.with_restart_policy(custom);
2827 assert_eq!(ext.restart_policy.max_attempts, 7);
2828 ext.shutdown().await;
2829 }
2830}
2831
2832#[cfg(test)]
2833mod capture_validator_tests {
2834 use super::*;
2835 use crate::extensions::permissions::{Permission, PermissionSet};
2836
2837 fn perms_with(grants: &[Permission]) -> PermissionSet {
2838 let mut p = PermissionSet::new();
2839 for g in grants {
2840 p.grant(*g);
2841 }
2842 p
2843 }
2844
2845 fn cap(kind: &str, name: &str, perms: &[&str]) -> CapabilityDeclaration {
2846 CapabilityDeclaration {
2847 kind: kind.to_string(),
2848 name: name.to_string(),
2849 permissions: perms.iter().map(|p| p.to_string()).collect(),
2850 params: serde_json::Value::Null,
2851 }
2852 }
2853
2854 #[test]
2855 fn capability_validator_rejects_empty_kind() {
2856 let d = cap(" ", "Sample", &["audio.input"]);
2857 let perms = perms_with(&[Permission::AudioInput]);
2858 let err = validate_capability(&d, &perms).unwrap_err();
2859 assert!(err.contains("kind"), "got: {}", err);
2860 }
2861
2862 #[test]
2863 fn capability_validator_rejects_empty_name() {
2864 let d = cap("capture", " ", &["audio.input"]);
2865 let perms = perms_with(&[Permission::AudioInput]);
2866 let err = validate_capability(&d, &perms).unwrap_err();
2867 assert!(err.contains("name"), "got: {}", err);
2868 }
2869
2870 #[test]
2871 fn capability_validator_rejects_unknown_permission_string() {
2872 let d = cap("capture", "Sample", &["audio.telepathy"]);
2873 let perms = perms_with(&[Permission::AudioInput, Permission::AudioOutput]);
2874 let err = validate_capability(&d, &perms).unwrap_err();
2875 assert!(
2876 err.contains("unknown permission") && err.contains("audio.telepathy"),
2877 "got: {}",
2878 err,
2879 );
2880 }
2881
2882 #[test]
2883 fn capability_validator_requires_every_declared_permission() {
2884 let d = cap("capture", "Sample", &["audio.input"]);
2885 let perms = perms_with(&[]);
2886 let err = validate_capability(&d, &perms).unwrap_err();
2887 assert!(
2888 err.contains("audio.input") && err.contains("not granted"),
2889 "got: {}",
2890 err,
2891 );
2892 }
2893
2894 #[test]
2895 fn capability_validator_accepts_when_all_permissions_granted() {
2896 let d = cap("capture", "Sample", &["audio.input", "audio.output"]);
2897 let perms = perms_with(&[Permission::AudioInput, Permission::AudioOutput]);
2898 validate_capability(&d, &perms).expect("should validate");
2899 }
2900
2901 #[test]
2902 fn capability_validator_accepts_no_permissions() {
2903 let d = cap("ocr", "Tesseract", &[]);
2907 let perms = perms_with(&[]);
2908 validate_capability(&d, &perms).expect("should validate");
2909 }
2910
2911 #[test]
2912 fn capability_validator_does_not_branch_on_kind() {
2913 let perms = perms_with(&[Permission::AudioInput]);
2917 for kind in ["capture", "ocr", "agent", "foot_pedal", "eeg"] {
2918 let d = cap(kind, "Anything", &["audio.input"]);
2919 validate_capability(&d, &perms).expect("should validate");
2920 }
2921 }
2922}
2923
2924#[cfg(test)]
2925mod invoke_command_dispatch_tests {
2926 use super::*;
2931 use crate::extensions::commands::CommandOutputEvent;
2932 use crate::extensions::invoke_output::{invoke_event_channel, InvokeOutputBudget};
2933 use crate::extensions::runtime::InvokeCommandEvent;
2934 use crate::extensions::tasks::{TaskEvent, TaskKind};
2935 use serde_json::json;
2936
2937 fn frame(method: &str, params: serde_json::Value) -> NotificationFrame {
2938 NotificationFrame {
2939 method: method.to_string(),
2940 params,
2941 }
2942 }
2943
2944 async fn dispatch(frames: Vec<NotificationFrame>) -> (Vec<InvokeCommandEvent>, bool) {
2949 let (tx, collector) = invoke_event_channel(InvokeOutputBudget::default());
2950 let mut open = true;
2951 let mut saw_done = false;
2952 for f in frames {
2953 saw_done |=
2954 ProcessExtension::forward_invoke_command_frame("ext-test", "r1", &tx, &mut open, f)
2955 .await;
2956 }
2957 drop(tx);
2958 let report = collector.collect().await;
2959 (report.events, saw_done)
2960 }
2961
2962 #[tokio::test]
2963 async fn forwards_mixed_event_stream_in_order() {
2964 let frames = vec![
2965 frame(
2966 "command.output",
2967 json!({"request_id":"r1","event":{"kind":"text","content":"A"}}),
2968 ),
2969 frame(
2970 "task.start",
2971 json!({"id":"dl","label":"Downloading","kind":"download"}),
2972 ),
2973 frame("task.update", json!({"id":"dl","current":50,"total":100})),
2974 frame(
2975 "command.output",
2976 json!({"request_id":"r1","event":{"kind":"system","content":"working"}}),
2977 ),
2978 frame("task.done", json!({"id":"dl"})),
2979 frame(
2980 "command.output",
2981 json!({"request_id":"r1","event":{"kind":"done"}}),
2982 ),
2983 ];
2984
2985 let (events, saw_done) = dispatch(frames).await;
2986 assert!(saw_done, "should have observed the command Done marker");
2987 assert_eq!(events.len(), 6);
2988 assert_eq!(
2989 events[0],
2990 InvokeCommandEvent::Output(CommandOutputEvent::Text {
2991 content: "A".into()
2992 })
2993 );
2994 assert!(matches!(
2995 events[1],
2996 InvokeCommandEvent::Task(TaskEvent::Start {
2997 kind: TaskKind::Download,
2998 ..
2999 })
3000 ));
3001 assert!(matches!(
3002 events[2],
3003 InvokeCommandEvent::Task(TaskEvent::Update { .. })
3004 ));
3005 assert!(matches!(
3006 events[3],
3007 InvokeCommandEvent::Output(CommandOutputEvent::System { .. })
3008 ));
3009 assert!(matches!(
3010 events[4],
3011 InvokeCommandEvent::Task(TaskEvent::Done { error: None, .. })
3012 ));
3013 assert_eq!(
3014 events[5],
3015 InvokeCommandEvent::Output(CommandOutputEvent::Done)
3016 );
3017 }
3018
3019 #[tokio::test]
3020 async fn ignores_command_output_for_unrelated_request_id() {
3021 let (events, saw_done) = dispatch(vec![frame(
3022 "command.output",
3023 json!({"request_id":"other","event":{"kind":"text","content":"x"}}),
3024 )])
3025 .await;
3026 assert!(!saw_done);
3027 assert!(events.is_empty());
3028 }
3029
3030 #[tokio::test]
3031 async fn skips_malformed_command_output_without_aborting() {
3032 let (events, saw_done) = dispatch(vec![
3033 frame("command.output", json!({"request_id":"r1","event":{}})),
3035 frame(
3037 "command.output",
3038 json!({"request_id":"r1","event":{"kind":"done"}}),
3039 ),
3040 ])
3041 .await;
3042 assert!(saw_done);
3043 assert_eq!(
3044 events,
3045 vec![InvokeCommandEvent::Output(CommandOutputEvent::Done)]
3046 );
3047 }
3048
3049 #[tokio::test]
3050 async fn task_events_pass_through_regardless_of_request_id() {
3051 let (events, _) = dispatch(vec![frame("task.log", json!({"id":"abc","line":"..."}))]).await;
3052 match &events[..] {
3053 [InvokeCommandEvent::Task(TaskEvent::Log { id, line })] => {
3054 assert_eq!(id, "abc");
3055 assert_eq!(line, "...");
3056 }
3057 other => panic!("unexpected: {other:?}"),
3058 }
3059 }
3060
3061 #[tokio::test]
3062 async fn unrelated_methods_are_dropped() {
3063 let (events, saw_done) = dispatch(vec![frame(
3064 "provider.stream.event",
3065 json!({"type":"text","delta":"x"}),
3066 )])
3067 .await;
3068 assert!(!saw_done);
3069 assert!(events.is_empty());
3070 }
3071
3072 #[tokio::test]
3075 async fn closed_sink_latches_open_false_without_blocking() {
3076 let (tx, collector) = invoke_event_channel(InvokeOutputBudget::default());
3077 drop(collector);
3078 let mut open = true;
3079 let saw_done = ProcessExtension::forward_invoke_command_frame(
3080 "ext-test",
3081 "r1",
3082 &tx,
3083 &mut open,
3084 frame(
3085 "command.output",
3086 json!({"request_id":"r1","event":{"kind":"done"}}),
3087 ),
3088 )
3089 .await;
3090 assert!(saw_done);
3092 assert!(!open);
3093 }
3094}
3095
3096#[cfg(all(test, unix))]
3097mod lifecycle_tests {
3098 use super::*;
3099 use std::time::Duration;
3100
3101 async fn fixture() -> ProcessExtension {
3104 let result = serde_json::json!({
3105 "protocol_version": CURRENT_EXTENSION_PROTOCOL_VERSION,
3106 "capabilities": {}
3107 })
3108 .to_string();
3109 let script = format!(
3110 r#"while IFS= read -r header; do
3111 length=${{header#Content-Length: }}
3112 length=${{length%?}}
3113 IFS= read -r blank || exit 1
3114 body=$(dd bs=1 count="$length" 2>/dev/null) || exit 1
3115 id=${{body##*\"id\":}}
3116 id=${{id%%[!0-9]*}}
3117 case "$body" in
3118 *'"method":"pending"'*) IFS= read -r hold; continue ;;
3119 esac
3120 reply='{{"jsonrpc":"2.0","id":'"$id"',"result":{result}}}'
3121 printf 'Content-Length: %s\r\n\r\n%s' "${{#reply}}" "$reply"
3122done"#
3123 );
3124 let ext = ProcessExtension::spawn("lifecycle-fixture", "/bin/sh", &["-c".into(), script])
3125 .await
3126 .expect("spawn offline lifecycle fixture");
3127 tokio::time::timeout(Duration::from_secs(3), ext.initialize_for_test(None))
3128 .await
3129 .expect("bounded initialize")
3130 .expect("initialize fixture");
3131 ext
3132 }
3133
3134 async fn unknown_method_fixture() -> ProcessExtension {
3138 let init = serde_json::json!({
3139 "protocol_version": CURRENT_EXTENSION_PROTOCOL_VERSION,
3140 "capabilities": {}
3141 })
3142 .to_string();
3143 let script = format!(
3144 r#"while IFS= read -r header; do
3145 length=${{header#Content-Length: }}
3146 length=${{length%?}}
3147 IFS= read -r blank || exit 1
3148 body=$(dd bs=1 count="$length" 2>/dev/null) || exit 1
3149 id=${{body##*\"id\":}}
3150 id=${{id%%[!0-9]*}}
3151 case "$body" in
3152 *'"method":"initialize"'*) reply='{{"jsonrpc":"2.0","id":'"$id"',"result":{init}}}' ;;
3153 *) reply='{{"jsonrpc":"2.0","id":'"$id"',"error":{{"code":-32601,"message":"unknown method"}}}}' ;;
3154 esac
3155 printf 'Content-Length: %s
3156
3157%s' "${{#reply}}" "$reply"
3158done"#
3159 );
3160 let ext =
3161 ProcessExtension::spawn("unknown-method-fixture", "/bin/sh", &["-c".into(), script])
3162 .await
3163 .expect("spawn offline unknown-method fixture");
3164 tokio::time::timeout(Duration::from_secs(3), ext.initialize_for_test(None))
3165 .await
3166 .expect("bounded initialize")
3167 .expect("initialize fixture");
3168 ext
3169 }
3170
3171 #[tokio::test]
3177 async fn application_error_reply_never_restarts_the_process() {
3178 let ext = unknown_method_fixture().await;
3179 let before = ext.lifecycle_snapshot().unwrap();
3180 for method in ["info.get", "sidecar.spawn_args", "no.such.method"] {
3181 let error = tokio::time::timeout(Duration::from_secs(3), ext.call(method, Value::Null))
3182 .await
3183 .expect("bounded call")
3184 .expect_err("fixture answers every non-initialize method with an error");
3185 assert!(is_application_error(&error), "{method}: {error}");
3186 assert!(error.contains("unknown method"), "{method}: {error}");
3187 assert!(
3188 !error.contains("retry after restart failed"),
3189 "{method}: must not restart on an application error: {error}"
3190 );
3191 }
3192 assert_eq!(ext.total_restarts.load(Ordering::Relaxed), 0);
3193 assert_eq!(ext.restart_count.load(Ordering::Relaxed), 0);
3194 let after = ext.lifecycle_snapshot().unwrap();
3195 assert_eq!(after.generation, before.generation);
3196 assert_eq!(after.health, before.health);
3197 assert!(!is_application_error(
3199 "transport closed: response channel dropped"
3200 ));
3201 assert!(!is_application_error("Write error: broken pipe"));
3202 ext.force_shutdown().await;
3203 }
3204
3205 async fn wait_failed(ext: &ProcessExtension) -> ExtensionLifecycle {
3206 tokio::time::timeout(Duration::from_secs(3), async {
3207 loop {
3208 if let Some(snapshot) = ext.lifecycle_snapshot() {
3209 if snapshot.health == ExtensionHealth::Failed {
3210 return snapshot;
3211 }
3212 }
3213 tokio::time::sleep(Duration::from_millis(5)).await;
3214 }
3215 })
3216 .await
3217 .expect("lifecycle must detect failure without another RPC")
3218 }
3219
3220 #[tokio::test]
3221 async fn actual_child_exit_is_detected_while_rpc_holds_state_and_reader_is_stopped() {
3222 let ext = fixture().await;
3223 let before = ext.lifecycle_snapshot().expect("live snapshot");
3224 assert_eq!(before.health, ExtensionHealth::Running);
3225 let call_guard = ext.call_lock.lock().await;
3226 let mut state_guard = ext.state.lock().await;
3227 let state = state_guard.as_mut().unwrap();
3228 let lifecycle = state.lifecycle.clone();
3229 state.reader_handle.abort();
3230 tokio::time::timeout(Duration::from_secs(3), async {
3232 while !state.reader_handle.is_finished() {
3233 tokio::task::yield_now().await;
3234 }
3235 })
3236 .await
3237 .expect("reader cancellation");
3238 {
3239 let rpc = ext.call_once_locked(state, "pending", Value::Null, 900);
3240 tokio::pin!(rpc);
3241 tokio::select! {
3242 biased;
3243 result = &mut rpc => panic!("RPC must remain pending: {result:?}"),
3244 _ = tokio::time::sleep(Duration::from_millis(20)) => {}
3245 }
3246 assert!(ext.inbox.pending.lock().await.contains_key(&900));
3247 assert_eq!(ext.lifecycle_snapshot(), Some(before));
3248 lifecycle.child.lock().await.start_kill().unwrap();
3249 let dead = wait_failed(&ext).await;
3250 assert_ne!(dead.generation, before.generation);
3251 assert_eq!(ext.lifecycle_snapshot(), Some(dead));
3252 assert!(!ext.inbox.closed.load(Ordering::Acquire));
3253 assert!(ext.inbox.pending.lock().await.contains_key(&900));
3254 }
3255 drop(state_guard);
3256 drop(call_guard);
3257 ext.force_shutdown().await;
3258 }
3259
3260 #[tokio::test]
3261 async fn shutdown_invalidates_generation_before_waiting_for_rpc_lock() {
3262 let ext = Arc::new(fixture().await);
3263 let before = ext.lifecycle_snapshot().unwrap();
3264 let call_guard = ext.call_lock.lock().await;
3265 let shutdown_ext = ext.clone();
3266 let shutdown = tokio::spawn(async move { shutdown_ext.shutdown().await });
3267 let dead = wait_failed(&ext).await;
3268 assert_ne!(dead.generation, before.generation);
3269 assert!(!shutdown.is_finished());
3270 drop(call_guard);
3271 tokio::time::timeout(Duration::from_secs(3), shutdown)
3272 .await
3273 .expect("bounded shutdown")
3274 .expect("shutdown task");
3275 assert_eq!(ext.lifecycle_snapshot(), Some(dead));
3276 }
3277
3278 #[tokio::test]
3282 async fn a_call_after_shutdown_fails_without_respawning() {
3283 let ext = fixture().await;
3284 ext.shutdown().await;
3285 let err = ext
3286 .call_inner("tool.call", serde_json::json!({}))
3287 .await
3288 .expect_err("a stopped handler must refuse the call");
3289 assert!(err.contains("was stopped"), "{err}");
3290 assert!(ext.state.lock().await.is_none(), "no process was respawned");
3291 assert_eq!(ext.total_restarts.load(Ordering::Relaxed), 0);
3292 }
3293
3294 #[tokio::test]
3295 async fn in_place_restart_changes_generation_and_old_invalidation_is_isolated() {
3296 let ext = fixture().await;
3297 let before = ext.lifecycle_snapshot().unwrap();
3298 let mut state = ext.state.lock().await;
3299 let old = state.as_ref().unwrap().lifecycle.clone();
3300 {
3301 let child_guard = old.child.lock().await;
3304 let restart = ext.restart_locked(&mut state);
3305 tokio::pin!(restart);
3306 tokio::select! {
3307 biased;
3308 result = &mut restart => panic!("restart must await child lock: {result:?}"),
3309 _ = tokio::time::sleep(Duration::from_millis(20)) => {}
3310 }
3311 let replacing = ext.lifecycle_snapshot().unwrap();
3312 assert_ne!(replacing.generation, before.generation);
3313 assert_eq!(replacing.health, ExtensionHealth::Failed);
3314 drop(child_guard);
3315 tokio::time::timeout(Duration::from_secs(5), restart)
3316 .await
3317 .expect("bounded restart")
3318 .expect("restart fixture");
3319 }
3320 let after = ext.lifecycle_snapshot().unwrap();
3321 assert_ne!(after.generation, before.generation);
3322 assert_eq!(after.health, ExtensionHealth::Degraded);
3323 assert_eq!(old.snapshot(false).unwrap().health, ExtensionHealth::Failed);
3324 old.invalidate();
3325 assert_eq!(ext.lifecycle_snapshot(), Some(after));
3326 drop(state);
3327 ext.force_shutdown().await;
3328 let stopped = ext.lifecycle_snapshot().unwrap();
3329 assert_ne!(stopped.generation, after.generation);
3330 assert_eq!(stopped.health, ExtensionHealth::Failed);
3331 }
3332
3333 #[tokio::test]
3334 async fn transport_eof_invalidates_generation_even_while_child_is_alive() {
3335 let ext = ProcessExtension::spawn(
3337 "lifecycle-eof",
3338 "/bin/sh",
3339 &[
3340 "-c".into(),
3341 "IFS= read -r trigger; exec 1>&-; IFS= read -r hold".into(),
3342 ],
3343 )
3344 .await
3345 .unwrap();
3346 let mut state_guard = ext.state.lock().await;
3347 let state = state_guard.as_mut().unwrap();
3348 state.lifecycle.initialized.store(true, Ordering::Release);
3349 let before = ext.lifecycle_snapshot().unwrap();
3350 state
3351 .stdin
3352 .lock()
3353 .await
3354 .write_all(b"close\n")
3355 .await
3356 .unwrap();
3357 let dead = wait_failed(&ext).await;
3358 assert_ne!(dead.generation, before.generation);
3359 assert!(state
3360 .lifecycle
3361 .child
3362 .lock()
3363 .await
3364 .try_wait()
3365 .unwrap()
3366 .is_none());
3367 drop(state_guard);
3368 ext.force_shutdown().await;
3369 }
3370
3371 #[tokio::test]
3372 async fn contended_child_probe_is_unavailable_without_changing_generation() {
3373 let ext = fixture().await;
3374 let before = ext.lifecycle_snapshot().unwrap();
3375 let lifecycle = ext.inbox.lifecycle.lock().await.clone().unwrap();
3376 let child_guard = lifecycle.child.lock().await;
3377 assert_eq!(ext.lifecycle_snapshot(), None);
3378 drop(child_guard);
3379 assert_eq!(ext.lifecycle_snapshot(), Some(before));
3380 let slot_guard = ext.inbox.lifecycle.lock().await;
3381 assert_eq!(ext.lifecycle_snapshot(), None);
3382 drop(slot_guard);
3383 assert_eq!(ext.lifecycle_snapshot(), Some(before));
3384 ext.force_shutdown().await;
3385 }
3386
3387 async fn replace_fixture() -> ProcessExtension {
3390 let init = serde_json::json!({
3391 "protocol_version": CURRENT_EXTENSION_PROTOCOL_VERSION,
3392 "capabilities": {}
3393 })
3394 .to_string();
3395 let script = format!(
3396 r#"while IFS= read -r header; do
3397 length=${{header#Content-Length: }}
3398 length=${{length%?}}
3399 IFS= read -r blank || exit 1
3400 body=$(dd bs=1 count="$length" 2>/dev/null) || exit 1
3401 id=${{body##*\"id\":}}
3402 id=${{id%%[!0-9]*}}
3403 case "$body" in
3404 *'"method":"initialize"'*) reply='{{"jsonrpc":"2.0","id":'"$id"',"result":{init}}}' ;;
3405 *'"method":"hook.handle"'*) reply='{{"jsonrpc":"2.0","id":'"$id"',"result":{{"action":"replace","output":"CRUSHED"}}}}' ;;
3406 *) reply='{{"jsonrpc":"2.0","id":'"$id"',"result":{{}}}}' ;;
3407 esac
3408 printf 'Content-Length: %s\r\n\r\n%s' "${{#reply}}" "$reply"
3409done"#
3410 );
3411 let ext = ProcessExtension::spawn("replace-fixture", "/bin/sh", &["-c".into(), script])
3412 .await
3413 .expect("spawn offline replace fixture");
3414 tokio::time::timeout(Duration::from_secs(3), ext.initialize_for_test(None))
3415 .await
3416 .expect("bounded initialize")
3417 .expect("initialize fixture");
3418 ext
3419 }
3420
3421 #[tokio::test]
3426 async fn a_worker_hook_never_respawns_a_crashed_or_stopped_extension() {
3427 let ext = replace_fixture().await;
3428 let event = HookEvent::after_tool_call(
3429 "bash",
3430 serde_json::json!({"command": "ls"}),
3431 "RAW".to_string(),
3432 );
3433 let crushed = HookResult::Replace { output: "CRUSHED".into() };
3434 async fn bounded<F: std::future::Future>(fut: F) -> F::Output {
3435 tokio::time::timeout(Duration::from_secs(3), fut)
3436 .await
3437 .expect("bounded: never hangs")
3438 }
3439 assert_eq!(bounded(ext.handle_no_spawn(&event)).await, crushed);
3440
3441 let lifecycle = ext.state.lock().await.as_ref().unwrap().lifecycle.clone();
3443 lifecycle.child.lock().await.start_kill().unwrap();
3444 tokio::time::timeout(Duration::from_secs(3), async {
3445 while !ext.inbox.closed.load(Ordering::Acquire) {
3446 tokio::time::sleep(Duration::from_millis(5)).await;
3447 }
3448 })
3449 .await
3450 .expect("reader observes the crash");
3451 for _ in 0..3 {
3452 assert_eq!(
3453 bounded(ext.handle_no_spawn(&event)).await,
3454 HookResult::Continue,
3455 "crashed: pass-through"
3456 );
3457 }
3458 assert_eq!(ext.total_restarts.load(Ordering::Relaxed), 0, "a worker never restarts it");
3459 assert!(ext.inbox.closed.load(Ordering::Acquire), "nothing was respawned");
3460
3461 assert_eq!(bounded(ext.handle(&event)).await, crushed);
3464 assert_eq!(ext.total_restarts.load(Ordering::Relaxed), 1);
3465 assert_eq!(bounded(ext.handle_no_spawn(&event)).await, crushed);
3466
3467 ext.shutdown().await;
3469 assert_eq!(bounded(ext.handle_no_spawn(&event)).await, HookResult::Continue);
3470 assert!(ext.state.lock().await.is_none(), "no process was respawned");
3471 assert_eq!(ext.total_restarts.load(Ordering::Relaxed), 1);
3472 }
3473}