Update Task Heartbeat to include Job_id (#594)

src/agent/onefuzz-agent/src/tasks/config.rs

@@ -47,7 +47,7 @@ impl CommonConfig {
     pub async fn init_heartbeat(&self) -> Result<Option<TaskHeartbeatClient>> {
         match &self.heartbeat_queue {
             Some(url) => {
-                let hb = init_task_heartbeat(url.clone(), self.task_id).await?;
+                let hb = init_task_heartbeat(url.clone(), self.task_id, self.job_id).await?;
                 Ok(Some(hb))
             }
             None => Ok(None),

src/agent/onefuzz-agent/src/tasks/heartbeat.rs

@@ -26,18 +26,24 @@ struct Heartbeat {
 #[derive(Clone)]
 pub struct TaskContext {
     task_id: Uuid,
+    job_id: Uuid,
     machine_id: Uuid,
     machine_name: String,
 }
 
 pub type TaskHeartbeatClient = HeartbeatClient<TaskContext, HeartbeatData>;
 
-pub async fn init_task_heartbeat(queue_url: Url, task_id: Uuid) -> Result<TaskHeartbeatClient> {
+pub async fn init_task_heartbeat(
+    queue_url: Url,
+    task_id: Uuid,
+    job_id: Uuid,
+) -> Result<TaskHeartbeatClient> {
     let machine_id = get_machine_id().await?;
     let machine_name = get_machine_name().await?;
     let hb = HeartbeatClient::init_heartbeat(
         TaskContext {
             task_id,
+            job_id,
             machine_id,
             machine_name,
         },

src/api-service/__app__/queue_task_heartbeat/__init__.py

@@ -30,7 +30,7 @@ def main(msg: func.QueueMessage, dashboard: func.Out[str]) -> None:
             task.heartbeat = datetime.utcnow()
             task.save()
     except ValidationError:
-        logging.error("invalid task heartbat: %s", raw)
+        logging.error("invalid task heartbeat: %s", raw)
 
     events = get_events()
     if events:

src/pytypes/onefuzztypes/models.py

@@ -516,6 +516,7 @@ class Job(BaseModel):
 
 class TaskHeartbeatEntry(BaseModel):
     task_id: UUID
+    job_id: Optional[UUID]
     machine_id: UUID
     data: List[Dict[str, HeartbeatType]]