use std::io::Write;
use std::path::PathBuf;
use std::time::Duration;
use rmcp::model::CallToolResult;
use shep_client::{Client, ConnectError, RequestError, SshOptions};
use shep_core::protocol::{HelloAck, ProcessInfo, Request, Response};
use crate::commands::rpc::deadline_for;
use crate::exit::ExitCode;
use crate::whistle::read::unexpected_response;
const HOST_CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
#[must_use]
pub fn host_ssh() -> SshOptions {
crate::commands::hosts::unattended(HOST_CONNECT_TIMEOUT)
}
#[derive(Debug, Clone)]
enum Target {
Local(PathBuf),
Host {
dest: String,
options: SshOptions,
},
}
impl Target {
fn host(&self) -> Option<&str> {
match self {
Self::Local(_) => None,
Self::Host { dest, .. } => Some(dest),
}
}
}
#[derive(Debug, Clone)]
pub struct Shepherd {
target: Target,
}
impl Shepherd {
#[must_use]
pub fn new(socket: PathBuf) -> Self {
Self {
target: Target::Local(socket),
}
}
#[must_use]
pub fn for_host(dest: &str, options: SshOptions) -> Self {
Self {
target: Target::Host {
dest: dest.to_owned(),
options,
},
}
}
pub async fn call(&self, request: Request) -> Result<Response, CallToolResult> {
self.call_with_ack(request)
.await
.map(|(_ack, response)| response)
}
pub async fn call_with_ack(
&self,
request: Request,
) -> Result<(HelloAck, Response), CallToolResult> {
let (client, ack) = self.connect().await?;
let deadline = deadline_for(&request);
let response = client
.request_with_deadline(request, deadline)
.await
.map_err(|err| refusal(&err));
let _ = client.close().await;
response.map(|response| (ack, response))
}
pub async fn list_flock_with_ack(
&self,
) -> Result<(HelloAck, Vec<ProcessInfo>), CallToolResult> {
let (client, ack) = self.connect().await?;
let flock = client.list_flock().await.map_err(|err| match err {
RequestError::UnexpectedReply { .. } => unexpected_response(),
other => refusal(&other),
});
let _ = client.close().await;
flock.map(|flock| (ack, flock))
}
async fn connect(&self) -> Result<(Client, HelloAck), CallToolResult> {
let dialled = match &self.target {
Target::Local(socket) => Client::connect(socket).await,
Target::Host { dest, options } => Client::connect_ssh_with(dest, options).await,
};
let client = dialled.map_err(|err| connect_refusal(&err, self.target.host()))?;
refuse_if_skewed(&client)?;
let ack = client.daemon().clone();
Ok((client, ack))
}
}
fn refuse_if_skewed(client: &Client) -> Result<(), CallToolResult> {
if client.host().is_some() {
return match crate::version_guard::host_skew(
client,
crate::version_guard::VersionGuard::Enforce,
) {
Ok(None) => Ok(()),
Ok(Some(notice)) => {
let _ = writeln!(std::io::stderr(), "whistle: {notice}");
Ok(())
}
Err(message) => Err(own_refusal(ExitCode::VersionSkew.code_str(), message)),
};
}
let mut sink = std::io::sink();
let mut err = std::io::stderr();
let mut streams = crate::output::Streams {
out: &mut sink,
err: &mut err,
style: crate::style::Presentation::BARE,
fmt: crate::cli::Format::Table,
};
crate::version_guard::refuse_version_skew(
&mut streams,
client,
crate::version_guard::VersionGuard::Enforce,
)
.map_err(|_code| {
CallToolResult::structured_error(serde_json::json!({
"code": crate::exit::ExitCode::VersionSkew.code_str(),
"message": format!(
"this shep is {}, the running shepherd is {}; \
`cargo install shep` replaced the binary without \
restarting it. Run `shep daemon reload`",
env!("CARGO_PKG_VERSION"),
client.daemon().daemon_version,
),
}))
})
}
fn connect_refusal(err: &ConnectError, host: Option<&str>) -> CallToolResult {
if matches!(err, ConnectError::ProtocolMismatch { .. }) {
return own_refusal(ExitCode::VersionSkew.code_str(), err.to_string());
}
let (code, message) = match (err, host) {
(ConnectError::Ssh { .. }, _) => ("no_shepherd", err.to_string()),
(ConnectError::Destination(_), Some(dest)) => (
ExitCode::InvalidConfig.code_str(),
format!("`{dest}` is a listed host ssh cannot take: {err}"),
),
(other, Some(dest)) => ("host_unreachable", format!("`{dest}`: {other}")),
(other, None) => ("no_shepherd", format!("no shepherd is running: {other}")),
};
own_refusal(code, message)
}
fn refusal(err: &RequestError) -> CallToolResult {
let (code, message) = match err {
RequestError::Rpc(rpc) => (
ExitCode::from(rpc.code).code_str().to_string(),
rpc.message.clone(),
),
other => ("transport".to_string(), other.to_string()),
};
CallToolResult::structured_error(serde_json::json!({
"code": code,
"message": message,
}))
}
pub fn own_refusal(code: &str, message: String) -> CallToolResult {
CallToolResult::structured_error(serde_json::json!({
"code": code,
"message": message,
}))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::whistle::matching_ack;
use shep_core::protocol::{RpcError, RpcErrorCode};
#[test]
fn a_daemon_refusal_is_an_in_band_error_keeping_its_own_message() {
let result = refusal(&RequestError::Rpc(RpcError {
code: RpcErrorCode::Internal,
message: "api is already being reloaded".to_string(),
daemon_version: None,
}));
assert_eq!(result.is_error, Some(true));
let structured = result
.structured_content
.expect("a refusal carries structured content a model can branch on");
assert_eq!(structured["message"], "api is already being reloaded");
assert_eq!(
structured["code"], "internal",
"and the code, so a model can tell a conflict from a not-found: {structured}"
);
}
#[test]
fn an_unreachable_shepherd_names_the_socket_once() {
let socket = std::path::Path::new("/nonexistent/shep/run/shep.sock");
let result = connect_refusal(
&ConnectError::Connect {
path: socket.to_path_buf(),
source: std::io::Error::from(std::io::ErrorKind::NotFound),
},
None,
);
assert_eq!(result.is_error, Some(true));
let message = result.structured_content.expect("structured")["message"]
.as_str()
.expect("a string")
.to_string();
assert!(message.contains("/nonexistent/shep/run/shep.sock"));
assert!(
message.contains("no shepherd"),
"and says what is missing, not just what failed: {message}"
);
assert_eq!(
message.matches("/nonexistent/shep/run/shep.sock").count(),
1,
"the socket path appears once, not once per layer: {message}"
);
}
#[test]
fn a_host_that_turns_down_the_handshake_is_refused_as_a_version_gap() {
let result = connect_refusal(
&ConnectError::ProtocolMismatch {
client: shep_core::protocol::PROTOCOL_VERSION,
daemon_version: Some("0.1.8".to_string()),
message: "this daemon speaks protocol 2".to_string(),
},
Some("amp2"),
);
assert_eq!(result.is_error, Some(true));
let structured = result
.structured_content
.expect("a refusal carries structured content a model can branch on");
assert_eq!(
structured["code"],
ExitCode::VersionSkew.code_str(),
"the same code this module gives a host that answers and agrees to disagree: {structured}"
);
let message = structured["message"]
.as_str()
.expect("a string")
.to_string();
assert!(message.contains("protocol mismatch"), "{message}");
assert!(
message.contains("older of the two"),
"and names which side of the gap to replace: {message}"
);
assert!(
!message.contains("no shepherd is running"),
"amp2 has a shepherd, which refused: {message}"
);
}
#[test]
fn a_host_that_never_answers_the_handshake_names_the_host() {
let result = connect_refusal(
&ConnectError::HandshakeTimeout {
after: HOST_CONNECT_TIMEOUT,
},
Some("amp2"),
);
assert_eq!(result.is_error, Some(true));
let structured = result
.structured_content
.expect("a refusal carries structured content a model can branch on");
assert_eq!(
structured["code"], "host_unreachable",
"a code that does not claim the shepherd is missing: {structured}"
);
let message = structured["message"]
.as_str()
.expect("a string")
.to_string();
assert!(
message.contains("amp2"),
"the refusal names the host it dialled: {message}"
);
assert!(message.contains("handshake"), "{message}");
assert!(
!message.contains("no shepherd is running"),
"the link was the slow part: {message}"
);
}
#[test]
fn a_listed_host_ssh_cannot_take_is_refused_as_a_config_problem() {
let listed = "amp 2";
let destination = shep_core::destination::Destination::parse(listed)
.expect_err("a space reads to ssh as a second argument");
let result = connect_refusal(&ConnectError::Destination(destination), Some(listed));
assert_eq!(result.is_error, Some(true));
let structured = result
.structured_content
.expect("a refusal carries structured content a model can branch on");
assert_eq!(
structured["code"],
ExitCode::InvalidConfig.code_str(),
"the entry is the problem, not the host: {structured}"
);
let message = structured["message"]
.as_str()
.expect("a string")
.to_string();
assert!(message.contains(listed), "{message}");
assert!(message.contains("destination"), "{message}");
}
#[tokio::test]
async fn two_calls_survive_a_shepherd_that_restarted_in_between() {
let dir = tempfile::tempdir().unwrap();
let socket = shep_client::testing::control_address(dir.path());
let (first, first_served) = shep_client::testing::fake_daemon_accepting_repeatedly_with_ack(
&socket,
matching_ack(),
Response::Pong,
);
let shepherd = Shepherd::new(socket.clone());
let one = tokio::time::timeout(
std::time::Duration::from_secs(10),
shepherd.call(Request::Ping),
)
.await
.expect("the first call finished within ten seconds");
assert!(one.is_ok());
assert_eq!(
first_served.load(std::sync::atomic::Ordering::SeqCst),
1,
"one call is one connection — not zero, and not a retry"
);
first.abort();
let _ = first.await;
#[cfg(unix)]
std::fs::remove_file(&socket).unwrap();
let (second, _second_served) =
shep_client::testing::fake_daemon_accepting_repeatedly_with_ack(
&socket,
matching_ack(),
Response::Pong,
);
let two = tokio::time::timeout(
std::time::Duration::from_secs(10),
shepherd.call(Request::Ping),
)
.await
.expect("the second call finished within ten seconds");
assert!(two.is_ok(), "a fresh connection per call needs no ladder");
second.abort();
}
#[tokio::test]
async fn a_version_skewed_shepherd_is_an_in_band_refusal() {
let dir = tempfile::tempdir().unwrap();
let addr = shep_client::testing::control_address(dir.path());
let ack = HelloAck {
daemon_version: "0.1.8".to_string(),
protocol: shep_core::protocol::PROTOCOL_VERSION,
pid: 4242,
min_supported: None,
};
let (client, _fake) = shep_client::testing::fake_client_with_ack(&addr, ack).await;
let result = super::refuse_if_skewed(&client).expect_err("a skew must be refused");
assert_eq!(result.is_error, Some(true));
let structured = result
.structured_content
.expect("a refusal carries structured content a model can branch on");
assert_eq!(structured["code"], "version_skew");
let message = structured["message"].as_str().expect("a string message");
assert!(message.contains(env!("CARGO_PKG_VERSION")), "{message}");
assert!(message.contains("0.1.8"), "{message}");
}
#[tokio::test]
async fn a_matching_version_proceeds() {
let dir = tempfile::tempdir().unwrap();
let addr = shep_client::testing::control_address(dir.path());
let ack = HelloAck {
daemon_version: env!("CARGO_PKG_VERSION").to_string(),
protocol: shep_core::protocol::PROTOCOL_VERSION,
pid: 4242,
min_supported: None,
};
let (client, _fake) = shep_client::testing::fake_client_with_ack(&addr, ack).await;
super::refuse_if_skewed(&client).expect("a matching version is not a skew");
}
}