Dispatch suite tasks to agents

netmito/src/api/agents.rs

@@ -1,7 +1,7 @@
 use axum::{
     extract::{Path, Query, State},
     middleware,
-    routing::{delete, post},
+    routing::{delete, get, post},
     Extension, Json, Router,
 };
 use uuid::Uuid;
@@ -11,8 +11,9 @@ use crate::{
     config::InfraPool,
     error::ApiError,
     schema::{
-        AgentHeartbeatReq, AgentHeartbeatResp, AgentShutdownReq, AgentsQueryReq, AgentsQueryResp,
-        RegisterAgentReq, RegisterAgentResp,
+        AcceptSuiteReq, AcceptSuiteResp, AgentHeartbeatReq, AgentHeartbeatResp, AgentShutdownReq,
+        AgentsQueryReq, AgentsQueryResp, FetchTasksReq, FetchTasksResp, JobReportOp, JobReportReq,
+        JobReportResp, RegisterAgentReq, RegisterAgentResp, StopAgentJobReq, StopAgentJobResp,
     },
     service::{
         self,
@@ -29,6 +30,7 @@ pub fn agents_router(st: InfraPool) -> Router<InfraPool> {
         .route("/", post(register_agent))
         .route("/query", post(query_agents))
         .route("/{uuid}", delete(shutdown_agent))
+        .route("/{uuid}/job", delete(stop_agent_job))
         .route_layer(middleware::from_fn_with_state(
             st.clone(),
             auth_middleware::<USER>,
@@ -38,6 +40,9 @@ pub fn agents_router(st: InfraPool) -> Router<InfraPool> {
     let agent_router = Router::new()
         .route("/", delete(agent_exit))
         .route("/heartbeat", post(heartbeat))
+        .route("/suite", post(accept_suite))
+        .route("/job", post(report_job))
+        .route("/tasks", get(fetch_tasks))
         .route_layer(middleware::from_fn_with_state(
             st.clone(),
             auth_middleware::<AGENT>,
@@ -84,6 +89,20 @@ async fn shutdown_agent(
     Ok(())
 }
 
+/// `DELETE /agents/{uuid}/job?op=graceful|force`: stop the job that the agent runs now. The agent
+/// stays up and picks a suite again.
+async fn stop_agent_job(
+    Extension(u): Extension<AuthUser>,
+    State(pool): State<InfraPool>,
+    Path(uuid): Path<Uuid>,
+    Query(req): Query<StopAgentJobReq>,
+) -> Result<Json<StopAgentJobResp>, ApiError> {
+    let resp = service::agent::user_stop_agent_job(u.id, uuid, req.op, &pool)
+        .await
+        .map_err(map_service_error)?;
+    Ok(Json(resp))
+}
+
 /// `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>,
@@ -106,3 +125,69 @@ async fn heartbeat(
         .map_err(map_service_error)?;
     Ok(Json(resp))
 }
+
+/// `POST /agents/suite`: pick a suite, claim it, and open a job.
+///
+/// The transaction picks the suite, thus the span records the suite and the job when they are
+/// known.
+#[tracing::instrument(
+    skip_all,
+    fields(
+        agent_uuid = %a.uuid,
+        suite_uuid = tracing::field::Empty,
+        job_id = tracing::field::Empty,
+    )
+)]
+async fn accept_suite(
+    Extension(a): Extension<AuthAgent>,
+    State(pool): State<InfraPool>,
+    Json(req): Json<AcceptSuiteReq>,
+) -> Result<Json<AcceptSuiteResp>, ApiError> {
+    let resp = service::agent::agent_accept_suite(a.id, a.uuid, &pool, req)
+        .await
+        .map_err(map_service_error)?;
+    Ok(Json(resp))
+}
+
+/// `POST /agents/job`: move the job to the next phase. Only `Complete` returns data.
+async fn report_job(
+    Extension(a): Extension<AuthAgent>,
+    State(pool): State<InfraPool>,
+    Json(req): Json<JobReportReq>,
+) -> Result<Json<JobReportResp>, ApiError> {
+    let next_suite_available = match req.op {
+        JobReportOp::Start => {
+            service::agent::agent_start_job(a.id, &pool, req.job)
+                .await
+                .map_err(map_service_error)?;
+            None
+        }
+        JobReportOp::EnterCleanup => {
+            service::agent::agent_enter_cleanup(a.id, &pool, req.job)
+                .await
+                .map_err(map_service_error)?;
+            None
+        }
+        JobReportOp::Complete { outcome } => Some(
+            service::agent::agent_complete_job(a.id, a.uuid, &pool, req.job, outcome)
+                .await
+                .map_err(map_service_error)?,
+        ),
+    };
+    Ok(Json(JobReportResp {
+        next_suite_available,
+    }))
+}
+
+/// `GET /agents/tasks?suite_uuid=…`: claim ready tasks of the suite.
+#[tracing::instrument(skip_all, fields(agent_uuid = %a.uuid, suite_uuid = %req.suite_uuid))]
+async fn fetch_tasks(
+    Extension(a): Extension<AuthAgent>,
+    State(pool): State<InfraPool>,
+    Query(req): Query<FetchTasksReq>,
+) -> Result<Json<FetchTasksResp>, ApiError> {
+    let resp = service::agent::task::agent_fetch_tasks(a.id, a.uuid, &pool, req.suite_uuid)
+        .await
+        .map_err(map_service_error)?;
+    Ok(Json(resp))
+}

netmito/src/api/groups.rs

@@ -8,13 +8,13 @@ use axum_extra::extract::Query;
 
 use crate::{
     config::InfraPool,
-    error::{ApiError, ApiResult, AuthError},
+    error::{ApiError, ApiResult},
     schema::*,
     service::{
         self,
         auth::{
             auth_middleware,
-            roles::{USER, WORKER},
+            roles::{AGENT, USER, WORKER},
             AuthIdentity, AuthUser,
         },
     },
@@ -28,7 +28,7 @@ pub fn groups_router(st: InfraPool) -> Router<InfraPool> {
         )
         .route_layer(middleware::from_fn_with_state(
             st.clone(),
-            auth_middleware::<{ USER | WORKER }>,
+            auth_middleware::<{ USER | WORKER | AGENT }>,
         ));
     let user_router = Router::new()
         .route("/", post(create_group))
@@ -195,9 +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()),
+        AuthIdentity::Agent(a) => {
+            service::s3::agent_get_attachment(&pool, a.id, group_name, key).await
+        }
     }
     .map_err(|e| match e {
         crate::error::Error::AuthError(err) => ApiError::AuthError(err),

netmito/src/api/suites.rs

@@ -12,8 +12,9 @@ use crate::{
     error::ApiError,
     schema::{
         CancelTaskSuiteParam, ChangeTaskSuiteReq, CreateTaskSuiteReq, CreateTaskSuiteResp,
-        SuiteAgentOverrideReq, SuiteAgentOverrideResp, TaskSuiteQueryResp, TaskSuitesQueryReq,
-        TaskSuitesQueryResp,
+        StopAgentJobReq, StopAgentJobResp, SuiteAgentOverrideReq, SuiteAgentOverrideResp,
+        SuiteJobQueryResp, SuiteJobsQueryReq, SuiteJobsQueryResp, TaskSuiteQueryResp,
+        TaskSuitesQueryReq, TaskSuitesQueryResp,
     },
     service::{
         self,
@@ -33,6 +34,11 @@ pub fn suites_router(st: InfraPool) -> Router<InfraPool> {
         )
         .route("/{uuid}/close", post(close_suite))
         .route("/{uuid}/agents/override", post(override_agents_for_suite))
+        .route("/{uuid}/jobs/query", post(query_suite_jobs))
+        .route(
+            "/{uuid}/jobs/{job_id}",
+            get(get_suite_job).delete(stop_suite_job),
+        )
         .route_layer(middleware::from_fn_with_state(
             st.clone(),
             auth_middleware::<USER>,
@@ -119,3 +125,42 @@ pub async fn override_agents_for_suite(
         .map_err(map_service_error)?;
     Ok(Json(resp))
 }
+
+/// `POST /suites/{uuid}/jobs/query`: the jobs of the suite, with a filter.
+pub async fn query_suite_jobs(
+    Extension(u): Extension<AuthUser>,
+    State(pool): State<InfraPool>,
+    Path(uuid): Path<Uuid>,
+    Json(req): Json<SuiteJobsQueryReq>,
+) -> Result<Json<SuiteJobsQueryResp>, ApiError> {
+    let resp = service::suite::user_query_suite_jobs(u.id, &pool, uuid, req)
+        .await
+        .map_err(map_service_error)?;
+    Ok(Json(resp))
+}
+
+/// `GET /suites/{uuid}/jobs/{job_id}`: one job with its hooks.
+pub async fn get_suite_job(
+    Extension(u): Extension<AuthUser>,
+    State(pool): State<InfraPool>,
+    Path((uuid, job_id)): Path<(Uuid, i32)>,
+) -> Result<Json<SuiteJobQueryResp>, ApiError> {
+    let resp = service::suite::user_get_suite_job(u.id, &pool, uuid, job_id)
+        .await
+        .map_err(map_service_error)?;
+    Ok(Json(resp))
+}
+
+/// `DELETE /suites/{uuid}/jobs/{job_id}?op=graceful|force`: stop the job. The agent stays up and
+/// picks a suite again, possibly this suite.
+pub async fn stop_suite_job(
+    Extension(u): Extension<AuthUser>,
+    State(pool): State<InfraPool>,
+    Path((uuid, job_id)): Path<(Uuid, i32)>,
+    Query(req): Query<StopAgentJobReq>,
+) -> Result<Json<StopAgentJobResp>, ApiError> {
+    let resp = service::suite::user_stop_suite_job(u.id, &pool, uuid, job_id, req.op)
+        .await
+        .map_err(map_service_error)?;
+    Ok(Json(resp))
+}

netmito/src/api/tasks.rs

@@ -9,14 +9,14 @@ use uuid::Uuid;
 use crate::{
     config::InfraPool,
     entity::content::ArtifactContentType,
-    error::ApiError,
+    error::{ApiError, AuthError},
     schema::*,
     service::{
         self,
         auth::{
             auth_middleware,
-            roles::{USER, WORKER},
-            AuthIdentity, AuthUser, AuthWorker,
+            roles::{AGENT, USER, WORKER},
+            AuthIdentity, AuthUser,
         },
     },
 };
@@ -30,13 +30,13 @@ pub fn tasks_router(st: InfraPool) -> Router<InfraPool> {
         )
         .route_layer(middleware::from_fn_with_state(
             st.clone(),
-            auth_middleware::<{ USER | WORKER }>,
+            auth_middleware::<{ USER | WORKER | AGENT }>,
         ));
-    let worker_router = Router::new()
+    let runner_router = Router::new()
         .route("/report", post(report_task))
         .route_layer(middleware::from_fn_with_state(
             st.clone(),
-            auth_middleware::<WORKER>,
+            auth_middleware::<{ WORKER | AGENT }>,
         ));
     let user_router = Router::new()
         .route("/", post(submit_task))
@@ -67,7 +67,7 @@ pub fn tasks_router(st: InfraPool) -> Router<InfraPool> {
         ));
     Router::new()
         .merge(shared_router)
-        .merge(worker_router)
+        .merge(runner_router)
         .merge(user_router)
         .with_state(st)
 }
@@ -164,12 +164,41 @@ pub async fn query_task(
     Ok(Json(task))
 }
 
+/// `POST /tasks/report`: a worker reports its task, an agent reports a task or a hook of its job.
 pub async fn report_task(
-    Extension(w): Extension<AuthWorker>,
+    Extension(p): Extension<AuthIdentity>,
     State(pool): State<InfraPool>,
-    Json(ReportTaskReq { id, op }): Json<ReportTaskReq>,
+    Json(ReportTaskReq { id, kind, job, op }): Json<ReportTaskReq>,
 ) -> Result<Json<ReportTaskResp>, ApiError> {
-    match service::worker::report_task(w.id, w.uuid, id, op, &pool).await {
+    let result = match (p, kind) {
+        (AuthIdentity::Worker(w), TaskKind::Task) => {
+            if job.is_some() {
+                return Err(ApiError::InvalidRequest(
+                    "A worker task has no job".to_string(),
+                ));
+            }
+            service::worker::report_task(w.id, w.uuid, id, op, &pool).await
+        }
+        (AuthIdentity::Worker(_), TaskKind::Hook) => {
+            return Err(ApiError::AuthError(AuthError::PermissionDenied));
+        }
+        (AuthIdentity::Agent(a), TaskKind::Task) => {
+            let Some(job) = job else {
+                return Err(ApiError::InvalidRequest(
+                    "An agent task report must give its job".to_string(),
+                ));
+            };
+            service::agent::task::agent_report_task(a.id, a.uuid, job, id, op, &pool).await
+        }
+        (AuthIdentity::Agent(a), TaskKind::Hook) => {
+            service::agent::hook::agent_report_hook(a.id, id, op, &pool).await
+        }
+        // The middleware of this route admits only workers and agents.
+        (AuthIdentity::User(_), _) => {
+            return Err(ApiError::AuthError(AuthError::PermissionDenied));
+        }
+    };
+    match result {
         Ok(url) => Ok(Json(ReportTaskResp { url })),
         Err(crate::error::Error::ApiError(e)) => Err(e),
         Err(e) => {