use std::collections::HashMap;
use std::process::{Command, Stdio};
use std::time::{Duration, Instant};
use serde::{Deserialize, Serialize};
use super::{
config::DiscoveredHost, host_key_verification_error, is_host_key_verification_failure,
strict_ssh_cli_tokens, wait_for_child_output_with_timeout,
};
pub const DEFAULT_PROBE_TIMEOUT: u64 = 10;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HostProbeResult {
pub host_name: String,
pub reachable: bool,
pub connection_time_ms: u64,
pub cass_status: CassStatus,
pub detected_agents: Vec<DetectedAgent>,
pub system_info: Option<SystemInfo>,
pub resources: Option<ResourceInfo>,
pub error: Option<String>,
}
impl HostProbeResult {
pub fn unreachable(host_name: &str, error: impl Into<String>) -> Self {
Self {
host_name: host_name.to_string(),
reachable: false,
connection_time_ms: 0,
cass_status: CassStatus::Unknown,
detected_agents: Vec::new(),
system_info: None,
resources: None,
error: Some(error.into()),
}
}
pub fn has_cass(&self) -> bool {
self.cass_status.is_installed()
}
pub fn has_agent_data(&self) -> bool {
!self.detected_agents.is_empty()
}
pub fn total_sessions(&self) -> u64 {
self.detected_agents
.iter()
.filter_map(|a| a.estimated_sessions)
.sum()
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "status", rename_all = "snake_case")]
pub enum CassStatus {
Indexed {
version: String,
session_count: u64,
last_indexed: Option<String>,
},
InstalledNotIndexed { version: String },
NotFound,
Unknown,
}
impl CassStatus {
pub fn is_installed(&self) -> bool {
matches!(
self,
CassStatus::Indexed { .. } | CassStatus::InstalledNotIndexed { .. }
)
}
pub fn version(&self) -> Option<&str> {
match self {
CassStatus::Indexed { version, .. } | CassStatus::InstalledNotIndexed { version } => {
Some(version)
}
_ => None,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DetectedAgent {
pub agent_type: String,
pub path: String,
pub estimated_sessions: Option<u64>,
pub estimated_size_mb: Option<u64>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SystemInfo {
pub os: String,
pub arch: String,
pub distro: Option<String>,
pub has_cargo: bool,
pub has_cargo_binstall: bool,
pub has_curl: bool,
pub has_wget: bool,
pub remote_home: String,
#[serde(default)]
pub machine_id: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ResourceInfo {
pub disk_available_mb: u64,
pub memory_total_mb: u64,
pub memory_available_mb: u64,
pub can_compile: bool,
}
impl ResourceInfo {
pub const MIN_DISK_MB: u64 = 1024;
pub const MIN_MEMORY_MB: u64 = 2048; }
fn shell_single_quote_arg(value: &str) -> String {
format!("'{}'", value.replace('\'', r#"'\''"#))
}
fn collect_probe_dirs(probe_paths: Vec<(&'static str, Vec<String>)>) -> Vec<String> {
let mut dir_list = Vec::new();
for (_slug, paths) in probe_paths {
for path in paths {
dir_list.push(path);
}
}
dir_list.sort();
dir_list.dedup();
dir_list
}
fn probe_dir_array_entries(dir_list: &[String]) -> String {
dir_list
.iter()
.map(|path| format!(" {}", shell_single_quote_arg(path)))
.collect::<Vec<_>>()
.join("\n")
}
fn build_probe_script() -> String {
let dir_list = collect_probe_dirs(franken_agent_detection::default_probe_paths_tilde());
build_probe_script_for_dirs(&dir_list)
}
fn build_probe_script_for_dirs(dir_list: &[String]) -> String {
let dirs_str = probe_dir_array_entries(dir_list);
format!(
r#"#!/bin/bash
echo "===PROBE_START==="
# System info
echo "OS=$(uname -s | tr '[:upper:]' '[:lower:]')"
echo "ARCH=$(uname -m)"
echo "HOME=$HOME"
# Distro detection (Linux only)
if [ -f /etc/os-release ]; then
. /etc/os-release
echo "DISTRO=$PRETTY_NAME"
fi
# Machine ID for deduplication of SSH aliases pointing to same host
# Linux: /etc/machine-id, macOS: IOPlatformUUID
if [ -f /etc/machine-id ]; then
MACHINE_ID=$(cat /etc/machine-id 2>/dev/null | tr -d '\n')
echo "MACHINE_ID=$MACHINE_ID"
elif command -v ioreg &> /dev/null; then
MACHINE_ID=$(ioreg -rd1 -c IOPlatformExpertDevice 2>/dev/null | awk -F'"' '/IOPlatformUUID/{{print $4}}')
echo "MACHINE_ID=$MACHINE_ID"
fi
# Cass status - check PATH and common install locations
# Non-interactive SSH doesn't source .bashrc, so user bin dirs may not be in PATH
CASS_BIN=""
if command -v cass &> /dev/null; then
CASS_BIN="cass"
elif [ -x "$HOME/.cargo/bin/cass" ]; then
CASS_BIN="$HOME/.cargo/bin/cass"
elif [ -x "$HOME/.local/bin/cass" ]; then
CASS_BIN="$HOME/.local/bin/cass"
elif [ -x "/usr/local/bin/cass" ]; then
CASS_BIN="/usr/local/bin/cass"
fi
if [ -n "$CASS_BIN" ]; then
CASS_VER=$("$CASS_BIN" --version 2>/dev/null | head -1 | awk '{{print $2}}')
if [ -z "$CASS_VER" ]; then
# Binary exists but version command failed - treat as not found
echo "CASS_VERSION=NOT_FOUND"
else
echo "CASS_VERSION=$CASS_VER"
# Get health status (JSON output) - only if version was detected
if "$CASS_BIN" health --json &>/dev/null; then
echo "CASS_HEALTH=OK"
# Try to get session count from stats
STATS=$("$CASS_BIN" stats --json 2>/dev/null)
if [ $? -eq 0 ] && [ -n "$STATS" ]; then
# Extract total conversations from JSON (allow whitespace/newlines)
SESSIONS=$(echo "$STATS" | tr -d '\n' | sed -n 's/.*"conversations"[[:space:]]*:[[:space:]]*\([0-9][0-9]*\).*/\1/p')
echo "CASS_SESSIONS=${{SESSIONS:-0}}"
else
echo "CASS_SESSIONS=0"
fi
else
echo "CASS_HEALTH=NOT_INDEXED"
fi
fi
else
echo "CASS_VERSION=NOT_FOUND"
fi
# Tool availability - also check ~/.cargo/bin for non-interactive SSH sessions
if command -v cargo &> /dev/null || [ -x "$HOME/.cargo/bin/cargo" ]; then
echo "HAS_CARGO=1"
else
echo "HAS_CARGO=0"
fi
if command -v cargo-binstall &> /dev/null || [ -x "$HOME/.cargo/bin/cargo-binstall" ]; then
echo "HAS_BINSTALL=1"
else
echo "HAS_BINSTALL=0"
fi
command -v curl &> /dev/null && echo "HAS_CURL=1" || echo "HAS_CURL=0"
command -v wget &> /dev/null && echo "HAS_WGET=1" || echo "HAS_WGET=0"
# Resource info - disk (in KB, converted later)
DISK_KB=$(df -k ~ 2>/dev/null | awk 'NR==2 {{print $4}}')
echo "DISK_AVAIL_KB=${{DISK_KB:-0}}"
# Memory info (Linux)
if [ -f /proc/meminfo ]; then
MEM_TOTAL=$(grep MemTotal /proc/meminfo 2>/dev/null | awk '{{print $2}}')
MEM_AVAIL=$(grep MemAvailable /proc/meminfo 2>/dev/null | awk '{{print $2}}')
echo "MEM_TOTAL_KB=${{MEM_TOTAL:-0}}"
echo "MEM_AVAIL_KB=${{MEM_AVAIL:-0}}"
else
# macOS - use sysctl
if command -v sysctl &> /dev/null; then
MEM_BYTES=$(sysctl -n hw.memsize 2>/dev/null)
MEM_KB=$((MEM_BYTES / 1024))
echo "MEM_TOTAL_KB=${{MEM_KB:-0}}"
echo "MEM_AVAIL_KB=${{MEM_KB:-0}}" # macOS doesn't have easy available mem
fi
fi
# Agent data detection (with sizes and file counts)
PROBE_DIRS=(
{dirs}
)
for dir in "${{PROBE_DIRS[@]}}"; do
# Expand only the leading tilde marker from our static probe list. Do not
# eval paths: connector-owned paths can contain shell metacharacters.
case "$dir" in
"~") expanded_dir="$HOME" ;;
"~/"*) expanded_dir="$HOME/${{dir#\~/}}" ;;
*) expanded_dir="$dir" ;;
esac
if [ -e "$expanded_dir" ]; then
SIZE=$(du -sm "$expanded_dir" 2>/dev/null | cut -f1)
# Count JSONL files for session estimate
if [ -d "$expanded_dir" ]; then
# Keep probe bounded for very large trees: depth-limit and timeout when available.
if command -v timeout &> /dev/null; then
COUNT=$(timeout 5s find "$expanded_dir" -maxdepth 8 \( -name "*.jsonl" -o -name "*.json" \) 2>/dev/null | wc -l | tr -d ' ')
elif command -v gtimeout &> /dev/null; then
COUNT=$(gtimeout 5s find "$expanded_dir" -maxdepth 8 \( -name "*.jsonl" -o -name "*.json" \) 2>/dev/null | wc -l | tr -d ' ')
else
COUNT=$(find "$expanded_dir" -maxdepth 8 \( -name "*.jsonl" -o -name "*.json" \) 2>/dev/null | wc -l | tr -d ' ')
fi
else
COUNT=1 # Single file
fi
echo "AGENT_DATA=$dir|${{SIZE:-0}}|${{COUNT:-0}}"
fi
done
echo "===PROBE_END==="
"#,
dirs = dirs_str
)
}
pub fn probe_host(host: &DiscoveredHost, timeout_secs: u64) -> HostProbeResult {
let start = Instant::now();
let timeout_secs = timeout_secs.max(1);
let command_timeout = Duration::from_secs(timeout_secs);
let mut cmd = Command::new("ssh");
cmd.args(strict_ssh_cli_tokens(timeout_secs))
.arg("--")
.arg(&host.name)
.arg("bash -s")
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped());
let mut child = match cmd.spawn() {
Ok(c) => c,
Err(e) => {
return HostProbeResult::unreachable(
&host.name,
format!("Failed to execute ssh: {}", e),
);
}
};
let probe_script = build_probe_script();
let write_error = if let Some(mut stdin) = child.stdin.take() {
use std::io::Write;
stdin.write_all(probe_script.as_bytes()).err()
} else {
None
};
let output = match wait_for_child_output_with_timeout(child, command_timeout) {
Ok(Some(o)) => o,
Ok(None) => {
return HostProbeResult::unreachable(
&host.name,
format!("Connection timed out after {timeout_secs} seconds"),
);
}
Err(e) => {
return HostProbeResult::unreachable(&host.name, format!("SSH command failed: {}", e));
}
};
let connection_time_ms = start.elapsed().as_millis() as u64;
if !output.status.success() {
let stderr = String::from_utf8_lossy(&output.stderr);
let error_msg = if stderr.contains("Connection refused") {
"Connection refused".to_string()
} else if stderr.contains("Connection timed out") || stderr.contains("timed out") {
"Connection timed out".to_string()
} else if stderr.contains("Permission denied") {
"Permission denied (key not loaded in ssh-agent?)".to_string()
} else if is_host_key_verification_failure(&stderr) {
host_key_verification_error(&host.name)
} else if stderr.contains("No route to host") {
"No route to host".to_string()
} else {
format!("SSH failed: {}", stderr.trim())
};
return HostProbeResult::unreachable(&host.name, error_msg);
}
if let Some(e) = write_error {
return HostProbeResult::unreachable(
&host.name,
format!("Failed to write probe script: {}", e),
);
}
let stdout = String::from_utf8_lossy(&output.stdout);
parse_probe_output(&host.name, &stdout, connection_time_ms)
}
fn parse_probe_output(host_name: &str, output: &str, connection_time_ms: u64) -> HostProbeResult {
let mut values: HashMap<String, String> = HashMap::new();
let mut agent_data: Vec<(String, u64, u64)> = Vec::new();
let mut inside_probe = false;
let mut saw_start = false;
let mut saw_end = false;
for line in output.lines() {
let line = line.trim();
if line == "===PROBE_START===" {
if saw_start {
return HostProbeResult::unreachable(host_name, "Probe script output malformed");
}
saw_start = true;
inside_probe = true;
continue;
}
if line == "===PROBE_END===" {
if !inside_probe {
return HostProbeResult::unreachable(host_name, "Probe script output malformed");
}
saw_end = true;
break;
}
if !inside_probe || line.is_empty() || line.starts_with("===") {
continue;
}
if line.starts_with("AGENT_DATA=") {
if let Some(data) = line.strip_prefix("AGENT_DATA=") {
let parts: Vec<&str> = data.rsplitn(3, '|').collect();
if parts.len() == 3 {
let count = parts[0].parse().unwrap_or(0);
let size = parts[1].parse().unwrap_or(0);
let path = parts[2].to_string();
agent_data.push((path, size, count));
}
}
} else if let Some((key, value)) = line.split_once('=') {
values.insert(key.to_string(), value.to_string());
}
}
if !saw_start || !saw_end {
return HostProbeResult::unreachable(host_name, "Probe script output malformed");
}
let cass_status = if let Some(version) = values.get("CASS_VERSION") {
if version == "NOT_FOUND" {
CassStatus::NotFound
} else {
let health = values.get("CASS_HEALTH").map(|s| s.as_str());
if health == Some("OK") {
let sessions = values
.get("CASS_SESSIONS")
.and_then(|s| s.parse().ok())
.unwrap_or(0);
CassStatus::Indexed {
version: version.clone(),
session_count: sessions,
last_indexed: None,
}
} else {
CassStatus::InstalledNotIndexed {
version: version.clone(),
}
}
}
} else {
CassStatus::Unknown
};
let system_info = values.get("OS").map(|os| SystemInfo {
os: os.clone(),
arch: values.get("ARCH").cloned().unwrap_or_default(),
distro: values.get("DISTRO").cloned(),
has_cargo: values.get("HAS_CARGO").map(|v| v == "1").unwrap_or(false),
has_cargo_binstall: values
.get("HAS_BINSTALL")
.map(|v| v == "1")
.unwrap_or(false),
has_curl: values.get("HAS_CURL").map(|v| v == "1").unwrap_or(false),
has_wget: values.get("HAS_WGET").map(|v| v == "1").unwrap_or(false),
remote_home: values.get("HOME").cloned().unwrap_or_default(),
machine_id: values.get("MACHINE_ID").cloned().filter(|s| !s.is_empty()),
});
let resources = {
let disk_kb = values
.get("DISK_AVAIL_KB")
.and_then(|s| s.parse::<u64>().ok())
.unwrap_or(0);
let mem_total_kb = values
.get("MEM_TOTAL_KB")
.and_then(|s| s.parse::<u64>().ok())
.unwrap_or(0);
let mem_avail_kb = values
.get("MEM_AVAIL_KB")
.and_then(|s| s.parse::<u64>().ok())
.unwrap_or(0);
if disk_kb > 0 || mem_total_kb > 0 {
let disk_mb = disk_kb / 1024;
let mem_total_mb = mem_total_kb / 1024;
let mem_avail_mb = mem_avail_kb / 1024;
Some(ResourceInfo {
disk_available_mb: disk_mb,
memory_total_mb: mem_total_mb,
memory_available_mb: mem_avail_mb,
can_compile: disk_mb >= ResourceInfo::MIN_DISK_MB
&& mem_total_mb >= ResourceInfo::MIN_MEMORY_MB,
})
} else {
None
}
};
let detected_agents: Vec<DetectedAgent> = agent_data
.into_iter()
.map(|(path, size_mb, count)| {
let agent_type = infer_agent_type(&path);
DetectedAgent {
agent_type,
path,
estimated_sessions: Some(count),
estimated_size_mb: Some(size_mb),
}
})
.collect();
HostProbeResult {
host_name: host_name.to_string(),
reachable: true,
connection_time_ms,
cass_status,
detected_agents,
system_info,
resources,
error: None,
}
}
fn infer_agent_type(path: &str) -> String {
if path.contains("saoudrizwan.claude-dev") || path.contains("rooveterinaryinc.roo-cline") {
"cline".to_string()
} else if path.contains(".claude") {
"claude_code".to_string()
} else if path.contains(".codex") {
"codex".to_string()
} else if path.contains(".cursor") || path.contains("Cursor") {
"cursor".to_string()
} else if path.contains(".gemini") {
"gemini".to_string()
} else if path.contains("/.pi/") || path.ends_with("/.pi") {
"pi_agent".to_string()
} else if path.contains(".aider") {
"aider".to_string()
} else if path.contains("opencode") {
"opencode".to_string()
} else if path.contains(".goose") {
"goose".to_string()
} else if path.contains("copilot-chat")
|| path.contains("gh-copilot")
|| path.contains("gh/copilot")
{
"copilot".to_string()
} else if path.contains(".continue") {
"continue".to_string()
} else if path.contains("sourcegraph.amp") || path.contains("/amp/") || path.ends_with("/amp") {
"amp".to_string()
} else if path.contains(".clawdbot") {
"clawdbot".to_string()
} else if path.contains(".factory") {
"factory".to_string()
} else if path.contains(".vibe") {
"vibe".to_string()
} else if path.contains(".windsurf") {
"windsurf".to_string()
} else {
"unknown".to_string()
}
}
pub fn probe_hosts_parallel<F>(
hosts: &[DiscoveredHost],
timeout_secs: u64,
on_progress: F,
) -> Vec<HostProbeResult>
where
F: Fn(usize, usize, &str) + Send + Sync,
{
use rayon::prelude::*;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
let total = hosts.len();
let completed = Arc::new(AtomicUsize::new(0));
let on_progress = Arc::new(on_progress);
hosts
.par_iter()
.map(|host| {
let result = probe_host(host, timeout_secs);
let done = completed.fetch_add(1, Ordering::SeqCst) + 1;
on_progress(done, total, &host.name);
result
})
.collect()
}
#[derive(Debug)]
pub struct ProbeCache {
results: HashMap<String, (HostProbeResult, std::time::Instant)>,
ttl_secs: u64,
}
impl Default for ProbeCache {
fn default() -> Self {
Self::new(300) }
}
impl ProbeCache {
pub fn new(ttl_secs: u64) -> Self {
Self {
results: HashMap::new(),
ttl_secs,
}
}
pub fn get(&self, host_name: &str) -> Option<&HostProbeResult> {
self.results.get(host_name).and_then(|(result, ts)| {
if ts.elapsed().as_secs() < self.ttl_secs {
Some(result)
} else {
None
}
})
}
pub fn insert(&mut self, result: HostProbeResult) {
self.results.insert(
result.host_name.clone(),
(result, std::time::Instant::now()),
);
}
pub fn clear_expired(&mut self) {
self.results
.retain(|_, (_, ts)| ts.elapsed().as_secs() < self.ttl_secs);
}
}
pub fn deduplicate_probe_results(
results: Vec<HostProbeResult>,
) -> (Vec<HostProbeResult>, HashMap<String, Vec<String>>) {
let mut by_machine_id: HashMap<String, Vec<HostProbeResult>> = HashMap::new();
let mut no_machine_id: Vec<HostProbeResult> = Vec::new();
for result in results {
if let Some(ref machine_id) = result
.system_info
.as_ref()
.and_then(|s| s.machine_id.clone())
{
by_machine_id
.entry(machine_id.clone())
.or_default()
.push(result);
} else {
no_machine_id.push(result);
}
}
let mut deduplicated: Vec<HostProbeResult> = Vec::new();
let mut merged_aliases: HashMap<String, Vec<String>> = HashMap::new();
for (_machine_id, mut group) in by_machine_id {
if group.len() == 1 {
deduplicated.push(group.remove(0));
} else {
group.sort_by(|a, b| {
let a_installed = a.cass_status.is_installed();
let b_installed = b.cass_status.is_installed();
if a_installed != b_installed {
return b_installed.cmp(&a_installed);
}
let a_sessions = match &a.cass_status {
CassStatus::Indexed { session_count, .. } => *session_count,
_ => 0,
};
let b_sessions = match &b.cass_status {
CassStatus::Indexed { session_count, .. } => *session_count,
_ => 0,
};
if a_sessions != b_sessions {
return b_sessions.cmp(&a_sessions);
}
a.host_name.cmp(&b.host_name)
});
let kept = group.remove(0);
let merged: Vec<String> = group.into_iter().map(|h| h.host_name).collect();
if !merged.is_empty() {
merged_aliases.insert(kept.host_name.clone(), merged);
}
deduplicated.push(kept);
}
}
deduplicated.extend(no_machine_id);
deduplicated.sort_by(|a, b| a.host_name.cmp(&b.host_name));
(deduplicated, merged_aliases)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_cass_status_is_installed() {
assert!(
CassStatus::Indexed {
version: "0.1.50".into(),
session_count: 100,
last_indexed: None
}
.is_installed()
);
assert!(
CassStatus::InstalledNotIndexed {
version: "0.1.50".into()
}
.is_installed()
);
assert!(!CassStatus::NotFound.is_installed());
assert!(!CassStatus::Unknown.is_installed());
}
#[test]
fn test_cass_status_version() {
assert_eq!(
CassStatus::Indexed {
version: "0.1.50".into(),
session_count: 0,
last_indexed: None
}
.version(),
Some("0.1.50")
);
assert_eq!(
CassStatus::InstalledNotIndexed {
version: "0.1.49".into()
}
.version(),
Some("0.1.49")
);
assert_eq!(CassStatus::NotFound.version(), None);
}
#[test]
fn test_infer_agent_type() {
assert_eq!(infer_agent_type("~/.claude/projects"), "claude_code");
assert_eq!(infer_agent_type("~/.codex/sessions"), "codex");
assert_eq!(infer_agent_type("~/.cursor"), "cursor");
assert_eq!(infer_agent_type("~/.gemini/tmp"), "gemini");
assert_eq!(
infer_agent_type("~/.config/Code/User/globalStorage/saoudrizwan.claude-dev"),
"cline"
);
assert_eq!(
infer_agent_type("~/.config/Code/User/globalStorage/github.copilot-chat"),
"copilot"
);
assert_eq!(infer_agent_type("~/.config/gh-copilot"), "copilot");
assert_eq!(infer_agent_type("/some/random/path"), "unknown");
}
#[test]
fn test_parse_probe_output_success() {
let output = r#"
===PROBE_START===
OS=linux
ARCH=x86_64
HOME=/home/user
DISTRO=Ubuntu 22.04
CASS_VERSION=0.1.50
CASS_HEALTH=OK
CASS_SESSIONS=1234
HAS_CARGO=1
HAS_BINSTALL=0
HAS_CURL=1
HAS_WGET=1
DISK_AVAIL_KB=52428800
MEM_TOTAL_KB=16777216
MEM_AVAIL_KB=8388608
AGENT_DATA=~/.claude/projects|150|42
AGENT_DATA=~/.codex/sessions|50|10
===PROBE_END===
"#;
let result = parse_probe_output("test-host", output, 100);
assert!(result.reachable);
assert_eq!(result.host_name, "test-host");
assert_eq!(result.connection_time_ms, 100);
assert!(
matches!(&result.cass_status, CassStatus::Indexed { .. }),
"expected Indexed status"
);
if let CassStatus::Indexed {
version,
session_count,
..
} = &result.cass_status
{
assert_eq!(version, "0.1.50");
assert_eq!(*session_count, 1234);
}
let sys = result.system_info.as_ref().unwrap();
assert_eq!(sys.os, "linux");
assert_eq!(sys.arch, "x86_64");
assert_eq!(sys.distro, Some("Ubuntu 22.04".into()));
assert!(sys.has_cargo);
assert!(!sys.has_cargo_binstall);
assert!(sys.has_curl);
let res = result.resources.as_ref().unwrap();
assert_eq!(res.disk_available_mb, 51200); assert_eq!(res.memory_total_mb, 16384); assert!(res.can_compile);
assert_eq!(result.detected_agents.len(), 2);
assert_eq!(result.detected_agents[0].agent_type, "claude_code");
assert_eq!(result.detected_agents[0].estimated_sessions, Some(42));
assert_eq!(result.detected_agents[1].agent_type, "codex");
}
#[test]
fn test_parse_probe_output_ignores_noise_outside_markers() {
let output = r#"
CASS_VERSION=NOT_FOUND
AGENT_DATA=/tmp/outside-before|999|999
===PROBE_START===
OS=linux
ARCH=x86_64
HOME=/home/user
CASS_VERSION=0.4.2
CASS_HEALTH=OK
CASS_SESSIONS=7
HAS_CARGO=1
HAS_BINSTALL=0
HAS_CURL=1
HAS_WGET=1
DISK_AVAIL_KB=2048000
MEM_TOTAL_KB=4096000
MEM_AVAIL_KB=1024000
===PROBE_END===
CASS_VERSION=NOT_FOUND
AGENT_DATA=/tmp/outside-after|999|999
"#;
let result = parse_probe_output("noisy-host", output, 42);
assert!(result.reachable);
assert!(result.detected_agents.is_empty());
assert!(matches!(
result.cass_status,
CassStatus::Indexed {
ref version,
session_count: 7,
..
} if version == "0.4.2"
));
}
#[test]
fn test_parse_probe_output_cass_not_found() {
let output = r#"
===PROBE_START===
OS=darwin
ARCH=arm64
HOME=/Users/user
CASS_VERSION=NOT_FOUND
HAS_CARGO=0
HAS_BINSTALL=0
HAS_CURL=1
HAS_WGET=0
DISK_AVAIL_KB=10240000
MEM_TOTAL_KB=8388608
MEM_AVAIL_KB=4194304
===PROBE_END===
"#;
let result = parse_probe_output("mac-host", output, 50);
assert!(result.reachable);
assert!(matches!(result.cass_status, CassStatus::NotFound));
let sys = result.system_info.as_ref().unwrap();
assert_eq!(sys.os, "darwin");
assert_eq!(sys.arch, "arm64");
assert!(!sys.has_cargo);
}
#[test]
fn test_parse_probe_output_malformed() {
let output = "random garbage";
let result = parse_probe_output("bad-host", output, 0);
assert!(!result.reachable);
assert!(result.error.is_some());
}
#[test]
fn test_parse_probe_output_rejects_out_of_order_markers() {
let output = r#"
===PROBE_END===
===PROBE_START===
OS=linux
CASS_VERSION=0.4.2
"#;
let result = parse_probe_output("bad-host", output, 0);
assert!(!result.reachable);
assert!(result.error.is_some());
}
#[test]
fn test_host_probe_result_unreachable() {
let result = HostProbeResult::unreachable("test", "Connection refused");
assert!(!result.reachable);
assert_eq!(result.error, Some("Connection refused".into()));
assert!(!result.has_cass());
assert!(!result.has_agent_data());
}
#[test]
fn test_probe_cache() {
let mut cache = ProbeCache::new(300);
let result = HostProbeResult {
host_name: "test".into(),
reachable: true,
connection_time_ms: 100,
cass_status: CassStatus::NotFound,
detected_agents: vec![],
system_info: None,
resources: None,
error: None,
};
cache.insert(result);
assert!(cache.get("test").is_some());
assert!(cache.get("nonexistent").is_none());
}
#[test]
fn test_resource_info_can_compile() {
let good = ResourceInfo {
disk_available_mb: 2000,
memory_total_mb: 4000,
memory_available_mb: 2000,
can_compile: true,
};
assert!(good.can_compile);
let low_disk = ResourceInfo {
disk_available_mb: 500,
memory_total_mb: 4000,
memory_available_mb: 2000,
can_compile: false,
};
assert!(!low_disk.can_compile);
}
fn run_probe_script_with_home(script: &str, home: Option<&std::path::Path>) -> String {
use std::io::Write;
let mut cmd = Command::new("bash");
cmd.arg("-s")
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped());
if let Some(home) = home {
cmd.env("HOME", home);
} else if dotenvy::var("HOME").is_err()
&& let Some(dirs) = directories::BaseDirs::new()
{
cmd.env("HOME", dirs.home_dir());
}
let mut child = cmd.spawn().expect("bash should be available");
if let Some(mut stdin) = child.stdin.take() {
stdin
.write_all(script.as_bytes())
.expect("write probe script");
}
let output = child
.wait_with_output()
.expect("probe script should finish");
assert!(
output.status.success(),
"probe script failed: {:?}",
String::from_utf8_lossy(&output.stderr)
);
String::from_utf8_lossy(&output.stdout).to_string()
}
fn run_probe_script_locally() -> String {
run_probe_script_with_home(&build_probe_script(), None)
}
#[test]
fn shell_single_quote_arg_quotes_shell_metacharacters() {
assert_eq!(shell_single_quote_arg("plain/path"), "'plain/path'");
assert_eq!(shell_single_quote_arg("can't"), "'can'\\''t'");
assert_eq!(
shell_single_quote_arg("$(touch /tmp/nope); `whoami`"),
"'$(touch /tmp/nope); `whoami`'"
);
}
#[test]
fn probe_script_uses_literal_array_without_eval() {
let script = build_probe_script();
assert!(script.contains("PROBE_DIRS=("));
assert!(script.contains("for dir in \"${PROBE_DIRS[@]}\""));
assert!(script.contains("expanded_dir=\"$HOME/${dir#\\~/}\""));
assert!(
!script.contains("eval echo"),
"probe paths must not be expanded through eval"
);
}
#[test]
fn probe_script_treats_special_probe_paths_as_literals() {
let home = tempfile::tempdir().expect("temp home");
let relative_path =
"Library/Application Support/Codex$(touch \"$HOME/SHOULD_NOT_EXIST\");can't";
std::fs::create_dir_all(home.path().join(relative_path)).expect("create special path");
let probe_path = format!("~/{relative_path}");
let script = build_probe_script_for_dirs(std::slice::from_ref(&probe_path));
let output = run_probe_script_with_home(&script, Some(home.path()));
assert!(
output.contains(&format!("AGENT_DATA={probe_path}|")),
"special probe path should be reported literally: {output}"
);
assert!(
!home.path().join("SHOULD_NOT_EXIST").exists(),
"probe path interpolation must not execute command substitutions"
);
let result = parse_probe_output("localhost", &output, 0);
assert!(
result
.detected_agents
.iter()
.any(|agent| agent.path == probe_path),
"parsed agent data should preserve literal path"
);
}
#[test]
fn real_probe_script_produces_valid_markers() {
let output = run_probe_script_locally();
assert!(
output.contains("===PROBE_START==="),
"missing PROBE_START marker"
);
assert!(
output.contains("===PROBE_END==="),
"missing PROBE_END marker"
);
}
#[test]
fn real_probe_script_parses_into_reachable_result() {
let output = run_probe_script_locally();
let result = parse_probe_output("localhost", &output, 0);
assert!(
result.reachable,
"local probe should be reachable: {:?}",
result.error
);
assert!(result.system_info.is_some(), "should have system info");
assert!(result.resources.is_some(), "should have resource info");
}
#[test]
fn real_probe_system_info_has_valid_os() {
let output = run_probe_script_locally();
let result = parse_probe_output("localhost", &output, 0);
let sys = result.system_info.as_ref().expect("system_info");
assert!(
sys.os == "linux" || sys.os == "darwin",
"OS should be linux or darwin, got: {}",
sys.os
);
}
#[test]
fn real_probe_system_info_has_valid_arch() {
let output = run_probe_script_locally();
let result = parse_probe_output("localhost", &output, 0);
let sys = result.system_info.as_ref().expect("system_info");
let valid_archs = [
"x86_64", "aarch64", "arm64", "armv7l", "i686", "s390x", "ppc64le",
];
assert!(
valid_archs.contains(&sys.arch.as_str()),
"arch should be a known value, got: {}",
sys.arch
);
}
#[test]
fn real_probe_system_info_has_nonempty_home() {
let output = run_probe_script_locally();
let result = parse_probe_output("localhost", &output, 0);
let sys = result.system_info.as_ref().expect("system_info");
assert!(!sys.remote_home.is_empty(), "home should not be empty");
assert!(
sys.remote_home.starts_with('/'),
"home should be absolute: {}",
sys.remote_home
);
}
#[test]
fn real_probe_resources_have_nonzero_disk() {
let output = run_probe_script_locally();
let result = parse_probe_output("localhost", &output, 0);
let res = result.resources.as_ref().expect("resources");
assert!(res.disk_available_mb > 0, "disk_available_mb should be > 0");
}
#[test]
fn real_probe_resources_have_nonzero_memory() {
let output = run_probe_script_locally();
let result = parse_probe_output("localhost", &output, 0);
let res = result.resources.as_ref().expect("resources");
assert!(res.memory_total_mb > 0, "memory_total_mb should be > 0");
assert!(
res.memory_available_mb > 0,
"memory_available_mb should be > 0"
);
}
#[test]
fn real_probe_resources_memory_invariant() {
let output = run_probe_script_locally();
let result = parse_probe_output("localhost", &output, 0);
let res = result.resources.as_ref().expect("resources");
assert!(
res.memory_available_mb <= res.memory_total_mb,
"available memory ({}) should not exceed total ({})",
res.memory_available_mb,
res.memory_total_mb
);
}
#[test]
fn real_probe_resources_can_compile_reflects_thresholds() {
let output = run_probe_script_locally();
let result = parse_probe_output("localhost", &output, 0);
let res = result.resources.as_ref().expect("resources");
let expected = res.disk_available_mb >= ResourceInfo::MIN_DISK_MB
&& res.memory_total_mb >= ResourceInfo::MIN_MEMORY_MB;
assert_eq!(
res.can_compile, expected,
"can_compile should match threshold check: disk={}MB mem={}MB",
res.disk_available_mb, res.memory_total_mb
);
}
#[test]
fn real_probe_tool_detection_is_consistent() {
let output = run_probe_script_locally();
let result = parse_probe_output("localhost", &output, 0);
let sys = result.system_info.as_ref().expect("system_info");
if sys.has_cargo_binstall {
assert!(sys.has_cargo, "binstall requires cargo");
}
assert!(
sys.has_curl || sys.has_wget,
"system should have at least curl or wget"
);
}
#[test]
fn probe_script_contains_all_franken_agent_detection_paths() {
let script = build_probe_script();
assert!(script.contains("~/.claude"), "missing claude paths");
assert!(script.contains("~/.codex/sessions"), "missing codex path");
assert!(script.contains("~/.gemini"), "missing gemini paths");
assert!(script.contains("~/.goose/sessions"), "missing goose path");
assert!(
script.contains("~/.continue/sessions"),
"missing continue path"
);
assert!(script.contains("~/.aider"), "missing aider path");
assert!(
script.contains("saoudrizwan.claude-dev"),
"missing cline path"
);
assert!(script.contains("copilot-chat"), "missing copilot path");
assert!(script.contains("~/.windsurf"), "missing windsurf path");
assert!(script.contains("~/.factory"), "missing factory path");
assert!(script.contains("~/.clawdbot"), "missing clawdbot path");
assert!(script.contains("~/.vibe"), "missing vibe path");
assert!(script.contains("sourcegraph.amp"), "missing amp path");
assert!(script.contains("===PROBE_START==="));
assert!(script.contains("===PROBE_END==="));
assert!(script.contains("for dir in \"${PROBE_DIRS[@]}\""));
}
#[test]
fn infer_agent_type_covers_all_dynamic_agents() {
assert_eq!(infer_agent_type("~/.goose/sessions"), "goose");
assert_eq!(infer_agent_type("~/.continue/sessions"), "continue");
assert_eq!(infer_agent_type("~/.clawdbot/sessions"), "clawdbot");
assert_eq!(infer_agent_type("~/.factory/sessions"), "factory");
assert_eq!(infer_agent_type("~/.vibe/logs/session"), "vibe");
assert_eq!(infer_agent_type("~/.windsurf"), "windsurf");
assert_eq!(
infer_agent_type("~/.config/Code/User/globalStorage/sourcegraph.amp"),
"amp"
);
assert_eq!(infer_agent_type("~/.pi/agent/sessions"), "pi_agent");
}
fn make_probe_result(
name: &str,
machine_id: Option<&str>,
sessions: Option<u64>,
) -> HostProbeResult {
HostProbeResult {
host_name: name.to_string(),
reachable: true,
connection_time_ms: 100,
cass_status: if let Some(s) = sessions {
CassStatus::Indexed {
version: "0.1.50".into(),
session_count: s,
last_indexed: None,
}
} else {
CassStatus::NotFound
},
detected_agents: vec![],
system_info: Some(SystemInfo {
os: "linux".into(),
arch: "x86_64".into(),
distro: Some("Ubuntu 25.10".into()),
has_cargo: true,
has_cargo_binstall: false,
has_curl: true,
has_wget: true,
remote_home: "/home/ubuntu".into(),
machine_id: machine_id.map(String::from),
}),
resources: Some(ResourceInfo {
disk_available_mb: 800_000,
memory_total_mb: 16_000,
memory_available_mb: 8_000,
can_compile: true,
}),
error: None,
}
}
#[test]
fn test_deduplicate_no_duplicates() {
let results = vec![
make_probe_result("host1", Some("machine-1"), Some(100)),
make_probe_result("host2", Some("machine-2"), Some(200)),
];
let (deduped, merged) = deduplicate_probe_results(results);
assert_eq!(deduped.len(), 2);
assert!(merged.is_empty());
}
#[test]
fn test_deduplicate_same_machine() {
let results = vec![
make_probe_result("jain", Some("abc123"), None),
make_probe_result("jain_ovh_box", Some("abc123"), None),
];
let (deduped, merged) = deduplicate_probe_results(results);
assert_eq!(deduped.len(), 1);
assert_eq!(deduped[0].host_name, "jain");
assert_eq!(
merged.get("jain").unwrap(),
&vec!["jain_ovh_box".to_string()]
);
}
#[test]
fn test_deduplicate_prefers_installed_cass() {
let results = vec![
make_probe_result("alias_a", Some("machine-x"), None), make_probe_result("alias_b", Some("machine-x"), Some(500)), ];
let (deduped, merged) = deduplicate_probe_results(results);
assert_eq!(deduped.len(), 1);
assert_eq!(deduped[0].host_name, "alias_b");
assert!(merged.contains_key("alias_b"));
}
#[test]
fn test_deduplicate_prefers_more_sessions() {
let results = vec![
make_probe_result("host_low", Some("machine-y"), Some(50)),
make_probe_result("host_high", Some("machine-y"), Some(500)),
];
let (deduped, merged) = deduplicate_probe_results(results);
assert_eq!(deduped.len(), 1);
assert_eq!(deduped[0].host_name, "host_high");
assert!(merged.contains_key("host_high"));
}
#[test]
fn test_deduplicate_no_machine_id_not_merged() {
let results = vec![
make_probe_result("host1", None, Some(100)),
make_probe_result("host2", None, Some(200)),
];
let (deduped, merged) = deduplicate_probe_results(results);
assert_eq!(deduped.len(), 2);
assert!(merged.is_empty());
}
#[test]
fn test_deduplicate_mixed_with_and_without_machine_id() {
let results = vec![
make_probe_result("aliasA", Some("same-machine"), Some(100)),
make_probe_result("aliasB", Some("same-machine"), Some(50)),
make_probe_result("standalone", None, Some(75)),
];
let (deduped, merged) = deduplicate_probe_results(results);
assert_eq!(deduped.len(), 2);
assert!(deduped.iter().any(|h| h.host_name == "aliasA"));
assert!(deduped.iter().any(|h| h.host_name == "standalone"));
assert_eq!(merged.len(), 1);
}
#[test]
fn test_deduplicate_three_aliases_same_machine() {
let results = vec![
make_probe_result("alias1", Some("same"), Some(100)),
make_probe_result("alias2", Some("same"), Some(200)),
make_probe_result("alias3", Some("same"), Some(150)),
];
let (deduped, merged) = deduplicate_probe_results(results);
assert_eq!(deduped.len(), 1);
assert_eq!(deduped[0].host_name, "alias2");
let merged_list = merged.get("alias2").unwrap();
assert_eq!(merged_list.len(), 2);
assert!(merged_list.contains(&"alias1".to_string()));
assert!(merged_list.contains(&"alias3".to_string()));
}
#[test]
fn real_probe_machine_id_present() {
let output = run_probe_script_locally();
let result = parse_probe_output("localhost", &output, 0);
let sys = result.system_info.as_ref().expect("system_info");
if sys.os == "linux" || sys.os == "darwin" {
assert!(
sys.machine_id.is_some(),
"machine_id should be present on {}",
sys.os
);
let mid = sys.machine_id.as_ref().unwrap();
assert!(!mid.is_empty(), "machine_id should not be empty");
assert!(
mid.len() >= 32,
"machine_id should be at least 32 chars, got: {}",
mid
);
}
}
}