@team-agent/installer 0.5.75 → 0.5.78

package/crates/team-agent/src/lifecycle/launch/pi_mcp.rs

@@ -16,6 +16,8 @@
 use std::collections::BTreeSet;
 use std::fs::OpenOptions;
 use std::io::{Read, Write};
+#[cfg(test)]
+use std::cell::RefCell;
 use std::os::unix::fs::{OpenOptionsExt, PermissionsExt};
 use std::path::{Path, PathBuf};
 use std::process::{Command, Stdio};
@@ -588,6 +590,116 @@ fn executable_file_type(path: &Path) -> Result<PiExecutableFileType, ProviderErr
     }
 }
 
+#[cfg(test)]
+#[derive(Debug, Default)]
+pub(crate) struct PiCatalogTestObservation {
+    pub(crate) spawn_count: usize,
+    pub(crate) argv: Vec<String>,
+    pub(crate) parent_pid: Option<u32>,
+    pub(crate) parent_exit_success: Option<bool>,
+    pub(crate) parent_exit_code: Option<i32>,
+    pub(crate) reader_timeout: bool,
+}
+
+#[cfg(test)]
+struct PiCatalogTestContext {
+    receipt: PathBuf,
+    helper_mode: Option<String>,
+    observation: PiCatalogTestObservation,
+}
+
+#[cfg(test)]
+thread_local! {
+    static PI_CATALOG_TEST_CONTEXT: RefCell<Option<PiCatalogTestContext>> = RefCell::new(None);
+}
+
+#[cfg(test)]
+pub(crate) fn with_pi_catalog_test_observation<T>(
+    receipt: &Path,
+    helper_mode: Option<&str>,
+    action: impl FnOnce() -> T,
+) -> (T, PiCatalogTestObservation) {
+    PI_CATALOG_TEST_CONTEXT.with(|slot| {
+        assert!(slot
+            .replace(Some(PiCatalogTestContext {
+                receipt: receipt.to_path_buf(),
+                helper_mode: helper_mode.map(str::to_owned),
+                observation: PiCatalogTestObservation::default(),
+            }))
+            .is_none());
+        let result = action();
+        let context = slot.take().expect("Pi catalog test context");
+        (result, context.observation)
+    })
+}
+
+#[cfg(test)]
+fn observe_pi_catalog_test(f: impl FnOnce(&mut PiCatalogTestContext)) {
+    PI_CATALOG_TEST_CONTEXT.with(|slot| {
+        if let Some(context) = slot.borrow_mut().as_mut() {
+            f(context);
+        }
+    });
+}
+
+fn spawn_catalog_reader(
+    stdout: std::process::ChildStdout,
+    max_bytes: u64,
+) -> mpsc::Receiver<(std::io::Result<usize>, Vec<u8>)> {
+    let (sender, receiver) = mpsc::channel();
+    std::thread::spawn(move || {
+        let mut bytes = Vec::new();
+        let result = stdout.take(max_bytes + 1).read_to_end(&mut bytes);
+        let _ = sender.send((result, bytes));
+    });
+    receiver
+}
+
+fn finish_catalog_reader(
+    receiver: mpsc::Receiver<(std::io::Result<usize>, Vec<u8>)>,
+    status: std::process::ExitStatus,
+    timeout: Duration,
+    max_bytes: u64,
+) -> Result<Vec<u8>, String> {
+    let (result, bytes) = match receiver.recv_timeout(timeout) {
+        Ok(value) => value,
+        Err(_) => {
+            #[cfg(test)]
+            observe_pi_catalog_test(|context| context.observation.reader_timeout = true);
+            return Err("Pi model catalog command timed out".to_string());
+        }
+    };
+    if !status.success() {
+        return Err("Pi model catalog command failed".into());
+    }
+    result.map_err(|_| "Pi model catalog could not be read".to_string())?;
+    if bytes.len() as u64 > max_bytes {
+        return Err("Pi model catalog exceeds the bounded output limit".into());
+    }
+    Ok(bytes)
+}
+
+#[cfg(test)]
+pub(crate) fn run_pi_catalog_after_parent_exit_for_test(
+    stdout: std::process::ChildStdout,
+    parent_pid: u32,
+    status: std::process::ExitStatus,
+    timeout: Duration,
+    max_bytes: u64,
+) -> Result<Vec<u8>, String> {
+    observe_pi_catalog_test(|context| {
+        context.observation.parent_pid = Some(parent_pid);
+        context.observation.parent_exit_success = Some(status.success());
+        context.observation.parent_exit_code = status.code();
+    });
+    finish_catalog_reader(
+        spawn_catalog_reader(stdout, max_bytes),
+        status,
+        timeout,
+        max_bytes,
+    )
+}
+
 /// Discover candidates with one PATH-first catalog observation. This deliberately
 /// does not resolve adapters or scan wrapper chains; callers use it only for
 /// rejecting an unqualified role model before lifecycle mutation.
