use axum::{Json, extract::State, http::StatusCode, response::IntoResponse};
use chrono::{DateTime, Duration, FixedOffset, Utc};
use serde::{Deserialize, Serialize};
use sqlx::PgPool;
use uuid::Uuid;
use crate::{app::AppState, auth::AuthenticatedAgent, error::ApiError};
pub const INTERVAL_SECONDS: i64 = 60;
pub const STALE_AFTER_SECONDS: i64 = INTERVAL_SECONDS * 3;
const MAX_CLOCK_SKEW_SECONDS: i64 = 60;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, sqlx::Type)]
#[sqlx(type_name = "agent_state", rename_all = "lowercase")]
#[serde(rename_all = "lowercase")]
pub enum AgentState {
Working,
Paused,
Idle,
}
#[derive(Debug, Deserialize)]
pub struct Pulse {
pub state: AgentState,
pub at: DateTime<FixedOffset>,
}
#[derive(Debug, Serialize)]
pub struct Accepted {
pub interval_seconds: i64,
pub stale_after_seconds: i64,
pub state: AgentState,
pub clock_skew_seconds: i64,
}
pub async fn beat(State(state): State<AppState>, agent: AuthenticatedAgent, Json(pulse): Json<Pulse>) -> Result<impl IntoResponse, ApiError> {
let now = Utc::now();
let at = pulse.at.with_timezone(&Utc);
let skew = (at - now).num_seconds();
if skew > MAX_CLOCK_SKEW_SECONDS {
return Err(ApiError::bad_request(format!(
"the pulse is stamped {skew} s ahead of this server; check the machine's clock"
)));
}
sqlx::query("UPDATE agents SET heartbeat_state = $2, heartbeat_at = $3, heartbeat_received_at = now() WHERE id = $1")
.bind(agent.agent_id)
.bind(pulse.state)
.bind(at)
.execute(&state.pool)
.await?;
Ok((
StatusCode::ACCEPTED,
Json(Accepted {
interval_seconds: INTERVAL_SECONDS,
stale_after_seconds: STALE_AFTER_SECONDS,
state: pulse.state,
clock_skew_seconds: skew,
}),
))
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
#[serde(rename_all = "lowercase")]
pub enum LiveStatus {
Working,
Paused,
Idle,
Offline,
Unknown,
}
impl LiveStatus {
pub fn resolve(state: Option<AgentState>, received: Option<DateTime<Utc>>, now: DateTime<Utc>) -> Self {
let (Some(state), Some(received)) = (state, received) else {
return Self::Unknown;
};
if now - received > Duration::seconds(STALE_AFTER_SECONDS) {
return Self::Offline;
}
match state {
AgentState::Working => Self::Working,
AgentState::Paused => Self::Paused,
AgentState::Idle => Self::Idle,
}
}
}
#[derive(Debug, Serialize)]
pub struct Live {
pub user_id: Uuid,
pub status: LiveStatus,
pub since_received: Option<i64>,
}
pub async fn load(pool: &PgPool, visible_users: &str, is_admin: bool, reader: Uuid) -> Result<Vec<Live>, ApiError> {
let rows: Vec<(Uuid, Option<AgentState>, Option<DateTime<Utc>>)> = sqlx::query_as(sqlx::AssertSqlSafe(format!(
"SELECT u.id, h.heartbeat_state, h.heartbeat_received_at
FROM users u
LEFT JOIN LATERAL (
-- The freshest pulse among this person's live agents. Someone
-- with a desktop and a laptop is working if either says so, and
-- a revoked agent's last words are not evidence of anything.
SELECT a.heartbeat_state, a.heartbeat_received_at
FROM agents a
WHERE a.user_id = u.id AND a.revoked_at IS NULL AND a.heartbeat_received_at IS NOT NULL
ORDER BY a.heartbeat_received_at DESC
LIMIT 1
) AS h ON true
WHERE u.active AND {visible_users}
ORDER BY u.display_name, u.email"
)))
.bind(is_admin)
.bind(reader)
.fetch_all(pool)
.await?;
let now = Utc::now();
Ok(rows
.into_iter()
.map(|(user_id, state, received)| Live {
user_id,
status: LiveStatus::resolve(state, received, now),
since_received: received.map(|received| (now - received).num_seconds().max(0)),
})
.collect())
}
#[cfg(test)]
mod tests {
use super::*;
fn seconds_ago(now: DateTime<Utc>, seconds: i64) -> Option<DateTime<Utc>> {
Some(now - Duration::seconds(seconds))
}
#[test]
fn the_wire_names_are_the_contract() {
assert_eq!(serde_json::to_string(&AgentState::Working).unwrap(), r#""working""#);
assert_eq!(serde_json::to_string(&AgentState::Paused).unwrap(), r#""paused""#);
assert_eq!(serde_json::to_string(&AgentState::Idle).unwrap(), r#""idle""#);
assert_eq!(serde_json::to_string(&LiveStatus::Offline).unwrap(), r#""offline""#);
assert_eq!(serde_json::to_string(&LiveStatus::Unknown).unwrap(), r#""unknown""#);
}
#[test]
fn a_state_round_trips_through_json() {
for state in [AgentState::Working, AgentState::Paused, AgentState::Idle] {
let json = serde_json::to_string(&state).unwrap();
assert_eq!(serde_json::from_str::<AgentState>(&json).unwrap(), state);
}
}
#[test]
fn a_fresh_pulse_is_shown_as_claimed() {
let now = Utc::now();
for (state, expected) in [
(AgentState::Working, LiveStatus::Working),
(AgentState::Paused, LiveStatus::Paused),
(AgentState::Idle, LiveStatus::Idle),
] {
assert_eq!(LiveStatus::resolve(Some(state), seconds_ago(now, 5), now), expected);
}
}
#[test]
fn a_pulse_survives_two_missed_intervals() {
let now = Utc::now();
let two_missed = INTERVAL_SECONDS * 2 + 5;
assert_eq!(
LiveStatus::resolve(Some(AgentState::Working), seconds_ago(now, two_missed), now),
LiveStatus::Working
);
}
#[test]
fn a_stale_pulse_is_offline_whatever_it_claimed() {
let now = Utc::now();
let stale = seconds_ago(now, STALE_AFTER_SECONDS + 1);
for state in [AgentState::Working, AgentState::Paused, AgentState::Idle] {
assert_eq!(LiveStatus::resolve(Some(state), stale, now), LiveStatus::Offline);
}
}
#[test]
fn silence_is_unknown_rather_than_idle() {
let now = Utc::now();
assert_eq!(LiveStatus::resolve(None, None, now), LiveStatus::Unknown);
assert_eq!(LiveStatus::resolve(Some(AgentState::Working), None, now), LiveStatus::Unknown);
assert_eq!(LiveStatus::resolve(None, seconds_ago(now, 1), now), LiveStatus::Unknown);
}
#[test]
fn the_interval_leaves_room_for_a_missed_pulse() {
const { assert!(STALE_AFTER_SECONDS > INTERVAL_SECONDS, "an agent gets no margin at all") };
}
}