Add agent registration, heartbeat and ws notification

netmito/Cargo.toml

@@ -18,7 +18,7 @@ argon2 = { version = "0.5.3", features = ["std"] }
 async-compression = { version = "0.4.32", features = ["gzip", "tokio"] }
 async-stream = "0.3.6"
 aws-sdk-s3 = { version = "1.110", features = ["behavior-version-latest"] }
-axum = { version = "0.8.6", features = ["http2"] }
+axum = { version = "0.8.6", features = ["http2", "ws"] }
 axum-extra = { version = "0.12", features = ["typed-header", "query"] }
 base64 = "0.22.1"
 clap = { workspace = true }
@@ -65,6 +65,7 @@ sea-orm = { version = "1.1.17", default-features = false, features = [
 sea-orm-migration = "1.1.17"
 serde = { workspace = true }
 serde_json = { workspace = true }
+speedy = { version = "0.8.7", features = ["uuid"] }
 thiserror = "2.0.17"
 time = { version = "0.3.44", features = ["serde-human-readable"] }
 tokio = { workspace = true }

netmito/src/api/agents.rs

@@ -0,0 +1,108 @@
+use axum::{
+    extract::{Path, Query, State},
+    middleware,
+    routing::{delete, post},
+    Extension, Json, Router,
+};
+use uuid::Uuid;
+
+use super::map_service_error;
+use crate::{
+    config::InfraPool,
+    error::ApiError,
+    schema::{
+        AgentHeartbeatReq, AgentHeartbeatResp, AgentShutdownReq, AgentsQueryReq, AgentsQueryResp,
+        RegisterAgentReq, RegisterAgentResp,
+    },
+    service::{
+        self,
+        auth::{
+            auth_middleware,
+            roles::{AGENT, USER},
+            AuthAgent, AuthUser,
+        },
+    },
+};
+
+pub fn agents_router(st: InfraPool) -> Router<InfraPool> {
+    let user_router = Router::new()
+        .route("/", post(register_agent))
+        .route("/query", post(query_agents))
+        .route("/{uuid}", delete(shutdown_agent))
+        .route_layer(middleware::from_fn_with_state(
+            st.clone(),
+            auth_middleware::<USER>,
+        ))
+        .with_state(st.clone());
+
+    let agent_router = Router::new()
+        .route("/", delete(agent_exit))
+        .route("/heartbeat", post(heartbeat))
+        .route_layer(middleware::from_fn_with_state(
+            st.clone(),
+            auth_middleware::<AGENT>,
+        ))
+        .with_state(st.clone());
+
+    Router::new().merge(user_router).merge(agent_router)
+}
+
+/// `POST /agents`: register an agent and issue its token.
+async fn register_agent(
+    Extension(u): Extension<AuthUser>,
+    State(pool): State<InfraPool>,
+    Json(req): Json<RegisterAgentReq>,
+) -> Result<Json<RegisterAgentResp>, ApiError> {
+    let resp = service::agent::user_register_agent(u.id, &pool, req)
+        .await
+        .map_err(map_service_error)?;
+    Ok(Json(resp))
+}
+
+/// `POST /agents/query`
+async fn query_agents(
+    Extension(u): Extension<AuthUser>,
+    State(pool): State<InfraPool>,
+    Json(req): Json<AgentsQueryReq>,
+) -> Result<Json<AgentsQueryResp>, ApiError> {
+    let resp = service::agent::user_query_agents(u.id, &pool, req)
+        .await
+        .map_err(map_service_error)?;
+    Ok(Json(resp))
+}
+
+/// `DELETE /agents/{uuid}?op=graceful|force`: shut the agent down. The agent row stays.
+async fn shutdown_agent(
+    Extension(u): Extension<AuthUser>,
+    State(pool): State<InfraPool>,
+    Path(uuid): Path<Uuid>,
+    Query(req): Query<AgentShutdownReq>,
+) -> Result<(), ApiError> {
+    service::agent::user_shutdown_agent_by_uuid(u.id, uuid, req.op, &pool)
+        .await
+        .map_err(map_service_error)?;
+    Ok(())
+}
+
+/// `DELETE /agents`: the agent exits. Retire it now, so that it does not wait for its heartbeat times out.
+async fn agent_exit(
+    Extension(a): Extension<AuthAgent>,
+    State(pool): State<InfraPool>,
+) -> Result<(), ApiError> {
+    service::agent::retire_agent(&pool, a.id, service::agent::RetireCause::SelfExit)
+        .await
+        .map_err(map_service_error)?;
+    Ok(())
+}
+
+/// `POST /agents/heartbeat`: record a heartbeat and return the notifications that the agent has not acknowledged.
+async fn heartbeat(
+    Extension(a): Extension<AuthAgent>,
+    State(pool): State<InfraPool>,
+    Json(req): Json<AgentHeartbeatReq>,
+) -> Result<Json<AgentHeartbeatResp>, ApiError> {
+    let resp = service::agent::agent_heartbeat(a.id, a.uuid, &pool, req)
+        .await
+        .map_err(map_service_error)?;
+    Ok(Json(resp))
+}

netmito/src/api/groups.rs

@@ -8,7 +8,7 @@ use axum_extra::extract::Query;
 
 use crate::{
     config::InfraPool,
-    error::{ApiError, ApiResult},
+    error::{ApiError, ApiResult, AuthError},
     schema::*,
     service::{
         self,
@@ -195,6 +195,9 @@ pub async fn download_attachment(
         AuthIdentity::Worker(w) => {
             service::s3::worker_get_attachment(&pool, w.uuid, group_name, key).await
         }
+        // TODO: let an agent read the attachments of a group whose suite it runs a job for.
+        // Do this when agents run jobs. Before that, no job rows exist and the check always fails.
+        AuthIdentity::Agent(_) => Err(AuthError::PermissionDenied.into()),
     }
     .map_err(|e| match e {
         crate::error::Error::AuthError(err) => ApiError::AuthError(err),

netmito/src/api/mod.rs

@@ -1,9 +1,11 @@
 pub mod admin;
+pub mod agents;
 pub mod groups;
 pub mod suites;
 pub mod tasks;
 pub mod users;
 pub mod workers;
+pub mod ws;
 
 #[cfg(feature = "debugging")]
 use std::net::SocketAddr;
@@ -77,6 +79,8 @@ pub fn router(st: InfraPool, cancel_token: CancellationToken) -> Router {
         .nest("/workers", workers::workers_router(st.clone()))
         .nest("/tasks", tasks::tasks_router(st.clone()))
         .nest("/suites", suites::suites_router(st.clone()))
+        .nest("/agents", agents::agents_router(st.clone()))
+        .nest("/ws", ws::ws_router(st.clone()))
         .with_state(st)
         .layer(CorsLayer::permissive())
         .layer(CatchPanicLayer::new());

netmito/src/api/workers.rs

@@ -16,7 +16,7 @@ use crate::{
         auth::{
             auth_middleware,
             roles::{USER, WORKER},
-            token::generate_worker_token,
+            token::{generate_runner_token, TokenRole},
             AuthUser, AuthWorker,
         },
     },
@@ -78,8 +78,8 @@ pub async fn user_register_worker(
             ApiError::InternalServerError
         }
     })?;
-    let token =
-        generate_worker_token(uuid.clone().to_string(), 0, req.lifetime).map_err(|e| match e {
+    let token = generate_runner_token(uuid.clone().to_string(), 0, req.lifetime, TokenRole::Worker)
+        .map_err(|e| match e {
             crate::error::Error::AuthError(err) => ApiError::AuthError(err),
             crate::error::Error::ApiError(e) => e,
             _ => {

netmito/src/api/ws.rs

@@ -0,0 +1,17 @@
+use axum::{middleware, routing::get, Router};
+
+use crate::{
+    config::InfraPool,
+    service::auth::{auth_middleware, roles::AGENT},
+    ws::websocket_handler,
+};
+
+pub fn ws_router(st: InfraPool) -> Router<InfraPool> {
+    Router::new()
+        .route("/agents", get(websocket_handler))
+        .route_layer(middleware::from_fn_with_state(
+            st.clone(),
+            auth_middleware::<AGENT>,
+        ))
+        .with_state(st)
+}