Add an RPC endpoint to report scout firmware upgrade status

Author: rahmonov

crates/api-model/src/machine/mod.rs

@@ -1576,6 +1576,22 @@ pub enum HostReprovisionState {
         reason: Option<String>,
     },
     WaitingForRackFirmwareUpgrade,
+    WaitingForScoutUpgrade {
+        component_type: FirmwareComponentType,
+        target_version: String,
+        started_at: DateTime<Utc>,
+        #[serde(default)]
+        result: Option<ScoutUpgradeResult>,
+    },
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)]
+pub struct ScoutUpgradeResult {
+    pub success: bool,
+    pub exit_code: i32,
+    pub stdout: String,
+    pub stderr: String,
+    pub error: String,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]

crates/api/src/api.rs

@@ -1145,6 +1145,14 @@ impl Forge for Api {
             .await
     }
 
+    async fn report_scout_firmware_upgrade_status(
+        &self,
+        request: Request<rpc::ScoutFirmwareUpgradeStatusRequest>,
+    ) -> Result<Response<()>, Status> {
+        crate::handlers::host_reprovisioning::report_scout_firmware_upgrade_status(self, request)
+            .await
+    }
+
     async fn list_hosts_waiting_for_reprovisioning(
         &self,
         request: Request<rpc::HostReprovisioningListRequest>,
@@ -3253,7 +3261,7 @@ pub(crate) fn log_tenant_organization_id(organization_id: &str) {
     tracing::Span::current().record("tenant.organization_id", organization_id);
 }
 
-fn truncate(mut s: String, len: usize) -> String {
+pub(crate) fn truncate(mut s: String, len: usize) -> String {
     if s.len() < len || len < 3 {
         return s;
     }

crates/api/src/auth/internal_rbac_rules.rs

@@ -168,6 +168,7 @@ impl InternalRBACRules {
         x.perm("DiscoveryCompleted", vec![Machineatron, Scout]);
         x.perm("CleanupMachineCompleted", vec![Machineatron, Scout]);
         x.perm("ReportForgeScoutError", vec![Scout]);
+        x.perm("ReportScoutFirmwareUpgradeStatus", vec![Scout]);
         x.perm("DiscoverDhcp", vec![Dhcp, Machineatron]);
         x.perm("ExpireDhcpLease", vec![Dhcp, Machineatron]);
         x.perm("AssignStaticAddress", vec![ForgeAdminCLI]);

crates/api/src/handlers/host_reprovisioning.rs

@@ -17,11 +17,13 @@
 use ::rpc::forge as rpc;
 use carbide_uuid::machine::MachineId;
 use itertools::Itertools;
-use model::machine::LoadSnapshotOptions;
+use model::machine::{
+    HostReprovisionState, LoadSnapshotOptions, ManagedHostState, ScoutUpgradeResult,
+};
 use tonic::{Request, Response, Status};
 
 use crate::CarbideError;
-use crate::api::{Api, log_request_data};
+use crate::api::{Api, log_request_data, truncate};
 use crate::handlers::utils::convert_and_log_machine_id;
 
 pub(crate) async fn reset_host_reprovisioning(
@@ -144,3 +146,72 @@ pub async fn mark_manual_firmware_upgrade_complete(
 
     Ok(Response::new(()))
 }
+
+pub async fn report_scout_firmware_upgrade_status(
+    api: &Api,
+    request: Request<rpc::ScoutFirmwareUpgradeStatusRequest>,
+) -> Result<Response<()>, Status> {
+    log_request_data(&request);
+
+    let req = request.into_inner();
+    let machine_id = convert_and_log_machine_id(req.machine_id.as_ref())?;
+
+    let (machine, mut txn) = api.load_machine(&machine_id, Default::default()).await?;
+
+    // Verify machine is in WaitingForScoutUpgrade state
+    let (component_type, target_version, started_at, retry_count) = match machine.current_state() {
+        ManagedHostState::HostReprovision {
+            reprovision_state:
+                HostReprovisionState::WaitingForScoutUpgrade {
+                    component_type,
+                    target_version,
+                    started_at,
+                    ..
+                },
+            retry_count,
+        } => (
+            component_type.clone(),
+            target_version.clone(),
+            *started_at,
+            *retry_count,
+        ),
+        _ => {
+            return Err(CarbideError::InvalidArgument(format!(
+                "Machine {machine_id} is not in WaitingForScoutUpgrade state"
+            ))
+            .into());
+        }
+    };
+
+    const MAX_STORED_OUTPUT_SIZE: usize = 1500;
+
+    let new_state = ManagedHostState::HostReprovision {
+        reprovision_state: HostReprovisionState::WaitingForScoutUpgrade {
+            component_type,
+            target_version,
+            started_at,
+            result: Some(ScoutUpgradeResult {
+                success: req.success,
+                exit_code: req.exit_code,
+                stdout: truncate(req.stdout, MAX_STORED_OUTPUT_SIZE),
+                stderr: truncate(req.stderr, MAX_STORED_OUTPUT_SIZE),
+                error: truncate(req.error, MAX_STORED_OUTPUT_SIZE),
+            }),
+        },
+        retry_count,
+    };
+
+    db::machine::advance(&machine, &mut txn, &new_state, None).await?;
+
+    txn.commit().await?;
+
+    if let Err(err) = api
+        .machine_state_handler_enqueuer
+        .enqueue_object(&machine_id)
+        .await
+    {
+        tracing::warn!(%err, %machine_id, "Failed to wake up state handler for machine");
+    }
+
+    Ok(Response::new(()))
+}

crates/api/src/state_controller/machine/handler.rs

@@ -6653,6 +6653,10 @@ impl HostUpgradeState {
                 self.host_new_firmware_reported_wait(state, ctx, details, machine_id, scenario)
                     .await
             }
+            HostReprovisionState::WaitingForScoutUpgrade { .. } => {
+                // TODO: will be implemented in a follow-up (@jrakhmonov)
+                Ok(StateHandlerOutcome::do_nothing())
+            }
             HostReprovisionState::FailedFirmwareUpgrade { report_time, .. } => {
                 let can_retry = retry_count < MAX_FIRMWARE_UPGRADE_RETRIES;
                 let waited_enough = Utc::now()

crates/rpc/proto/forge.proto

@@ -315,6 +315,8 @@ service Forge {
   // TODO: Remove when manual upgrade feature is removed
   // Mark host as having completed manual firmware upgrade
   rpc MarkManualFirmwareUpgradeComplete(common.MachineId) returns (google.protobuf.Empty);
+  // Report the result of a scout-based firmware upgrade
+  rpc ReportScoutFirmwareUpgradeStatus(ScoutFirmwareUpgradeStatusRequest) returns (google.protobuf.Empty);
 
   rpc GetDpuInfoList(GetDpuInfoListRequest) returns (GetDpuInfoListResponse);
 
@@ -5162,6 +5164,16 @@ message MachineRebootCompletedResponse {
 message MachineRebootCompletedRequest {
   common.MachineId machine_id = 1;
 }
+
+message ScoutFirmwareUpgradeStatusRequest {
+  common.MachineId machine_id = 1;
+  bool success = 2;
+  int32 exit_code = 3;
+  string stdout = 4;
+  string stderr = 5;
+  string error = 6;
+}
+
 // enum MachineValidationState {
 //     Started = 0;
 //     InProgress = 1;