fix(suite): resolve PR comments

Author: un-lock-able

netmito/src/api/suites.rs

@@ -8,12 +8,10 @@ use uuid::Uuid;
 
 use crate::{
     config::InfraPool,
-    entity::task_suite_agent::SuiteAgentSelectionType,
     error::ApiError,
     schema::{
-        CancelSuiteResp, CancelTaskSuiteParam, CreateTaskSuiteReq, CreateTaskSuiteResp,
-        ResetSuiteAgentResp, SuiteAgentReq, SuiteAgentResp, TaskSuiteQueryResp, TaskSuitesQueryReq,
-        TaskSuitesQueryResp,
+        CancelTaskSuiteParam, CreateTaskSuiteReq, CreateTaskSuiteResp, SuiteAgentSelectionReq,
+        SuiteAgentSelectionResp, TaskSuiteQueryResp, TaskSuitesQueryReq, TaskSuitesQueryResp,
     },
     service::{
         self,
@@ -27,9 +25,7 @@ pub fn suites_router(st: InfraPool) -> Router<InfraPool> {
         .route("/query", post(query_suites))
         .route("/{uuid}", get(get_suite_details).delete(cancel_suite))
         .route("/{uuid}/close", post(close_suite))
-        .route("/{uuid}/agents/include", post(include_suite_agents))
-        .route("/{uuid}/agents/exclude", post(exclude_suite_agents))
-        .route("/{uuid}/agents/reset", post(reset_suite_agent))
+        .route("/{uuid}/agents/selection", post(select_suite_agents))
         .route_layer(middleware::from_fn_with_state(
             st.clone(),
             user_auth_middleware,
@@ -98,57 +94,20 @@ pub async fn cancel_suite(
     State(pool): State<InfraPool>,
     Path(uuid): Path<Uuid>,
     Query(param): Query<CancelTaskSuiteParam>,
-) -> Result<Json<CancelSuiteResp>, ApiError> {
-    let resp =
-        service::suite::user_cancel_task_suite(u.id, &pool, uuid, param.op.unwrap_or_default())
-            .await
-            .map_err(map_service_error)?;
-    Ok(Json(resp))
-}
-
-pub async fn include_suite_agents(
-    Extension(u): Extension<AuthUser>,
-    State(pool): State<InfraPool>,
-    Path(uuid): Path<Uuid>,
-    Json(req): Json<SuiteAgentReq>,
-) -> Result<Json<SuiteAgentResp>, ApiError> {
-    let resp = service::suite::user_set_suite_agent(
-        u.id,
-        &pool,
-        uuid,
-        req.agent_uuid,
-        SuiteAgentSelectionType::UserIncluded,
-    )
-    .await
-    .map_err(map_service_error)?;
-    Ok(Json(resp))
-}
-
-pub async fn exclude_suite_agents(
-    Extension(u): Extension<AuthUser>,
-    State(pool): State<InfraPool>,
-    Path(uuid): Path<Uuid>,
-    Json(req): Json<SuiteAgentReq>,
-) -> Result<Json<SuiteAgentResp>, ApiError> {
-    let resp = service::suite::user_set_suite_agent(
-        u.id,
-        &pool,
-        uuid,
-        req.agent_uuid,
-        SuiteAgentSelectionType::UserExcluded,
-    )
-    .await
-    .map_err(map_service_error)?;
-    Ok(Json(resp))
+) -> Result<(), ApiError> {
+    service::suite::user_cancel_task_suite(u.id, &pool, uuid, param.op.unwrap_or_default())
+        .await
+        .map_err(map_service_error)?;
+    Ok(())
 }
 
-pub async fn reset_suite_agent(
+pub async fn select_suite_agents(
     Extension(u): Extension<AuthUser>,
     State(pool): State<InfraPool>,
     Path(uuid): Path<Uuid>,
-    Json(req): Json<SuiteAgentReq>,
-) -> Result<Json<ResetSuiteAgentResp>, ApiError> {
-    let resp = service::suite::user_reset_suite_agent(u.id, &pool, uuid, req.agent_uuid)
+    Json(req): Json<SuiteAgentSelectionReq>,
+) -> Result<Json<SuiteAgentSelectionResp>, ApiError> {
+    let resp = service::suite::user_add_agents_to_suite(u.id, &pool, uuid, req)
         .await
         .map_err(map_service_error)?;
     Ok(Json(resp))

netmito/src/client/http.rs

@@ -1,3 +1,4 @@
+use std::collections::HashMap;
 use std::path::PathBuf;
 
 use figment::value::magic::RelativePathBuf;
@@ -1394,7 +1395,7 @@ impl MitoHttpClient {
         &mut self,
         uuid: Uuid,
         force: bool,
-    ) -> crate::error::Result<CancelSuiteResp> {
+    ) -> crate::error::Result<()> {
         self.url.set_path(&format!("suites/{uuid}"));
         if force {
             self.url.set_query(Some("op=force"));
@@ -1409,61 +1410,32 @@ impl MitoHttpClient {
         // Clear the query so it does not leak into later requests reusing self.url.
         self.url.set_query(None);
         if resp.status().is_success() {
-            let resp = resp
-                .json::<CancelSuiteResp>()
-                .await
-                .map_err(RequestError::from)?;
-            Ok(resp)
-        } else {
-            Err(get_error_from_resp(resp).await.into())
-        }
-    }
-
-    /// Include (`include == true`) or exclude an agent on a suite.
-    pub async fn set_suite_agent(
-        &mut self,
-        uuid: Uuid,
-        include: bool,
-        req: SuiteAgentReq,
-    ) -> crate::error::Result<SuiteAgentResp> {
-        let op = if include { "include" } else { "exclude" };
-        self.url.set_path(&format!("suites/{uuid}/agents/{op}"));
-        let resp = self
-            .http_client
-            .post(self.url.as_str())
-            .bearer_auth(&self.credential)
-            .json(&req)
-            .send()
-            .await
-            .map_err(map_reqwest_err)?;
-        if resp.status().is_success() {
-            let resp = resp
-                .json::<SuiteAgentResp>()
-                .await
-                .map_err(RequestError::from)?;
-            Ok(resp)
+            Ok(())
         } else {
             Err(get_error_from_resp(resp).await.into())
         }
     }
 
-    pub async fn reset_suite_agent(
+    /// Batch-set agent selection overrides on a suite. Each entry pins (`Include`),
+    /// blocks (`Exclude`), or clears the override for (`Match`) one agent.
+    pub async fn select_suite_agents(
         &mut self,
         uuid: Uuid,
-        req: SuiteAgentReq,
-    ) -> crate::error::Result<ResetSuiteAgentResp> {
-        self.url.set_path(&format!("suites/{uuid}/agents/reset"));
+        selection: HashMap<Uuid, SuiteAgentSelectionAction>,
+    ) -> crate::error::Result<SuiteAgentSelectionResp> {
+        self.url
+            .set_path(&format!("suites/{uuid}/agents/selection"));
         let resp = self
             .http_client
             .post(self.url.as_str())
             .bearer_auth(&self.credential)
-            .json(&req)
+            .json(&SuiteAgentSelectionReq { selection })
             .send()
             .await
             .map_err(map_reqwest_err)?;
         if resp.status().is_success() {
             let resp = resp
-                .json::<ResetSuiteAgentResp>()
+                .json::<SuiteAgentSelectionResp>()
                 .await
                 .map_err(RequestError::from)?;
             Ok(resp)

netmito/src/client/interactive.rs

@@ -229,7 +229,7 @@ pub(crate) fn output_group_info(info: &GroupQueryInfo) {
     }
 }
 
-pub(crate) fn output_suite_list_info(info: &TaskSuiteInfo) {
+pub(crate) fn output_suite_info(info: &TaskSuiteInfo) {
     tracing::info!("Suite UUID: {}", info.uuid);
     if let Some(ref name) = info.name {
         tracing::info!("Name: {}", name);

netmito/src/client/mod.rs

@@ -1,4 +1,4 @@
-use std::{io::Write, process::Stdio};
+use std::{collections::HashMap, io::Write, process::Stdio};
 
 use clap_repl::ReadCommandOutput;
 use http::MitoHttpClient;
@@ -824,36 +824,32 @@ impl MitoClient {
         self.http_client.close_task_suite(args.uuid).await
     }
 
-    pub async fn suites_cancel(
-        &mut self,
-        args: CancelSuiteArgs,
-    ) -> crate::error::Result<CancelSuiteResp> {
+    pub async fn suites_cancel(&mut self, args: CancelSuiteArgs) -> crate::error::Result<()> {
         self.http_client
             .cancel_task_suite(args.uuid, args.force)
             .await
     }
 
-    pub async fn suites_set_agent(
+    /// Apply a single-agent selection override, reporting the per-agent outcome. Backed
+    /// by the batch endpoint with a one-entry request.
+    async fn suites_select_agent(
         &mut self,
         args: SuiteAgentArgs,
-        include: bool,
-    ) -> crate::error::Result<SuiteAgentResp> {
-        let req = SuiteAgentReq {
-            agent_uuid: args.agent,
-        };
-        self.http_client
-            .set_suite_agent(args.uuid, include, req)
-            .await
-    }
-
-    pub async fn suites_reset_agent(
-        &mut self,
-        args: SuiteAgentArgs,
-    ) -> crate::error::Result<ResetSuiteAgentResp> {
-        let req = SuiteAgentReq {
-            agent_uuid: args.agent,
-        };
-        self.http_client.reset_suite_agent(args.uuid, req).await
+        action: SuiteAgentSelectionAction,
+    ) {
+        let (suite, agent) = (args.uuid, args.agent);
+        let selection = HashMap::from([(agent, action)]);
+        match self.http_client.select_suite_agents(suite, selection).await {
+            Ok(resp) => match resp.failed.get(&agent) {
+                Some(err) => {
+                    tracing::error!(
+                        "Failed to apply {action:?} to agent {agent} on suite {suite}: {err:?}"
+                    )
+                }
+                None => tracing::info!("Applied {action:?} to agent {agent} on suite {suite}"),
+            },
+            Err(e) => tracing::error!("{}", e),
+        }
     }
 
     pub async fn tasks_batch_cancel(
@@ -2016,9 +2012,16 @@ impl MitoClient {
                             if !counted {
                                 for suite in resp.suites {
                                     if verbose {
-                                        output_suite_list_info(&suite);
+                                        output_suite_info(&suite);
                                     } else {
-                                        tracing::info!("{} ({})", suite.uuid, suite.state);
+                                        tracing::info!(
+                                            "{}, {} ({}), Tasks (incomplete/total): {} / {}",
+                                            suite.uuid,
+                                            suite.name.unwrap_or("[no name]".to_string()),
+                                            suite.state,
+                                            suite.incomplete_tasks,
+                                            suite.total_tasks
+                                        );
                                     }
                                 }
                             }
@@ -2030,7 +2033,7 @@ impl MitoClient {
                 }
                 SuitesCommands::Get(args) => match self.suites_get(args).await {
                     Ok(resp) => {
-                        output_parsed_suite_info(&resp.info, &resp.assigned_agents);
+                        output_parsed_suite_info(&resp.info, &resp.eligible_agents);
                     }
                     Err(e) => {
                         tracing::error!("{}", e);
@@ -2045,58 +2048,25 @@ impl MitoClient {
                     }
                 },
                 SuitesCommands::Cancel(args) => match self.suites_cancel(args).await {
-                    Ok(resp) => {
-                        tracing::info!(
-                            "Suite cancelled; {} tasks cancelled",
-                            resp.cancelled_task_count
-                        );
+                    Ok(_) => {
+                        tracing::info!("Suite cancelled");
                     }
                     Err(e) => {
                         tracing::error!("{}", e);
                     }
                 },
                 SuitesCommands::IncludeAgent(args) => {
-                    match self.suites_set_agent(args, true).await {
-                        Ok(resp) => {
-                            tracing::info!(
-                                "Agent {} set to {:?} on suite {}",
-                                resp.agent_uuid,
-                                resp.selection,
-                                resp.suite_uuid
-                            );
-                        }
-                        Err(e) => {
-                            tracing::error!("{}", e);
-                        }
-                    }
+                    self.suites_select_agent(args, SuiteAgentSelectionAction::Include)
+                        .await;
                 }
                 SuitesCommands::ExcludeAgent(args) => {
-                    match self.suites_set_agent(args, false).await {
-                        Ok(resp) => {
-                            tracing::info!(
-                                "Agent {} set to {:?} on suite {}",
-                                resp.agent_uuid,
-                                resp.selection,
-                                resp.suite_uuid
-                            );
-                        }
-                        Err(e) => {
-                            tracing::error!("{}", e);
-                        }
-                    }
+                    self.suites_select_agent(args, SuiteAgentSelectionAction::Exclude)
+                        .await;
+                }
+                SuitesCommands::ResetAgent(args) => {
+                    self.suites_select_agent(args, SuiteAgentSelectionAction::Match)
+                        .await;
                 }
-                SuitesCommands::ResetAgent(args) => match self.suites_reset_agent(args).await {
-                    Ok(resp) => {
-                        if resp.reset {
-                            tracing::info!("Reset the agent to the tag-match default");
-                        } else {
-                            tracing::info!("No manual override to reset");
-                        }
-                    }
-                    Err(e) => {
-                        tracing::error!("{}", e);
-                    }
-                },
             },
             ClientCommand::Quit => {
                 return false;

netmito/src/config/client/tasks.rs

@@ -50,7 +50,7 @@ pub struct SubmitTaskArgs {
     #[arg(short = 'g', long = "group")]
     pub group_name: Option<String>,
     /// The UUID of the task suite to submit this task to. When set, the task is executed
-    /// by the suite's agents instead of traditional workers, and inherits the suite's tags.
+    /// by the suite's agents instead of workers
     #[arg(long = "suite")]
     pub suite_uuid: Option<Uuid>,
     /// The tags of the task, used to filter workers to execute the task