hyperhive/swarm-controller/src/agent_control.rs

//! `POST /api/agents/{name}/cancel` — stop the turn an agent is in the middle
//! of. A command, not text: it travels on its own subject
//! ([`swarm_queue_client::agent_control_subject`]), never the one a turn's
//! message body does, and no hive is in the path - the agent's harness
//! subscribes itself.
//!
//! **Sent as a NATS request, and unlike input it is answered**: the harness
//! replies with whether a turn was actually running. That reply is the only
//! thing that distinguishes "cancelled" from "was idle", and it is the only
//! confirmation a command can have - so a request nobody answers within
//! [`swarm_queue_client::AGENT_CONTROL_REPLY_DEADLINE`] is a 504, never a
//! success. No subscriber at all is a 503, as for input.
//!
//! The sender is the `Remote-User` header, for the same reasons as
//! [`crate::agent_input`]; a command is refused without one.

use axum::Json;
use axum::extract::{Path, State};
use axum::http::{HeaderMap, StatusCode};
use problem_details::ProblemDetails;
use serde::Serialize;
use swarm_queue_client::{AgentControl, AgentControlReply};
use utoipa::ToSchema;

use crate::{AppState, error_problem};

/// What became of one command.
#[derive(Debug, PartialEq, Eq)]
pub(crate) enum Outcome {
    /// The agent answered.
    Answered(AgentControlReply),
    /// Nothing subscribes to the agent's control subject.
    NotRunning,
    /// A subscriber held the command and did not answer in time.
    Unconfirmed,
    /// The queue refused or failed the request, or the answer was unreadable.
    Failed(String),
}

/// Where a command goes: the queue in production, a fake in the tests.
pub(crate) trait ControlSink {
    fn request(&self, subject: String, payload: Vec<u8>) -> impl Future<Output = Outcome> + Send;
}

impl ControlSink for async_nats::Client {
    async fn request(&self, subject: String, payload: Vec<u8>) -> Outcome {
        use async_nats::RequestErrorKind;

        let request = async_nats::Request::new()
            .payload(payload.into())
            .timeout(Some(swarm_queue_client::AGENT_CONTROL_REPLY_DEADLINE));
        match self.send_request(subject, request).await {
            Ok(reply) => match serde_json::from_slice(&reply.payload) {
                Ok(reply) => Outcome::Answered(reply),
                Err(e) => Outcome::Failed(format!("the agent's answer was unreadable: {e}")),
            },
            Err(e) => match e.kind() {
                RequestErrorKind::TimedOut => Outcome::Unconfirmed,
                RequestErrorKind::NoResponders => Outcome::NotRunning,
                _ => Outcome::Failed(e.to_string()),
            },
        }
    }
}

/// The agent's answer to a cancel.
#[derive(Serialize, ToSchema)]
pub struct CancelResponse {
    /// Whether a turn was running and was stopped. `false` means the agent
    /// heard the command and had nothing to stop.
    acted: bool,
}

/// Stop the turn an agent is running.
#[utoipa::path(
    post,
    path = "/api/agents/{name}/cancel",
    params(("name" = String, Path, description = "agent whose turn to cancel")),
    responses(
        (status = 200, description = "the agent answered: `acted` says whether a turn was \
            running and was stopped", body = CancelResponse),
        (status = 400, description = "the agent name is not shaped like an identifier \
            (problem+json)", body = String),
        (status = 401, description = "the request carries no `Remote-User` (problem+json)",
            body = String),
        (status = 503, description = "the agent's harness is not subscribed - not running, \
            or not connected to the swarm queue - or no swarm queue is configured on this \
            host or it could not be reached (problem+json)", body = String),
        (status = 504, description = "the harness held the command and did not answer in \
            time, so whether anything was cancelled is unknown (problem+json)", body = String),
        (status = 500, description = "the queue refused or failed the request (problem+json)",
            body = String),
    ),
    tag = "agents"
)]
pub(crate) async fn post_agent_cancel(
    State(state): State<AppState>,
    Path(agent): Path<String>,
    headers: HeaderMap,
) -> Result<Json<CancelResponse>, ProblemDetails> {
    let (agent, from) = target(&headers, &agent).map_err(|(s, d)| error_problem(s, d))?;
    let Some(status) = state.status.as_ref() else {
        return Err(error_problem(
            StatusCode::SERVICE_UNAVAILABLE,
            "no swarm queue is configured on this host",
        ));
    };
    let client = status.queue_client();
    swarm_queue_client::ensure_connected(&client).map_err(|e| {
        error_problem(
            StatusCode::SERVICE_UNAVAILABLE,
            &swarm_queue_client::chain(&e),
        )
    })?;
    send(&client, &agent, &from, AgentControl::Cancel).await
}

/// The agent and who is asking, or the refusal. Everything decided before the
/// queue is touched.
fn target(
    headers: &HeaderMap,
    agent: &str,
) -> Result<(String, String), (StatusCode, &'static str)> {
    let from = crate::agent_input::sender(headers)?;
    let agent = hive_types::Ident::parse(agent)
        .map_err(|reason| (StatusCode::BAD_REQUEST, reason))?
        .into_string();
    Ok((agent, from))
}

