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