@@ -596,8 +708,18 @@ pub(crate) fn run_pi_catalog(
     timeout: Duration,
     max_bytes: u64,
 ) -> Result<Vec<u8>, String> {
-    let mut child = Command::new(executable)
-        .arg("--list-models")
+    let mut command = Command::new(executable);
+    command.arg("--list-models");
+    #[cfg(test)]
+    observe_pi_catalog_test(|context| {
+        context.observation.spawn_count += 1;
+        context.observation.argv.push("--list-models".to_string());
+        command.env("TEAM_AGENT_MODELS_TIMEOUT_RECEIPT", &context.receipt);
+        if let Some(mode) = context.helper_mode.as_deref() {
+            command.env("TEAM_AGENT_MODELS_TIMEOUT_MODE", mode);
+        }
+    });
+    let mut child = command
         .stdin(Stdio::null())
         .stdout(Stdio::piped())
         .stderr(Stdio::null())
@@ -614,27 +736,29 @@ pub(crate) fn run_pi_catalog(
             let _ = error;
             "Pi executable is unavailable on PATH".to_string()
         })?;
+    #[cfg(test)]
+    observe_pi_catalog_test(|context| context.observation.parent_pid = Some(child.id()));
     let stdout = child.stdout.take().ok_or_else(|| {
         let _ = child.kill();
         let _ = child.wait();
         "Pi model catalog output unavailable".to_string()
     })?;
-    let (sender, receiver) = mpsc::channel();
-    std::thread::spawn(move || {
-        let mut bytes = Vec::new();
-        let result = stdout.take(max_bytes + 1).read_to_end(&mut bytes);
-        let _ = sender.send((result, bytes));
-    });
+    let receiver = spawn_catalog_reader(stdout, max_bytes);
     let start = Instant::now();
     loop {
         match child.try_wait() {
             Ok(Some(status)) => {
-                let (result, bytes) = receiver.recv_timeout(timeout.saturating_sub(start.elapsed()))
-                    .map_err(|_| "Pi model catalog command timed out".to_string())?;
-                if !status.success() { return Err("Pi model catalog command failed".into()); }
-                result.map_err(|_| "Pi model catalog could not be read".to_string())?;
-                if bytes.len() as u64 > max_bytes { return Err("Pi model catalog exceeds the bounded output limit".into()); }
-                return Ok(bytes);
+                #[cfg(test)]
+                observe_pi_catalog_test(|context| {
+                    context.observation.parent_exit_success = Some(status.success());
+                    context.observation.parent_exit_code = status.code();
+                });
+                return finish_catalog_reader(
+                    receiver,
+                    status,
+                    timeout.saturating_sub(start.elapsed()),
+                    max_bytes,
+                );
             }
             Ok(None) if start.elapsed() < timeout => std::thread::sleep(Duration::from_millis(20)),
             Ok(None) => {