async fn send(
    sink: &impl ControlSink,
    agent: &str,
    from: &str,
    command: AgentControl,
) -> Result<Json<CancelResponse>, ProblemDetails> {
    let subject = swarm_queue_client::agent_control_subject(agent);
    let payload = serde_json::to_vec(&command).map_err(|e| {
        error_problem(
            StatusCode::INTERNAL_SERVER_ERROR,
            &format!("encoding the command failed: {e}"),
        )
    })?;
    match sink.request(subject.clone(), payload).await {
        Outcome::Answered(reply) => {
            tracing::info!(%agent, %from, ?command, acted = reply.acted, "agent command answered");
            Ok(Json(CancelResponse { acted: reply.acted }))
        }
        Outcome::NotRunning => Err(error_problem(
            StatusCode::SERVICE_UNAVAILABLE,
            &format!(
                "{agent} is not listening for commands: its harness is not running or not connected to the swarm queue"
            ),
        )),
        Outcome::Unconfirmed => {
            tracing::warn!(%agent, %from, ?command, "agent command unanswered");
            Err(error_problem(
                StatusCode::GATEWAY_TIMEOUT,
                &format!(
                    "{agent} did not answer in time: whether anything was cancelled is unknown"
                ),
            ))
        }
        Outcome::Failed(e) => {
            tracing::warn!(%agent, %subject, error = %e, "agent command failed");
            Err(error_problem(
                StatusCode::INTERNAL_SERVER_ERROR,
                &format!("sending to {subject} failed: {e}"),
            ))
        }
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Mutex;

    use axum::http::{HeaderMap, HeaderValue, StatusCode};
    use swarm_queue_client::{AgentControl, AgentControlReply};

    use super::{ControlSink, Outcome, send, target};

    /// Answers every command with one fixed outcome and records what it got.
    struct FakeSink {
        answer: fn() -> Outcome,
        got: Mutex<Vec<(String, Vec<u8>)>>,
    }

    impl FakeSink {
        fn new(answer: fn() -> Outcome) -> Self {
            Self {
                answer,
                got: Mutex::default(),
            }
        }
    }

    impl ControlSink for FakeSink {
        async fn request(&self, subject: String, payload: Vec<u8>) -> Outcome {
            self.got
                .lock()
                .expect("unpoisoned")
                .push((subject, payload));
            (self.answer)()
        }
    }

    fn from(user: &str) -> HeaderMap {
        let mut h = HeaderMap::new();
        h.insert("Remote-User", HeaderValue::from_str(user).expect("ascii"));
        h
    }

    #[test]
    fn a_command_needs_a_sender_and_an_identifier() {
        assert_eq!(
            target(&from("mara"), "atlas").expect("valid"),
            ("atlas".to_owned(), "mara".to_owned())
        );
        assert_eq!(
            target(&HeaderMap::new(), "../x").expect_err("no sender").0,
            StatusCode::UNAUTHORIZED
        );
        assert_eq!(
            target(&from("mara"), "../x").expect_err("bad agent").0,
            StatusCode::BAD_REQUEST
        );
    }

    #[tokio::test]
    async fn the_answer_is_relayed_from_the_agent_s_own_control_subject() {
        let answers: [(bool, fn() -> Outcome); 2] = [
            (true, || {
                Outcome::Answered(AgentControlReply { acted: true })
            }),
            (false, || {
                Outcome::Answered(AgentControlReply { acted: false })
            }),
        ];
        for (acted, answer) in answers {
            let sink = FakeSink::new(answer);
            let body = send(&sink, "atlas", "mara", AgentControl::Cancel)
                .await
                .expect("answered");
            assert_eq!(body.0.acted, acted);
            let got = sink.got.lock().expect("unpoisoned");
            assert_eq!(got[0].0, "$SWARM.control.atlas");
            let sent: AgentControl = serde_json::from_slice(&got[0].1).expect("a command");
            assert_eq!(sent, AgentControl::Cancel);
        }
    }

    #[tokio::test]
    async fn unanswered_is_504_not_success_and_the_rest_are_503_and_500() {
        let status = |answer: fn() -> Outcome| async move {
            send(
                &FakeSink::new(answer),
                "atlas",
                "mara",
                AgentControl::Cancel,
            )
            .await
            .err()
            .expect("not a success")
            .status
        };
        assert_eq!(
            status(|| Outcome::Unconfirmed).await,
            Some(StatusCode::GATEWAY_TIMEOUT)
        );
        assert_eq!(
            status(|| Outcome::NotRunning).await,
            Some(StatusCode::SERVICE_UNAVAILABLE)
        );
        assert_eq!(
            status(|| Outcome::Failed("refused".to_owned())).await,
            Some(StatusCode::INTERNAL_SERVER_ERROR)
        );
    }
}